Policy Gradient Methods › Summary

Policy gradient methods provide the mathematical foundation for optimizing language models using reward signals. The key insights from this chapter are:

• Language models are policies. The autoregressive generation process naturally maps to the reinforcement learning framework: states are contexts, actions are tokens, and the policy $\pi_\theta(a_t|s_t)$ is the model's next-token probability distribution. This framing connects language model alignment directly to decades of reinforcement learning research.

• The policy gradient theorem enables optimization. By applying the log-derivative trick, we rewrite the gradient of expected reward as $\mathbb{E}[\nabla_\theta \log \pi_\theta(y|x) \cdot R(y)]$. This turns an intractable sum over all possible sequences into a Monte Carlo estimate we can compute from samples.

• REINFORCE is simple but high-variance. The basic algorithm samples sequences, computes their rewards, and updates parameters in proportion to $\log \pi(y) \cdot R$. This works in principle but requires many samples for stable estimates, which makes it impractical at scale for large language models.

• Baselines reduce variance without introducing bias. Subtracting a baseline $b$ from rewards yields the same expected gradient but lower variance. The score function has zero mean in expectation, which is the key mathematical fact that makes baselines work. The optimal baseline is the expected reward, leading to the advantage formulation $A = R - V$.

• Reward-to-go exploits causal structure. Actions cannot affect past rewards, only future ones. Using only future rewards in the gradient estimate removes noise from past outcomes while preserving the unbiasedness of the estimator.

• Additional techniques help in practice. Reward normalization, entropy regularization, and gradient clipping all contribute to more stable training. These techniques address different failure modes: scale variation, policy collapse, and catastrophic updates respectively.

The variance problem with REINFORCE motivates the development of more sophisticated algorithms. In the next chapter, we will see how Proximal Policy Optimization (PPO) addresses these issues by constraining policy updates and using careful advantage estimation with generalized advantage estimation (GAE), making RL-based alignment practical for production language models.