Understanding Policy Gradients in LLMs

I’ve tried to learn RL for LLMs several times, but the connection between the intuition and math was not clear to me. Recently, I came across Nathan Lambert’s RLHF Book, which helped me connect the pieces. I’m truly thankful for this great material. In this blog, I’ll discuss what a policy gradient is and why it is useful.

RL for LLMs#

In general, there are two major types of methods in RL.

• The first one is value-based, meaning that you want to study how much return you get when you take a certain action in a given state.
• The second one is policy-based, meaning that you want to learn a policy (an LLM) that gets the maximum expected return in the environment.

Most of the RL methods for LLMs we have heard about are policy-based, which means we improve the model itself directly. For LLMs, policy-based RL feels like a natural continuation of pretraining: we already have a model that knows how to generate, and now we use rewards to improve what it generates. Though value-based methods are more fundamental, policy-based methods take a direct route: using rewards to improve the policy without requiring a separate value-learning step.

In LLM generation, the state is the current token sequence, and an action appends one token to it.

Before we start, let’s introduce the notation we will use.

• $\tau=(s_0,a_0,r_0,s_1,a_1,r_1,\ldots)$: a sampled trajectory.
• $s_t$: the current token sequence, with $s_0=\text{prompt}$ and $s_{t+1}=\operatorname{concat}(s_t,a_t)$.
• $\pi_\theta(a_t\mid s_t)$: the next-token policy.
• $R(\tau)$: the sequence reward, with no explicit dependence on $\theta$.

The LLM RL Objective#

Our objective is to find a policy that maximizes the expected return, starting from an initial state drawn from some distribution. Ideally, this prompt distribution should reflect the tasks we want the model to perform.

In an RL environment, taking an action can bring you to a new state and give you an immediate reward. So we can write our objective as

$$ J(\theta) =\mathbb{E}_{\tau\sim p_\theta(\cdot)}[R(\tau)] =\sum_\tau p_\theta(\tau)R(\tau) $$

with

$$ p_\theta(\tau) =p(s_0)\prod_{t=0}^{T-1}\pi_\theta(a_t\mid s_t) $$

Here, $p(s_0)$ is the initial-state distribution and does not depend on $\theta$. With deterministic transitions, the remaining trajectory probability is the product of the action probabilities.

Computing this expectation exactly is impossible because the space of possible responses is too large. At each step, there are $|V|$ possible actions. For a fixed response length of $T$ tokens, there are $|V|^{T}$ possible sequences. So how do we find the optimal policy that maximizes our objective, the expected return?

Deriving the Policy Gradient#

How can we use gradient ascent to maximize this objective? We start by differentiating it with respect to the policy parameters:

$$ \begin{aligned} \nabla_\theta J(\theta) &=\nabla_\theta \mathbb{E}_{\tau\sim p_\theta(\cdot)}[R(\tau)] \\ &=\nabla_\theta\sum_\tau p_\theta(\tau)R(\tau) \\ &=\sum_\tau \nabla_\theta p_\theta(\tau)R(\tau) \\ &=\sum_\tau p_\theta(\tau) R(\tau) \nabla_\theta \log p_\theta(\tau) \\ &\boldsymbol{=\mathbb{E}_{\tau\sim p_\theta(\cdot)}[R(\tau)\nabla_\theta\log p_\theta(\tau)]} \\ &\boldsymbol{=\mathbb{E}_{\tau\sim p_\theta(\cdot)}\left[ R(\tau)\sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t\mid s_t) \right]}. \end{aligned} $$

The last step expands the trajectory log probability into a sum over tokens; $p(s_0)$ drops out because it does not depend on $\theta$.

From the derivation above, we can see that the gradient of the objective is the expectation of the reward-weighted gradient of the trajectory log probability. Notice that this is an exact transformation!

Since the gradient is now written as an expectation over trajectories, we can use sampling to estimate it.

$$ \nabla_\theta J(\theta) \approx\frac{1}{B}\sum_{i=1}^{B} R(\tau_i)\nabla_\theta\log p_\theta(\tau_i), \qquad \tau_i\sim p_\theta. $$

The approximation comes from replacing the expectation with an average over $B$ sampled rollouts.

But why estimate the gradient rather than the objective itself? Think about what happens if we estimate the objective directly by sampling:

$$ \hat J=\frac{1}{B}\sum_{i=1}^{B}R(\tau_i). $$

Once the rollouts are fixed, this is just an average of fixed rewards, with no differentiable path back to $\theta$. In the policy-gradient estimate, however, $\log p_\theta(\tau_i)$ still depends on the policy parameters, giving us a way to do backpropagation.

We now have the gradient estimate we need. To compute it with automatic differentiation, we write a loss whose derivative gives that estimate with a negative sign:

$$ L_{\mathrm{PG}}(\theta) =-\frac{1}{B}\sum_{i=1}^{B} R(\tau_i)\log p_\theta(\tau_i). $$

With the sampled responses and rewards held fixed, calling backward on this loss computes the gradient for us. Minimizing the loss then moves the policy in the direction of higher expected reward.

References#

1. Nathan Lambert. RLHF Book: Policy Gradients and Implementation.