Policy Gradient Methods › Additional Variance Reduction Techniques
Language AI HandbookAll chapters Policy Gradient Methods › Additional Variance Reduction TechniquesOne section · Sections in this chapter11 / 14 ↓
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
Beyond baselines and reward-to-go, several additional techniques help control variance in practice. Together, these techniques form a toolkit that practitioners reach for when raw REINFORCE proves too unstable.
Reward Normalization
Reward normalization standardizes rewards across a batch to have zero mean and unit variance:
def normalize_rewards(rewards, eps=1e-8):
"""
Normalize rewards to zero mean and unit variance within a batch.
This prevents large reward magnitudes from causing gradient explosions
while preserving the relative ordering of rewards within the batch.
The eps term prevents division by zero when all rewards are identical.
"""
mean = rewards.mean()
std = rewards.std()
return (rewards - mean) / (std + eps)This technique prevents large reward magnitudes from causing gradient explosions while preserving the relative ordering of rewards. It is particularly useful when the reward scale varies across training: perhaps early in training rewards cluster near zero, while later in training they spread across a wider range. Normalization ensures the learning rate remains appropriate throughout.
The key insight is that what matters for policy improvement is not the absolute scale of rewards but their relative ordering. Whether responses score 5, 7, and 9, or 500, 700, and 900, the model should learn the same qualitative lessons. Normalization makes the gradient magnitudes consistent regardless of the reward scale.
Entropy Regularization
Entropy regularization adds a bonus for maintaining diversity in the policy's probability distribution, preventing premature convergence to deterministic behavior:
import torch
def policy_entropy(probs):
"""
Compute the Shannon entropy of an action distribution.
Higher entropy means more uniform distribution (more exploration).
Lower entropy means more concentrated distribution (more exploitation).
"""
return -(probs * torch.log(probs + 1e-10)).sum(dim=-1)
def entropy_regularized_loss(policy_loss, probs, entropy_coef=0.01):
"""
Add entropy bonus to encourage continued exploration.
The entropy coefficient controls the exploration-exploitation tradeoff:
larger values encourage more diversity in generated responses.
"""
entropy = policy_entropy(probs).mean()
return policy_loss - entropy_coef * entropyHigh entropy means the policy is exploring diverse actions, generating varied responses. Low entropy means it is becoming deterministic, repeatedly generating the same responses. The entropy bonus counteracts the natural tendency of policy gradient training to become more and more confident (lower entropy) over time.
Without entropy regularization, a policy gradient algorithm can get stuck in a local optimum: it finds a response type that consistently receives decent rewards, concentrates all probability on that response type, and never explores alternative approaches that might receive higher rewards. The entropy bonus keeps the policy somewhat diffuse, maintaining the exploration needed to discover better response strategies.
Gradient Clipping
Gradient clipping limits the magnitude of gradient updates by rescaling the gradient vector when its norm exceeds a threshold:
import torch
def clip_gradients(model, max_norm=1.0):
"""
Clip gradients to prevent excessively large parameter updates.
When the gradient norm exceeds max_norm, the gradient vector is
rescaled to have norm exactly max_norm. This preserves the direction
of the gradient while bounding its magnitude.
"""
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)As we discussed in the chapter on gradient clipping, this technique prevents individual updates from being too large, which is especially important given REINFORCE's high variance. A single high-reward sample can produce an enormous gradient that catastrophically overwrites learned behavior. Clipping ensures no single update can move parameters too far from their current values.
The choice of max_norm involves a tradeoff: too small and learning slows unnecessarily; too large and you lose the protection against large updates. Values between 0.5 and 5.0 are common in practice, with 1.0 being a typical starting point.
Putting the Techniques Together
In production RLHF systems, these variance reduction techniques are applied together. A typical implementation would:
- Generate a batch of responses (multiple responses per prompt for better variance reduction)
- Score each response with the reward model
- Normalize rewards within the batch to zero mean and unit variance
- Compute baseline predictions from the value network
- Compute per-sequence advantages (normalized reward minus baseline)
- Compute the REINFORCE loss using advantages instead of raw rewards
- Backpropagate and clip gradients before applying the update
- Add entropy bonus to prevent policy collapse
Each technique addresses a different source of instability: normalization handles reward scale variation, baselines address absolute reward level, clipping prevents catastrophic updates, and entropy regularization prevents premature convergence.