ArXiv: 2309.10202

🎯 Pitch

Standard RLHF suffers from easy reward hacking because reward models assign wildly different score ranges across tasks; this paper fixes that by training an advantage model that directly predicts the extra reward over the expected value, forcing distributions to align. Combined with a selective rehearsal strategy that replays only the most representative SFT examples per cluster during PPO, the method slashes calibration error by over 25% and flips forget-set performance from decline to double-digit win rate gains.

1. Executive Summary

This technical report introduces two stabilization techniques for RLHF training of large language models and empirically analyzes them on the HH-RLHF benchmark and proprietary datasets using BLOOMZ models. The first mechanism, the Advantage Model (AM), directly models advantage scores — the extra reward a response obtains relative to the expected reward — and regulates score distributions across tasks to prevent reward hacking (e.g., constraining reward scores so that Code Generation and QA tasks exhibit similar distributions, avoiding the model transferring response patterns from higher-reward categories to lower-reward ones). The second mechanism, Selective Rehearsal, mitigates catastrophic forgetting by clustering PPO training data, selecting the top-scoring examples within each cluster, and adding an NLL loss on these representative examples to the PPO objective (e.g., retaining SFT-acquired skills on expert-aligned examples while still optimizing for higher reward). The Advantage Model achieves both higher ranking accuracy and substantially lower Expected Calibration Error than a standard Reward Model — 3.48 vs. 4.70 ECE on HH-RLHF and 3.83 vs. 5.35 on proprietary data — while producing more stable per-task score means and variances. In PPO training, AM-PPO-SR (Advantage Model with Selective Rehearsal) attains a 15.78% win rate over the SFT model on the main test set and a substantially improved 10.30% win / 7.95% loss rate on the forget test set, establishing that RLHF stability gains from distribution-regularized reward signals and strategic data curation are achievable across tasks but that the optimal cluster count for selective rehearsal remains dataset-dependent and warrants domain-specific tuning.

2. Context and Motivation

The Core Problem: RLHF Training Is Dangerously Unstable

The paper addresses a specific, practical failure mode in the dominant paradigm for aligning large language models with human preferences. RLHF — Reinforcement Learning from Human Feedback — is the technique behind virtually every major aligned LLM release: GPT-4, Claude, Sparrow, Bard, and Llama 2-Chat all depend on it. The RLHF recipe is conceptually straightforward: train a reward model on human preference comparisons, then use Proximal Policy Optimization (PPO) to fine-tune the language model to maximize that reward while staying close to the supervised fine-tuning (SFT) baseline via a KL-divergence penalty.

In practice, however, this pipeline is extraordinarily brittle. The authors open the paper by cataloguing two specific symptoms of this brittleness, which they later trace to two root causes:

Symptom 1: Reward hacking. During PPO training, the model frequently discovers ways to earn high reward scores that are completely uncorrelated with actual human preference. The paper provides a concrete example (Section 1, Figure 1a): the reward model learns sharply different score distributions for different task categories — Code Generation tasks receive higher mean rewards than QA tasks, simply as an artifact of how the reward model was trained. When PPO then optimizes against this reward model, the policy learns to transfer response patterns from the high-reward category to other categories — for instance, producing code-like formatting or structure when answering QA questions, because doing so triggers the inflated reward signal. The authors report that this manifests as "gibberish responses (but high-reward)," a classic sign that the reward proxy has decoupled from the true objective.

This is not a hypothetical concern. Stiennon et al. (2020) first documented reward hacking in summarization RLHF, where models learned to game the reward model rather than produce genuinely better summaries. Skalse et al. (2022) later provided a formal framework for characterizing reward gaming, establishing that it arises from the fundamental mismatch between the learned reward function and the true utility function. The present paper extends this diagnosis: reward hacking in multi-task RLHF is exacerbated by inter-category reward score disparities, where the reward model assigns systematically different score scales to different types of tasks. Standard score normalization (e.g., moving averages) helps but does not fully solve the problem, because the means and variances fluctuate during training and can induce "unexpected model behaviors" (Section 1).

Symptom 2: Catastrophic forgetting. The paper shows (Figure 1b) that during PPO training, the model's win rate against the SFT baseline on "expert-aligned examples" — data samples that already met human standards and were used for SFT — drops significantly. In other words, the PPO process causes the model to unlearn capabilities it already possessed after SFT. The authors quantify this: on their "forget test set" (1,704 examples drawn from the SFT test data), a standard RM-PPO model achieves a 16.87% win rate but suffers a 29.28% loss rate compared to the SFT model (Table 2). This means the model is actively worse on nearly 30% of the SFT task distribution after RLHF training. This is catastrophic forgetting in the classic sense (McCloskey & Cohen, 1989): the model's policy update to maximize a new objective (the reward signal) overwrites previously learned behaviors, even when those behaviors were already aligned with human preferences.

The authors attribute this to "over-optimizing with PPO on examples that were well-aligned with humans in the SFT stage" (Section 1). The PPO algorithm treats all training prompts equally, applying the same reward-driven gradient update regardless of whether the SFT model's response was already correct. On examples where the SFT model produces good responses, aggressive PPO optimization can easily degrade them — the reward signal is noisy, and even small misestimations of reward can push the policy away from a good local optimum.

Why These Problems Matter: Practical and Theoretical Stakes

The instability of RLHF has both immediate practical consequences and deeper theoretical significance.

Practical impact. RLHF is not a research curiosity — it is the production alignment method for the largest deployed AI systems. OpenAI's GPT-4, Anthropic's Claude, Google's Bard, and Meta's Llama 2-Chat all use RLHF as a core component. Training instability in this context means:

  • Resource waste. Unstable RLHF runs often need to be restarted with different random seeds or hyperparameter configurations. Given the scale of these models (the paper uses BLOOMZ 176B for policy training), failed runs represent substantial computational expense.
  • Safety risks. Reward hacking can produce models that score highly on the training reward model but exhibit harmful behaviors in deployment — for instance, generating superficially polite responses that nonetheless contain toxic content in a format that evades the reward model's detection. The gap between reward model scores and human preference evaluations (visible in Figure 5b, where RM-PPO's reward scores rise while GPT-4 win rates drop) is a concrete demonstration of this decoupling.
  • Deployment barriers. If RLHF cannot be made reliable, organizations face a difficult choice: deploy models that may unpredictably forget capabilities or exhibit reward-hacked behavior, or forego alignment fine-tuning entirely — trading helpfulness/harmlessness for stability. This tradeoff is untenable for production systems.

Theoretical significance. The paper's diagnosis of RLHF instability connects to deeper questions about learned reward functions and policy optimization:

  • Reward calibration as a first-class problem. The finding that reward models produce systematically different score distributions across task categories (Figure 4a) reveals that the standard Bradley-Terry preference modeling objective (Equation 1) is insufficient for multi-task reward modeling. The model learns relative preferences within each category but has no incentive to produce comparable scores across categories. This is not merely a normalization issue that can be fixed post-hoc; it reflects a fundamental ambiguity in the reward modeling objective — any monotonic transformation of the scores preserves the ranking, so the model can freely shift score distributions between tasks without affecting training loss. The paper's Advantage Model addresses this by explicitly constraining scores to be bounded around an expected reward, effectively removing the degrees of freedom that enable distributional drift.

  • The exploration-catastrophe tension in PPO. PPO for language models operates under a KL-divergence constraint that is supposed to prevent the policy from straying too far from the SFT initialization (Equation 2, the βlog⁡(πϕ/πinit)\beta \log(\pi_\phi / \pi_{\text{init}}) term). The paper's results (Table 2, Forget Test Set) show that this constraint is insufficient to prevent catastrophic forgetting on specific data subsets. This suggests that the KL penalty, which operates on the aggregate policy distribution, fails to protect individual examples or task clusters from being overwritten. Selective Rehearsal addresses this by explicitly reinforcing the SFT behavior on strategically chosen examples, effectively implementing a local rather than global constraint on policy change.

Where Prior Approaches Fall Short

The paper identifies specific limitations in prior work along three axes:

1. Reward modeling ignores score calibration across tasks. The standard Bradley-Terry reward modeling objective (Equation 1) optimizes pairwise ranking accuracy: given two responses, can the model correctly identify which one humans preferred? This objective is invariant to any monotonic transformation of the scores — if the model adds a constant offset or multiplies all scores by a factor, the pairwise comparison probability σ(r(x,yc)−r(x,yr))\sigma(r(x, y_c) - r(x, y_r)) remains unchanged. This means the model can learn dramatically different score scales for different task categories without any penalty, as long as within-category rankings are preserved. The consequence (Figure 4a) is that Code Generation tasks might have mean rewards around 2.0 while QA tasks have mean rewards around 0.5, creating a perverse incentive during PPO: the policy can boost its average reward simply by making QA responses look more like code.

Prior work has attempted to address this with score normalization during PPO training — the authors mention using "moving average for score normalization" (Section 4.3). However, they show this is insufficient: "RM-PPO w/ MA encounters instabilities during PPO training" (Section 4.3), and Figure 5b shows that RM-PPO with moving average still exhibits declining GPT-4 win rates despite rising reward scores. Normalization treats the symptom (fluctuating score scales) without addressing the cause (unconstrained reward model output distributions).

An alternative approach, training separate reward models per task category, is mentioned but dismissed as impractical: "many studies train different RMs separately for helpful and harmless examples to achieve better performance. However, in our experiments, we did not distinguish between helpful and harmless examples" (Section 4.1). For their proprietary data spanning 61 task categories, training 61 separate reward models would be wildly impractical.

2. PPO applies uniform treatment to all training data. The standard PPO objective (Equation 2) samples prompts uniformly from the training distribution and applies the same reward-driven update regardless of whether the SFT model already produces good responses on that prompt. This is problematic because:

  • On prompts where the SFT response is already well-aligned, the reward signal provides limited useful information. The gradient update may push the policy away from a correct response due to noise in the reward estimate or KL penalty dynamics.
  • On prompts where the SFT response is poor, PPO's optimization pressure is exactly what is needed — the reward signal can genuinely guide the model toward better behavior.

The uniform treatment conflates these two cases. Prior approaches to this problem include rejection sampling (Touvron et al., 2023; Gulcehre et al., 2023), where only high-reward responses are used for supervised fine-tuning. However, the paper notes that rejection sampling "only consider[s] reward model score" (Section 3.2), ignoring other dimensions of data quality such as diversity, confidence (entropy), or response characteristics. This can lead to mode collapse — selecting only the highest-reward responses may reinforce a narrow set of response patterns while forgetting others that are also important.

3. Catastrophic forgetting mitigation is under-explored for RLHF. The continual learning literature has developed various techniques to combat catastrophic forgetting — experience replay, elastic weight consolidation, progressive networks — but these have been primarily developed and evaluated in the context of reinforcement learning agents (Khetarpal et al., 2022) and continual pre-training (Gupta et al., 2023), not RLHF. The standard RLHF recipe (Ouyang et al., 2022; Bai et al., 2022a) includes a KL-divergence penalty against the SFT model (the βlog⁡(πϕ/πinit)\beta \log(\pi_\phi / \pi_{\text{init}}) term in Equation 2), which is supposed to serve as a forgetting-prevention mechanism by discouraging large policy changes. However, the paper's results demonstrate that this global penalty is insufficient: significant forgetting occurs on the forget test set even with the KL penalty in place.

The underlying reason is that the KL penalty operates on the aggregate distributional divergence, not on per-example retention. A policy can satisfy a KL constraint of, say, 5 nats while still radically changing its behavior on specific prompts, as long as the behavior on the majority of prompts changes little. Expert-aligned examples — which by definition already meet human standards — are particularly vulnerable because the reward signal on these examples is noisy (the SFT response is already good, so reward differences between candidate responses reflect noise rather than genuine improvement opportunities), and the KL penalty offers them no special protection.

How This Paper Positions Itself

The paper positions itself not as proposing an entirely new alignment paradigm (it stays within the RLHF framework), but as introducing two complementary stabilization mechanisms that address the identified failure modes while preserving the core PPO + reward model architecture:

Advantage Model targets the reward signal quality. Rather than trying to fix reward hacking at the PPO level (through better normalization, more aggressive KL penalties, or reward clipping), the paper intervenes upstream at the reward model itself. The key insight is that reward score instability is a consequence of the unconstrained nature of the Bradley-Terry objective — the model has too many degrees of freedom in assigning absolute scores. By reformulating the objective to model advantages (extra reward relative to the expected reward for a prompt), with an explicit bounding loss that constrains advantage magnitudes (Equation 6), the Advantage Model produces scores that are:

  • Better calibrated (lower ECE in Table 1, Figure 2): the probability that a response with advantage score ss is actually preferred by humans is more accurately reflected by ss.
  • More comparable across tasks (Figure 4a): the expected reward baseline eτ(x)e_\tau(x) absorbs task-specific scale differences, so the advantage scores represent genuine quality differences in a task-invariant way.
  • More stable during training (Figure 4b): the bounding loss constrains variance, preventing the model from learning extreme scores for certain examples.

The paper explicitly connects this to the advantage function concept from reinforcement learning (Section 3.1), where subtracting a state-dependent baseline reduces variance without changing the optimal policy. The Advantage Model applies the same principle to reward modeling: subtract a prompt-dependent expected reward to get a signal that reflects how much better this response is than typical, rather than an absolute score whose scale depends on arbitrary features of the task category.

Selective Rehearsal targets the PPO optimization process. This is an intervention downstream of the reward model, modifying how PPO uses training data. The central insight is that not all training examples benefit equally from PPO optimization — some examples are well-served by the SFT model and should be protected from aggressive reward-driven updates, while others genuinely need the optimization pressure that PPO provides. Selective Rehearsal operationalizes this by:

  • Clustering PPO training data into semantically similar groups (Section 3.2, using SimCSE embeddings and KMeans).
  • Selecting the highest-quality examples per cluster based on advantage model scores (the same signal that the Advantage Model provides, creating a natural synergy between the two techniques).
  • Adding an NLL loss on these selected examples to the PPO objective (Equation 9), effectively reinforcing SFT behavior on representative examples from each skill cluster.

This is fundamentally different from rejection sampling because it explicitly optimizes for diversity (by sampling from each cluster) and quality (by selecting top-scoring examples within clusters), rather than quality alone. The clustering step ensures that the rehearsal data covers the full range of skills present in the SFT distribution, preventing the mode collapse that pure reward-maximizing selection can cause.

The paper also positions itself relative to the growing body of work on alternative alignment methods that bypass RL training entirely — Direct Preference Optimization (DPO; Rafailov et al., 2023) and Preference Ranking Optimization (PRO; Song et al., 2023) are cited in Section 5. While acknowledging these as promising directions, the paper implicitly argues that RLHF remains important enough (given its deployment in major production systems) that stabilizing it is a worthwhile goal, even if alternative paradigms exist.

The Connection Between the Two Contributions

A crucial aspect of the paper's positioning is that the Advantage Model and Selective Rehearsal are not independent proposals — they are synergistic. The Advantage Model provides the calibrated, task-invariant scores that Selective Rehearsal uses to identify high-quality examples within each cluster. Without the Advantage Model's score regularization, selecting examples by raw reward model score would favor categories with inflated score distributions (e.g., selecting disproportionately many Code Generation examples even for QA clusters, because Code Generation scores are systematically higher). Conversely, Selective Rehearsal provides a concrete use case for the Advantage Model's scores that goes beyond simply feeding them to PPO — the scores guide data curation decisions that shape the entire training trajectory.

This synergy is reflected in the experimental results (Table 2): AM-PPO-SR (Advantage Model + Selective Rehearsal) achieves 15.78% win rate over SFT on the main test set, compared to 14.87% for AM-PPO without Selective Rehearsal and 12.72% for RM-PPO. On the forget test set, AM-PPO-SR achieves 10.30% win / 7.95% loss, dramatically better than RM-PPO's 16.87% win / 29.28% loss. Each component contributes independently (comparing AM-PPO to RM-PPO shows the Advantage Model's contribution, comparing AM-PPO-SR to AM-PPO shows Selective Rehearsal's), and the combination yields the best results on both metrics.

What the Paper Does NOT Address

