Policy Gradient Methods › From Sequences to Tokens

For language models, we need to express the log-probability of an entire sequence in terms of individual token log-probabilities. This connection is needed because our models compute token-level probabilities, not sequence-level probabilities directly. Fortunately, the autoregressive structure of language models makes this conversion straightforward.

Using the autoregressive factorization from the previous section, we can apply the logarithm to convert the product of probabilities into a sum of log-probabilities:

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

where:

• $\sum_{t=1}^{T}$: sum over all timesteps in the sequence
• $\log \pi_\theta(y_t | x, y_{<t})$: log-probability of the specific token $y_t$ given context

The logarithm converts the product of probabilities into a sum of log-probabilities. This is where the log-derivative trick pays off: instead of dealing with products of many small numbers (which can underflow numerically when sequence lengths are in the hundreds), we work with sums of log-probabilities, which remain in a numerically stable range.

Each term $\log \pi_\theta(y_t | s_t)$ is simply the log-softmax output of our model at position $t$, evaluated at the selected token. This is exactly what we compute during the forward pass when training language models with cross-entropy loss. When the model processes a sequence, it produces logits for every position, and applying log-softmax gives us log-probabilities. To compute the log-probability of a specific token choice, we simply index into this log-probability vector at the chosen token's index.

The gradient of the full sequence log-probability then follows from the linearity of differentiation. The gradient of a sum is the sum of gradients:

$$\nabla_\theta \log \pi_\theta(y|x) = \sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(y_t | s_t)$$

where:

• $s_t$: the state (context) at timestep $t$, equivalent to $(x, y_{<t})$
• $\nabla_\theta \log \pi_\theta(y_t | s_t)$: gradient of the log-probability for a single token decision

Each term in the sum corresponds to one token's contribution to the overall sequence probability. We compute each term by running a standard backward pass through the model at that timestep.

Putting this together, the policy gradient becomes a sum over token-level contributions, all weighted by the same sequence-level reward:

$$\nabla_\theta J(\theta) = \mathbb{E}_{y \sim \pi_\theta}\left[\sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(y_t | s_t) \cdot R(x, y)\right]$$

where:

• $\sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(y_t | s_t)$: the accumulated gradient from all tokens in the sequence
• $R(x, y)$: the final reward, which scales the gradient of every token choice equally

The same reward $R(x, y)$ multiplies every token's gradient. This is the basic credit assignment mechanism of REINFORCE: even though we only receive reward at the end, it propagates back to influence all token choices. Every token that contributed to the response receives the same credit or blame for the final outcome. This is both a strength and a weakness of the basic approach. It allows learning from sparse rewards, but it does not distinguish between tokens that were important for the reward and tokens that were irrelevant to it.