Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
The paper introduces GRADE, a novel method for aligning LLMs with human preferences that replaces policy gradient methods with direct backpropagation. GRADE utilizes the Gumbel-Softmax reparameterization with straight-through estimation (GRADE-STE) to enable end-to-end gradient flow from reward signals through generated tokens to model parameters. Experiments on sentiment-controlled text generation using the IMDB dataset demonstrate that GRADE-STE achieves a 50% relative improvement over PPO, exhibits significantly lower gradient variance, and maintains stable training dynamics.
Ditch finicky RLHF: GRADE offers a backpropagation alternative that boosts reward by 50% and slashes gradient variance for LLM alignment.
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.