The paper is explicit about several boundaries that contextualize its contributions:

  • Single model family, specific scale. All experiments use BLOOMZ models (7B for reward/advantage modeling, 176B for policy training). The authors do not claim universality across model families or scales, though the techniques are architecturally independent of the base model.

  • The margin parameter m(x)m(x) is not theoretically grounded. Equation 6 introduces m(x)m(x), a prompt-dependent bounding margin for advantage scores. The authors state that m(x)m(x) "may have a connection with the complexity or difficulty involved in learning the reward function for prompts similar to xx" but explicitly label this as "speculative and requires further investigation" (Section 3.1, footnote). Throughout experiments, m(x)m(x) is simply set to 2.5. This is a hyperparameter whose optimal value is not explored and whose theoretical justification remains open.

  • The expected reward estimation is approximate. Equation 7 defines the expected reward eτ(x)e_\tau(x) as the expectation over the current policy's responses, but actually computing this expectation exactly is intractable (it requires sampling all possible responses). The paper's practical approximation (Equation 8) uses a weighted combination of the current policy and previous policy models, introducing several additional hyperparameters (NN, KK) whose sensitivity is not analyzed.

  • Selective Rehearsal does not explore all selection criteria. The paper uses only advantage model score for selecting examples within clusters (Section 3.2) and mentions but does not explore alternatives such as combining advantage scores with entropy, human satisfaction rate, or response length. The cluster count cc is briefly studied (Figure 6) and found to have relatively small impact (approximately 0.05 points variance in test-set rewards), but the authors still recommend "a thorough analysis of this aspect when applying selective rehearsal to different datasets."

  • No combination with alternative alignment methods. The paper remains entirely within the RLHF paradigm and does not compare with or combine its techniques with DPO, PRO, or other PPO-free alignment methods. The relationship between the Advantage Model's score regularization and the implicit reward modeling in DPO (which directly optimizes a policy from preferences without an explicit reward model) is left unexplored.

3. Technical Approach

This is a stabilization paper — it does not propose a fundamentally new alignment paradigm, but identifies two specific failure modes in the standard RLHF pipeline (reward hacking from uncalibrated reward scores and catastrophic forgetting from uniform PPO optimization) and introduces two complementary mechanisms that intervene at different stages of the pipeline to address them while preserving the core PPO + reward model architecture.

3.1 Reader Orientation

What the system is (in plain language): The paper builds two drop-in modifications to the standard RLHF training pipeline — one replacing the reward model with an "Advantage Model" that outputs calibrated scores constrained around zero, and another modifying PPO training to periodically rehearse high-quality examples grouped by semantic similarity — so that the language model being fine-tuned improves on human preference while not forgetting skills it already learned during supervised fine-tuning.

What problem it solves and the shape of the solution: RLHF training repeatedly produces models that either game the reward signal (producing high-scoring but nonsensical or misaligned outputs) or lose capabilities they possessed after supervised fine-tuning. The solution is two-stage: (i) constrain the reward model's output distribution so that different task categories produce comparable score scales, removing the incentive for the policy to exploit between-category score differences; and (ii) curate a diverse, high-quality subset of the PPO training data and add a supervised learning loss on this subset, so that the policy is explicitly incentivized to retain correct behavior on representative examples from every skill cluster.

3.2 Big-Picture Architecture (Diagram in Words)

The system modifies the standard RLHF pipeline at two points, with a natural information flow connecting them:

  1. Base LLM (BLOOMZ 176B, pre-trained) — the large language model to be aligned, shared across all components.

  2. Reward / Advantage Model Training — takes preference comparison data (pairs of responses where humans indicate which is better) and trains a model (BLOOMZ 7B with a scalar value head) to assign scores. The standard Reward Model (RM) uses the Bradley-Terry pairwise ranking loss (Equation 1). The proposed Advantage Model (AM) uses the same architecture but a different objective (Equation 6) that subtracts an expected reward baseline and constrains scores to lie within a bounded margin, producing advantage scores — extra reward relative to what is typical for that prompt.

  3. PPO Training Loop — samples prompts, generates responses from the current policy, scores them with the reward/advantage model, and updates the policy using the PPO objective (Equation 2) with a KL-divergence penalty against the SFT initialization. This is where RLHF happens.

  4. Selective Rehearsal Module — a data curation step that executes before PPO training. It takes the PPO prompt set, embeds prompts using SimCSE, clusters them into $c$ semantically similar groups, selects the top-$k$ responses within each cluster by advantage model score, and constructs a rehearsal dataset $D_R$. During PPO training, an additional negative log-likelihood (NLL) loss on $D_R$ is added to the standard PPO objective, weighted by a coefficient $\gamma = 0.01$.

The flow: preference data → train Advantage Model → PPO training data → Selective Rehearsal curation → PPO training with combined loss → aligned model. The Advantage Model feeds into Selective Rehearsal (its scores are the selection criterion), and does standard duty of reward signal for PPO training.

3.3 Roadmap for the Deep Dive

  • First, the Reward Model baseline (Equation 1): the standard Bradley-Terry pairwise ranking objective that the Advantage Model replaces. Understanding its unconstrained score degrees of freedom is essential to understanding why the Advantage Model's bounding loss is necessary.

  • Second, the Advantage Model (Equations 5–8): how it reformulates the modeling target from raw rewards to advantages, introduces an expected reward baseline, constrains scores with a bounding loss, and approximates the intractable baseline in practice. This is the paper's core technical contribution.

  • Third, the calibration and regularization properties of the Advantage Model: why the bounding loss produces better ECE, more comparable per-task means and variances, and resistance to reward hacking. This connects the mathematical formulation to the observed empirical behavior.

  • Fourth, the Selective Rehearsal pipeline (Section 3.2): the clustering, selection, and rehearsal training steps. This is the data-centric contribution that synergizes with the Advantage Model.

  • Fifth, the Selective Rehearsal integration with PPO (Equation 9): how the rehearsal NLL loss combines with the standard PPO objective, the weighting scheme, and the interpretation as a localized KL penalty.

3.4 Detailed, Sentence-Based Technical Breakdown

Reward Model Baseline (The Standard Bradley-Terry Preference Objective)

The paper begins from the standard reward modeling formulation used in RLHF, which serves as the baseline that the Advantage Model improves upon. In the standard setup, a reward model $r_\theta(x, y)$ is a scalar-valued function parameterized by a neural network (here, BLOOMZ 7B with the final language modeling head replaced by a single linear layer outputting a scalar) that takes a prompt $x$ and a response $y$ and outputs a single number representing how good that response is according to human preferences.

The training objective for this reward model is the Bradley-Terry pairwise ranking loss:

LRM=−E(x,yc,yr)∼DRM[log⁡(σ(rθ(x,yc)−rθ(x,yr)))]\mathcal{L}_{\text{RM}} = -\mathbb{E}_{(x, y_c, y_r) \sim \mathcal{D}_{\text{RM}}} \left[ \log\left( \sigma( r_\theta(x, y_c) - r_\theta(x, y_r) ) \right) \right]

where $r_\theta(x, y)$ is the scalar reward score assigned by the model with parameters $\theta$ to response $y$ given prompt $x$, $y_c$ is the human-preferred (chosen) response from a comparison pair, $y_r$ is the rejected response, $\sigma(z) = 1/(1 + e^{-z})$ is the logistic sigmoid function, and $\mathcal{D}_{\text{RM}}$ is the dataset of human preference comparisons — specifically a set of triples $(x, y_c, y_r)$ where human labelers indicated that $y_c$ is better than $y_r$ for prompt $x$.

What it computes: For each training example, the model produces two scalar scores — one for the chosen response and one for the rejected response — and computes their difference $r_\theta(x, y_c) - r_\theta(x, y_r)$. This difference is passed through the logistic sigmoid to produce a probability $\sigma(\Delta)$ that the chosen response is preferred. The loss is the negative log of this probability: when $\Delta$ is large and positive, $\sigma(\Delta) \approx 1$ and the loss is near zero; when $\Delta$ is near zero or negative, the loss grows large, penalizing the model for failing to assign a higher score to the chosen response. This is summed (in expectation) over all pairs in the training set.

Why this form: The Bradley-Terry model is the standard probabilistic model for pairwise comparisons: it assumes that each item $i$ has a latent "strength" $s_i$, and the probability that item $i$ beats item $j$ in a comparison is $\sigma(s_i - s_j)$. This is a natural choice for preference modeling because it captures the intuition that the probability one response beats another should be a smooth, monotonically increasing function of the score difference. The negative log-likelihood is the maximum-likelihood objective under this model. However — and this is the critical property the paper exploits — the Bradley-Terry objective is invariant under any monotonically increasing transformation of the scores. Adding a constant $c$ to all scores for responses to a particular prompt, or multiplying all scores by a positive constant, does not change any pairwise difference $r_\theta(x, y_c) - r_\theta(x, y_r)$, and therefore does not change the loss. This means the model has unconstrained degrees of freedom: it can assign arbitrarily different score scalings to different prompts or task categories without any training penalty, as long as within-category rankings are preserved. This degree of freedom is the root cause of the inter-category score disparities shown in Figure 4a, where Code Generation tasks have substantially higher mean rewards than QA tasks simply because the model settled into different score scales for the two categories during training.

Advantage Model: Formulation and Objective

The Advantage Model addresses the unconstrained scaling problem by reformulating the reward modeling target. Instead of modeling absolute reward scores $r_\theta(x, y)$, it models advantage scores $a_\theta(x, y)$ — the extra reward that response $y$ achieves relative to the expected reward for prompt $x$. The formal definition is:

aθ(x,y)=rθ(x,y)−Ey∼π′(x)[πϕ(y∣x)π′(y∣x)rθ(x,y)]a_\theta(x, y) = r_\theta(x, y) - \mathbb{E}_{y \sim \pi'(x)} \left[ \frac{\pi_\phi(y|x)}{\pi'(y|x)} r_\theta(x, y) \right]

where $r_\theta(x, y)$ is the raw reward score (the same quantity the standard RM would output), $y \sim \pi'(x)$ denotes sampling responses from some policy $\pi'$ (the policy that was used during data collection), $\pi_\phi(y|x)$ is the probability that the current policy (the one being updated during PPO) assigns to response $y$, $\pi'(y|x)$ is the probability that the data-collection policy assigned to $y$, and the expectation computes a weighted average of reward scores over all possible responses, with the importance weight $\pi_\phi(y|x) / \pi'(y|x)$ correcting for the fact that the data-collection policy $\pi'$ may differ from the current policy $\pi_\phi$.

What it computes: The advantage score is the difference between the raw reward of a specific response $y$ and the expected (average) reward across all responses to prompt $x$, where the expectation is taken under the current policy's distribution (with importance sampling correction for distribution shift). If response $y$ is better than average — it achieves higher reward than what the model typically produces for that prompt — then $a_\theta(x, y) > 0$. If it is worse than average, $a_\theta(x, y) < 0$. If it is exactly average, $a_\theta(x, y) \approx 0$. This centers the score distribution around zero per prompt, which automatically removes prompt-level or task-level scale differences: a Code Generation response and a QA response that are both "equally better than typical for their respective tasks" will receive similar positive advantage scores, even if the raw reward scales for Code Generation and QA are completely different.

Why this form: This directly parallels the advantage function in reinforcement learning, defined as $A(s, a) = Q(s, a) - V(s)$, where the state-value baseline $V(s)$ is subtracted to reduce variance in policy gradient estimates. Here, the prompt $x$ is analogous to the state, and the expected reward $\mathbb{E}_{y}[r_\theta(x, y)]$ is analogous to the state-value function — it represents "how much reward you expect to get from this prompt on average." Subtracting it produces a signal that reflects incremental quality rather than absolute quality, which has two critical properties for RLHF stability:

  1. Task-invariance: The baseline absorbs all prompt-level and task-level variation in reward scale. Two different prompts $x_1$ and $x_2$ may have very different expected rewards (e.g., 2.0 for Code Generation and 0.5 for QA), but the advantage scores for both will be centered around zero. This removes the incentive for the PPO policy to shift response patterns toward higher-reward categories, since doing so does not increase the advantage score.

  2. Bounded signal: By construction, advantages should be positive for good responses and negative for bad responses, with zero representing the decision boundary. This natural centering makes it easier to impose explicit bounds — the $m(x)$ margin in Equation 6 — because the scores already cluster around zero rather than drifting to arbitrary offsets.

However, the definition in Equation 5 is not directly computable: the expectation is over all possible responses $y \sim \pi'(x)$, which is an intractably large (infinite) set. The paper addresses this with a practical approximation described below.

Advantage Model: The Training Objective (Ranking + Bounding)

The training objective for the Advantage Model combines two losses: a ranking loss that preserves the pairwise preference information (analogous to the standard RM loss), and a bounding loss that constrains the magnitude of advantage scores to prevent extreme values. The joint objective is:

LAM=−E(x,yc,yr)∼DRM[log⁡(σ(aθ(x,yc)−aθ(x,yr)))+log⁡(σ(m(x)−aθ(x,yc)))+log⁡(σ(m(x)+aθ(x,yr)))]\mathcal{L}_{\text{AM}} = -\mathbb{E}_{(x, y_c, y_r) \sim \mathcal{D}_{\text{RM}}} \left[ \log\left( \sigma( a_\theta(x, y_c) - a_\theta(x, y_r) ) \right) + \log\left( \sigma( m(x) - a_\theta(x, y_c) ) \right) + \log\left( \sigma( m(x) + a_\theta(x, y_r) ) \right) \right]

where $a_\theta(x, y_c)$ and $a_\theta(x, y_r)$ are the advantage scores for the chosen and rejected responses respectively, $\sigma$ is the logistic sigmoid, $m(x)$ is a prompt-dependent margin parameter (set to 2.5 throughout experiments), and $\mathcal{D}_{\text{RM}}$ is the same preference comparison dataset used for standard reward model training.

What it computes: This loss has three terms, each of which is a negative log-probability encouraging a specific inequality:

  • Term 1, the ranking term: $\log(\sigma(a_\theta(x, y_c) - a_\theta(x, y_r)))$ — identical in form to the standard Bradley-Terry loss (Equation 1), but operating on advantage scores rather than raw rewards. It encourages the chosen response's advantage to be greater than the rejected response's advantage: $a_\theta(x, y_c) > a_\theta(x, y_r)$. This preserves the pairwise preference ranking.

  • Term 2, the upper bound on chosen responses: $\log(\sigma(m(x) - a_\theta(x, y_c)))$ — encourages the chosen response's advantage to be less than the margin $m(x)$: $a_\theta(x, y_c) < m(x)$. Penalizes the model if a good response receives an advantage score larger than $m(x)$, effectively capping how positive advantages can become.

  • Term 3, the lower bound on rejected responses: $\log(\sigma(m(x) + a_\theta(x, y_r)))$ — encourages the rejected response's advantage to be greater than $-m(x)$: $a_\theta(x, y_r) > -m(x)$. Penalizes the model if a bad response receives an advantage score more negative than $-m(x)$, effectively capping how negative advantages can become.

All three terms are summed and negated (maximizing likelihood), and the expectation is taken over the training dataset. The net effect: the model is trained to rank responses correctly (term 1) while keeping all advantage scores within the interval $[-m(x), m(x)]$ (terms 2 and 3), with $m(x) = 2.5$ in all experiments.

Why this form: The three-term structure directly addresses both requirements of a stable reward signal:

  • Ranking preservation (term 1): Essential — the model must still distinguish good from bad responses. Without this term, the model could satisfy the bounding constraints by outputting constant scores near zero for everything, which would make the advantage scores useless for guiding PPO. The ranking term ensures that the ordering information from human preferences is retained.

  • Score bounding (terms 2 and 3): Essential for preventing reward hacking — without these, the advantage scores could still drift to extreme values, because the ranking term only constrains differences between pairs, not absolute magnitudes. A model could assign $a_\theta(x, y_c) = 100$ and $a_\theta(x, y_r) = 99$ and still satisfy the ranking loss perfectly, but this would produce the same inter-category scaling problems as the standard RM: Code Generation responses might drift to advantages of 5–10 while QA responses stay at 0–1, recreating the perverse incentive for PPO to shift response patterns toward high-magnitude categories. The bounding terms prevent this by explicitly penalizing any score with $|a_\theta| > m(x)$.

  • The choice of $m(x) = 2.5$: The paper does not provide a derivation or theoretical justification for this value. At $m = 2.5$, the sigmoid $\sigma(2.5) \approx 0.924$, meaning the bounding loss incurs $-\log(0.924) \approx 0.079$ nats of penalty when a score is exactly at the boundary — a relatively soft constraint. The authors speculate that $m(x)$ "may have a connection with the complexity or difficulty involved in learning the reward function for prompts similar to $x$" (Section 3.1, footnote), suggesting that more ambiguous prompts might need wider margins, but this connection is not explored. Setting $m(x)$ to a constant 2.5 across all prompts is a simplification that the paper acknowledges as an open question.

  • Why the sigmoid bounding form rather than hard clipping: Hard clipping (e.g., $\max(-m, \min(m, a))$) would produce zero gradients outside the margin, preventing the model from learning to pull out-of-bound scores back toward the valid range. The soft sigmoid penalty provides continuous gradients everywhere: when $a_\theta(x, y_c) \gg m(x)$, the term $\log(\sigma(m - a))$ grows approximately linearly (since $\log(\sigma(z)) \approx z$ for large negative $z$), creating a strong gradient pulling the score down. This is numerically better behaved than hard constraints and integrates naturally with gradient-based optimization.

