Policy Gradient Methods › The Objective Function

To optimize our policy, we need a clear mathematical objective. In RLHF, we want to maximize the expected reward over all possible responses the model might generate. This expectation is important because language generation is inherently stochastic. The same prompt can lead to many different responses depending on which tokens get sampled at each step. Some responses will be excellent, others will be mediocre, and still others will be poor. We want to tune the model so that, on average, it produces responses that receive high rewards.

Think of the objective function as a report card for a policy, but one computed by averaging over all possible exams the policy might face. A policy scores well on this report card if, when we sample random responses from it, those responses tend to be good. The objective does not ask that every response be perfect; it asks that the distribution over responses be skewed toward high-quality outcomes.

Given a prompt $x$, the model generates a response $y = (y_1, y_2, \ldots, y_T)$ by sampling tokens according to its policy. The reward model then scores the complete response, giving us $R(x, y)$. Our objective is to find parameters $\theta$ that maximize the expected reward:

$$J(\theta) = \mathbb{E}_{y \sim \pi_\theta(\cdot|x)}[R(x, y)]$$

where:

• $J(\theta)$: the objective function we want to maximize
• $\theta$: the parameters of the language model
• $\mathbb{E}$: the expectation over sequences sampled from the policy
• $\pi_\theta(\cdot|x)$: the policy (language model) distribution conditioned on prompt $x$
• $R(x, y)$: the reward scalar for the generated response $y$

This expectation is taken over all possible responses the model could generate. Some responses will be excellent (high reward), others mediocre, others poor. The expectation weights each by its probability under the current policy. This weighting matters: if the model rarely generates a particular response, that response contributes little to the expected value, even if it happens to have high reward. A policy that occasionally produces brilliant responses but usually produces mediocre ones will score lower than a policy that consistently produces good responses.

We can expand this expectation using the probability of generating the complete sequence:

$$J(\theta) = \sum_{y} \pi_\theta(y|x) R(x, y)$$

where:

• $\sum_{y}$: summation over all possible output sequences in the vocabulary
• $\pi_\theta(y|x)$: probability of generating sequence $y$ given prompt $x$
• $R(x, y)$: reward for the specific sequence $y$

This summation form makes the objective's structure transparent, but it also reveals a computational challenge. The number of possible sequences is astronomically large: for a vocabulary of 50,000 tokens and sequences of length 100, there are $50000^{100}$ possible sequences. We cannot enumerate them all. We need a way to optimize this objective without explicitly summing over all sequences.

The probability $\pi_\theta(y|x)$ represents the likelihood of generating the entire response $y$. It is computed as the product of individual token probabilities:

$$\pi_\theta(y|x) = \prod_{t=1}^{T} \pi_\theta(y_t | x, y_{<t})$$

where:

• $T$: the length of the generated sequence
• $y_t$: the token generated at timestep $t$
• $y_{<t}$: the history of tokens generated before step $t$, that is $(y_1, y_2, \ldots, y_{t-1})$
• $\pi_\theta(y_t | x, y_{<t})$: the probability the model assigns to token $y_t$ given the full context

This autoregressive factorization, which we have seen throughout our discussion of language models in earlier chapters, connects the policy formulation directly to how transformers and LSTMs generate text. Each factor in the product corresponds to one step of the generation process, one call to the model's forward function, one decision about which token to emit next. The full sequence probability is simply the product of all these individual decisions.

You might wonder why we cannot just directly maximize $J(\theta)$ by computing its gradient. The obstacle is that $J(\theta)$ involves an expectation over sequences sampled from $\pi_\theta$, and sampling is not a differentiable operation. When we sample a token from a probability distribution, we make a discrete choice, and discrete choices are not differentiable with respect to the parameters that shape that distribution.

Think of it this way: if you tweak the parameters of the model slightly so that the probability of token "excellent" increases by 0.001, you cannot directly measure how that change affects the expected reward. The reward depends on which specific tokens get sampled, and sampling is random. We need a mathematical trick to convert the gradient of an expectation into an expectation of gradients that we can estimate from samples.