PaperlayerView exact revision on arXiv

Source-linked AI summary

SAFE: Stable Alignment Finetuning with Entropy-Aware Predictive Control for Reinforcement Learning from Human Feedback (RLHF)

Dipan Maity

arXiv:2602.04651v2cs.LG

TL;DR

RLHF optimization can suffer reward volatility, entropy collapse, value drift, and late-stage policy divergence, while KL control alone does not regulate these dynamics. SAFE combines pessimistic value estimation with entropy-aware KL control and reward-driven threshold adaptation. Experiments report improved stability and reward performance, while limitations include higher value loss and remaining KL spikes.

  • Problem

    RLHF training remains unstable in long-horizon optimization, with reward oscillations, entropy collapse, value drift, and policy divergence that KL control alone does not fully address.

  • Method

    SAFE combines pessimistic double soft-min value estimation, entropy-aware KL regulation, and reward-driven PID-based adaptive thresholds.

  • Results

    SAFE improves training stability and reports a 5.2% reward improvement with 2.8× lower reward variance than PPO while avoiding catastrophic collapses.

  • Takeaways & Limitations

    SAFE’s primary benefit is reduced reward volatility and more consistent convergence through joint regulation of divergence, entropy dynamics, and value estimation.

  • Takeaways & Limitations

    Value loss remains higher than PPO, and KL spike frequency remains non-negligible despite reduced reward volatility.

Abstract

from arXiv · show

Proximal Policy Optimization (PPO) has been positioned by recent literature as the canonical method for the RL part of Reinforcement Learning from Human Feedback (RLHF). PPO performs well empirically but has a heuristic motivation and handles the KL-divergence constraint used in LM-RLHF in an ad-hoc manner and suffers form reward oscillations, entropy collapse, value function drift, and sudden policy divergence that require frequent restarts and extensive hyperparameter tuning. In this paper, we develop a new pure on policy actor-critic RL method for the LM-RLHF setting. We present SAFE (Stable Alignment Finetuning with Entropy-aware control),a novel RLHF algorithm that combines a Double Soft-Min Critic for pessimistic value estimation with a new multi-layer stabilization framework combining entropy-gated KL regulation, and PID-controlled adaptive thresholds. Unlike standard PPO's symmetric KL penalties, SAFE distinguishes high-entropy exploration from low-entropy mode collapse and adjusts penalties dynamically based on reward velocity. Experiments on a 3B parameter model show SAFE achieves +5.15\% training-average reward than PPO (0.725 vs 0.689), negligible reward crashes, and superior KL control than ppo . Our method adds minimal computational overhead and provides an interpretable, crash-resistant RLHF framework that maintains aggressive learning speed while ensuring stable long-horizon optimization suitable for production deployment. Code is available at https://github.com/ryyzn9/SAFE

1 Introduction

RLHF training can become unstable because KL control alone does not regulate value estimation, divergence dynamics, and exploration behavior together. SAFE addresses these coupled modes with a coordinated multi-layer stabilization framework.

  • Motivation: RLHF exhibits reward oscillations, entropy collapse, value function drift, and sudden policy divergence during long-horizon on-policy optimization.These failures can require frequent restarts and extensive hyperparameter tuning.
  • Motivation: KL regularization limits policy drift but may miss determinization, gradient sensitivity, and unstable value estimation within acceptable KL bounds.Static or symmetric divergence control therefore does not fully capture internal policy dynamics.
  • Motivation: Overestimating critics can amplify outlier rewards through inflated advantages, driving aggressive updates and cascading instability.Bootstrapping and function-approximation error make high-variance regions especially vulnerable.
  • Motivation: Late-stage instabilities can emerge after hundreds or thousands of optimization steps and evade instantaneous metrics or static regularization schedules.The paper motivates coordinated intervention across multiple optimization levels.
  • SAFE framework: SAFE integrates pessimistic value aggregation, entropy-conditioned divergence control, and predictive threshold adaptation across value estimation, divergence regulation, and training dynamics.The framework is designed as a unified alternative to relying on a single stabilizer.
  • Contributions: SAFE’s contributions target interacting instability modes and report improved stability and robustness without sacrificing reward performance.The paper also presents theoretical analysis suggesting partial resistance to reward hacking through the pessimistic critic mechanism.

