Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Policy Gradient Methods › The Variance Problem

Michael BrenndoerferDecember 26, 20255 min read

Language AI HandbookAll chapters

Part 37 of 67Alignment and RLHF

Chapter 277Policy Gradient Methods › The Variance Problem

Reading section 7 of 14 · The Variance Problem · 5 minSection 7 of 14 · 5 min

Section 7 of 14 · 5 min read

Policy Gradient Methods › The Variance ProblemOne section · Sections in this chapter7 / 14 ↓

Sections in this chapter14

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

REINFORCE has a fundamental weakness that prevents it from being directly used for training large language models: extremely high variance in its gradient estimates. This variance explains why every practical RLHF system uses more sophisticated algorithms. It directly motivates PPO, which we cover in the next chapter.

To understand where the variance comes from, recall our Monte Carlo gradient estimate:

\nabla_\theta J(\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \nabla_\theta \log \pi_\theta(y^{(i)}|x) \cdot R(x, y^{(i)})

where:

$N$ : the batch size (number of sampled sequences)
$y^{(i)}$ : the $i$ -th sampled sequence in the batch
$R(x, y^{(i)})$ : the reward for the $i$ -th sequence

We are estimating the true gradient by averaging over $N$ samples. If $N = 1$ , our estimate equals the product of a random gradient direction and a random reward. Both quantities vary across samples, and their product varies even more. Several sources contribute to the overall variance:

Sampling variance. We estimate an expectation using a small batch of samples. Different samples from the same policy can have wildly different rewards. Consider a prompt where some responses are brilliant (reward +10) and others are terrible (reward -10), but most are mediocre (reward near 0). If we happen to sample one of the rare brilliant responses, we get a huge positive gradient. If we sample a terrible one, we get a huge negative gradient. The expected gradient might be small and stable, but individual estimates can be orders of magnitude larger.

Reward magnitude. Large rewards amplify gradient noise. A reward of 100 versus 0.1 changes gradient magnitudes by a factor of 1000. If the reward scale varies across prompts or across training, the gradient scale varies too. This makes it hard to choose a stable learning rate that works well across all prompts.

Long sequences. With more tokens, there is more opportunity for randomness in the sampling process. Each token choice is stochastic, and these random choices compound. The same prompt can lead to vastly different sequences purely due to sampling randomness early in generation. A small difference in the first few tokens can cascade into completely different responses.

To see why the variance matters practically, imagine training a model with batch size 1 on a task where rewards range from -10 to +10. One batch might produce a sample with reward +8, causing a large positive gradient update. The next batch produces a sample with reward -7, causing a large negative gradient update. The model is being pushed back and forth, making little net progress. It is like trying to navigate by compass in a room where someone randomly rotates the compass between readings.

In[11]:

Code

import numpy as np

# Simulate variance in REINFORCE estimates
np.random.seed(42)


def simulate_reinforce_variance(
    n_samples_per_estimate, true_gradient=1.0, reward_std=5.0
):
    """Simulate variance in gradient estimates."""
    # Simulate rewards with high variance
    rewards = np.random.normal(
        loc=true_gradient, scale=reward_std, size=n_samples_per_estimate
    )
    # Each estimate is the mean of samples
    estimate = np.mean(rewards)
    return estimate


# Compare variance with different batch sizes
batch_sizes = [1, 4, 16, 64, 256]
n_estimates = 200

variance_by_batch = {}
for batch_size in batch_sizes:
    estimates = [
        simulate_reinforce_variance(batch_size) for _ in range(n_estimates)
    ]
    variance_by_batch[batch_size] = np.var(estimates)

Out[12]:

Visualization

Bar chart showing gradient estimate variance decreasing from about 25 to near 0.1 as batch size increases. — REINFORCE gradient estimate variance as a function of batch size. Variance decreases following a $1/N$ scaling law, requiring exponentially more samples to achieve linear improvements in stability. Even at batch size 256, variance remains substantial compared to supervised learning objectives, showing the basic sample inefficiency of Monte Carlo policy gradient estimation.

The variance decreases at rate $O(1/N)$ with batch size $N$ , but this means we need exponentially more samples to achieve linear improvements in stability. To reduce variance by a factor of 10, we need 10 times more samples. To reduce it by a factor of 100, we need 100 times more samples. For large language models where each sample requires generating a complete response (potentially hundreds of tokens) through a billion-parameter model, and where each generated response must be scored by a separate reward model, this sample inefficiency becomes a serious practical barrier.

The practical consequence is that naive REINFORCE is essentially unusable for training large language models. The gradient noise dominates the signal, training progress is erratic, and the compute cost is prohibitive. This motivates variance reduction techniques, which we develop in the following sections.