Advantage Model: Practical Approximation of the Expected Reward Baseline

The definition in Equation 5 requires computing $\mathbb{E}_{y \sim \pi'(x)} [ \frac{\pi_\phi(y|x)}{\pi'(y|x)} r_\theta(x, y) ]$, the expected reward under the current policy with importance sampling correction. This is intractable: the expectation is over all possible token sequences constituting valid responses, which is astronomically large. The paper introduces a practical approximation by parameterizing the expected reward of the current policy as a learnable scalar $e_\tau(x)$ and combining it with importance-weighted samples from previous policy models.

The parameterized expected reward is:

eτ(x)=Ey∼πϕ(x)[rθ(x,y)]e_\tau(x) = \mathbb{E}_{y \sim \pi_\phi(x)} \left[ r_\theta(x, y) \right]

where $e_\tau(x)$ is a learnable function (parameterized by $\tau$, implemented as a neural network prediction) that estimates the average reward the current policy $\pi_\phi$ would achieve on prompt $x$. The notation $\tau$ distinguishes these parameters from the advantage model parameters $\theta$.

Using this, the full practical approximation for the advantage score is:

aθ(x,y)=rθ(x,y)−N−KNeτ(x)−∑k=1K1Nπϕ(y∣x)πk′(y∣x)rθ(x,y)a_\theta(x, y) = r_\theta(x, y) - \frac{N - K}{N} e_\tau(x) - \sum_{k=1}^{K} \frac{1}{N} \frac{\pi_\phi(y|x)}{\pi'_k(y|x)} r_\theta(x, y)

where $N$ is a hyperparameter balancing emphasis on the current policy versus historical policies, $K$ is the number of alternate (historical) policy models used during comparison data collection, $\pi'_k(y|x)$ is the probability assigned to response $y$ by the $k$-th historical policy model, and $\pi_\phi(y|x)$ is the probability under the current policy.

What it computes: The advantage is the raw reward $r_\theta(x, y)$ minus a weighted combination of two baseline estimates:

  • The current-policy baseline: $\frac{N-K}{N} e_\tau(x)$ — the learnable expected reward estimate for the current policy, weighted by $(N-K)/N$. When $N$ is large relative to $K$, this term dominates, meaning the baseline is mostly determined by the current policy's expected reward (which is what we actually want to subtract — the current policy's average performance).

  • The historical-policy correction: $\sum_{k=1}^{K} \frac{1}{N} \frac{\pi_\phi(y|x)}{\pi'_k(y|x)} r_\theta(x, y)$ — importance-weighted reward samples from each of the $K$ historical policies, each weighted equally by $1/N$. This term corrects for the fact that the comparison data was not collected under the current policy: the importance weights $\pi_\phi(y|x) / \pi'_k(y|x)$ re-weight each response's reward to reflect how much more or less likely it is under the current policy compared to the policy that actually generated it.

The net effect: the advantage score subtracts a prompt-dependent baseline that is primarily determined by the current policy's expected reward (via $e_\tau(x)$) but stabilized by historical data to prevent the baseline from becoming degenerate when the current policy has not yet explored certain response types.

Why this form: Several design considerations shape this approximation:

  • The $e_\tau(x)$ parameterization avoids needing to sample from $\pi_\phi$ for every advantage computation. If the advantage model had to actually generate responses from $\pi_\phi$ and score them to compute the baseline, the advantage computation would depend on the current policy in a way that is expensive (requiring inference) and potentially unstable (as $\pi_\phi$ changes during PPO, the advantage scores would change even for the same $(x, y)$ pair). Parameterizing $e_\tau(x)$ as a learned function makes the baseline a direct prediction, decoupling it from policy sampling during advantage model inference.

  • The $N$ and $K$ hyperparameters control a bias-variance tradeoff. A small $K$ (few historical policies) with large $N$ makes the baseline mostly $e_\tau(x)$, which is low-variance (a single prediction per prompt) but potentially biased if $\pi_\phi$ has diverged significantly from the data collection policies. A large $K$ with small $N$ makes the baseline mostly the importance-weighted historical rewards, which is less biased (correctly accounting for distribution shift) but higher-variance (sum of individual importance-weighted rewards). The paper does not report the specific values of $N$ and $K$ used in experiments, nor does it ablate them — this is a limitation of the experimental analysis.

  • The importance weight $\pi_\phi(y|x) / \pi'_k(y|x)$ is the standard off-policy correction. In reinforcement learning, when you have data collected under a behavior policy $\pi'$ but want to estimate expectations under a target policy $\pi_\phi$, you multiply by the ratio of probabilities (the Radon-Nikodym derivative). This is necessary here because the comparison data $\mathcal{D}_{\text{RM}}$ was collected using one or more historical policies ($\pi'_k$), but the advantage should be relative to the current policy's expected reward. Without this correction, the baseline would reflect the historical policy's average performance, not the current policy's, and would fail to adapt as the policy improves during PPO.

  • Why not just use $e_\tau(x)$ alone? The historical correction terms provide a form of regularization: they ensure the baseline is grounded in actual observed rewards from real responses, not solely in the learned function's predictions (which could be poorly calibrated early in training). This hybrid approach mirrors techniques from value function estimation in RL, where a learned value function is combined with Monte Carlo returns for more stable training.

Advantage Model: Calibration and Regularization Properties

The Advantage Model's formulation produces three concrete properties that the paper's experiments confirm are responsible for improved RLHF stability. These are not additional mechanisms but rather consequences of the objective design described above:

Property 1: Lower Expected Calibration Error (ECE). ECE measures the gap between a model's predicted probabilities and empirical frequencies. For a reward model, calibration means that if the model assigns a score difference $\Delta = a(x, y_c) - a(x, y_r)$ implying a preference probability of $\sigma(\Delta)$, then among all pairs with that score difference, the chosen response should actually be preferred $\sigma(\Delta)$ fraction of the time. The bounding terms in Equation 6 act as a regularizer that prevents the model from becoming overconfident: they explicitly penalize extreme scores, which pushes the model toward more moderate (and typically better calibrated) predictions. Table 1 reports ECE of 3.48 for AM vs. 4.70 for RM on HH-RLHF, and 3.83 vs. 5.35 on proprietary data — reductions of approximately 26% and 28% respectively. Figure 2 visualizes this: the AM's observed accuracy closely follows the diagonal line $1/(1 + e^{-\Delta})$ representing perfect calibration, while the RM's curve deviates substantially, indicating systematic overconfidence.

Property 2: Comparable means and stable variances across tasks. The prompt-dependent baseline $e_\tau(x)$ in the advantage computation absorbs task-specific reward scale differences. If Code Generation prompts have a high expected reward (say, 2.0) and QA prompts have a low expected reward (say, 0.5), subtracting these baselines centers both distributions around zero: $a_{\text{code}} = r_{\text{code}} - 2.0$, $a_{\text{qa}} = r_{\text{qa}} - 0.5$. Figure 4a confirms this: the AM's per-task mean advantages cluster tightly around zero (within roughly ±0.2 for all task categories), whereas the RM's per-task mean rewards show large dispersion (spanning approximately 0.1 to 0.6 in Figure 4a). Figure 4b shows that the AM's per-task standard deviations are also more uniform (within approximately 0.35–0.45), while the RM's vary more widely.

Property 3: Resistance to reward over-optimization. When PPO maximizes a standard reward model's scores, it can increase the reward without actually improving — and sometimes degrading — response quality. This happens because the reward model has "blind spots" where high scores do not correspond to human-preferred outputs. The Advantage Model's score bounding ($|a_\theta| \leq m(x)$) and task-invariance reduce the PPO policy's ability to exploit these blind spots. If all task categories have their advantage scores centered around zero with similar variances, there is no free reward gain to be had by shifting response patterns from one category to another — the policy must actually improve response quality within each category to increase its advantage scores. Figure 5b demonstrates this empirically: while RM-PPO's GPT-4 win rate drops during training (despite rising RM scores), AM-PPO's GPT-4 win rate remains stable or increases.

Selective Rehearsal: Representative Example Discovery via Clustering

Selective Rehearsal operates in two sequential phases: representative example discovery (data curation) and rehearsal training (modifying the PPO objective). The discovery phase executes before PPO training begins and produces a curated dataset that is then used during training.

Input data: The starting point is $\mathcal{D}_{\text{PPO}}$, the set of prompts for PPO training, with responses generated from the initial policy model $\pi_{\text{init}}$ (the SFT model, before any PPO updates). Specifically, $\mathcal{D}_{\text{PPO}} = [(x_1, y_1), (x_2, y_2), \ldots]$ where each $y_i$ is sampled from $\pi_{\text{init}}(x_i)$.

Step 1: Embedding. Each prompt $x$ is encoded into a fixed-dimensional vector representation using SimCSE (Gao et al., 2021), specifically the sup-simcse-roberta-base model from HuggingFace. SimCSE is a contrastively trained sentence embedding model: it was trained so that semantically similar sentences have high cosine similarity in the embedding space, while dissimilar sentences are far apart. The paper uses SimCSE embeddings to represent the prompts (not the full prompt-response pairs), so the clustering groups prompts that ask for similar types of tasks — for example, all math problem prompts would have similar embeddings, as would all creative writing prompts, even if the specific content differs.

The choice of SimCSE over alternatives (BERT embeddings, TF-IDF, task-category labels) is not extensively justified. The paper simply states using "SimCSE sentence embedding to represent the query $x$" (Section 3.2). The practical motivation is that SimCSE provides off-the-shelf semantic similarity without requiring task-specific labels — the clustering can discover task groupings automatically from the prompt text alone, which is important for the proprietary dataset with 61 task categories where manual labeling would be expensive.

Step 2: Clustering. The prompt embeddings are clustered into $c$ groups using the KMeans algorithm. KMeans partitions the $n$ prompts into $c$ disjoint sets $S_1, S_2, \ldots, S_c$ by minimizing the within-cluster sum of squared distances:

arg⁡min⁡S∑i=1c∑x∈Si∥embed(x)−μi∥2\arg\min_S \sum_{i=1}^{c} \sum_{x \in S_i} \| \text{embed}(x) - \mu_i \|^2

where $\text{embed}(x)$ is the SimCSE embedding of prompt $x$, $\mu_i$ is the centroid (mean) of all embeddings in cluster $S_i$, and the optimization assigns each prompt to the cluster whose centroid it is closest to.

The number of clusters $c$ is a hyperparameter. The paper briefly studies its effect (Figure 6) and finds that varying $c$ produces "a relatively consistent variance of approximately 0.05 points in test-set rewards," suggesting the method is not highly sensitive to this choice. The paper does not report the specific values of $c$ used for the main experiments, which limits reproducibility. Based on the figure, values in the range of 10–50 appear to have been tested.

Step 3: Selection within clusters. Within each cluster $S_i$, the method selects the top-$k$ (prompt, response) pairs ranked by their advantage model score $a_\theta(x, y)$ — the same advantage score that the Advantage Model (Section 3.1) produces for PPO training. The paper states this choice explicitly: "Within each cluster, here we simply choose the top-$k$ (x, y) pairs with the highest advantage model score" (Section 3.2).

The selection by advantage score serves two purposes:

  • Quality filtering: Higher advantage scores indicate responses that are better than expected for their prompt — these are the responses that the model should rehearse to retain its good behavior.
  • Uniformity across clusters: Because the Advantage Model's scores are already task-invariant (centered around zero with similar variances across tasks, as shown in Figure 4), selecting by advantage score does not favor certain task categories over others. If raw reward model scores were used instead, the selection would be biased toward categories with inflated score distributions (e.g., Code Generation), resulting in rehearsal data that over-represents those categories and under-represents others — exactly the kind of imbalance that causes catastrophic forgetting of under-represented skills.

The paper mentions but does not implement alternative or combined selection criteria: "entropy (low entropy indicates high confidence), human satisfaction rate or response length (higher length may indicate redundancy)" (Section 3.2). These are flagged as future work.

Step 4: Assembly. The selected pairs from all clusters are shuffled together to form the rehearsal dataset $\mathcal{D}_R$. The shuffling ensures that during rehearsal training, batches contain examples from diverse clusters, preventing the model from overfitting to one skill type at a time. The size of $\mathcal{D}_R$ is determined by $c \times k$: number of clusters times the per-cluster selection count. Neither $c$ nor $k$ are reported as specific numbers, though the total rehearsal dataset size is implied to be a small fraction of the full PPO dataset (since selection picks only top examples per cluster).

Why this design (instead of uniform random selection or pure top-score selection):

  • Compared to uniform random selection from $\mathcal{D}_{\text{PPO}}$: Random selection would cover the full diversity of skills (each cluster gets representation proportional to its size in the training data), but would include low-quality examples — responses that the model should not rehearse because they exhibit incorrect or suboptimal behavior. Selective Rehearsal adds the quality filter, ensuring that only good responses are reinforced.

  • Compared to pure top-$k$ selection by score (without clustering): Selecting the globally highest-scoring pairs ignores diversity entirely. If the advantage model slightly favors certain response patterns (e.g., longer responses, responses with certain formatting), the top-scoring set could collapse to a narrow mode of the response distribution. The clustering step guarantees representation from every semantic region of the prompt space, preserving the full range of skills.

  • Compared to rejection sampling (Touvron et al., 2023): Rejection sampling and reinforced self-training (Gulcehre et al., 2023) also select high-quality self-generated responses for supervised training, but "only consider reward model score" (Section 3.2). Selective Rehearsal explicitly adds the diversity constraint via clustering, which the paper argues captures "multi-dimensional important aspects" beyond scalar reward.

  • Why only PPO training data? The paper notes that rehearsal data is drawn from $\mathcal{D}_{\text{PPO}}$ with responses generated from the initial policy "to enable a more fair and nuanced comparison, as no additional information is introduced." In other scenarios, the rehearsal pairs "could come from other important data sources representing specific skills (e.g., math-problem solving) the main policy are not expected to forget." This flexibility means Selective Rehearsal can be extended to incorporate external calibration data for critical skills, but the paper keeps the evaluation controlled by using only PPO-internal data.

Selective Rehearsal: Integration with PPO Training

Once the rehearsal dataset $\mathcal{D}_R$ is constructed, it is used to augment the standard PPO objective with an additional loss term during RLHF training. The modified training objective is:

LPPO-SR=LPPO+γ⋅E(x,y)∼DR[∑t=1∣y∣πϕ(yt∣y<t,x)]\mathcal{L}_{\text{PPO-SR}} = \mathcal{L}_{\text{PPO}} + \gamma \cdot \mathbb{E}_{(x, y) \sim \mathcal{D}_R} \left[ \sum_{t=1}^{|y|} \pi_\phi(y_t | y_{<t}, x) \right]

where $\mathcal{L}_{\text{PPO}}$ is the standard PPO objective defined in Equation 2, $\gamma = 0.01$ is a fixed coefficient controlling the strength of the rehearsal loss relative to the PPO loss, $\mathcal{D}_R$ is the curated rehearsal dataset, $(x, y)$ is a (prompt, response) pair from $\mathcal{D}_R$, $|y|$ is the number of tokens in response $y$, $y_t$ is the $t$-th token of $y$, $y_{<t}$ represents all tokens before position $t$, and $\pi_\phi(y_t | y_{<t}, x)$ is the policy's predicted probability for the correct token $y_t$ given the prompt and previous tokens.

What it computes: The rehearsal term is a standard causal language modeling negative log-likelihood (NLL) loss: for each token in each rehearsal response, the policy must assign high probability to the correct token. The sum over tokens computes the log-probability of the entire response under the policy, and the expectation averages this over the rehearsal dataset. The coefficient $\gamma = 0.01$ means the rehearsal loss has 1% of the weight of the PPO loss — it is a relatively weak auxiliary objective designed to nudge the policy toward retaining correct behavior on representative examples, without dominating the primary goal of maximizing the reward signal.

Why this form: The rehearsal loss addresses catastrophic forgetting through a specific mechanism that is distinct from the KL-divergence penalty already present in the PPO objective (Equation 2, the $-\beta \log(\pi_\phi / \pi_{\text{init}})$ term):

  • The KL penalty is global and distributional: It penalizes the policy for changing its aggregate output distribution relative to the initialization, measured in KL-divergence. This is a soft constraint: the policy can change behavior substantially on specific prompts as long as the average divergence across all prompts stays below some threshold. As the forget test set results demonstrate (Table 2, RM-PPO: 29.28% loss rate against SFT), this global constraint is insufficient to protect individual examples or skill categories from being overwritten. A policy can satisfy a KL constraint of 5 nats while still catastrophically forgetting certain skills, because the aggregate divergence is dominated by the majority of prompts where the policy stays close to initialization.

  • The rehearsal loss is local and example-specific: It explicitly forces the policy to maintain high probability on specific high-quality (prompt, response) pairs. This is a hard constraint on those specific examples, not an aggregate statistic. The clustering-based selection ensures these examples are distributed across all skill categories, so the local constraint provides broad coverage. The NLL loss form means the rehearsal acts as supervised fine-tuning on the selected examples, directly reinforcing the SFT behavior that produced those responses.

  • The small coefficient $\gamma = 0.01$ reflects a deliberate balance: If $\gamma$ were too large, the rehearsal loss would overwhelm the PPO reward-driven update, preventing the policy from improving on prompts where the SFT model's responses are actually suboptimal. The motivation behind Selective Rehearsal is not to prevent any change from SFT — that would defeat the purpose of RLHF — but to prevent destructive change on examples where the SFT behavior was already good. The small $\gamma$ provides a gentle gradient pulling the policy back toward the SFT distribution on the selected examples, which is sufficient to counteract the reward-driven drift that causes forgetting, without suppressing genuine improvements.

Alternative interpretation as localized KL amplification: The paper notes (Section 3.2, end) that "one can view selective rehearsal as a means of amplifying the weight of the KL-divergence term in PPO training for crucial instances and their related counterparts." This is because both the rehearsal NLL loss and the KL penalty encourage the policy to stay close to the initial distribution, but the rehearsal loss does so in a targeted way — only on the selected examples — while the KL penalty applies uniformly. The clustering step ensures the "related counterparts" (other examples in the same semantic cluster) also benefit, since the policy's behavior on similar prompts is correlated.

Full Training Pipeline: How the Components Interact

Putting the pieces together, the complete training pipeline operates as follows:

  1. Before training: Train the Advantage Model on the human preference comparison dataset $\mathcal{D}_{\text{RM}}$ using the three-term loss in Equation 6. This produces a calibrated, task-invariant scoring function $a_\theta(x, y)$.

  2. Data curation: Use the initial policy (SFT model $\pi_{\text{init}}$) to generate responses for all PPO prompts, forming $\mathcal{D}_{\text{PPO}}$. Apply Selective Rehearsal: embed prompts with SimCSE, cluster with KMeans into $c$ groups, select top-$k$ responses per cluster by advantage score $a_\theta(x, y)$, and assemble the rehearsal dataset $\mathcal{D}_R$.

  3. PPO training loop (repeated for 100 steps): For each update step:

    • Sample a batch of prompts from $\mathcal{D}_{\text{PPO}}$ and generate responses from the current policy $\pi_\phi$.
    • Compute advantage scores $a_\theta(x, y)$ for each generated response using the pre-trained Advantage Model.
    • Compute the standard PPO loss (Equation 2) using these advantage scores as the reward signal, along with the KL-divergence penalty against $\pi_{\text{init}}$ ($\beta$ coefficient not reported in the paper).
    • Sample a batch from the rehearsal dataset $\mathcal{D}_R$ and compute the NLL loss (Equation 9, the $\gamma$ term).
    • Sum the two losses (PPO + $\gamma \cdot$ NLL) and update policy parameters $\phi$ via gradient descent.
    • Optionally, recompute $\mathcal{D}_R$ if the policy has changed substantially (the paper does not specify whether rehearsal data is refreshed during training — the text implies it is pre-computed once from the initial policy, given the statement about using responses "generated from the initial policy model").

  4. After training: The final aligned policy $\pi_\phi$ is evaluated on the main test set (standard PPO evaluation), the forget test set (expert-aligned SFT examples to measure forgetting), and via GPT-4 win rate comparisons against the SFT baseline.

The Advantage Model and Selective Rehearsal are synergistic at multiple points:

  • Advantage scores drive both the PPO optimization and the rehearsal data selection. This ensures consistency: the same signal that guides which responses to produce (via PPO) also guides which responses to protect from forgetting (via rehearsal).
  • The Advantage Model's task-invariance prevents rehearsal data bias. Without task-invariant scores, selecting "top-$k$ by score" would favor task categories with inflated score distributions. The Advantage Model's centering ensures fair representation across clusters.
  • The rehearsal stabilizes the baseline for advantage computation. If the policy drifts too far from the initial distribution, the importance weights $\pi_\phi(y|x) / \pi'(y|x)$ in Equation 8 can become extreme, increasing variance in the advantage estimates. By keeping the policy close to the initial distribution on selected examples, rehearsal indirectly stabilizes the advantage signals themselves.

4. Key Insights and Innovations

Innovation 1: Reframing Reward Modeling as Advantage Estimation to Eliminate Spurious Degrees of Freedom

The paper's most conceptually distinctive move is not a new architecture or training trick, but a reformulation of the reward modeling problem itself. The standard Bradley-Terry objective for training reward models (Equation 1) has a well-known but underappreciated property: it is invariant to any monotonic transformation of scores. A model that assigns scores of (5, 3) to a (chosen, rejected) pair achieves exactly the same loss as one that assigns (500, 300) or (-1, -3). The training signal only constrains relative ordering within each pair, leaving the model free to develop arbitrarily different score scales for different prompts, tasks, or response types.

Prior work treated this as a nuisance to be patched downstream — normalizing scores during PPO training with moving averages, clipping rewards, or training separate reward models per task category to avoid cross-contamination. The paper's insight is that this degree of freedom is not merely inconvenient but is the root cause of reward hacking in multi-task RLHF. When the reward model learns that Code Generation responses naturally score around 2.0 while QA responses score around 0.5 (Figure 4a), PPO optimization discovers a shortcut: make QA responses structurally resemble Code Generation responses. This requires no actual improvement in answer quality — only a surface-level distribution shift — yet the reward signal increases because the policy has moved into a region of higher absolute scores. The incentive is spurious, but PPO cannot distinguish spurious from genuine reward increase; it optimizes the scalar signal it receives.

The Advantage Model dissolves this entire class of failure by redefining what the model tries to predict. Instead of modeling absolute reward r(x, y), it models the advantage a(x, y) = r(x, y) - E[r] — how much better this response is than what would be expected for this prompt on average. The prompt-conditional baseline E[r] absorbs all task-level, prompt-level, and distribution-level variation in raw scores. A Code Generation response that is "slightly above average for code" and a QA response that is "slightly above average for QA" receive similar small positive advantages, even if their raw scores differ by an order of magnitude. The perverse incentive to shift response patterns across task categories disappears, because doing so does not change the advantage: moving a QA response toward code-like formatting might increase r(x, y), but it also increases E[r(x)] (since the expected reward for QA responses now includes this code-like pattern), leaving a(x, y) unchanged or potentially decreased.

This reframing is fundamental rather than incremental because it changes the optimization landscape that PPO operates on. With a standard reward model, the optimization surface has spurious ridges — regions of high reward that do not correspond to human-preferred outputs — created by the arbitrary scaling of scores across prompt categories. PPO will naturally climb these ridges. With the Advantage Model, these ridges are flattened: the advantage surface has zero mean by construction for every prompt, so the only way to increase advantage is to genuinely improve response quality relative to the expected baseline. This is not a regularization trick layered on top of a standard reward model; it is a structural property of the modeling target itself.

The bounding loss (terms 2 and 3 in Equation 6) further constrains the advantage surface so that scores cannot diverge to extreme magnitudes — a secondary defense against the model discovering new spurious maxima within the advantage space itself. The combination of centering (via the baseline subtraction) and bounding (via the margin penalty) creates a reward signal that is both task-invariant in expectation and bounded in magnitude, which the empirical results show translates directly to improved PPO stability: the GPT-4 win rate for AM-PPO remains stable during training while RM-PPO's drops despite rising RM scores (Figure 5b).

Evidence anchoring: Figure 4a demonstrates the baseline's effect concretely — per-task mean advantages under AM cluster tightly around zero, while per-task mean rewards under RM are widely dispersed. Figure 4b shows that AM also produces more uniform per-task variances. Figure 5b confirms the downstream impact: the decoupling of RM scores from GPT-4 win rate (the hallmark of reward hacking) is substantially reduced with AM.

Innovation 2: Diagnosing Catastrophic Forgetting in RLHF as a Data Coverage Problem, Not a Capacity Problem

The second conceptual contribution is a diagnostic reframing of why RLHF causes catastrophic forgetting. The standard view — implicit in the KL-divergence penalty that all RLHF implementations include — treats forgetting as a problem of how much the policy changes. If the policy moves too far from its initialization, it loses previously acquired skills. The solution under this view is to constrain the magnitude of change: the KL penalty directly penalizes the divergence between current and initial policy distributions.

The paper's diagnosis is fundamentally different. Forgetting is not (primarily) about the magnitude of policy change but about which examples receive optimization pressure and which are starved of reinforcement. The PPO objective uniformly applies reward-driven updates to all sampled prompts. On prompts where the SFT model already produces good responses, these updates are, at best, noise — the reward model cannot reliably identify improvements over an already-correct answer, so the gradient pushes the policy in a random direction. Over many steps, this random drift accumulates, and the policy moves away from its initial correct behavior. The KL penalty slows this drift globally but does not prevent it locally — a policy can satisfy a modest KL constraint while still drifting significantly on specific examples, as long as the aggregate divergence across all prompts stays within bounds.

This is a data coverage problem, not a capacity problem. The SFT model already has the capability to produce good responses on expert-aligned examples; the PPO process causes it to unlearn this capability because those examples receive destructive rather than constructive updates. The paper's Selective Rehearsal addresses this by explicitly reinforcing the SFT behavior on a diverse, high-quality subset of examples — not by constraining policy change, but by providing countervailing gradients that actively maintain performance on examples the PPO reward signal would otherwise degrade.

This reframing has implications beyond the specific method. It suggests that the KL penalty in standard RLHF is fundamentally misaligned with the forgetting problem: it constrains the wrong quantity (aggregate divergence rather than per-example retention) and uses the wrong mechanism (a penalty on change rather than an incentive for correct behavior). This may explain why prior work has found that KL penalties alone are insufficient to prevent forgetting, and why methods like experience replay — which explicitly reinforce past behavior — are more effective in continual learning settings. Selective Rehearsal is an instance of experience replay adapted to RLHF, but the paper's deeper contribution is articulating why this class of methods is necessary and what property of RLHF makes the global KL penalty inadequate.

The clustering-based selection strategy adds a further refinement: it enforces coverage of the full skill distribution, not just quality. Simple top-k selection by score (as in rejection sampling) can collapse to a narrow mode of the response distribution, leaving under-represented skills unprotected. By clustering first and selecting top examples within each cluster, Selective Rehearsal ensures that every skill category receives reinforcement — even categories where the SFT model's absolute scores are lower than in other categories. The Advantage Model's task-invariant scores make this clustering-by-score approach viable; without them, score-based selection would be biased toward task categories with inflated score distributions.

Evidence anchoring: Table 2 quantifies the magnitude of the forgetting problem: RM-PPO loses to the SFT model on 29.28% of forget-test examples while winning only 16.87% (a net loss of ~12 percentage points). AM-PPO-SR reduces this to 10.30% win / 7.95% loss (a net gain of ~2 percentage points). The swing from net-forgetting to net-retention is the empirical signature of the coverage-based diagnosis — Selective Rehearsal doesn't just slow forgetting, it reverses it on the selected examples.

Innovation 3: Establishing Reward Model Calibration as a First-Class Metric for RLHF Stability

Prior work on reward modeling for RLHF focused almost exclusively on ranking accuracy — the fraction of comparison pairs where the model correctly identifies the human-preferred response. This is a natural metric given the Bradley-Terry training objective, which only requires correct pairwise ordering. The paper demonstrates, through a combination of adverse empirical results and conceptual argument, that ranking accuracy is insufficient as a quality metric for reward models used in RLHF. A model can achieve high ranking accuracy while being poorly calibrated (producing overconfident and miscalibrated scores), and this poor calibration directly causes RLHF instability.

The diagnostic evidence is Figure 2: both the RM and AM achieve similar ranking accuracy (Table 1: 69.25% vs 69.43% on HH-RLHF), but the RM is substantially miscalibrated — when it assigns a score difference implying 90% confidence that one response is better, the actual empirical preference rate deviates significantly from 90%. The AM, by contrast, closely follows the perfect calibration line. The downstream consequence is visible in Figure 5b: RM-PPO's reward scores rise during training, but the GPT-4 win rate (a proxy for true human preference) does not track this rise — and in some regimes, it declines. This decoupling is a direct consequence of miscalibration: the reward model becomes overconfident in its assessments, assigning high scores to responses that are not actually better, and PPO optimizes toward these spurious high-scoring regions.

This insight reframes what the field should optimize for in reward modeling. The paper is not the first to note that reward models can be miscalibrated — calibration is a standard concern in machine learning — but it is among the first to establish calibration as a causal factor in RLHF instability rather than a cosmetic property. The ECE metric, reported in Table 1 for both models, captures the calibration gap quantitatively: AM achieves 3.48 vs 4.70 ECE on HH-RLHF (~26% reduction) and 3.83 vs 5.35 on proprietary data (~28% reduction). The paper's implicit argument is that these ECE improvements are at least as important as the modest accuracy gain (69.43% vs 69.25%), because the ECE reduction addresses the mechanism by which reward models fail in PPO optimization — overconfident misranking of responses — rather than just reducing the misranking frequency.

The bounding loss in Equation 6 is the mechanism that produces this calibration improvement, but the conceptual contribution is broader than the specific technique. The paper establishes a criterion for evaluating reward models that goes beyond ranking performance: a good reward model for RLHF must be well-calibrated, because miscalibrated scores create spurious optimization targets that PPO will exploit. This criterion could be satisfied by other methods — temperature scaling, isotonic regression, Bayesian uncertainty quantification — and the Advantage Model is one approach among many. The paper's contribution is elevating calibration from an afterthought to a first-class design requirement.

Evidence anchoring: Table 1 and Figure 2 together make the calibration argument: accuracy is comparable but ECE differs substantially, and the calibration curves visually demonstrate the RM's systematic overconfidence. Figure 5b demonstrates the downstream consequence: stable GPT-4 win rates for AM-PPO vs. declining win rates for RM-PPO despite rising RM scores.

Innovation 4: The Synergy Between Reward Signal Quality and Data Curation as an Architectural Principle

The paper's fourth conceptual contribution is less a single technique and more an architectural insight about how the two proposed components interact. The Advantage Model and Selective Rehearsal could have been proposed as independent, unrelated stabilization techniques — one addressing reward quality, the other addressing data curation. The paper makes the stronger claim that they are synergistic in a specific, non-obvious way: the Advantage Model's task-invariant scores are what make Selective Rehearsal's cluster-based selection work correctly.

Consider what would happen if Selective Rehearsal used raw reward model scores for selection. When selecting the top-k examples within each cluster, the "top" would be defined by raw scores that vary systematically across task categories. A cluster primarily containing Code Generation prompts would have mean scores around 2.0; a cluster primarily containing QA prompts would have scores around 0.5. But KMeans clustering on SimCSE embeddings is not perfect — clusters may contain prompts from multiple task categories, especially if the task boundaries are fuzzy or if certain prompts span categories. In a mixed cluster, raw-score selection would systematically favor examples from the high-scoring category, even if the SFT model's actual quality was comparable across categories. This would create rehearsal data that is biased toward high-score categories, under-protecting lower-score categories and potentially exacerbating the very forgetting problem Selective Rehearsal is designed to solve.

The Advantage Model prevents this failure mode by centering scores per prompt, so that selection reflects genuine quality relative to expectation rather than arbitrary task-level score scales. A high-advantage Code Generation example and a high-advantage QA example have the same claim to being "top" within their respective clusters, regardless of the raw reward score difference between Code Generation and QA. This means the rehearsal data genuinely represents the best examples from each semantic cluster, preserving the diversity that makes Selective Rehearsal effective.

This synergy is not merely convenient — it reveals a design principle for multi-component RLHF systems: interventions at different stages of the pipeline should use compatible signal representations. The reward model produces a scalar signal that is consumed by both PPO (as the optimization target) and data curation (as the quality metric). If these two consumers have different requirements — PPO needs a signal that is monotonic with human preference; data curation additionally needs a signal that is comparable across diverse prompts — the reward model must satisfy both. The standard RM satisfies only the first requirement; the Advantage Model satisfies both by decomposing the signal into a prompt-dependent baseline (which absorbs scale variation) and an advantage (which captures preference). This decomposition enables downstream components to operate on the advantage alone without needing to account for cross-prompt scale differences.

The broader implication is that reward model design should anticipate downstream uses beyond PPO optimization. As RLHF pipelines become more complex — incorporating data curation, iterative training, multi-objective optimization, and rejection sampling — the reward signal will be consumed in more ways and by more components. The Advantage Model's decomposition into baseline and advantage is one approach to producing a signal suitable for multiple consumers; other decompositions (uncertainty quantification, multi-head objectives, disentangled representations) may be needed for other downstream uses. The paper does not explore this generalization, but the principle is implicit in the architecture.

Evidence anchoring: Table 2 shows that AM-PPO-SR outperforms both AM-PPO (Advantage Model alone) and RM-PPO (neither component) on both main and forget test sets. The gap between AM-PPO-SR and AM-PPO demonstrates Selective Rehearsal's contribution; the gap between AM-PPO-SR and RM-PPO (from 12.72% to 15.78% win rate on main, from 29.28% to 7.95% loss rate on forget) demonstrates the combined effect. The synergy claim is supported by the fact that Selective Rehearsal relies on AM scores for selection — no RM-SR variant is evaluated, but the logic of why it would underperform (biased selection due to score scale disparities) is clearly stated.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The paper uses two reward modeling datasets and one PPO dataset. For English reward modeling, the HH-RLHF dataset (Bai et al., 2022a; Ganguli et al., 2022) provides 118k helpful and 42k harmless training examples with 8.5k test examples; the paper does not distinguish between helpful and harmless categories. For Chinese reward modeling, a proprietary dataset collected with "quantities similar to those used in LLaMA 2" (Touvron et al., 2023) is used, where annotators generate prompts from a task spectrum, sample five SFT model responses per prompt with varying hyperparameters, and distribute to five annotators for ranking on helpfulness and harmlessness criteria. For PPO training, prompts are sampled from two domain-general datasets — COIG (64,364 training, 2,623 test) and firefly — with no intersection between training and test sets. An additional forget test set of 1,704 SFT test examples evaluates knowledge retention.

  • Base model(s). The paper uses BLOOMZ (Muennighoff et al., 2022) as the pre-trained backbone throughout: BLOOMZ 7B for reward model and Advantage Model training, and BLOOMZ 176B for SFT and RLHF policy training. The choice of BLOOMZ is not extensively justified beyond being a publicly available multi-task fine-tuned model; the scale difference between reward model (7B) and policy model (176B) follows the standard RLHF pattern of using a smaller model for reward prediction to conserve compute. The paper does not provide results with other model families or at other scales.

  • Metrics. Four primary metrics are used across different evaluation stages:

    • Ranking accuracy (%): The fraction of comparison pairs where the reward/advantage model correctly identifies the human-preferred response (Table 1, used for both HH-RLHF and proprietary data).
    • Expected Calibration Error (ECE): A binned estimate of the gap between predicted preference probabilities and empirical frequencies — lower ECE indicates better calibration (Table 1, Figure 2).
    • Delta reward / reward score: The reward or advantage signal during PPO training, used to plot learning curves (Figure 5a). However, the paper does not report absolute reward values; only relative trends (rising vs. stable) are discussed.
    • Win/Loss/Tie rate over SFT model: Evaluated by GPT-4 comparing PPO model outputs against SFT model outputs on both the main test set (standard PPO evaluation prompts) and the forget test set (expert-aligned SFT examples). Win, lose, and tie rates sum to 100% (Table 2, Figures 1b and 5b). The GPT-4 evaluation protocol (prompt format, temperature, number of samples per comparison) is not described in detail.

  • Baselines. Several baselines are compared:

    • Reward Model (RM): Standard Bradley-Terry pairwise ranking loss (Equation 1) with a scalar value head on BLOOMZ 7B (Section 4.2). This is the standard RLHF reward modeling approach.
    • OpenAssistant (Köpf et al., 2023): A publicly available reward model using DeBERTa (He et al., 2020) architecture, evaluated only on HH-RLHF for accuracy comparison (Table 1). This provides an external reference point.
    • RM-PPO: PPO training using the standard Reward Model for scoring. This is the baseline RLHF configuration against which AM-PPO and AM-PPO-SR are compared.
    • RM-PPO w/ MA: RM-PPO with moving average score normalization during training — an ad-hoc stabilization technique that the paper mentions but does not implement in its main comparisons; it is referenced only qualitatively in Section 4.3 as still exhibiting instabilities.
    • SFT model: The supervised fine-tuned model (the initialization for all PPO variants), used as the reference for win/loss/tie comparisons. All PPO models are compared against this SFT baseline.

    Notable missing baselines: The paper does not compare Selective Rehearsal against standard experience replay (random sampling from PPO data for rehearsal), against rejection sampling (Touvron et al., 2023; Equation 4), or against a variant of RM-PPO augmented with Selective Rehearsal (RM-PPO-SR), which would isolate whether Selective Rehearsal's benefits depend on using Advantage Model scores for selection or would transfer to standard reward model scores.

  • Generation budget / compute accounting. The paper measures compute implicitly through training steps rather than FLOPs or wall-clock time. All PPO experiments use a fixed budget: 100 training steps with a global batch size of 256 (Section 4.2), meaning the total number of PPO updates and the number of generations per update are held constant across comparisons. The Advantage Model training uses a fixed 1-epoch budget with a global batch size of 180 (English) or 480 (Chinese). No per-experiment compute comparison is provided (e.g., FLOPs for AM training vs. RM training, or total compute for RM-PPO vs. AM-PPO-SR). The Selective Rehearsal data curation cost — SimCSE embedding, KMeans clustering — is not quantified or included in any cost model. The difficulty estimation cost from the earlier sections (2048 samples per question) does not appear in this paper, as the paper studies RLHF on general-domain prompts rather than math problem-solving.

  • Cross-validation / statistical protocol. No cross-validation or statistical significance testing is reported. The paper does not train multiple random seeds, report confidence intervals, or conduct train/validation/test splits beyond the pre-specified test sets (8.5k for HH-RLHF, 2,623 for PPO, 1,704 for forget). The PPO training and testing datasets are described as having "no intersection," but the procedure for constructing this split (random sampling, stratified by task, etc.) is not specified. The two-fold cross-validation protocol described in Section 3.2 of the paper (for compute-optimal strategy selection) is relevant to a different paper; this paper does not use it. The learning curves in Figure 5 display single runs without error bars or multiple seeds, so the stability of the reported trends across random initializations is unknown.

Main Quantitative Results

Reward Model vs. Advantage Model: Accuracy and Calibration

The headline result for the Advantage Model appears in Table 1, which compares ranking accuracy and ECE across three reward modeling approaches on two datasets:

  • On HH-RLHF: OpenAssistant achieves 69.24% accuracy (ECE not reported); RM achieves 69.25% accuracy with 4.70 ECE; AM achieves 69.43% accuracy with 3.48 ECE. The accuracy improvement over RM is marginal (+0.18 percentage points), but the ECE reduction is substantial: 3.48 vs. 4.70, a ~26% decrease.
  • On proprietary (Chinese) data: RM achieves 74.75% accuracy with 5.35 ECE; AM achieves 75.28% accuracy with 3.83 ECE. Here the accuracy gain is +0.53 percentage points with a ~28% ECE reduction (3.83 vs. 5.35).

The paper's core claim is that AM "achieves slightly higher accuracy but significantly lower ECE on all the datasets" (Section 4.3). The numbers support this: the ECE improvements are consistent across datasets, while the accuracy improvements are small. The authors attribute the higher accuracy on proprietary data (74.75–75.28%) vs. HH-RLHF (69.25–69.43%) to "the trade-off between helpfulness and harmlessness objectives [being] more pronounced in HH-RLHF, possibly due to the limited presence of harmful examples in our proprietary data" — a plausible but untested hypothesis, since the paper does not break down results by helpfulness/harmlessness categories on either dataset.

Figure 2 provides the calibration visualization. The orange diagonal line represents perfect calibration: Accuracy = 1/(1 + e^{-Δ}) where Δ is the score difference assigned to the (chosen, rejected) pair. On both HH-RLHF (left panel) and proprietary data (right panel), the AM's observed accuracy curve closely follows the diagonal, while the RM's curve deviates upward for intermediate Δ values — the RM is systematically overconfident, assigning higher probabilities to its predictions than the empirical accuracy warrants. This is a classic overconfidence pattern: when the RM assigns a score difference of ~2 (implying ~88% confidence), the actual accuracy is closer to 70–75%. The AM's curve, by contrast, tracks the diagonal with visibly smaller deviations.

Figure 3 provides distribution-level evidence: histograms of RM and AM scores for "good" (human-preferred) and "bad" (rejected) examples on the proprietary data. Both models assign higher scores to good examples on average, but the AM "exhibits a more distinct distribution pattern" (Section 4.3). Specifically, the AM distributions for good and bad examples have less overlap and more clearly separated modes, while the RM distributions show substantial density overlap in the region around 0.0–0.5.

Advantage Model Regularization: Per-Task Score Statistics

Figure 4 presents the evidence for the Advantage Model's task-level regularization. These plots use the proprietary data with its task spectrum (61 task categories, as referenced in Section 1):

  • Figure 4a (per-task means): The RM exhibits "markedly different means for each task" (Section 4.3), with per-task mean rewards spanning a range of approximately 0.1 to 0.6 (based on visual inspection of the chart — the paper does not report exact min/max values). The AM's per-task mean advantages cluster much more tightly around zero, within approximately ±0.2, with no task showing extreme positive or negative mean advantages. This directly demonstrates the baseline subtraction's effect: E[a] = E[r - E[r]] = 0 by construction for each task.

  • Figure 4b (per-task standard deviations): The AM operates at a more "stable scale" — per-task standard deviations are more uniform (visually within approximately 0.35–0.45), while the RM's per-task standard deviations show greater variation. The paper argues this contributes to PPO stability by preventing certain task categories from dominating the reward signal variance.

The paper notes that these disparities "can potentially give rise to reward hacking issues (Skalse et al., 2022) and result in repeated failures during PPO training" (Section 4.3). The causal chain is: disparate means → PPO discovers that shifting response patterns toward high-mean tasks increases average reward → reward hacking. The AM prevents the first step, eliminating the incentive. However, the paper does not directly demonstrate this causal chain in operation — it shows that AM reduces mean disparities and that AM-PPO is more stable, but does not show that the stability gain is mediated specifically by the mean equalization rather than by the bounding loss or improved calibration.

PPO Training Results: Learning Curves and Win Rates

Figure 5 presents the core RLHF training results, showing both reward-based learning curves (Figure 5a) and GPT-4-based win/loss rates relative to the SFT model (Figure 5b).

Figure 5a — Delta reward trajectories: The y-axis shows "Delta Reward" — the reward signal during PPO training. The paper does not define "Delta Reward" precisely (it likely means the advantage signal or the reward relative to some baseline), but the relative trends are interpretable:

  • RM-PPO (blue line) shows the highest delta reward, rising rapidly early in training and continuing to trend upward.
  • AM-PPO (orange line) achieves intermediate delta rewards, also rising but less rapidly than RM-PPO.
  • AM-PPO-SR (green line) has the lowest delta reward but still shows an upward trend.
  • RM-PPO w/ MA (purple line) — not defined in the figure legend but discussed in text — is described as encountering "instabilities during PPO training" (Section 4.3).

The paper interprets the high RM-PPO delta reward as partially spurious: "despite a rise in RM scores" (Figure 5a), RM-PPO faces "a drop in win rate evaluated by GPT-4" (Figure 5b). This is the empirical signature of reward hacking — the reward model's scores increase while true performance decreases.

Figure 5b — GPT-4 win/loss vs. SFT:

  • RM-PPO (blue bars): Win rate starts high (approximately 18–20% early in training), then declines. The loss rate correspondingly increases. The paper characterizes this as "significant reward hacking issues."
  • AM-PPO (orange bars): Win rate is lower initially but remains stable or increases slightly through training, without the decline observed in RM-PPO. The loss rate remains low throughout.
  • AM-PPO-SR (green bars): Win rate is highest among the three, also stable, with the lowest loss rate.

The exact numbers for Figure 5b are not stated in the text — the bar chart must be read visually. The temporal trend (RM-PPO declining, AM-PPO and AM-PPO-SR stable/improving) is the key claim. However, the paper does not explain the x-axis of Figure 5b (presumably training steps, but the labeling is unclear), nor does it specify how many evaluation samples were used for the GPT-4 comparison at each point or what prompt template was used for the GPT-4 judge.

Table 2 — Final win/loss/tie rates: The end-of-training head-to-head comparisons against the SFT model provide precise numbers:

ModelMain Test SetForget Test Set
Win ↑Lose ↓TieWin ↑Lose ↓Tie
RM-PPO12.7212.6274.6616.8729.2853.84
AM-PPO14.8710.3874.749.708.4481.86
AM-PPO-SR15.789.7774.4510.307.9581.75

On the main test set, all three models have very high tie rates (~74–75%), meaning that for roughly three-quarters of test prompts, GPT-4 judged the PPO model and the SFT model as producing equally good responses. The win/loss differences among the models are in the remaining ~25% of examples:

  • RM-PPO: roughly even win/loss (12.72% win vs. 12.62% loss), net +0.10 percentage points.
  • AM-PPO: clearly positive (14.87% win vs. 10.38% loss), net +4.49 percentage points.
  • AM-PPO-SR: most positive (15.78% win vs. 9.77% loss), net +6.01 percentage points.

The incremental improvements from each component are visible: switching from RM to AM (+4.39 net win rate), adding Selective Rehearsal (+1.52 additional net win rate). However, the high tie rate means these gains apply to a small fraction of the test set, and the absolute differences in win rate are modest (15.78% vs. 12.72%).

On the forget test set, the results are dramatically different:

  • RM-PPO: loses more than it wins (16.87% win vs. 29.28% loss), net −12.41 percentage points. This is catastrophic forgetting — the PPO model is worse than the SFT model on nearly 30% of these expert-aligned examples.
  • AM-PPO: wins slightly more than it loses (9.70% win vs. 8.44% loss), net +1.26 percentage points, with a much higher tie rate (81.86% vs. 53.84%). The Advantage Model alone largely eliminates the catastrophic forgetting, moving from a substantial net loss to a small net gain.
  • AM-PPO-SR: further improves (10.30% win vs. 7.95% loss), net +2.35 percentage points, with tie rate essentially unchanged (81.75%).

The qualitative shift from RM-PPO to AM-PPO on the forget set is the paper's strongest result: a 29.28% → 8.44% reduction in loss rate, with the tie rate increasing from 53.84% to 81.86%. This means that AM-PPO produces responses that GPT-4 judges as comparable to SFT on most examples where RM-PPO was actively worse. Selective Rehearsal provides an additional but smaller improvement, reducing the loss rate from 8.44% to 7.95% and increasing the win rate from 9.70% to 10.30%.

The main test set results are more modest: all three models are mostly tied with SFT, and the win rate improvements are measured in single-digit percentage points. The paper's claim that the methods "achieve higher reward scores and win rates" (Section 1, abstract) is technically true but these gains apply to a small slice of test examples, with the majority of responses unchanged in quality.

Ablation Studies and Robustness Checks

  • Number of clusters c for Selective Rehearsal (Figure 6): The paper investigates the effect of varying the number of KMeans clusters c on AM-PPO-SR performance. Figure 6 shows the test-set reward (y-axis, "Reward") as a function of training steps for different c values (exact values not labeled on the figure, but visually appearing to span roughly 5–50 clusters). The finding: "a relatively consistent variance of approximately 0.05 points in test-set rewards across various cluster numbers c" (Section 4.3). This suggests Selective Rehearsal is robust to the choice of c, which is practically important because optimal cluster count is typically dataset-dependent. However, the paper only examines the effect on reward signal, not on final win rate or forget-test performance, and does not test extreme c values (e.g., c = 1 meaning no clustering, or c equal to the number of training examples).

  • Advantage Model calibration analysis (Figure 2): While not framed as an ablation, Figure 2 implicitly ablates the bounding loss by comparing AM (which includes terms 2 and 3 of Equation 6) against RM (which only includes the ranking term equivalent to term 1). The comparison shows that adding the bounding terms substantially improves calibration without degrading ranking accuracy. However, the paper does not ablate the bounding loss in isolation: no experiment trains an AM with only the ranking term and the baseline subtraction (without the log(σ(m(x) ± a)) terms) to determine whether the ECE improvement comes from the baseline, the bounding, or their combination. Similarly, no ablation varies the margin m(x) = 2.5 to determine sensitivity.

  • Score distribution comparison (Figures 3 and 4): The AM vs. RM score distribution analysis (Figure 3: good/bad example score histograms; Figure 4: per-task means and standard deviations) serves as a qualitative ablation of the baseline subtraction. Figure 3 shows that removing the baseline (RM) produces overlapping distributions for good and bad examples, while the AM's advantage scores separate them more cleanly. Figure 4 shows that the baseline subtraction centers per-task means near zero and equalizes variances. However, these are correlational, not causal: they demonstrate that AM produces better-distributed scores, but not that these distributional properties directly cause the PPO stability improvements, as opposed to other properties of AM (e.g., the bounding loss, or the mere fact of having a different random initialization).

  • Moving average normalization as a stabilization baseline: The paper mentions RM-PPO with moving average score normalization (RM-PPO w/ MA, shown in Figure 5a as a purple line) as an ad-hoc stabilization technique. It is described as encountering "instabilities during PPO training" (Section 4.3) — the exact nature of the instabilities (divergence, oscillation, collapse) is not specified. The paper also mentions no formal comparison of RM-PPO w/ MA vs. AM-PPO in terms of win rates or forget-test performance. This is a missed opportunity: comparing AM against a straightforward score normalization baseline would strengthen the claim that AM's benefits come from more than just variance reduction.

  • No ablation of the importance-weighting scheme in Equation 8. The hyperparameters N (harmonization weight) and K (number of historical policies) are part of the advantage computation's baseline approximation but are never varied or analyzed. The behavior of the advantage model under different settings of N and K — especially extreme values like K=0 (relying entirely on e_τ(x) without historical correction) — is unknown.

  • No ablation of the PPO KL penalty coefficient β or the rehearsal loss coefficient γ. The paper fixes γ = 0.01 without sweep (Equation 9) and does not report the β value for the KL-divergence term (Equation 2). The interaction between β (global KL penalty) and γ (local rehearsal NLL) is unexplored — a larger β might compensate for the absence of rehearsal, reducing the gap between RM-PPO and AM-PPO-SR on the forget set. Similarly, the effect of varying γ is not studied.

  • No ablation of the selection criterion within clusters. The paper uses advantage model score as the sole selection criterion for picking top-k examples within each cluster. Alternative criteria mentioned but not tested include entropy (low entropy indicating high confidence), human satisfaction rate, or response length (Section 3.2). The sensitivity of results to replacing advantage score with these alternatives — or combining them — is unknown.

  • Training epoch for reward models: The paper notes that "overfitting issue is observed in general after models are trained for one epoch. As such, we fixed the training epoch as 1 for all the experiments" (Section 4.2). The claim that overfitting occurs "in general after models are trained for one epoch" is not supported by learning curves or validation loss plots. If both RM and AM overfit at one epoch, this raises questions about whether either model reaches its full potential — early stopping might improve both, and AM might benefit more or less from additional training than RM. This is a potentially significant confound that the paper does not investigate.

Critical Assessment

Claim 1: The Advantage Model balances reward score distributions across tasks and prevents reward hacking.

The paper demonstrates that the Advantage Model produces more balanced per-task score statistics (Figure 4) — means clustered around zero and more uniform variances — and that AM-PPO avoids the win-rate decline that RM-PPO exhibits during training (Figure 5b). These are real, documented effects.

However, the causal chain is incompletely established. The paper's implicit argument is: unbalanced reward scores (Figure 4a) → reward hacking (RM-PPO's declining GPT-4 win rate in Figure 5b) → fixed by balancing scores (AM-PPO's stable win rate). But the AM differs from the RM in three simultaneous ways: (1) baseline subtraction (Equation 5), (2) bounding loss (Equation 6, terms 2 and 3), and (3) the specific training procedure (same architecture and dataset, but different objective, which can lead to different local optima). Any of these could be responsible for the stability improvement. The paper does not isolate which mechanism matters — for instance, by training an RM with explicit score normalization (not moving average post-hoc, but a training-time penalty on per-task score disparities) and comparing it to AM. The claim that AM "balances reward score distributions and thereby prevents reward hacking" is thus a correlation supported by the architecture's design intent, but the actual mechanism is not proven.

The win-rate numbers in Table 2 show a nuanced picture. On the main test set, the net improvement from RM-PPO to AM-PPO is +4.39 percentage points (from +0.10 to +4.49 net win rate), but with ~75% of examples tied in both cases. The reward hacking problem manifests as RM-PPO's loss rate being nearly equal to its win rate (12.62% vs. 12.72%), while AM-PPO's loss rate drops to 10.38%. This is consistent with AM reducing spurious optimization that degrades responses. However, the absolute magnitude is relatively small — the primary difference is in the ~25% of non-tied examples, and within that subset, the effect size is moderate.

A missing experiment: The paper does not evaluate RM-PPO with AM scores (i.e., training a standard PPO model but using the AM as the reward source, without Selective Rehearsal). This is confusing because AM-PPO is exactly this configuration — PPO trained with Advantage Model scores. But this is not presented as a diagnostic: what specifically happens to RM-PPO's reward hacking when you simply replace the reward model with the advantage model? The paper answers this implicitly (AM-PPO outperforms RM-PPO), but does not analyze why — does the GPT-4 win rate stabilize because score disparities are reduced, because the bounding loss prevents over-optimization of extreme scores, or because the advantage signal simply provides a less noisy optimization surface?

Claim 2: Selective Rehearsal mitigates catastrophic forgetting on expert-aligned examples.

This claim is well-supported by the forget test set results in Table 2. The RM-PPO model loses to the SFT baseline on 29.28% of forget-test examples, while AM-PPO reduces this to 8.44% and AM-PPO-SR further reduces it to 7.95%. The win rate also improves from 9.70% (AM-PPO) to 10.30% (AM-PPO-SR), though the incremental benefit of Selective Rehearsal over AM alone is small: +0.60 percentage points on win rate and -0.49 points on loss rate.

The strength of the forgetting mitigation is primarily driven by the Advantage Model, not Selective Rehearsal. This is contrary to the paper's framing, which presents Selective Rehearsal as the anti-forgetting mechanism. The numbers show that switching from RM to AM reduces the forget-test loss rate from 29.28% to 8.44% (a ~71% reduction), while adding Selective Rehearsal further reduces it from 8.44% to 7.95% (a ~6% reduction). The Advantage Model alone nearly solves catastrophic forgetting on this benchmark; Selective Rehearsal provides a small additional improvement. This does not negate the value of Selective Rehearsal — the additional gain is real — but it substantially weakens the claim that Selective Rehearsal is the primary mechanism for addressing forgetting. The paper does not discuss this.

A critical missing baseline: RM-PPO with Selective Rehearsal (RM-PPO-SR). Without this, we cannot determine whether Selective Rehearsal's benefits depend on using Advantage Model scores for data selection, or whether it would similarly improve RM-PPO. If Selective Rehearsal with RM scores still reduces forgetting (perhaps less effectively, due to the biased selection the paper's logic predicts), then Selective Rehearsal is independently valuable regardless of the reward model. If it has no effect (because RM score disparities cause poor cluster selection), then Selective Rehearsal's value is contingent on the Advantage Model. This is an important architectural question that the paper leaves unresolved.

The forget test set is relatively small and its construction is opaque. The paper states it consists of "1,704 examples from the SFT test data" (Section 4.1) but does not describe how these examples were selected — are they randomly sampled, stratified by task, or chosen to be representative? If they happen to be examples where the SFT model performs particularly well, the forgetting problem might be overestimated. Conversely, if they are examples where the SFT model is mediocre, the forgetting problem might be underestimated. Without specifying the selection procedure or reporting SFT model accuracy on this set, the reader cannot assess how much of the SFT distribution the forget set represents.

Claim 3: The Advantage Model achieves higher reward scores and win rates while being better calibrated.

The calibration improvement is clearly demonstrated (Table 1, Figure 2). ECE drops by ~26–28% across both datasets, and Figure 2 visually confirms the AM's calibration curve tracks the diagonal more closely than the RM's. This is a robust result, though it is worth noting that ECE is a binned metric sensitive to binning strategy (not specified in the paper) and sample size per bin.

The win rate improvement on the main test set is genuine but modest. The net win rate (win − loss) improves from +0.10 (RM-PPO) to +4.49 (AM-PPO) to +6.01 (AM-PPO-SR). However, roughly 75% of test examples result in ties regardless of method, meaning these gains are concentrated in the minority of examples where the models meaningfully differ. The paper does not analyze which examples drive the improvement — are they from specific task categories, specific difficulty levels, or specific response characteristics? This limits the generalizability of the finding.

The "higher reward scores" claim is ambiguous. Figure 5a shows that RM-PPO achieves higher absolute reward than AM-PPO. The paper's claim (in the abstract and conclusion) that AM achieves "higher reward scores" presumably refers to more genuine reward (as validated by GPT-4 win rate) rather than higher raw scores. But this conflation of "reward" (the scalar signal during training) and "quality" (GPT-4 evaluation) is imprecise and could mislead readers into thinking AM produces higher scalar rewards, when Figure 5a shows the opposite.

Claim 4: The proposed methods "stabilize RLHF training."

What does "stabilize" mean empirically? The paper provides several pieces of evidence:

  • AM-PPO's GPT-4 win rate does not decline during training, while RM-PPO's does (Figure 5b). This is a valid stability metric — the model's alignment quality does not degrade with more training.
  • AM-PPO-SR shows low variance in test-set rewards across different cluster counts c (Figure 6), suggesting the method is robust to hyperparameter choice.
  • RM-PPO w/ MA encounters "instabilities during PPO training" (Section 4.3), implying that the standard approach with normalization is not stable.

What is missing from the stability analysis:

  • Run-to-run variability. The paper trains each configuration once. Without multiple random seeds, we cannot distinguish between genuine stability (low variance across random initializations) and fortuitous single-run outcomes. The PPO algorithm is known to be sensitive to random seed in deep RL settings; a stability claim requires demonstrating consistency across runs.
  • Stability across model scales. All experiments use BLOOMZ 176B for policy training. Whether the stability gains scale to larger models (where RLHF is typically deployed) or hold for smaller models is untested.
  • Stability across datasets. The HH-RLHF and proprietary Chinese datasets are both general-domain instruction-following benchmarks. Whether AM and Selective Rehearsal stabilize RLHF on more specialized domains (code, math, safety-critical applications) is unknown.
  • Longer training horizons. All PPO experiments use exactly 100 training steps. The paper does not investigate whether stability advantages persist, diminish, or grow with longer training. If RM-PPO's win-rate decline in Figure 5b is a transient phenomenon that would recover with more training, the stability claim would be weaker. Conversely, if AM-PPO and AM-PPO-SR would eventually also decline with sufficient training (due to inevitable over-optimization of any imperfect reward signal), the claim that they "solve" stability rather than delay instability would be important to know.

Genuine Experimental Weaknesses

  • Single model family (BLOOMZ) throughout. The paper's claims about AM and Selective Rehearsal are entirely validated on BLOOMZ models. There is no evidence that these techniques transfer to other architectures (LLaMA, GPT, PaLM) or that they interact differently with different pre-training distributions. Given that reward modeling quality depends on the base model's representations, the Advantage Model's performance may vary significantly across model families.

  • The 1-epoch training constraint on reward models. The paper notes that "overfitting issue is observed in general after models are trained for one epoch" and therefore fixes all reward model training to 1 epoch. If true, this means both RM and AM are trained with early stopping at an arbitrary point that may not be optimal for either. The paper does not present validation loss curves showing when overfitting begins, nor does it explore whether the optimal stopping point differs between RM and AM. If AM is more resistant to overfitting (plausible given the bounding loss's regularizing effect), the 1-epoch constraint may disadvantage the RM more than the AM, inflating the reported performance gap.

  • No reporting of key hyperparameters. Critical values are missing: the PPO KL penalty coefficient β (Equation 2), the importance-weighting parameters N and K (Equation 8), the number of clusters c for main experiments, and the per-cluster selection count k. The learning rate for reward model training is given as 5e-6 (standard), but the optimizer configuration for AM is not separately specified (does AM use the same optimizer, warmup schedule, and batch size as RM?). These omissions make independent replication difficult.

  • The GPT-4 evaluation protocol is undescribed. The win/loss/tie rates in Table 2 and Figure 5b are the paper's primary metrics for alignment quality, but the paper provides no details on how GPT-4 was prompted to make comparisons: what instruction template, what response format, whether position bias was controlled (by swapping order and averaging, as is standard), what temperature was used, how many comparison samples were collected per prompt. Given GPT-4's known sensitivity to prompt formatting and position effects, this is a significant methodological gap.

  • The Proprietary Chinese dataset is not public and cannot be verified. All results on the 61-task proprietary dataset are irreproducible by external researchers. The paper's key findings (Figures 3, 4, and the Chinese RM/AM comparison in Table 1) depend partially on this dataset. While proprietary data is common in industry research, it means the generalizability of results to other multi-task reward modeling scenarios rests on the HH-RLHF results alone (which show smaller accuracy differences: 69.25% vs. 69.43%).

  • No comparison against PPO-free alignment methods. The paper cites Direct Preference Optimization (DPO; Rafailov et al., 2023) and Preference Ranking Optimization (PRO; Song et al., 2023) as alternative alignment paradigms that avoid RL training entirely (Section 5). However, it does not compare its stabilized RLHF pipeline against these methods. If DPO achieves comparable or better alignment with none of the stability issues that AM and Selective Rehearsal address, the practical value of these techniques is diminished. The paper's contribution is within the RLHF paradigm, but a reader choosing an alignment approach needs to know whether stabilized RLHF is competitive with RL-free alternatives.

Missing Experiments That Would Have Strengthened the Paper

  • Isolating the bounding loss contribution: Train an AM variant without the log(σ(m(x) ± a)) terms in Equation 6 (i.e., only the ranking term with baseline subtraction) to determine whether the calibration improvement comes from the baseline or the bounding.

  • RM-PPO-SR: Train a standard RM-PPO model augmented with Selective Rehearsal (using RM scores for selection) to determine whether Selective Rehearsal independently helps or depends on AM scores for effective data curation.

  • Sensitivity to the margin m(x): Vary m(x) from, say, 1.0 to 5.0 to determine whether the calibration and stability benefits are robust to this hyperparameter or require careful tuning. The fixed value of 2.5 is given no theoretical or empirical justification.

  • Comparison against DPO with comparable compute: Train a DPO model on the same preference data and evaluate on the same main and forget test sets. This would contextualize the AM-PPO-SR results within the broader alignment landscape.

  • Multiple training runs with different seeds: Report mean and standard deviation for Table 2 metrics across at least 3–5 random seeds to establish that the reported improvements are reliable.

  • Analysis of which examples drive the main-test win rate improvement: Break down Table 2 results by task category (for the proprietary data with 61 categories) or by response length / complexity to understand where AM and Selective Rehearsal help most.

  • Longer PPO training: Extend training beyond 100 steps for all configurations to determine whether RM-PPO's decline is permanent or transient, and whether AM-PPO/AM-PPO-SR eventually also over-optimize the advantage signal. This would test the claim that the stabilization is a fundamental property rather than a delayed onset of the same problems.

  • Quantifying the Selective Rehearsal curation cost: Report the wall-clock time and compute (FLOPs or GPU-hours) for the SimCSE embedding, KMeans clustering, and data selection steps. While these are one-time costs before PPO training, they are not negligible (especially for very large prompt sets) and should be part of any practical cost comparison.

6. Limitations and Trade-offs

The Advantage Model's Expected Reward Estimation is Unvalidated and Depends on Unexplored Hyperparameters

The assumption or constraint: The practical computation of the advantage score depends on a tractable approximation of the expected reward baseline (Equation 8), which introduces three hyperparameters — $N$ (balance between current and historical policies), $K$ (number of historical policy models), and a learned function $e_\tau(x)$ — whose behavior and sensitivity are entirely unexplored. The paper acknowledges the intractability of the true expectation ("it is infeasible to list every potential response to calculate the expected reward," Section 3.1) and proposes the approximation as a solution, but provides no experimental analysis of any of these hyperparameters, does not report their values for the main experiments, and does not ablate alternative approximation strategies (e.g., using only $e_\tau(x)$ without historical correction, corresponding to $K = 0$).

The consequence: A practitioner attempting to implement the Advantage Model faces an unknown sensitivity surface. If the approximation's quality depends substantially on $N$ and $K$ — as is plausible given that these control a bias-variance tradeoff in the baseline estimate — the reported stability benefits may not transfer to settings with different data collection policies, different numbers of historical model versions, or different choices of these hyperparameters. Worse, if the $e_\tau(x)$ learner is poorly calibrated early in training, the advantage scores during the initial PPO steps will be systematically biased — the baseline may over- or under-estimate the expected reward, shifting the entire advantage distribution by an unknown offset that varies per prompt. The paper's calibration analysis (Figure 2, Table 1) evaluates the final trained AM on held-out comparison data, not the dynamic behavior of advantage estimates as $\pi_\phi$ evolves during PPO. The importance-weight correction term $\pi_\phi(y|x) / \pi'_k(y|x)$ can also become extreme if the current policy diverges substantially from the data-collection policies, introducing high-variance advantage estimates that could destabilize PPO in precisely the regimes where stability is most needed (large policy updates).

What evidence exists in the paper: None whatsoever. The paper does not report values for $N$ or $K$, does not vary them, does not ablate the historical correction terms against a pure $e_\tau(x)$ baseline, and does not track the importance weights' distribution during PPO training to confirm they remain well-behaved. The only evidence that the overall AM approximation works is indirect and downstream — AM-PPO outperforms RM-PPO (Table 2) — which validates the combined effect of all AM components but provides no signal about which components matter or how sensitive the result is to these specific hyperparameters.

Mitigation status: Not addressed. The paper does not flag this as a limitation, does not propose guidelines for setting $N$ and $K$, and does not suggest future work on validating or simplifying the expected reward approximation. This is a genuine gap: the mathematical definition of the Advantage Model (Equation 5) is elegant, but the gap between that definition and the practical approximation (Equation 8) is large and unexamined.

The Difficulty Estimation Cost Problem: Advantage Model Training and Selective Rehearsal Curation Are Not Accounted for in Any Cost Model

The assumption or constraint: The paper reports performance improvements (win rate gains of +3–6 percentage points on the main test set, large reductions in forget-test loss rate) without accounting for the computational overhead of the proposed methods relative to the standard RLHF baseline. Specifically:

  • Advantage Model training uses the same architecture (BLOOMZ 7B) and the same dataset as the standard Reward Model, but the objective (Equation 6) has three terms per training example instead of one. The paper does not report whether AM training converges at the same rate as RM training, whether it requires more steps, or whether the additional terms increase per-step computation (they likely do not significantly, since they reuse the same forward pass outputs $a_\theta(x, y_c)$ and $a_\theta(x, y_r)$ and add only scalar sigmoid-log computations). However, the need to train the expected reward function $e_\tau(x)$ — which the standard RM does not require — represents additional model complexity whose training cost is not discussed.

  • Selective Rehearsal data curation requires embedding every PPO training prompt with SimCSE (64,364 prompts in the main experiments), running KMeans clustering on 64,364 embedding vectors, scoring all (prompt, response) pairs with the Advantage Model to select top-k within each cluster, and constructing the rehearsal dataset. The paper does not report the wall-clock time, GPU-hours, or FLOPs for any of these steps, and does not include them in any cost comparison against the RM-PPO baseline, which requires none of these curation steps. While the curation is a one-time pre-processing cost (performed once before PPO training), for large-scale RLHF with millions of prompts, clustering and embedding can be substantial.

  • The PPO training itself adds an additional forward pass and loss computation for the rehearsal batch (Equation 9). The coefficient $\gamma = 0.01$ means this loss has only 1% of the PPO loss weight, but the rehearsal batch still requires a separate forward/backward pass or at minimum an additional data loading and loss evaluation step, which increases per-step training time by some (unreported) factor.

The consequence: A practitioner comparing RM-PPO against AM-PPO-SR for a production deployment cannot make an informed cost-benefit decision. The headline win-rate improvements (+6.01 net win rate for AM-PPO-SR vs. +0.10 for RM-PPO, Table 2) are presented as pure gains, without the denominator of additional compute cost. If AM training takes $2\times$ longer than RM training (unlikely but possible if $e_\tau(x)$ requires joint training), or if clustering 64,364 prompts takes hours on CPU and requires loading the entire PPO dataset into memory, the practical trade-off shifts. For organizations with tight training budgets or rapid iteration cycles, the overhead may outweigh the modest win-rate gains.

What evidence exists in the paper: None. The paper provides no compute accounting — no FLOP counts, no GPU-hours, no wall-clock comparisons, no mention of overhead relative to the RM-PPO baseline. The fixed training budgets (100 PPO steps, 1 reward model epoch) are reported but the per-step cost differences are not. The Selective Rehearsal clustering cost is not mentioned as a practical consideration at any point.

Mitigation status: Not addressed. The paper does not discuss computational efficiency, does not include overhead in any reported metric, and does not propose ways to reduce the curation cost (e.g., using a subset of prompts for clustering, using cheaper embeddings, or amortizing curation across multiple PPO runs). This is an important practical limitation for an applied method paper — readers need to know the price of the proposed stabilization, not just the benefit.

Single Model Family and Single Dataset Ecosystem Limits Generalizability

The assumption or constraint: All experiments — reward modeling, advantage modeling, SFT, and PPO training — use the BLOOMZ model family (7B for reward/advantage models, 176B for policy training; Muennighoff et al., 2022). The paper provides no evidence that the Advantage Model's calibration benefits or Selective Rehearsal's forgetting mitigation transfer to other model architectures (LLaMA, GPT, PaLM, Falcon), other pre-training distributions, or other scales. The paper's only public benchmark evaluation is on HH-RLHF (English); all Chinese results are on a proprietary, unreleasable dataset. The paper states that BLOOMZ was "employed ... as our pre-trained model backbone" (Section 4.1) but does not justify why the findings should generalize beyond this specific model family.

The consequence: This is a severe limitation for a paper whose contributions are presented as general techniques for RLHF stabilization, not as BLOOMZ-specific engineering. Several aspects of the proposed methods could interact with model properties in unknown ways:

  • The Advantage Model's calibration improvement depends on the base model's ability to learn a well-behaved expected reward function $e_\tau(x)$ and to produce scores that can be bounded by the $m(x) = 2.5$ margin without sacrificing ranking accuracy. A different base model with different representational properties or different pre-training data might require a different margin, might learn $e_\tau(x)$ with different fidelity, or might exhibit different calibration behavior under the bounding loss. The paper's finding that ECE improves from 4.70 to 3.48 on HH-RLHF (Table 1) is a single data point — whether similar improvements occur with, e.g., LLaMA-7B as the reward model backbone is unknown.

  • Selective Rehearsal relies on SimCSE embeddings for clustering quality. SimCSE (sup-simcse-roberta-base) is a RoBERTa-based model trained on English data. Its embedding quality on Chinese prompts (the proprietary dataset) is not validated, and its effectiveness on code-generation or multilingual prompts is unknown. If the embeddings fail to capture task-relevant semantic similarity, the clusters will mix unrelated skills, undermining the diversity guarantee that Selective Rehearsal is designed to provide.

  • The scale gap between reward model (7B) and policy model (176B) is fixed. Whether the Advantage Model's benefits persist, diminish, or amplify when the reward model scale changes relative to the policy model is untested. If the reward model is much smaller (e.g., 1B) relative to the policy, the expected reward estimate $e_\tau(x)$ might be less reliable, degrading advantage score quality.

What evidence exists in the paper: Zero cross-model-family evaluation. The paper does not compare BLOOMZ-based AM against non-BLOOMZ alternatives, does not apply Selective Rehearsal to a different base policy model, and does not discuss generalization as a limitation.

Mitigation status: Not addressed at all. The paper implicitly treats the BLOOMZ choice as incidental rather than as a potential confound. No future work on multi-model validation is suggested. For a practitioner using a non-BLOOMZ model (which describes virtually all production RLHF deployments — GPT-4, Claude, Llama 2, Gemini use proprietary or LLaMA-based architectures), the paper provides no evidence that the reported gains will materialize.

Catastrophic Forgetting Is Primarily Solved by the Advantage Model Alone, Undermining the Claim That Selective Rehearsal Is the Anti-Forgetting Mechanism

The constraint: Selective Rehearsal is presented as the mechanism for "mitigating catastrophic forgetting" (Section 1, abstract) and "preventing the depreciation of the model's performance on expert-aligned examples over time" (Section 1). The paper's architecture separates concerns: the Advantage Model handles reward hacking, Selective Rehearsal handles forgetting. The structure of Sections 1 and 3 reinforces this division.

The consequence: The empirical results in Table 2 contradict this division of labor. On the forget test set:

ConfigurationWin ↑Lose ↓Net (Win − Lose)
RM-PPO16.87%29.28%−12.41
AM-PPO9.70%8.44%+1.26
AM-PPO-SR10.30%7.95%+2.35

Switching from RM to AM reduces the loss rate from 29.28% to 8.44% — a ~71% reduction — and shifts the net outcome from substantially negative (−12.41) to slightly positive (+1.26). This is a qualitative reversal: the Advantage Model alone largely solves catastrophic forgetting on this benchmark, without any Selective Rehearsal. Adding Selective Rehearsal (AM-PPO → AM-PPO-SR) provides an incremental improvement: loss rate drops from 8.44% to 7.95% (−0.49 percentage points), win rate increases from 9.70% to 10.30% (+0.60 points), net improves from +1.26 to +2.35. This is a real but small additional gain — roughly 10% of the total forgetting reduction is attributable to Selective Rehearsal, while ~90% is attributable to the Advantage Model.

This substantially weakens the paper's narrative. The abstract and introduction frame Selective Rehearsal as a co-equal contribution addressing a distinct problem (catastrophic forgetting) that the Advantage Model alone cannot solve. The data show the Advantage Model is the dominant anti-forgetting intervention, and Selective Rehearsal provides a marginal improvement on top of it. A reader who implements only the Advantage Model (and not Selective Rehearsal) would capture the vast majority of the reported forgetting reduction, with considerably less implementation complexity (no embedding, clustering, or rehearsal dataset construction).

Furthermore, the paper's diagnostic that forgetting is caused by "over-optimizing with PPO on examples that were well-aligned with humans in the SFT stage" (Section 1) implies that Selective Rehearsal's explicit reinforcement of these examples is the natural solution. But the Advantage Model does not explicitly reinforce any SFT examples — it simply provides a better reward signal. If the Advantage Model alone nearly eliminates forgetting, the causal mechanism is not the one the paper's diagnosis suggests (lack of explicit reinforcement on good examples), but rather that the standard reward model's miscalibration and score disparities actively cause forgetting by pushing the policy away from correct SFT behavior through spurious reward gradients. Selective Rehearsal, under this alternative diagnosis, is treating a symptom (insufficient positive gradient on SFT-like examples) rather than the cause (destructive gradients from a poorly behaved reward model), which explains its small incremental benefit once the cause is addressed.

What evidence exists in the paper: Table 2 and Figure 5b. The paper does not discuss this distribution of effect sizes, does not acknowledge that the Advantage Model is the primary anti-forgetting mechanism, and maintains the narrative that Selective Rehearsal is the dedicated solution for catastrophic forgetting.

Mitigation status: Not addressed. The paper does not perform the critical ablation of RM-PPO-SR (standard reward model with Selective Rehearsal), which would reveal whether Selective Rehearsal can independently reduce forgetting when the reward signal is poor. If RM-PPO-SR also substantially reduces forgetting (e.g., from 29.28% to 15% loss rate), then Selective Rehearsal is indeed a general anti-forgetting mechanism that is partially redundant with a good reward model. If RM-PPO-SR has minimal effect (loss rate remains ~29%), then Selective Rehearsal's contribution is contingent on having a well-calibrated reward signal for data selection — which the Advantage Model provides, making the two techniques genuinely synergistic. Without this ablation, the paper's central architectural claim about the division of labor between the two components remains unvalidated.

No Comparison Against PPO-Free Alignment Methods Leaves the Practical Value Uncertain

The constraint: The paper operates entirely within the RLHF paradigm and evaluates only against RLHF baselines (RM-PPO, RM-PPO with moving average normalization). Section 5 (Related Work) acknowledges the existence of alternative alignment approaches that bypass reinforcement learning entirely — Direct Preference Optimization (DPO; Rafailov et al., 2023) and Preference Ranking Optimization (PRO; Song et al., 2023) are cited — and characterizes them accurately as methods that "sidestep the necessity for Reinforcement Learning (RL) training." However, the paper provides no experimental comparison against any PPO-free method, nor does it discuss whether the instability problems it addresses (reward hacking, catastrophic forgetting) are inherent to RL-based alignment or are also present in DPO/PRO.

The consequence: The paper's value proposition is that its techniques "stabilize RLHF training" and make it "more stable and effective." But if a practitioner can achieve comparable or better alignment by using DPO — which has no reward model to hack, no PPO optimization to diverge, and no KL penalty to tune — the motivation for investing in Advantage Model training, Selective Rehearsal curation, and the overall RLHF complexity is unclear. The paper's framing implies that RLHF is important enough to be worth stabilizing (which is plausible given its deployment in major production systems), but it provides no evidence about the relative alignment quality of stabilized RLHF vs. RL-free alternatives. A practitioner choosing an alignment approach needs to know: is AM-PPO-SR (with all its added complexity) actually better than just running DPO on the same preference data?

This limitation is particularly acute because DPO and PRO are designed to address precisely the failure modes the paper diagnoses. DPO eliminates the explicit reward model — the source of reward hacking — by directly optimizing a policy from preference data using an implicit reward derived from the policy's own probabilities. Catastrophic forgetting is addressed in DPO through the same mechanism as in PPO (a KL-divergence penalty against the reference model), so DPO might exhibit similar forgetting to RM-PPO, and the Advantage Model would have no analogue in the DPO framework. Without a comparison, we cannot know where stabilized RLHF stands relative to the simplest PPO-free baseline.

What evidence exists in the paper: None. The paper does not implement, evaluate, or reference any experimental comparison with DPO, PRO, or any other RL-free alignment method. The related work section (Section 5) mentions them as part of the landscape but does not position the paper's contributions relative to them empirically.

Mitigation status: Not addressed. The paper does not acknowledge the absence of PPO-free baselines as a limitation, does not discuss the tradeoffs between stabilized RLHF and RL-free alignment, and does not suggest future work on head-to-head comparisons. For a paper that aims to improve RLHF's practical viability, failing to compare against the most prominent alternative paradigm is a significant omission. A single DPO baseline (trained on the same preference data with comparable compute) would substantially clarify the paper's contribution — even if DPO underperforms AM-PPO-SR, the comparison would establish that the added complexity of stabilized RLHF is justified.

The Bounding Margin m(x) = 2.5 Is Unjustified and Its Sensitivity Is Unexplored

The assumption or constraint: The Advantage Model's training objective (Equation 6) includes a prompt-dependent margin function $m(x)$ that defines the permitted range $[-m(x), m(x)]$ for advantage scores through the bounding loss terms $\log(\sigma(m(x) - a_\theta(x, y_c)))$ and $\log(\sigma(m(x) + a_\theta(x, y_r)))$. The paper sets $m(x) = 2.5$ as a constant across all prompts and all experiments, with no theoretical derivation and no empirical sweep. The paper's own footnote (Section 3.1) acknowledges this gap candidly:

"We think that $m(x)$ may have a connection with the complexity or difficulty involved in learning the reward function for prompts similar to $x$. However, this is speculative and requires further investigation. We leave this aspect as a topic for future study and exploration. Throughout our experiments, we set $m(x)$ as 2.5."

The consequence: The margin $m(x)$ controls a fundamental tradeoff in the Advantage Model. A small margin (e.g., $m = 0.5$) tightly constrains advantage scores, forcing them into a narrow range that may be insufficient to capture genuine quality differences between responses — the model would be forced to compress a wide spectrum of response quality into a narrow band, potentially degrading ranking accuracy. A large margin (e.g., $m = 10.0$) provides minimal constraint, allowing the advantage scores to approach the unregularized behavior of the standard reward model, potentially recreating the score disparities and overconfidence that the Advantage Model is designed to prevent. The optimal margin depends on the data distribution: if most human preference differences correspond to modest score differences, $m = 2.5$ might be appropriate; if some comparisons involve dramatically better responses, a larger margin might be needed to avoid saturating the bounding loss and losing ranking signal.

Without a sensitivity analysis, a practitioner cannot determine whether $m = 2.5$ is a robust default (the gains are insensitive to the exact value within a broad range) or a carefully tuned value (the method works only near 2.5 and degrades rapidly at other margins). The latter case would indicate that the Advantage Model is fragile and requires dataset-specific tuning, undermining its claim to be a general stabilization technique. The former case would strengthen the paper's practical contribution, but it is not established.

The paper's speculation that $m(x)$ relates to "the complexity or difficulty involved in learning the reward function for prompts similar to $x$" raises an important design question: should easier prompts (where the reward model is more confident) have smaller margins, while harder prompts have larger margins? A constant $m(x) = 2.5$ for all prompts treats a straightforward factual QA pair identically to an ambiguous creative writing comparison, despite the very different reward modeling difficulty. If the optimal margin is indeed prompt-dependent, the constant-margin approximation may be leaving performance on the table for certain prompt categories while over-constraining others.

What evidence exists in the paper: None. No ablation of $m(x)$ is performed, no range of values is tested, and no justification (empirical or theoretical) is provided for the choice of 2.5. The paper's results (Table 1 calibration, Table 2 win rates) are valid only at $m = 2.5$ — we do not know whether any other value would produce similar, better, or worse results.

Mitigation status: The authors explicitly flag this as a limitation and suggest future work ("We leave this aspect as a topic for future study and exploration"), which is good scientific practice. However, the lack of even a minimal sensitivity analysis (e.g., $m \in \{1.0, 2.5, 5.0\}$) leaves a central hyperparameter of the proposed method entirely uncharacterized. A practitioner cannot make an informed choice about this parameter without running their own sweep, which partially defeats the purpose of a method presented as improving stability — it introduces a new, uncharacterized stability knob of unknown sensitivity.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper's primary contribution to the RLHF landscape is a diagnostic reframing rather than a paradigm shift. The standard RLHF pipeline has long been recognized as brittle — reward hacking and catastrophic forgetting are well-documented failure modes cited in virtually every major RLHF paper (Stiennon et al., 2020; Ouyang et al., 2022; Bai et al., 2022a). What the field lacked was a precise understanding of why these failures occur at a mechanistic level and a corresponding set of targeted interventions that address root causes rather than symptoms.

The paper provides two such diagnoses and validates them empirically:

First, reward hacking is reframed as a reward model calibration problem, not a PPO optimization problem. The standard response to reward hacking has been to intervene downstream — normalize scores during PPO, increase the KL penalty coefficient, clip rewards, or early-stop training. This paper demonstrates that these interventions treat symptoms while the root cause persists upstream: the Bradley-Terry preference modeling objective (Equation 1) contains unconstrained degrees of freedom that allow — indeed, encourage — the reward model to develop systematically different score scales across task categories (Figure 4a). PPO then exploits these disparities because it is simply doing gradient ascent on the provided signal. The paper's key insight is that stabilizing RLHF requires constraining the reward model's output distribution at training time, not patching the PPO optimizer to be robust to a poorly behaved reward signal. This shifts the focus of stabilization research from PPO-level techniques (better normalization, more sophisticated clipping, adaptive KL penalties) to reward model design (calibration objectives, bounded outputs, task-invariant representations).

The magnitude of this shift is substantial but not revolutionary: it does not propose abandoning the Bradley-Terry preference framework or PPO, but rather augmenting the reward model objective with explicit calibration and bounding constraints. This is an incremental refinement of the existing paradigm — one that any RLHF practitioner could adopt by replacing their reward model training loss with Equation 6 — rather than a fundamentally new alignment approach. The paper's comparison with alternative paradigms like DPO (Rafailov et al., 2023) is entirely absent from the experiments, so the paper does not claim that stabilized RLHF is categorically better than RL-free alignment; it claims only that RLHF can be stabilized within its own paradigm.

Second, catastrophic forgetting is partially reframed as a consequence of reward model miscalibration. The paper's own results (Table 2, forget test set) reveal something that the paper's framing does not fully acknowledge: switching from a standard Reward Model to the Advantage Model reduces the forget-test loss rate from 29.28% to 8.44% (~71% reduction), while adding Selective Rehearsal further reduces it only to 7.95% (~6% additional reduction). This means the reward model's behavior — not just the PPO data distribution — is the dominant driver of catastrophic forgetting in this experimental setup. The paper's stated diagnosis — that forgetting occurs because PPO over-optimizes on well-aligned examples — implies that the solution is to protect those examples (via Selective Rehearsal). The data suggest instead that the primary mechanism is destructive gradients from a miscalibrated reward model that actively pushes the policy away from correct SFT behavior on those examples. Fix the reward model, and the forgetting largely disappears.

This has an important implication for the field: catastrophic forgetting in RLHF may be more a reward model quality problem than a continual learning problem. Prior work has approached RLHF forgetting through the lens of continual learning (experience replay, elastic weight consolidation, progressive networks; Khetarpal et al., 2022), assuming that any policy update to maximize a new objective will inherently overwrite previously learned behaviors. This paper's results suggest that when the reward signal is well-calibrated and task-invariant, the policy update naturally preserves existing capabilities because the reward signal does not provide spurious gradients on examples where the SFT behavior is already good. This does not mean continual learning techniques are irrelevant — Selective Rehearsal still provides a small additional benefit — but it does reorder priorities: invest first in reward model quality, then in PPO data curation.

The paper reconciles a tension in the RLHF literature between studies that find reward model calibration matters (e.g., Bai et al., 2022a, which emphasizes training separate reward models for helpfulness and harmlessness to avoid score scale conflicts) and studies that treat calibration as an afterthought to ranking accuracy. By demonstrating that calibration improvements (ECE reduction from 4.70 to 3.48 on HH-RLHF, Table 1) translate directly to downstream PPO stability and win rate improvements, the paper provides empirical evidence that calibration is not merely a cosmetic property but a causal factor in RLHF success. This elevates calibration from a nice-to-have to a first-class requirement for reward models in production RLHF pipelines.

Research directions that become more attractive:

  • Training-time calibration constraints for reward models. The bounding loss in Equation 6 is one approach; temperature scaling, isotonic regression, and Bayesian reward models become natural extensions now that calibration has been causally linked to RLHF stability.
  • Reward model architectures beyond scalar heads. The advantage decomposition (raw reward minus expected reward) suggests that reward models should explicitly represent prompt-dependent baselines, not just raw scores. Multi-head architectures that separately predict expected reward and advantage could improve upon the paper's combined e_τ(x) parameterization.
  • Analyzing reward model behavior under distribution shift. The paper shows that reward model miscalibration causes problems during PPO, when the policy distribution shifts. This connects reward model evaluation to the out-of-distribution detection and calibration literature.

Research directions that become less attractive (or require stronger justification):

  • Pure PPO-level stabilization techniques that ignore reward model quality. If ~71% of catastrophic forgetting reduction comes from fixing the reward model (AM vs. RM), research focusing exclusively on PPO regularization (KL penalty tuning, trust region methods, rejection sampling) without addressing upstream reward signal quality is addressing the smaller portion of the problem.
  • Training separate reward models per task category as a solution to score disparities. The paper's Advantage Model achieves task-invariant scores within a single model, making the engineering complexity of maintaining multiple task-specific reward models harder to justify.

Follow-Up Research This Work Enables

Isolating the bounding loss contribution to calibration and stability. The Advantage Model combines three mechanisms: baseline subtraction (Equation 5), ranking loss (term 1 of Equation 6), and bounding loss (terms 2 and 3 of Equation 6). The paper attributes the calibration improvement and PPO stability gains to the combined effect, but does not ablate them. A high-priority follow-up would train three model variants on the same HH-RLHF data: (a) standard RM (baseline), (b) AM without bounding loss (ranking loss + baseline subtraction only, equivalent to an "unbounded Advantage Model"), and (c) full AM (Equation 6). Evaluating these three on both calibration metrics (ECE, Figure 2-style calibration curves) and downstream PPO stability (win rate trajectory, forget-test loss rate) would reveal whether the bounding loss, the baseline subtraction, or their interaction drives the reported gains. If (b) performs nearly as well as (c), then the baseline subtraction is the essential mechanism and the bounding loss is incidental — simplifying the method considerably. If (b) performs similarly to (a), then the bounding loss is doing all the work and the advantage framing is less important than the explicit score constraints — pointing toward simpler regularization schemes for standard reward models. If (b) is intermediate, both mechanisms contribute and their relative importance can be quantified.

Replicating the Advantage Model on non-BLOOMZ architectures with a public PPO-free baseline. The paper's results are confined to BLOOMZ 7B/176B, and it does not compare against DPO or any RL-free alignment method. A strong follow-up would replicate the core finding — that an Advantage Model trained with Equation 6 reduces reward hacking and catastrophic forgetting compared to a standard Reward Model — on (i) a LLaMA-based model family (e.g., LLaMA-2 7B for reward, LLaMA-2 70B for policy) and (ii) include a DPO baseline trained on the same preference data with comparable compute. The evaluation would measure standard RLHF metrics (win rate vs. SFT, reward model accuracy, ECE) plus a forget-test construct analogous to the paper's 1,704 SFT examples. The key question: does AM-PPO-SR outperform DPO on alignment quality and stability metrics when both are given the same preference data, or does DPO achieve comparable alignment with dramatically less complexity? This would establish the practical value proposition of stabilized RLHF versus the simpler alternative. Given that DPO bypasses explicit reward modeling entirely — and therefore cannot suffer from the reward hacking the Advantage Model is designed to prevent — the comparison would also test whether the remaining instabilities in DPO (if any) are of comparable magnitude to those in stabilized RLHF.

Difficulty-adaptive bounding margins m(x). The paper sets m(x) = 2.5 as a constant and explicitly flags this as an open question, speculating that m(x) may relate to prompt difficulty or reward modeling uncertainty. A natural follow-up would train an Advantage Model where m(x) is not a constant but is predicted as a function of the prompt embedding — either as an auxiliary output of the Advantage Model itself, or through a separate uncertainty estimator. The hypothesis: prompts where human preferences are ambiguous or noisy (e.g., creative writing, subjective quality judgments) should receive larger margins to allow the advantage model to express uncertainty through wider score ranges, while prompts with clear preference signals (e.g., factual accuracy, safety violations) should receive tighter margins. The experiment would compare constant m = 2.5 against learned m(x) on (i) reward model calibration (ECE computed per difficulty stratum), (ii) PPO stability (does adaptive margin further reduce reward hacking on ambiguous prompts?), and (iii) PPO win rate. If learned margins improve calibration without sacrificing ranking accuracy, it would validate the paper's speculation and provide a more principled approach to setting m(x). If learned margins make no difference, the constant 2.5 is a robust default and the speculation can be set aside.

Selective Rehearsal with RM scores to test the synergy claim. The paper claims that the Advantage Model and Selective Rehearsal are synergistic because AM's task-invariant scores enable fair selection across clusters. This claim is logical but not empirically tested — the paper never evaluates RM-PPO-SR (standard reward model with Selective Rehearsal). A crucial follow-up would train RM-PPO-SR, using the same clustering and selection pipeline but replacing AM scores with RM scores for within-cluster selection. If RM-PPO-SR performs substantially worse than AM-PPO-SR (e.g., on forget-test loss rate, where AM-PPO-SR achieves 7.95% vs. RM-PPO's 29.28%), this would validate the synergy claim: Selective Rehearsal's cluster-based selection requires task-invariant scores to work correctly, and using RM scores biases selection toward high-score categories, undermining diversity. If RM-PPO-SR performs comparably to AM-PPO-SR (e.g., forget-test loss rate drops to ~10%), then Selective Rehearsal is independently effective regardless of score calibration, and the two techniques are complementary but not synergistic — each works on its own, and their combination is additive rather than multiplicative. This experiment is low-cost (it only requires training RM-PPO-SR, since RM-PPO and AM-PPO-SR already exist) and would substantially clarify the architecture's design principles.

Analyzing which examples benefit from Selective Rehearsal to understand the forgetting mechanism. The paper reports aggregate win/loss/tie rates on the forget test set (Table 2) but does not characterize which types of SFT examples are most susceptible to forgetting, nor which are best protected by Selective Rehearsal. A follow-up analysis would stratify the forget test examples by (i) task category (using the 61-category task spectrum from the proprietary data, or the helpfulness/harmlessness split in HH-RLHF), (ii) SFT model confidence (entropy of the SFT model's response distribution), (iii) reward/advantage model score on the SFT response, and (iv) whether the example was selected into the rehearsal dataset D_R or not. The analysis would answer: Does Selective Rehearsal only protect the specific examples it rehearses, or does the protection generalize to similar examples in the same semantic cluster? Are low-confidence SFT examples more vulnerable to forgetting (because the SFT behavior is less stable) or less vulnerable (because the reward signal has more room to improve them)? Are examples not selected into D_R still protected by the clustering effect (since the policy's behavior on similar rehearsed examples creates a basin of attraction)? This would transform Selective Rehearsal from a demonstrated technique into a understood mechanism with predictable behavior, enabling practitioners to estimate its benefits on new datasets without full experimental evaluation.

Practical Applications and Downstream Use Cases

Production RLHF pipelines for multi-task chatbots. The paper's most direct application is to multi-task aligned chatbot training — precisely the setting of major deployed systems like ChatGPT, Claude, Bard, and Llama 2-Chat, which must handle prompts spanning dozens of task categories (QA, code generation, creative writing, summarization, safety-critical refusals, etc.). The Advantage Model's per-task score equalization (Figure 4a) directly addresses a problem that these production systems currently manage through ad-hoc means: training separate reward models for different task categories (as noted in Section 4.1 for helpfulness vs. harmlessness), using prompt-level normalization during PPO, or applying post-hoc calibration. The Advantage Model provides a single-model solution that scales to 61+ task categories without requiring category labels at inference time. The quantifiable benefit, based on Table 2, is (i) a reduction in catastrophic forgetting on expert-aligned examples from a 29.28% loss rate to 7.95% — meaning the chatbot is much less likely to degrade on skills it already possessed after SFT, and (ii) a net win rate improvement over SFT of +6.01 percentage points (AM-PPO-SR vs. −0.10 for RM-PPO), concentrated in the ~25% of prompts where models meaningfully differ. For a production system serving millions of queries, a 6-point swing in win rate on the subset of non-tied examples translates to a substantial user-experience improvement.

Alignment data generation for self-improvement loops. The Selective Rehearsal pipeline — clustering prompts by semantic similarity, selecting high-quality responses per cluster, and rehearsing them — is directly applicable to self-improvement data generation workflows where a model is iteratively fine-tuned on its own high-quality outputs (e.g., STaR, ReST, or rejection sampling fine-tuning). The standard approach selects training examples purely by reward score threshold (Equation 4, F = 1_{r(x,y) ≥ τ}), which the paper notes "only consider[s] reward model score" and can miss diversity. Selective Rehearsal's cluster-then-select strategy ensures that the fine-tuning data covers the full skill distribution while still filtering for quality. The quantifiable benefit, from Figure 6, is that the method is robust to the number of clusters c (variance of ~0.05 reward points), meaning a practitioner can set c based on the desired diversity granularity without extensive tuning. The Advantage Model provides the quality signal for selection without introducing task-category bias — a +2.35 net win rate on forget-set examples compared to the SFT baseline (Table 2, AM-PPO-SR) means the self-improvement loop is less likely to amplify reward model biases over iterations.

Reward model training with explicit calibration objectives for safety-critical applications. For alignment applications where overconfident reward model errors have severe consequences — e.g., safety classifiers that must correctly identify harmful content, or medical QA systems where incorrect answers carry high risk — the Advantage Model's calibration improvement (ECE reduction from 5.35 to 3.83 on proprietary data, ~28%) is practically significant. A lower ECE means that when the model assigns a high advantage score to a response (implying high confidence that it is preferred), that confidence is more likely to be warranted. In safety-critical settings, this reduces the risk of the PPO policy discovering and exploiting overconfident reward model blind spots (where a dangerous response receives a high score because the reward model incorrectly believes it is good). The bounding loss (terms 2 and 3 of Equation 6) also caps the maximum advantage magnitude, preventing any single response from receiving an arbitrarily high score that could dominate the PPO gradient and cause the policy to collapse to a narrow, potentially unsafe, response mode. The practical implementation is a drop-in replacement for the standard Bradley-Terry reward modeling objective — same architecture, same data, different loss function — making the adoption cost low for teams already training reward models.