2 Background

RLHF optimizes a policy against learned preference rewards while constraining deviation from a supervised reference, but PPO and static entropy control remain vulnerable to instability. Background mechanisms therefore motivate entropy-aware divergence regulation and pessimistic value estimation.

  • RLHF pipeline: RLHF uses supervised fine-tuning, reward-model training, and reinforcement-learning optimization as three sequential stages.The policy is optimized for reward-model outputs while remaining close to the supervised reference policy.
  • RLHF objective: The RLHF objective combines learned reward-model scores with KL regularization relative to a frozen supervised policy.The coefficient β controls divergence-regularization strength.
  • RLHF objective: KL control limits large policy updates and preserves linguistic quality by anchoring the learned policy to the pretrained distribution.This gives the constraint both an optimization and a quality-preservation role.
  • PPO limitations: PPO-based RLHF remains sensitive to entropy collapse, reward hacking, and late-stage instability.These limitations motivate more structured control mechanisms.
  • KL estimation: KL divergence in practical RLHF is estimated from sampled tokens using Monte Carlo log-probability ratios rather than computed exactly over the full vocabulary.Finite-sample estimates can be negative because of sampling variance and estimator bias.
  • KL estimation: Negative log-ratio estimates typically accompany high-entropy exploration, whereas positive estimates reflect increasing confidence and concentration on high-reward modes.Sustained positive estimates with decreasing entropy are observed before exploitative behavior and reward hacking.
  • Pessimistic critics: Value overestimation from bootstrapped updates and function-approximation error can make policy updates overly aggressive and lead to divergence.Double Q-learning addresses this by maintaining two critics and using pessimistic aggregation.
  • Pessimistic critics: Pessimistic aggregation reduces positive bias at the cost of a small underestimation bias and can improve stability and robustness.Twin-critic architectures are used in high-variance and nonstationary actor–critic settings.

3 Method:

SAFE combines pessimistic value estimation with asymmetric, momentum-aware divergence control to stabilize RLHF policy optimization. The reported asymmetric-controller results show stronger KL regulation but substantial value-learning and reward-performance trade-offs, motivating further stabilization.

  • Value-Level Stabilization via Pessimistic Aggregation: SAFE uses two critics with differentiable soft-min aggregation to suppress optimistic value outliers while preserving gradient continuity.Layer normalization and Polyak-averaged target networks further stabilize critic learning.
  • Entropy-Aware Predictive Control: The entropy-aware predictive controller adapts KL thresholds using reward velocity, entropy-conditioned penalty scaling, training-phase modulation, and gradient preview checks.These mechanisms target changing exploration and convergence regimes without hard entropy targets or rigid divergence constraints.
  • Asymmetric Divergence Control: SAFE selectively penalizes large positive KL divergence, tolerates small positive deviations, and applies zero penalty to negative KL estimates.A momentum term also enables intervention when divergence rises rapidly, even at moderate absolute KL.
  • Empirical Behavior: 0.241 mean absolute KL divergence under Asymmetric KL exceeded PPO’s 0.131, while only 28.6% versus 55.9% of steps exceeded KL > 0.2.KL volatility was also reduced by 44%, with rolling standard deviation 0.294 versus 0.526.
  • Empirical Behavior: 0.672 mean reward under Double Soft-Min Q with Asymmetric KL was below PPO’s 0.689, while both methods exhibited 2 reward crashes exceeding 20% loss.The asymmetric controller also produced 1256 value-loss spikes above 0.1 versus 28 under PPO, with mean value loss 0.226 versus 0.006.
  • Empirical Behavior: These results indicate that asymmetric KL control improves divergence regulation but, combined with pessimistic critics, can destabilize value learning and sacrifice reward performance.The findings motivate the entropy-aware predictive controller as a coordinated policy-and-critic stabilization layer.

3.4 Entropy-Aware Predictive Controller: Motivation and Design

