Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

GRADE: Replacing Policy Gradients with Backpropagation for LLM Alignment

Lukas Abrie Nel

TL;DR

GRADE (Gumbel-softmax Relaxation for Alignment 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, offers a simpler, more stable, and more effective alternative to reinforcement learning for LLM alignment.

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.

GRADE: Replacing Policy Gradients with Backpropagation for LLM Alignment

TL;DR

GRADE (Gumbel-softmax Relaxation for Alignment 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, offers a simpler, more stable, and more effective alternative to reinforcement learning for LLM alignment.

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.
Paper Structure (52 sections, 2 theorems, 19 equations, 3 figures, 3 tables, 1 algorithm)

This paper contains 52 sections, 2 theorems, 19 equations, 3 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Reinforcement Learning from Human Feedback.
  4. Policy Gradient Methods and Variance Reduction.
  5. Direct Alignment Methods.
  6. Gumbel-Softmax and Differentiable Relaxations.
  7. Differentiable Text Generation.
  8. Prompt Tuning and Soft Tokens.
  9. Background
  10. Problem Setup: LLM Alignment
  11. Policy Gradient Methods
  12. Gumbel-Softmax Distribution
  13. Straight-Through Estimator
  14. Method: GRADE
  15. Differentiable Token Generation
  16. ...and 37 more sections

Key Result

Proposition 1

Let $\hat{g}_{\text{PG}}$ be the REINFORCE policy gradient estimator and $\hat{g}_{\text{GS}}$ be the Gumbel-Softmax gradient estimator for the same objective. Under mild regularity conditions on the reward function $r$, we have: with equality only when the policy is deterministic.

Figures (3)

  • Figure 1: Training Dynamics. (a) Training reward over steps: GRADE-STE shows steady improvement to $\sim$0.75, while PPO plateaus around 0.5 and vanilla GRADE stays near 0.35. (b) Loss curves. (c) KL divergence from reference model: GRADE-STE maintains higher KL, indicating more substantial policy updates enabled by stable gradients.
  • Figure 2: Validation Performance and Generalization. (a) Validation reward during training, with best checkpoints marked. GRADE-STE achieves highest validation reward. (b) Generalization gap (train $-$ val reward): positive values indicate overfitting. GRADE variants show negative gaps (good generalization), while PPO shows increasing positive gap (overfitting).
  • Figure 3: Temperature Schedule Analysis. (a) Linear annealing from $\tau=2.0$ to $\tau\approx 1.82$ over 250 steps. (b) Reward vs temperature colored by training step, showing reward improvement correlates with both lower temperature and training progress.

Theorems & Definitions (3)

  • Proposition 1: Variance Reduction
  • proof : Proof Sketch
  • Proposition 2: Bias-Variance Tradeoff