GRADE: Replacing Policy Gradients with Backpropagation for LLM Alignment

TL;DR

GRADE introduces a differentiable relaxation technique for LLM alignment, replacing high-variance policy gradients with backpropagation, leading to more stable training and improved reward performance.

Contribution

The paper presents GRADE, a novel method that replaces policy gradient estimation with backpropagation using Gumbel-Softmax, enabling more stable and efficient LLM alignment.

Findings

GRADE-STE achieves higher test rewards than PPO and REINFORCE.

GRADE-STE exhibits over 14 times lower gradient variance.

The method generalizes well to held-out data.

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.