The Entropy-Aware Predictive Controller treats KL regulation as a dynamic control problem shaped by entropy, reward velocity, and training phase. It combines multi-timescale tracking, PID threshold adaptation, entropy-scaled penalties, and preview-based safeguards.

  • Motivation: KL magnitude alone may miss policy determinization, gradient sensitivity, and value instability, even when divergence remains within acceptable bounds.Entropy is therefore used as a modulating signal rather than assumed to directly cause reward outcomes.
  • Multi-timescale Tracking: The controller maintains short- and long-timescale KL exponential moving averages to distinguish transient spikes from slower distributional drift.This separation reduces overreaction to single-minibatch noise.
  • Reward-Driven Predictive Threshold Adaptation: PID control adjusts divergence tolerance from reward improvement velocity rather than absolute reward magnitude.The target improvement rate is vtarget = 0.001, or 0.1% per step, using exponentially smoothed rewards with decay factor 0.95.
  • Reward-Driven Predictive Threshold Adaptation: When improvement exceeds target, the threshold increases to permit exploration; when improvement falls below target, it decreases to tighten KL constraints.The resulting feedback loop relaxes constraints during productive improvement and tightens them during plateaus or instability.
  • Entropy-Based Penalty Modulation: Entropy-conditioned scaling applies stronger corrective pressure in low-entropy regimes and gentler regulation to higher-entropy policies.The mechanism modulates penalty strength without directly predicting entropy collapse.
  • Phase Awareness and Preview Control: A phase detector applies threshold multipliers for early exploration, stagnation, and convergence, while gradient preview scales down updates that would exceed a maximum KL limit.The framework reports adaptive regulation, robustness to noisy estimates, stabilized confidence–divergence interaction, and protection against post-update KL spikes.

4 Components

SAFE combines pessimistic value estimation, entropy-aware KL regulation, and adaptive training-dynamics control in a unified on-policy RLHF framework. Its phase-aware PID controller adjusts divergence tolerance using reward progress while additional critic and optimization safeguards improve robustness.

  • Framework: SAFE extends PPO-style optimization with stabilization at the value, policy-divergence, and training-dynamics levels.The framework combines pessimistic value aggregation, entropy-conditioned divergence control, and predictive threshold adaptation.
  • Value estimation: Twin critics use differentiable pessimistic aggregation to suppress optimistic value-estimation errors and stabilize advantage estimation.The soft-min operator preserves smooth gradients while reducing overestimation bias.
  • Policy control: Entropy-gated KL regulation amplifies penalties below the entropy floor and attenuates them when entropy remains high.This adapts divergence-control strength to the current exploration regime.
  • Adaptive thresholds: PID control adjusts the KL threshold from reward velocity, relaxing constraints during improvement and tightening them during stagnation or decline.The proportional, integral, and derivative terms respond to current gaps, persistent under-performance, and acceleration changes.
  • Phase detection: Phase-aware threshold modulation improves controller robustness without task-specific heuristics or manual scheduling.A sliding-window detector applies multiplicative phase adjustments across exploration, climbing, plateau, and convergence regimes.
  • Additional safeguards: Clipped Huber value loss, gradient clipping, synchronized learning rates, entropy regularization, Polyak targets, and dual-timescale KL tracking provide additional stabilization.These mechanisms limit outlier-driven critic gradients, pathological updates, bootstrap noise, and misleading short-term KL signals.

5 Training Algorithm

SAFE training follows an on-policy PPO-style loop that evaluates twin critics, computes pessimistic standardized advantages, and updates policy and critics using adaptive control signals.

  • Rollout and estimation: Each training step samples prompts, collects on-policy data, and evaluates the two critics and target critics.The algorithm then computes pessimistic values before policy optimization.
  • Policy update: Pessimistic critic estimates produce standardized advantages used across multiple clipped PPO epochs.Importance ratios and the clipped PPO loss are computed for each update epoch.
  • Control and optimization: The controller updates PID state from reward velocity before computing the adaptive KL threshold and entropy-gated penalty.The same step computes the clipped Huber value loss and updates policy and critic parameters.
  • Stabilization: Gradient norm clipping and Polyak averaging are applied after optimization to limit update magnitude and stabilize target critics.These operations complete the training step.

