Abstract
Reinforcement learning from human feedback (RLHF) has become the dominant paradigm for aligning large language models with human preferences. However, policy gradient methods such as PPO suffer from high variance gradient estimates, requiring careful hyperparameter tuning and extensive computational resources. We introduce GRADE (Gumbel-softmax Relaxation for Alignment via Differentiable Estimation), a method that replaces high-variance policy gradient estimation with direct backpropagation through a differentiable relaxation of the discrete token sampling process. Using the Gumbel-Softmax reparameterization with straight-through estimation (GRADE-STE), we enable end-to-end gradient flow from reward signals through generated tokens to model parameters. On sentiment-controlled text generation using the IMDB dataset, GRADE-STE achieves a test reward of 0.763 +- 0.344 compared to PPO's 0.510 +- 0.313 and REINFORCE's 0.617 +- 0.378, representing a 50% relative improvement over PPO. Critically, GRADE-STE exhibits gradient variance over 14 times lower than REINFORCE and maintains stable training dynamics throughout optimization. Our rigorous evaluation with proper train/validation/test splits demonstrates that these improvements generalize to held-out data, with GRADE-STE showing the best generalization characteristics among all methods tested. GRADE offers a simpler, more stable, and more effective alternative to reinforcement learning for LLM alignment.
Abstract (translated)
从人类反馈中的强化学习(RLHF)已成为将大型语言模型与人类偏好对齐的主导方法。然而,策略梯度方法如PPO(近端策略优化)由于高方差的梯度估计而需要仔细调整超参数和大量的计算资源。 我们提出了一种名为GRADE(通过可微估计进行对齐的Gumbel-softmax松弛)的方法,该方法通过直接反向传播穿过离散标记采样过程的一个可微分松弛来替代高方差策略梯度估计。使用带有直通估计(GRADE-STE)的Gumbel-Softmax参数化,我们使从奖励信号到生成令牌再到模型参数的端到端梯度流动成为可能。 在使用IMDB数据集进行情感控制文本生成时,GRADE-STE实现了0.763 ± 0.344的测试奖励值,相比之下PPO为0.510 ± 0.313,REINFORCE为0.617 ± 0.378,这代表了与PPO相比约50%的相对改进。更重要的是,GRADE-STE表现出比REINFORCE低14倍以上的梯度方差,并在整个优化过程中保持稳定的训练动态。 我们通过适当的训练/验证/测试拆分进行了严格的评估,这些结果表明这些改进可以推广到未见过的数据上,在所有测试的方法中,GRADE-STE显示出了最佳的泛化特性。GRADE为大型语言模型对齐提供了一种更为简单、稳定和有效的替代强化学习方法。
URL
https://arxiv.org/abs/2601.11574