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

TL;DR

SAFE introduces a novel entropy-aware RLHF algorithm that improves training stability and reward performance over PPO by dynamically regulating KL divergence and employing a double critic for pessimistic value estimation.

Contribution

The paper presents SAFE, a new on-policy RLHF method with entropy-aware control, combining a double soft-min critic and adaptive KL regulation for stable, efficient alignment finetuning.

Findings

SAFE outperforms PPO with +5.15% reward improvement

SAFE exhibits negligible reward crashes during training

SAFE achieves superior KL control and stability

Abstract

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