6 Overall Objective Function

SAFE optimizes a composite objective that balances PPO policy improvement, robust value learning, adaptive divergence control, and entropy-preserving exploration.

  • Objective composition: The composite loss combines policy improvement, value learning, divergence control, and exploration terms.This extends PPO with pessimistic value aggregation, Huber robustness, and entropy-aware KL control.
  • Policy loss: The PPO policy term uses an importance ratio and advantages computed from the pessimistic soft-min critic.The surrogate objective remains clipped as in PPO.
  • Value loss: The value term uses clipped Huber regression to cap gradient effects from reward-model outliers.This supports robust critic learning under heavy-tailed rewards.
  • Divergence and entropy: The KL term applies entropy-gated, adaptively thresholded divergence control with a PID-adjusted threshold and entropy-dependent scaling.The entropy bonus maintains baseline exploration.

7 Experiments

Experiments evaluate SAFE against PPO on a 3B-parameter Qwen model using matched training conditions. SAFE reports higher reward and substantially improved stability while maintaining comparable KL and near-identical computational resource usage, but the gains are demonstrated on one benchmark configuration.

  • Experimental setup: SAFE is evaluated with identical computational budgets, datasets, and optimization settings to PPO.The setup uses Qwen2.5-3B with LoRA, a frozen reference policy, ArmoRM-Llama3-8B-v0.1, and 5,000 training prompts plus 500 held-out samples.
  • Experimental setup: The experiment runs 2,000 on-policy update steps with batch size 16, gradient accumulation 2, synchronized learning rates of 1 × 10^-5, and gradient clipping.SAFE includes double soft-min critics, PID-controlled KL thresholds, and phase-aware modulation.
  • Overall results: SAFE achieves higher average reward than PPO while maintaining comparable KL divergence and substantially reducing reward variance.Aggregate statistics are reported over the full 2,000-step run.
  • Training stability: The reward coefficient of variation decreases from 0.114 under PPO to 0.040 under SAFE, while rolling reward standard deviation decreases from 0.0208 to 0.0123.PPO has two reward-collapse events; SAFE has none.
  • Training stability: SAFE reduces rolling KL standard deviation to 0.306 versus 0.526 for PPO while keeping average KL differences statistically nonsignificant.The reported KL comparison has p = 0.10.
  • Convergence behavior: SAFE increases average reward from 0.711 in the early phase to 0.731 in the late phase, while PPO retains higher variability and increasing late-stage divergence.SAFE’s KL trends downward across phases.
  • Convergence behavior: SAFE maintains entropy above its configured floor, stabilizes completion lengths earlier, and produces less oscillatory reward curves than PPO.The adaptive threshold tightens during plateaus and relaxes during early exploration.
  • Discussion: SAFE’s reported robustness improvements occur without substantially increasing average KL divergence.The paper attributes the result to combining pessimistic value estimation with entropy-aware predictive KL regulation.

8 Experimental Analysis

SAFE improves reward stability and divergence control through adaptive, entropy-aware regulation during training. Its benefits primarily appear as lower volatility, smoother convergence, and controlled entropy decay rather than substantially higher mean KL.

  • Reward progression: 0.7249 average reward versus 0.6894 for PPO, with reward standard deviation 0.0291 versus 0.0788.The lower variance indicates that SAFE primarily reduces training noise rather than merely increasing peak reward.
  • Divergence control: KL volatility decreases by approximately 42%, with rolling standard deviation 0.306 versus 0.526 for PPO.Mean KL magnitude remains comparable to PPO while divergence stays within a bounded operating region.
  • Entropy regulation: Entropy remains above the configured floor as entropy-aware gating amplifies KL penalties when entropy decreases.This allows entropy to decay gradually instead of collapsing abruptly.
  • Entropy regulation: Completion length stabilizes earlier and remains more consistent across training steps, indicating controlled policy determinization.SAFE achieves this without adding explicit entropy-maximization terms beyond the standard PPO entropy coefficient.
  • Value estimation: The double critic produces bounded value loss with smoother temporal behavior and fewer abrupt divergence events than PPO.Absolute value loss remains higher than PPO, reflecting the conservative value-estimation behavior.
  • Statistical analysis: Welch’s t-test and Mann–Whitney U test reject equal reward distributions with p < 10^-50, while KL differences are not statistically significant at p = 0.10.The reward effect size is Cohen’s d = 0.60, indicating a moderate practical improvement.
  • Overall assessment: SAFE improves stability through jointly regulating divergence, entropy dynamics, and value estimation, with reduced reward volatility and more consistent convergence.Adaptive modulation adjusts controller behavior to training dynamics rather than static scheduling rules.

