Analyzing Policy Shift During RL Tuning

Understanding how the language model's policy, denoted as $\pi_{\text{RL}}$, evolves during the reinforcement learning phase is fundamental for successful alignment. Simply maximizing the reward signal from the Reward Model (RM) isn't enough; we need to ensure this optimization process leads to genuine improvements in helpfulness, harmlessness, and honesty, without unintended consequences. Analyzing the policy shift provides insights into the training dynamics, helps detect potential issues early, and validates that the alignment process is effective.

The primary motivations for tracking policy changes include:

1. Monitoring Alignment Progress: Verifying that the policy is indeed learning behaviors preferred by the RM and, by proxy, human labelers.
2. Detecting Reward Hacking: Identifying situations where the policy finds shortcuts to maximize the RM score without producing genuinely better outputs. This often manifests as unnatural, repetitive, or overly verbose responses that exploit loopholes in the RM's understanding.
3. Managing Policy Stability: The KL divergence term in PPO aims to keep the RL policy close to the initial SFT policy ($\pi_{\text{SFT}}$). Monitoring this helps ensure the model doesn't forget its base knowledge or language fluency while optimizing for the reward.
4. Diagnosing Training Issues: Sudden jumps or collapses in reward, KL divergence, or output quality can indicate problems like exploding gradients, incorrect hyperparameter settings, or issues with the RM itself.

Several metrics and techniques are employed to analyze this shift:

KL Divergence Tracking

The Kullback-Leibler (KL) divergence measures the difference between the probability distributions of tokens predicted by the current RL policy ($\pi_{\text{RL}}$) and the reference policy (usually the initial SFT model, $\pi_{\text{SFT}}$), given the same prompt. In PPO for RLHF, we explicitly penalize high KL divergence to prevent the policy from straying too far from the SFT model. Tracking the per-token KL divergence averaged over batches is a standard practice.

A plot showing the mean per-token KL divergence between the RL policy and the SFT policy over training steps. Stable training usually shows a controlled increase, while rapid, unbounded growth might signal instability or excessive deviation from the base model.

A consistently low KL divergence might indicate that the policy isn't learning much or the KL penalty coefficient is too high. Conversely, a rapidly increasing KL divergence suggests the policy is changing significantly, which could be good if rewards are also increasing, but warrants investigation for potential reward hacking or loss of capabilities if rewards stagnate or outputs degrade.

Reward Score Analysis

Monitoring the average reward score assigned by the RM to the policy's generations is essential. We expect this score to increase over time as the policy learns to generate responses favored by the RM. Plotting the distribution of reward scores can also be informative.

Histograms showing the distribution of reward scores assigned by the RM to policy generations at different stages of RL training. A shift towards higher scores indicates successful optimization according to the RM.

A tightening distribution around high reward values generally indicates convergence. However, if the distribution shifts dramatically or develops unusual shapes, it might point towards the policy exploiting specific aspects of the RM, potentially indicating reward hacking.

Qualitative Evaluation

Quantitative metrics alone don't tell the whole story. Regularly sampling and manually inspecting the model's outputs for a fixed set of prompts throughout training is indispensable. Compare generations from:

• The initial SFT model ($\pi_{\text{SFT}}$)
• Early RL policy checkpoints (e.g., after 10% of training)
• Mid-training RL policy checkpoints
• Late or final RL policy ($\pi_{\text{RL}}$)

Look for desired improvements (e.g., increased helpfulness, better instruction following, reduced harmfulness) but also be vigilant for regressions (e.g., decreased coherence, repetition, sycophancy, emergence of new failure modes). This qualitative feedback loop is critical for understanding how the policy is changing.

Additional Quantitative Metrics

• Perplexity: Calculate the perplexity of the RL policy on hold-out datasets (either general text corpora or the SFT dataset) to monitor potential degradation in language modeling fundamentals. A significant increase in perplexity might suggest the model is losing fluency or general knowledge.
• Automated Evaluation Benchmarks: If computationally feasible, run periodic evaluations on established alignment benchmarks (like HELM subsets or specific instruction-following tests) to get an objective measure of performance on tasks.
• Win Rate vs. Reference: Calculate the win rate of the current RL policy's generations against the SFT policy's generations, using the RM as the judge. This provides a direct measure of improvement according to the learned preference model.

Interpreting Combined Signals

Analyzing these metrics together provides a more complete picture:

• Increasing Reward, Controlled KL, Good Qualitative Examples: The ideal scenario. The policy is learning aligned behaviors without drastic, potentially harmful shifts.
• Increasing Reward, High/Rapidly Increasing KL: Investigate closely. This could be effective learning of complex preferences, but also risks reward hacking or forgetting SFT knowledge. Qualitative checks are essential.
• Stagnant/Decreasing Reward, Increasing KL: Problematic. The policy is changing but not improving according to the RM. Could indicate training instability, poor hyperparameter choices, or an inadequate RM.
• Stagnant Reward, Stagnant KL: Training may have converged, or the learning rate/KL penalty might be preventing further progress. Consider adjusting hyperparameters or checking gradients.

Tooling

Standard ML experiment tracking tools like Weights & Biases or TensorBoard are invaluable for logging and visualizing these metrics (KL divergence, reward scores, evaluation metrics) over time. Libraries like Hugging Face's TRL often provide built-in utilities for calculating and logging KL divergence and reward statistics during PPO training, simplifying the analysis process.

By systematically analyzing how the policy shifts during RL tuning using a combination of quantitative metrics and qualitative inspection, you can gain confidence that the RLHF process is achieving meaningful alignment and catch potential problems before they derail training or lead to poorly behaved models.