9 Conclusion

SAFE combines pessimistic value estimation, entropy-gated KL control, and PID-based adaptive thresholds in a multi-layer RLHF stabilization framework. On a 3B model, it improves reward stability over PPO, although broader applicability remains unproven.

  • Contribution: SAFE combines pessimistic value estimation, entropy-gated KL control, and PID-based adaptive thresholds.The framework coordinates control across value, divergence, and temporal dynamics.
  • Empirical findings: SAFE reduces reward variance by 2.8× compared to PPO and achieves a statistically significant 5.2% reward improvement.The experiments are conducted on a 3B-parameter model.
  • Scope: Broader claims about production readiness or general applicability require validation across multiple scales, tasks, and longer training horizons.The reported evidence supports layered control in a specific experimental configuration.

10 Limitations

SAFE’s reported stability improvements are constrained by limited evaluation scope, persistent value instability, manual tuning, and the absence of component-isolation experiments. Future work therefore requires broader testing and systematic analysis.

  • Generalizability: Evaluation is restricted to 3B parameters, 2,000 steps, and one dataset-reward-model combination.Scaling beyond 7B parameters, dynamics beyond 5,000 steps, and other domains remain unvalidated.
  • Value stability: SAFE still exhibits 435 value-loss spikes versus PPO’s 28, indicating that soft-min aggregation does not eliminate value instability.The conservative critic introduces bias while reducing, rather than removing, instability.
  • Tuning: The method requires manual tuning of 8+ hyperparameters without systematic ablation or sensitivity analysis.The listed parameters include the entropy floor, PID gains, and phase thresholds.
  • Component analysis: The relative contributions of pessimistic critics, entropy gating, and PID control remain unclear because components were not isolated.This limits attribution of the observed improvements to individual mechanisms.
  • Future work: Future work requires multi-scale, long-horizon, diverse-task evaluation, systematic ablations, automated tuning, and direct reward-hacking tests.The proposed evaluation scope extends to 1B–70B models and 10,000+ training steps.

A Theoretical Analysis: Hypothesized Reward Hacking Resistance

SAFE’s potential resistance to reward hacking is argued mechanistically through pessimistic values, directional and entropy-gated penalties, and tighter reference anchoring. However, the paper explicitly treats this as a hypothesis because it lacks direct reward-hacking experiments.

  • Caveat: The proposed reward-hacking resistance remains theoretical because SAFE was not evaluated on deliberately hackable reward models.Rigorous assessment requires adversarial, out-of-distribution, and long-horizon experiments with explicitly hackable reward functions.
  • Pessimistic value estimation: Pessimistic double-soft-min values reduce the attractiveness of high-variance reward outliers during advantage computation.This may make policy updates less likely to chase exploitation spikes than a single overestimating critic.
  • Directional KL control: Directional KL control is intended to detect positive divergence combined with decreasing entropy before catastrophic exploitative drift.The asymmetric component penalizes positive deviations above a threshold while imposing zero penalty on negative estimates.
  • Entropy gating: Entropy-gated penalties amplify corrective pressure as policy entropy decreases, targeting deterministic behaviors associated with reward-hacking strategies.The mechanism allows relatively unpenalized exploration when entropy remains high.
  • Distributional anchoring: Lower KL volatility and tighter anchoring to the supervised reference may reduce the policy’s capacity to exploit distant reward-model artifacts.The reference policy has not been optimized against the reward model.
Loading 2602.04651v2…