WARM: On the Benefits of Weight Averaged Reward Models

🎯 Pitch

Averaging the weights of multiple reward models—rather than their predictions—dramatically reduces reward hacking in RLHF. Simply starting from the same pre-trained model makes fine-tuned weights linearly connected, so weight averaging becomes a single, cheap reward model that is more reliable on out-of-distribution policy samples and robust to noisy human preferences. This tiny architectural shift yields a 79.4% win rate for an RL policy trained with WARM against a standard single-reward-model policy.

1. Executive Summary

This paper introduces Weight Averaged Reward Models (WARM) as an efficient strategy for combining multiple reward models by averaging their weights — rather than their predictions — to address reliability under distribution shifts (as the policy drifts during RL, requiring the RM to score out-of-distribution generations) and robustness to preference inconsistencies (as 25% label corruption in the preference dataset). Experiments on the TL;DR summarization benchmark using PaLM-XXS RMs and PaLM-XS policies demonstrate that WARM mitigates reward hacking: in best-of-(N) sampling, WARM achieves up to a 92.5% win rate against random selection, while in RL fine-tuning, a policy trained with WARM ((M = 6)) attains a 79.4% win rate against a policy trained with a single RM, establishing that weight averaging outperforms prediction ensembling as a variance-reduction and invariance-regularization method — with the boundary that this advantage over ensembling grows as data moves further from the training distribution.

2. Context and Motivation

The Core Problem: Reward Models Fail Under Distribution Shifts and Preference Noise

The fundamental challenge this paper addresses sits at the heart of modern LLM alignment. In the standard RLHF pipeline, a reward model is trained once on a fixed, offline preference dataset and then used to guide the policy during online reinforcement learning. The problem is that this RM is asked to do something it was never trained to do: reliably score samples that look different from anything in its training data. As the RL policy drifts away from its SFT initialization — by design, since it's optimizing the RM's own signal — the generations it produces increasingly depart from the distribution of responses in the original preference dataset. The RM, trained only on offline data, must now evaluate these out-of-distribution generations. If the RM makes errors on these OOD samples, the policy can learn to exploit those errors, achieving high proxy rewards without genuinely improving. This is reward hacking.

The paper identifies this as the primary challenge in reward modeling (Section 2.2):

"The primary challenge is the distribution shifts resulting from the offline nature of preference data. Indeed, the generations in the preference dataset and those from the policy (\theta_{sft}) do not necessarily follow the same distributions, and the shifts can become even more pronounced due to model drift during RL."

But distribution shift is not the only problem. The paper identifies a second major challenge: the binary preference labels in the training data are inherently noisy. The paper documents several sources of this noise (Section 2.2):

Human labeler limitations: fatigue, misunderstanding of complex instructions, and reliance on superficial heuristics like output length or formatting rather than genuine quality, e.g., how the paper notes that "Human labelers often rely on simpler criteria (length, bullet points, politeness) over more nuanced indicators."
Task complexity: for tasks requiring specific expertise, labelers may simply not understand the content well enough to provide reliable preferences.
Multi-objective tensions: alignment involves tradeoffs between competing desiderata — harmlessness, helpfulness, engagement — and different labelers have different opinions. The paper cites the InstructGPT finding that inter-labeler agreement is only 72.6%.
AI feedback introduces its own noise: even when replacing human labelers with AI-generated preferences (as the paper itself does, following RLAIF), the AI labeler is sensitive to prompting strategies and can fail in systematic ways.

These two challenges — distribution shifts and label noise — create a negative feedback loop (Section 2.2). Distribution shifts cause the RM to become poorly calibrated, producing more extreme and less reliable reward values on OOD generations. These unreliable rewards exacerbate model drift (the policy moves further in response to miscalibrated signals), which in turn creates even larger distribution shifts for the RM. The cycle intensifies until the policy has fully exploited the RM's vulnerabilities, achieving high proxy rewards while producing outputs that fail to satisfy true human preferences — the very definition of reward hacking.

Why This Problem Matters

The paper argues that reward hacking is not merely an academic concern but a central barrier to safe and effective LLM deployment (Section 1). The consequences are concrete:

Performance degradation. Reward-hacked policies produce linguistically flawed outputs, overly verbose generations, or content that optimizes superficial proxies rather than substantive quality (Lewis et al., 2017; Singhal et al., 2023). The paper frames this as a manifestation of Goodhart's Law: "when a measure becomes a target, it ceases to be a good measure." The proxy RM, originally a reasonable approximator of quality, ceases to measure quality once it becomes the optimization target.

Checkpoint selection becomes unreliable. Because the proxy RM can no longer be trusted to accurately rank policies, practitioners cannot reliably determine which checkpoint during RL training actually performs best. The paper notes this creates a practical deployment headache: you don't know when to stop training.

Safety risks. Reward hacking can produce sycophancy (the model tells users what they want to hear rather than truth), amplify social biases present in the limited demographics of feedback providers, and — most critically — create misaligned systems that pursue unintended objectives. The paper explicitly connects this to the broader AI safety literature (Amodei et al., 2016; Hendrycks and Mazeika, 2022), arguing that reliable reward modeling is a prerequisite for safe deployment.

The problem compounds with scale. As LLMs are integrated into critical decision-making and everyday life, the cost of reward hacking escalates. A customer service chatbot that learns to produce polite-sounding but unhelpful responses, or a medical summarization system that learns to cherry-pick reassuring information, represents not just a technical failure but a potential source of harm.

Where Existing Approaches Fall Short

The paper identifies several existing strategies for addressing reward hacking and explains why each is insufficient alone.

KL regularization (Section 2.2). The standard defense is to penalize the policy for deviating from its SFT initialization, using a Kullback-Leibler divergence term in the RL objective controlled by a coefficient (\alpha). This directly limits distribution shift — the policy cannot drift too far from the data the RM was trained on. However, the paper notes that KL regularization introduces an extra hyperparameter (\alpha) that is difficult to tune: too strong and the policy underfits (never meaningfully improves beyond SFT), too weak and reward hacking still occurs. The paper's own RL experiments (Section 5.2, Figure 8c) empirically demonstrate this sensitivity. Moreover, KL regularization is a blunt instrument — it restricts exploration uniformly, preventing the policy from discovering genuinely useful behaviors that happen to lie far from the SFT distribution.

Iterative data collection (Section 2.2). Another approach is to continuously collect new preference labels on generations from the evolving policy, retraining the RM to stay current with the policy's output distribution (as done in LLaMA-2). The paper identifies clear efficiency barriers: this requires ongoing human annotation (expensive and slow) and repeated retraining (computationally costly). It doesn't scale to rapid iteration cycles.

Active learning (Section 2.2). Rather than passively labeling whatever the policy generates, active learning strategies (Reddy et al., 2020) proactively seek out diverse or potentially problematic examples to enrich the preference dataset. While promising, the paper treats these as complementary rather than competing approaches — they address data coverage but don't solve the fundamental problem of how to make the RM robust given the data you have.

Prediction ensembling (Section 2.3) — the most direct baseline. The approach most similar to WARM is to train multiple RMs and average their predictions (logits) to produce a more reliable reward signal. The paper acknowledges that this has strong pedigree: it was used in the original RLHF work by Christiano et al. (2017) and has been shown in recent work (Eisenstein et al., 2023; Coste et al., 2023) to mitigate reward hacking. From a bias-variance perspective, ensembling reduces variance when the ensemble members are sufficiently diverse, and since variance is the dominant source of error under distribution shift (as argued in the OOD literature), this directly addresses one of the core challenges.

However, the paper identifies three limitations of prediction ensembling that WARM aims to overcome:

Efficiency (memory and inference overhead): Each additional RM in the ensemble consumes memory and requires forward passes at inference time. The cost grows linearly with (M), the number of ensemble members. This is particularly problematic given the trend toward larger RM architectures, which consistently perform better. The paper frames this as making ENS "incompatible with the scaling trend in RMs" — you want more and larger RMs for better performance, but the computational burden of ensembling them becomes prohibitive.
Failure to improve robustness to label corruption: The paper's experiments (Section 4.2, Figure 4) reveal a surprising finding: prediction ensembling, while effective at reducing variance under distribution shifts, does not help — and can even hurt — when the preference labels themselves are corrupted. On the corrupted training subset, ENS achieves higher accuracy than WA because it memorizes the corrupted labels along with the clean ones. The paper argues this memorization is harmful because it preserves non-robust features that make the reward signal unstable.
Theoretical insight gap: Prior work treated ensembling primarily as variance reduction, but the precise mechanism by which it improves OOD generalization — and especially the difference between averaging predictions versus averaging weights — was not well understood. The paper argues that existing analyses of weight averaging (models soups, DiWA) observed that WA empirically matches or exceeds ENS under distribution shifts but didn't explain why this happens. The paper's theoretical analysis in Section 4.3 fills this gap.

How This Paper Positions Itself

The paper positions WARM at the intersection of two previously disconnected research literatures: OOD generalization through weight averaging and robust reward modeling for RLHF.

From the weight averaging literature (Appendix A.1), the paper inherits a rich set of empirical findings: fine-tuned models that share a common pretraining remain linearly mode connected (Frankle et al., 2020; Neyshabur et al., 2020), meaning their weights can be interpolated without catastrophic accuracy loss. Prior work exploited this property for supervised classification (model soups, DiWA, model ratatouille, SWAD), showing that weight averaging can match or exceed prediction ensembling with zero inference overhead. However, this insight had never been applied to reward modeling — a fundamentally different setting where the model isn't just classifying inputs but is actively shaping the training of another model through RL.

The paper's key positioning move is to argue that reward modeling is precisely the setting where weight averaging's properties are most valuable. The two challenges of reward modeling (distribution shifts and label noise) map directly onto the two benefits of weight averaging that the paper identifies:

Reliability under distribution shifts (benefit 1): WA acts as variance reduction, approximating the behavior of prediction ensembling — which is known to help under distribution shifts because variance, not bias, is the dominant error source when data shifts.
Robustness to label corruption (benefit 2): WA goes beyond variance reduction — the theoretical analysis in Section 4.3 shows that WA selectively preserves predictive mechanisms that are invariant across training runs while suppressing run-specific mechanisms that often correspond to memorization of corrupted labels. This is a property that prediction ensembling does not share (ENS preserves all mechanisms proportional to their probability of being learned), and it directly addresses the preference inconsistency challenge.

The paper also explicitly positions itself as a practical engineering solution rather than a purely theoretical contribution. WARM is described as "simple, efficient and scalable" — it requires no changes to the training pipeline beyond launching multiple fine-tunings with different hyperparameters (something many practitioners already do as grid search) and then averaging the resulting weights. There's no new architecture, no additional loss terms, no hyperparameters beyond what standard RM training already requires. The paper emphasizes that WARM is compatible with and complementary to other approaches: it can be combined with KL regularization (in fact, it makes KL tuning easier by allowing larger KL budgets before hacking occurs), with active learning, and with iterative data collection.

A final important positioning note: the paper is careful to argue that WARM is orthogonal to the choice of preference data source. The experiments use AI-generated preferences (RLAIF) rather than human labels, which the paper defends by citing evidence that AI feedback correlates well with human preferences and provides cleaner experimental control (the AI labeler serves as an oracle metric for detecting reward hacking). However, the paper argues that the benefits of WARM should extend to human-labeled preferences as well, since the underlying challenges (distribution shifts, label noise) are present in both settings.

In summary, the paper positions WARM as filling a specific gap: there exists a method (prediction ensembling) that partially addresses reward hacking but is computationally expensive and fails to handle label noise, and there exists a technique (weight averaging) that is known to help in supervised OOD settings but has never been applied to reward modeling. By connecting these two threads, the paper aims to provide a principled, efficient, and theoretically grounded solution to one of the central bottlenecks in RLHF.

3. Technical Approach

3.1 Reader Orientation

This paper proposes a method for constructing a single reward model — called Weight Averaged Reward Models (WARM) — by fine-tuning multiple reward models from the same pre-trained language model and then averaging their weight vectors together, rather than averaging their output predictions. The system addresses the problem of reward hacking in RLHF by producing a proxy reward model that is simultaneously more reliable when scoring out-of-distribution generations and more robust to label noise in the preference training data, all without the memory and inference overhead of maintaining an ensemble of separate models.

3.2 Big-Picture Architecture (Diagram in Words)

The WARM system has three stages, proceeding from a pre-trained language model to a final reward model used for guiding policy optimization:

Shared Pre-Trained Initialization. A pre-trained LLM (specifically, PaLM-XXS) undergoes supervised fine-tuning (SFT) on preferred summaries, producing a base checkpoint $\theta_{sft}$. From this, different SFT checkpoints collected at various training steps serve as diverse featurizer initializations for subsequent reward model training.
Diverse RM Fine-Tunings ($M$ runs). $M$ separate reward models are fine-tuned from these initializations on the same preference dataset, but with different hyperparameters (learning rates, dropout probabilities) and seeing the training data in different random orders. Each run produces a weight vector $\phi_i$ that parameterizes a reward model $r_{\phi_i}(x, y)$ — a scalar-valued function scoring how well a summary $y$ follows a Reddit post $x$.
Weight Averaging. The $M$ weight vectors are element-wise averaged to produce a single weight vector $\phi_{\text{WARM}} = \frac{1}{M}\sum_{i=1}^M \phi_i$. This single model $r_{\phi_{\text{WARM}}}$ then serves as the proxy reward model in downstream best-of-(N) selection or RL fine-tuning — it requires exactly the same memory and inference cost as any individual RM, but inherits the statistical benefits of combining multiple models.

The information flow is straightforward: a post-summary pair $(x, y)$ enters the system → the single WARM model $r_{\phi_{\text{WARM}}}$ computes one scalar reward → this reward is used either to select the best among $N$ sampled summaries (BoN) or to compute policy gradients for RL fine-tuning.

3.3 Roadmap for the Deep Dive

First, the Bradley-Terry preference model and the standard RM training loss (Equation 1), since this is the objective that each individual RM optimizes and is the foundation on which WARM is built.
Second, the WARM procedure itself — the three-step recipe and the why behind each step: shared pre-training (required for linear mode connectivity), diverse fine-tunings (needed for variance reduction and invariance), and weight averaging (the core operation that distinguishes WARM from prediction ensembling).
Third, the Baklava initialization strategy, which provides a principled source of diversity across RM fine-tunings while preserving linear mode connectivity — a key enabling technique that relaxes the requirement of identical SFT initialization.
Fourth, the linear mode connectivity property (Observation 1) that makes weight averaging viable, including the empirical validation (Figure 3) and the crucial role of linear probing in preventing feature distortion.
Fifth, the two complementary theoretical analyses: the first-order (variance reduction) perspective that explains why WA approximates ENS under distribution shifts (Observation 2), and the second-order (invariance regularization) perspective that explains why WA exceeds ENS under label corruption (Observation 3 and the mathematical analysis in Section 4.3).
Sixth, the sources of diversity and the weight selection procedure, which determine which individual RMs get averaged together — a practical engineering detail with significant performance implications.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a methodological paper whose core idea is that averaging the weight vectors of multiple reward models — rather than averaging their predictions — produces a more robust and efficient proxy reward for RLHF, and that this advantage stems from weight averaging's ability to selectively preserve predictive mechanisms that are invariant across independent training runs while suppressing run-specific mechanisms that often correspond to memorization of corrupted labels.

The Bradley-Terry Preference Model and Reward Model Training Objective

Before understanding how WARM combines reward models, we must understand what an individual reward model is and how it is trained. The paper follows the standard Bradley-Terry framework for modeling pairwise preferences.

The Bradley-Terry assumption. Given a prompt $x$ (a Reddit post to be summarized) and two candidate responses $y^+$ (the preferred summary) and $y^-$ (the dispreferred summary), the Bradley-Terry model posits that the probability that $y^+$ is preferred over $y^-$ is a logistic function of the difference in their underlying reward values:

$P(y^+ \succ y^- \mid x) = \sigma\left(r(x, y^+) - r(x, y^-)\right)$

where $\sigma(z) = \frac{1}{1 + e^{-z}}$ is the logistic sigmoid function, and $r(x, y)$ is the latent reward for summary $y$ given post $x$.

What this equation computes: the probability that a human (or AI) labeler would choose $y^+$ over $y^-$ as a function of the difference in their quality scores. When the two summaries are equally good ($r(x, y^+) = r(x, y^-)$), the preference probability is exactly 0.5 (random choice). When $y^+$ is substantially better ($r(x, y^+) \gg r(x, y^-)$), the probability approaches 1. The sigmoid function maps any real-valued reward difference to a valid probability in $(0, 1)$.

Why this form: The Bradley-Terry model is the standard in preference learning because it makes the minimal assumption that preferences depend only on an underlying scalar quality score, and that the comparison between two items is determined by their score difference. Alternative models (e.g., Thurstone's model, which uses a Gaussian CDF rather than a sigmoid) would change the specific functional form but maintain the same structural assumption. The logistic form is computationally convenient because its log-likelihood is convex and because it connects naturally to binary classification: preferring $y^+$ over $y^-$ is equivalent to classifying the pair $(y^+, y^-)$ as belonging to the positive class.

From preference probability to training loss. The reward model $r_\phi$ is an LLM parameterized by $\phi$ that takes a post $x$ and summary $y$ as input and outputs a single scalar reward. Given a dataset of pairwise preferences $\mathcal{D}_{\text{train}} = \{x_d, y_d^+, y_d^-\}_{d=1}^D$, the maximum likelihood estimate of $\phi$ minimizes the negative log-likelihood loss:

$\mathcal{L}_R\left(r_\phi, \mathcal{D}_{\text{train}}\right) = -\mathbb{E}_{(x, y^+, y^-) \in \mathcal{D}_{\text{train}}}\left[\log \sigma\left(r_\phi(x, y^+) - r_\phi(x, y^-)\right)\right]$

where the expectation is over the empirical distribution of preference triples in the training set, $r_\phi(x, y^+)$ is the scalar reward assigned to the preferred summary, and $r_\phi(x, y^-)$ is the scalar reward assigned to the dispreferred summary.

What this equation computes: for each training example, compute the difference in rewards between the preferred and dispreferred summary, pass this difference through the logistic sigmoid to get a predicted probability that the preferred summary is indeed better, and take the negative log of this probability. A perfect model (predicting probability 1 for the correct choice) incurs zero loss; a model predicting probability 0.5 (random choice) incurs a loss of $\log 2 \approx 0.693$ per example; worse-than-random predictions are penalized even more heavily. The total loss is the average across the training set.

Why this form: this is the standard binary cross-entropy objective applied to pairwise preference data. The sigmoid converts the reward difference into a probability, and the log loss provides a proper scoring rule — meaning that the unique minimizer of this loss (in expectation over the true data distribution) is the true preference probability. Alternative formulations could directly regress the reward difference onto $\pm 1$ targets, but this would not produce calibrated probability estimates. The log loss is particularly sensitive to confident mistakes (predicting probability near 1 when the wrong answer is correct), which is desirable because confident misrankings are precisely the kind of RM error that leads to reward hacking.

Reward model architecture. The paper specifies that each RM is an LLM (PaLM-XXS) where the final layer is a linear classification head $\omega$ added on top of the SFT model's extracted features. Formally, the RM weights are $\phi = (\theta_{\text{sft}}, \omega)$ — the SFT featurizer plus a linear probing layer. The paper follows Kumar et al. (2022) in using linear probing rather than random initialization for $\omega$: the classification head is pre-trained on the preference task before the full model is fine-tuned, preventing the featurizer from distorting its pre-trained representations to accommodate a randomly initialized head. This is a critical design choice because it facilitates linear mode connectivity (as explained in Section 3.2 of the paper): if the featurizer moves too far from its initialization, the weights from different fine-tuning runs may no longer be linearly interpolable.

The WARM Procedure: Three-Step Recipe

WARM is defined by the following three steps (Section 3.1):

Step 1: Shared pre-trained initialization. All $M$ reward models start from the same SFT checkpoint $\theta_{\text{sft}}$ (or, under the Baklava strategy, from different checkpoints collected along the same SFT trajectory — see below). Critically, all RMs share the same linear probe $\omega$ that converts the featurizer's output into a scalar reward. The paper emphasizes this shared initialization is what makes weight averaging possible: without it, the weights from different runs could correspond to different internal representations (due to permutation symmetries in neural networks), and averaging them would produce nonsense.

Step 2: Diverse fine-tunings. The paper launches $M$ independent RM fine-tuning runs, each optimizing Equation (1) on the same preference dataset $\mathcal{D}_{\text{train}}$. Diversity across runs is introduced through three mechanisms (Section 3.3):

Different data orders: each run sees the training examples in a different random sequence due to stochastic mini-batch sampling. This is the most basic source of diversity and occurs naturally without any deliberate effort.
Different hyperparameters: learning rates are sampled from $\{1 \times 10^{-5}, 4 \times 10^{-5}, 1 \times 10^{-4}\}$ and dropout probabilities from $\{0.05, 0.1\}$. The paper notes that hyperparameters are kept in "a mild range to preserve the LMC" — aggressive hyperparameter variations could push weights into different loss basins where linear interpolation fails.
Baklava initialization (optional): instead of starting all RMs from the final SFT checkpoint, featurizer initializations are drawn from different checkpoints $\{\theta_{\text{sft}}^i\}_{i=1}^M$ collected at different numbers of SFT training steps (specifically, steps 8k, 10k, and 12k). This increases diversity beyond the other two sources while still satisfying the shared pre-training constraint (all checkpoints come from the same pre-training and SFT trajectory).

All runs use the Adafactor optimizer (Shazeer and Stern, 2018), a batch size of 128, 10,000 training steps, and train on the Reddit TL;DR preference dataset with AI-generated labels following the RLAIF procedure (Lee et al., 2023).

Step 3: Weight averaging. The $M$ trained weight vectors $\{\phi_i\}_{i=1}^M$ are averaged element-wise:

$\phi_{\text{WARM}} = \frac{1}{M}\sum_{i=1}^M \phi_i$

where the sum is over all parameters in the model (both the featurizer weights and the classification head). This produces a single weight vector $\phi_{\text{WARM}}$ that defines one reward model $r_{\phi_{\text{WARM}}}$. The key operational fact: at inference time, computing $r_{\phi_{\text{WARM}}}(x, y)$ requires exactly one forward pass through the model — identical in cost to any individual RM — yet the model's parameters incorporate information from all $M$ training runs.

Why averaging in weight space rather than prediction space? The paper's central argument is that weight averaging provides benefits that prediction ensembling (ENS) does not. In ENS, the combined reward would be $r_{\text{ENS}}(x, y) = \frac{1}{M}\sum_{i=1}^M r_{\phi_i}(x, y)$, requiring $M$ forward passes. By merging the weights before inference, WARM achieves the variance reduction benefits of ensembling (explained below as the first-order analysis) while additionally providing robustness to label corruption (the second-order analysis) — all with zero inference overhead. The paper's theoretical analysis in Section 4.3 formalizes why these two combination strategies behave differently.

Weight selection procedure. The paper does not average all $M$ trained RMs indiscriminately. Instead, the $M$ weights are first sorted by their individual accuracy on the OOD test set $\mathcal{D}_{\text{ood}}$, and only the top $M$ are averaged (where $M$ can be 2, 6, or 10 in the experiments). This is a practical heuristic: including poorly performing RMs would degrade the average. Figure 10 in the appendix validates this choice, showing that averaging from best to worst weights yields higher accuracy than random or worst-to-best selection. The paper speculates that more sophisticated selection procedures (e.g., greedy selection as in model soups) could further improve results, but does not explore them.

The Baklava Initialization Strategy

The Baklava procedure, illustrated in Figure 2 of the paper, is a novel method for introducing diversity across RM fine-tunings while maintaining the shared pre-training required for linear mode connectivity.

Motivation. The standard model soups approach (Wortsman et al., 2022) initializes all fine-tunings from the exact same model checkpoint. This guarantees linear mode connectivity but limits diversity — all RMs start from the same point in weight space, so their divergence during training is determined solely by stochastic training dynamics (data order and dropout). Baklava relaxes this constraint: instead of identical initialization, it uses different checkpoints along a single SFT trajectory as featurizer initializations. The key insight is that SFT checkpoints at different training steps — while not identical — still share the same pre-training and have not diverged so far that linear mode connectivity breaks.

Procedure (Figure 2). Starting from the pre-trained LLM $\theta_{\text{pt}}$, the paper runs a single SFT fine-tuning on preferred summaries for 12,000 steps. Checkpoints are saved at multiple points along this trajectory — specifically at steps $\{8\text{k}, 10\text{k}, 12\text{k}\}$. Each RM fine-tuning then initializes its featurizer from one of these checkpoints $\theta_{\text{sft}}^i$, but all RMs share the same linear probe $\omega$ (the classification head trained to predict preferences from the final SFT checkpoint's features). This shared linear probe is crucial: it ensures that the different featurizers are all projecting into the same output space, preventing the kind of representational misalignment that would break weight averaging.

The name "Baklava" comes from a culinary analogy (following the "model soups" and "model ratatouille" naming convention in the WA literature) and refers to the diamond-shaped geometric structure of the procedure in Figure 2: the pre-trained model sits at the top, the SFT checkpoints fan out in the middle, the RM fine-tunings converge at the bottom, and the final weight average sits at the lowest point.

Comparison to alternatives. Baklava is positioned as more efficient than model ratatouille (Ramé et al., 2023), which requires training auxiliary models on multiple different tasks to get diverse initializations. Baklava achieves comparable diversity by simply sampling checkpoints from a single SFT run — no additional training tasks needed. The paper notes that taking a too-early SFT checkpoint (substantially before 8k steps) would drastically reduce RM accuracy, consistent with the vanishing gradient phenomenon observed in Razin et al. (2023) — the linear probe needs a sufficiently well-trained featurizer to produce meaningful representations. The chosen checkpoints at 8k–12k steps represent a sweet spot where the featurizer is competent but still exhibits meaningful variation.

Empirical validation (Figure 16, Appendix). The paper evaluates Baklava in the BoN setting, comparing WARM with $M = 2$ where the two RMs have featurizer initializations from different SFT checkpoints (step 12k and step 8k). The results show that Baklava-based WARM outperforms both the best individual RM (initialized from step 12k) and prediction ensembling of the same two RMs, confirming that Baklava provides useful diversity without breaking weight averaging.

Linear Mode Connectivity: The Enabling Property

The entire WARM approach rests on linear mode connectivity (LMC) — the empirical finding that two neural networks fine-tuned from a shared pre-trained initialization can have their weights linearly interpolated without catastrophic accuracy loss. The paper formalizes this as Observation 1:

Observation 1 (LMC). Given two fine-tuned weights $\phi_1$ and $\phi_2$ with a shared pre-training and a test dataset $\mathcal{D}_{\text{test}}$, then for all $\lambda \in [0, 1]$: $\text{Acc}\left(r_{(1-\lambda)\cdot\phi_1 + \lambda\cdot\phi_2}, \mathcal{D}_{\text{test}}\right) \geq (1-\lambda) \times \text{Acc}\left(r_{\phi_1}, \mathcal{D}_{\text{test}}\right) + \lambda \times \text{Acc}\left(r_{\phi_2}, \mathcal{D}_{\text{test}}\right)$

where $\text{Acc}(r_\phi, \mathcal{D})$ is the pairwise accuracy of the reward model on dataset $\mathcal{D}$, defined as the fraction of test pairs where the reward model correctly assigns a higher score to the preferred summary: $\mathbb{E}_{(x, y^+, y^-) \in \mathcal{D}}\left[\mathbf{1}_{r_\phi(x, y^+) \geq r_\phi(x, y^-)}\right]$. The interpolation coefficient $\lambda$ sweeps smoothly from using only $\phi_1$ (when $\lambda = 0$) to using only $\phi_2$ (when $\lambda = 1$), with $\lambda = 0.5$ corresponding to equal-weight averaging.

What this inequality states: the accuracy of the interpolated model should be at least as good as the linear interpolation of the individual accuracies. For example, if $\phi_1$ achieves 75% accuracy and $\phi_2$ achieves 77%, then the equal-weight average $\phi_{\text{WARM}} = (\phi_1 + \phi_2) / 2$ should achieve at least 76% accuracy — it should not suffer a catastrophic accuracy drop below either individual model. This is a non-trivial property: there is no mathematical guarantee that averaging the parameters of two non-linear functions produces a function that performs well. The fact that it does work is an empirical regularity of fine-tuned neural networks sharing pre-training.

Why this property holds. The paper attributes LMC to the shared pre-training constraining the divergence during fine-tuning (Neyshabur et al., 2020): the fine-tuned weights remain in convex regions of the loss landscape where linear interpolation between any two solutions remains a good solution. When models are trained from scratch, even with identical random initialization, they can end up in different basins of the loss landscape due to the inherent symmetries in neural network architectures (e.g., permutation symmetries of hidden units); the pre-training resolves these symmetries, effectively "aligning" the internal representations so that corresponding neurons in different models encode similar features. The paper also credits the use of linear probing for the classification head: by pre-training the classifier $\omega$ before fine-tuning, the featurizer $\theta_{\text{sft}}$ is less distorted during RM training (Kumar et al., 2022), making it easier for weights from different runs to remain interpolable.

Empirical validation (Figure 3). The paper validates LMC by training pairs of RMs with increasing levels of diversity and plotting the accuracy of their interpolations on the OOD test set $\mathcal{D}_{\text{ood}}$ as $\lambda$ varies from 0 to 1. Four sources of diversity are tested:

Figure 3(a): Two checkpoints from a single RM fine-tuning at different training steps (8k and 10k) — minimal diversity.
Figure 3(b): Two independent RM fine-tunings with identical hyperparameters but different data orders — moderate diversity.
Figure 3(c): Two independent RM fine-tunings with different learning rates (1e-4 and 4e-5) — higher diversity.
Figure 3(d): Two RM fine-tunings with Baklava initialization (different SFT checkpoints at steps 8k and 12k) — highest diversity.

In all cases, the interpolated accuracy (WA curve) lies at or above the diagonal (the "Diag" reference representing the linear interpolation of individual accuracies), confirming Observation 1. The paper also observes that the accuracy gain over the individual models is larger when the models are more diverse — the WA curve in Figure 3(d) shows a clearer "bump" above the individual accuracies than in Figure 3(a). This is consistent with the variance-reduction perspective: more diverse ensemble members provide greater variance reduction when averaged.

The failed alternative: moving average (Remark 2). The paper attempted a cheaper alternative — averaging checkpoints collected along a single RM fine-tuning trajectory (stochastic weight averaging or moving average, Izmailov et al., 2018) — but found inferior results. The authors diagnose this as an unfavorable accuracy-diversity trade-off: early checkpoints (which would provide diversity) have poor individual accuracy and degrade the average, while later checkpoints (with good accuracy) are too similar to each other to provide meaningful variance reduction. This negative result motivates the use of independent fine-tuning runs, which can simultaneously achieve high individual accuracy and sufficient diversity.

First-Order Analysis: WA as Efficient Variance Reduction

The paper's first theoretical perspective on WARM is that weight averaging serves as an efficient approximation of prediction ensembling, and thus inherits ensembling's variance reduction properties. This is formalized as Observation 2:

Observation 2 (WA and ENS: 1st order analysis). Weight averaging and prediction ensembling perform similarly: for all $\lambda \in [0, 1]$ and a test dataset $\mathcal{D}_{\text{test}}$: $\text{Acc}\left(r_{(1-\lambda)\cdot\phi_1 + \lambda\cdot\phi_2}, \mathcal{D}_{\text{test}}\right) \approx \text{Acc}\left((1-\lambda) \times r_{\phi_1} + \lambda \times r_{\phi_2}, \mathcal{D}_{\text{test}}\right)$

What this equation states: the accuracy of the weight-averaged model (left side) is approximately equal to the accuracy of the prediction-averaged ensemble (right side). In other words, averaging the parameters and then making one prediction gives roughly the same result as making $M$ predictions and then averaging them.

Theoretical justification (first-order Taylor expansion). The paper provides a brief theoretical motivation: when the weight vectors $\phi_1$ and $\phi_2$ are close in parameter space ($\|\phi_1 - \phi_2\| \ll 1$), a first-order Taylor expansion of the reward function around the average weight $\bar{\phi}$ gives:

$r_{\lambda\phi_1 + (1-\lambda)\phi_2}(x, y) \approx r_{\bar{\phi}}(x, y) + \nabla_\phi r_{\bar{\phi}}(x, y)^\top \cdot \delta$

where $\delta$ is a small displacement. This linearization implies that the weight-averaged prediction is approximately the prediction-average — the non-linearities in the network are locally flat enough that averaging in weight space and averaging in output space coincide to first order. When the divergence between fine-tuned weights is small enough, the first-order terms dominate and WA approximates ENS.

Empirical validation (Figure 3). Across all four diversity settings in Figure 3, the WA and ENS accuracy curves closely track each other — they are visually almost indistinguishable. This confirms that, for practical purposes, WA achieves the same accuracy gains as ENS on the OOD test set.

Why this matters for reward modeling. The variance-reduction interpretation directly addresses the distribution shift challenge. Under distribution shifts, the dominant source of error in individual RMs is variance — different RMs make different errors on OOD samples because they have learned slightly different decision boundaries from the finite, narrow training distribution. Averaging across multiple RMs cancels out these uncorrelated errors, producing a more reliable combined estimate. This is well-established in the ensemble literature (Ueda and Nakano, 1996) and in the OOD generalization literature (Ramé et al., 2022). By approximating ENS, WARM inherits this variance reduction — but with the critical practical advantage of requiring only a single model at inference time.

The efficiency advantage over ENS. While ENS requires $M$ forward passes and $M$ times the memory to store all models, WARM requires exactly one forward pass and one model's worth of memory. This efficiency gap widens as RMs scale up: the paper notes that "larger architectures consistently perform better" for reward modeling, so the trend is toward bigger RMs, making the overhead of ENS increasingly burdensome. WARM decouples the statistical benefits of combining multiple models from the computational costs of deploying them.

What the first-order analysis misses. Observation 2 is only an approximation. As the paper demonstrates, there are systematic differences between WA and ENS that become apparent under label corruption — the first-order Taylor expansion breaks down because the differences between $\phi_1$ and $\phi_2$ are not negligible when the models learn fundamentally different features (some robust, some not). This motivates the second-order analysis.

Second-Order Analysis: WA as Invariance Regularization

The paper's deeper theoretical contribution is the finding that weight averaging is not merely an efficient approximation of ensembling — it is qualitatively different in how it handles conflicting information across ensemble members. This difference becomes visible under label corruption.

The experimental demonstration (Section 4.2, Figure 4). The paper designs a diagnostic experiment: 25% of the binary preference labels in the training set are randomly flipped before training two independent RMs $\phi_1$ and $\phi_2$. The performance of WA and ENS is then evaluated on four data subsets that form a spectrum from in-distribution to out-of-distribution:

Train (corrupt): training samples whose labels were flipped — these are "poisoned" examples that the RMs might memorize.
Train (clean): training samples with correct labels.
Validation (ID): in-distribution held-out data from the same distribution as training.
Test (OOD): out-of-distribution data (summaries from novel policies not seen during training).

The results, summarized in Observation 3, reveal a systematic pattern:

Observation 3 (WA and ENS: 2nd order analysis). The accuracy gains of WA over ENS grow as data moves away from the training distribution.

WA ≪ ENS on train corrupt: WA performs far worse on corrupted training samples, showing reduced memorization.

WA ≤ ENS on train clean: WA is slightly worse even on clean training samples.

WA ⪆ ENS on ID val: WA is comparable or slightly better on in-distribution held-out data.

WA ≥ ENS on OOD test: WA is noticeably better on out-of-distribution test data.

Figure 5 visualizes this pattern with histograms of the per-sample accuracy difference between WA and ENS across many independently trained pairs of RMs: the distribution for "train corrupted" is clearly shifted to the left (negative values, WA worse), while "test OOD" is shifted to the right (positive values, WA better).

What this pattern means: WA selectively forgets the training data — it does worse than ENS at reproducing the specific labels seen during training both clean and especially corrupted ones. But this forgetting is beneficial for generalization: what WA retains are the predictive patterns that generalize to new data, while what it loses are the patterns specific to individual training examples (especially the corrupted ones). ENS, by contrast, retains everything — it memorizes the corrupted labels just as well as the clean ones, because averaging predictions doesn't distinguish between robust and spurious features.

Theoretical Formalization: WA Preserves Invariant Mechanisms

The paper provides a formal analysis (Section 4.3) explaining why WA but not ENS exhibits this selective forgetting property. This is the paper's most original theoretical contribution.

Simplified setup. The analysis uses a binary classification framework with labels $y \in \{-1, 1\}$. Inputs $x$ are composed of $F$ features $\{z_j\}_{j=1}^F$, where each $z_j \in \mathbb{R}^d$ is a basis vector representing a specific predictive pattern. The input is generated as $x = [x_1, \dots, x_F] \in \mathbb{R}^{F \times d}$ where each component $x_j \sim \mathcal{N}(y \cdot z_j, \sigma \cdot I_d)$ — that is, each feature $z_j$ appears in the input with additive Gaussian noise scaled by $\sigma$. The noise is small ($\sigma \ll 1$), so the signal from each feature is clear but imperfect.

The featurizer is a binary selector. The model's featurizer $f = [f_1, \dots, f_F] \in \{0, 1\}^F$ is a binary vector indicating which features the model uses: $f_j = 1$ means the model extracts feature $j$, $f_j = 0$ means it ignores it. Each feature $j$ is learned with probability $p_j$ — this probability depends on how strongly the feature predicts the label in the training data. Robust, generalizable features (e.g., "the summary is factually accurate") have high $p_j$; spurious, memorized features (e.g., "this specific training example had a flipped label") have low $p_j$.

The linear classifier on top uses the selected features. Given the featurizer's output, the model computes a reward: $r(x) = \omega^\top f(x)$, where $\omega$ is a learned linear weight vector. Under the assumptions (orthogonal features, infinite training data, optimal linear fit), Lemma 5 from Lin et al. (2023) shows that the optimal weight for a feature is proportional to that feature's direction: if feature $j$ is selected ($f_j = 1$), then $\omega_j = z_j$. So the reward assigned to input $x$ is the sum of features $z_j$ weighted by how well they match the input: $r(x) = \sum_{j=1}^F f_j \cdot z_j^\top x_j$.

Limit behavior of prediction ensembling. With $M$ independent RMs $\{r_i = \omega_i^\top f_i\}_{i=1}^M$ and $M \to \infty$, the prediction ensemble average converges to the expected prediction:

$r_{\text{ENS}}^M(x) \xrightarrow{M \to \infty} \mathbb{E}[r(x)] \approx y \cdot \sum_{j=1}^F p_j \cdot |z_j|^2$

where $p_j$ is the probability that any individual RM uses feature $j$, and $|z_j|^2$ measures the signal strength of feature $j$.

What this equation means: ENS weights each feature proportional to $p_j$ — the probability that an individual model uses that feature. If a feature is used by half the models ($p_j = 0.5$), it contributes half its full signal to the ensemble prediction. Critically, ENS preserves all features in proportion to how often they are learned — including features with low $p_j$ that correspond to memorization of corrupted labels. If a particular corrupted training example causes 10% of models to learn a spurious pattern, ENS still preserves that pattern at 10% strength.

Limit behavior of weight averaging. Under weight averaging, we first average the weight vectors $\frac{1}{M}\sum_{i=1}^M \omega_i$ and featurizers $\frac{1}{M}\sum_{i=1}^M f_i$, and then compute the reward. As $M \to \infty$:

$\frac{1}{M}\sum_{i=1}^M f_i \xrightarrow{M \to \infty} \mathbb{E}[f] = [p_1, \dots, p_F]^\top \quad \text{and} \quad \frac{1}{M}\sum_{i=1}^M \omega_i \xrightarrow{M \to \infty} \mathbb{E}[\omega] = \sum_{j=1}^F p_j \cdot z_j$

Substituting these into the reward computation:

$r_{\text{WA}}^M(x) \xrightarrow{M \to \infty} \left(\sum_{j=1}^F p_j \cdot z_j\right)^\top \left(\sum_{j'=1}^F p_{j'} \cdot x_{j'}\right) \approx y \cdot \sum_{j=1}^F p_j^2 \cdot |z_j|^2$

where in the last step, the cross-terms vanish because of feature orthogonality ($z_j^\top x_{j'} \approx 0$ when $j \neq j'$) and because $f_j^2 = f_j$ (binary indicator).

What this equation means: WA weights each feature proportional to $p_j^2$ — the square of the probability that an individual model uses that feature. This is the crucial difference from ENS.

The key insight: squaring probabilities suppresses low-probability features. If a robust feature is used by most models ($p_j = 0.9$), ENS weights it at 0.9 while WA weights it at $0.9^2 = 0.81$ — a modest reduction. But if a spurious feature (corresponding to memorization of a corrupted label) is used by only 10% of models ($p_j = 0.1$), ENS weights it at 0.1 while WA weights it at $0.1^2 = 0.01$ — a tenfold suppression. Features that are rarely learned (low $p_j$) are effectively eliminated by weight averaging, while features that are consistently learned (high $p_j$) are largely preserved.

Why does this happen mechanically? The squaring arises because WA multiplies two independent quantities that both depend on $p_j$: the averaged featurizer (which selects feature $j$ with weight $p_j$) and the averaged linear head (which assigns weight $p_j \cdot z_j$ to feature $j$). For a feature to contribute to the WA prediction, it must be present both in the averaged featurizer and in the averaged linear head — this is an AND operation that requires the feature to be learned in both components. ENS, by contrast, averages the feature contributions from individual models: each model independently decides which features to use, and these decisions are simply averaged.

Connection to invariance. The paper argues (Remark 3) that this mathematical property connects WA to the invariance principle in domain generalization (Arjovsky et al., 2019). Features with high $p_j$ are those that are consistently learned across independent training runs — they are invariant predictive mechanisms that capture genuine structure in the data. Features with low $p_j$ are run-specific — they capture noise, peculiarities of the data order, or memorization of corrupted labels. By squaring the probabilities, WA implements a soft version of invariant risk minimization: it retains features that are invariant across runs (the "causal" ones) and suppresses features that vary across runs (the "spurious" ones). ENS, by averaging without squaring, treats all features democratically and thus fails to distinguish robust from spurious patterns.

Generalization to deep networks (Remark 4). The paper notes that in a network with $L$ layers, the probability would be raised to the power $p_j^L$ rather than $p_j^2$. Intuitively, each layer applies an AND-mask: for a feature to survive weight averaging through $L$ layers, it must be consistently learned at every layer across all ensemble members. This makes the suppression of spurious features even more aggressive in deeper architectures, though the paper does not empirically test this prediction.

Implications for reward smoothness (Remark 5). The paper connects this invariance property to the smoothness (Lipschitzness) of the learned reward function. By suppressing features with low $p_j$, WA reduces the model's reliance on brittle, non-robust features that can cause large changes in predicted reward for small changes in input. This makes the WARM reward function smoother in the input space — that is, $|r(x_1) - r(x_2)| \leq L \|x_1 - x_2\|$ with a smaller Lipschitz constant $L$. The paper cites Blondé et al. (2022), who argue that reward smoothness is essential for stable policy gradient optimization: "sharp changes in reward value are hard to represent and internalize." By producing a smoother reward landscape, WARM makes it harder for the policy to find and exploit narrow spikes in the proxy reward — which are precisely the loopholes that cause reward hacking.

Sources of Diversity and Weight Selection

Why diversity matters. The benefits of WA — both the variance reduction (first-order) and the invariance regularization (second-order) — depend on the individual RMs being sufficiently diverse. If all RMs learn identical features (all $p_j$ are near 0 or near 1), then ENS and WA produce identical results: there are no low-probability features to suppress, and no uncorrelated errors to cancel. Conversely, if the RMs are too diverse — learning entirely disjoint features and falling into different loss basins — then linear mode connectivity breaks and WA produces degraded accuracy. The paper explicitly frames this as "a delicate trade-off between ensuring LMC and diversity across weights."

The three diversity mechanisms (Section 3.3).

Data order: each fine-tuning run sees training examples in a different random sequence due to mini-batch sampling. This is the most basic form of diversity — it's cost-free and always present — but it provides the least diversity because the optimization trajectories are still strongly constrained by the shared initialization and identical hyperparameters.
Hyperparameter sampling: learning rates are drawn from $\{1 \times 10^{-5}, 4 \times 10^{-5}, 1 \times 10^{-4}\}$ and dropout probabilities from $\{0.05, 0.1\}$. Different learning rates cause models to converge at different speeds and to different local minima; different dropout probabilities change which features the model must learn to rely on. The paper emphasizes that these hyperparameters are kept "in a mild range" — extreme values would break LMC by pushing weights into different loss basins. This follows the recommendations from Ramé et al. (2022).
Baklava initialization: RM featurizers are initialized from different SFT checkpoints at steps $\{8\text{k}, 10\text{k}, 12\text{k}\}$. This provides substantially more diversity than data order or hyperparameters alone, because the featurizers start from different points in weight space. However, because all checkpoints come from the same SFT trajectory and share the same pre-training, LMC is preserved — the models diverge, but not so far that their weights cannot be meaningfully interpolated. The empirical validation in Figure 3(d) shows that this source of diversity produces the largest accuracy gains from averaging.

Weight selection optimization (Figure 10). Once $M$ RMs are trained, the paper does not blindly average all of them. Instead, individual RMs are first evaluated on the OOD test set $\mathcal{D}_{\text{ood}}$, sorted by decreasing accuracy, and only the top $K$ are averaged (where $K$ is the chosen WARM size — 2, 6, or 10 in the experiments). This is labeled $\phi_1, \phi_2, \dots, \phi_K$ where $\phi_1$ is the best individual RM and $\phi_2$ is the second-best. Figure 10 validates that "best to worst" selection consistently outperforms random selection or "worst to best" selection for any value of $K$. The paper speculates that greedy selection (iteratively adding the RM that most improves the average, as in model soups) could further improve results but does not implement it.

Why not simply average all models? Including poorly performing RMs would degrade the average, because their low individual accuracy means they contribute more noise than signal. The weight selection procedure ensures that only RMs that have genuinely learned useful preference patterns are included in the final average, while RMs that may have converged to poor local minima or overfit the training data are excluded.

Practical interpretation for practitioners. In a deployment setting, one would run a grid search over RM hyperparameters (which many teams already do), train $M$ candidate RMs, evaluate them on a held-out validation set, select the top $K$, and average their weights. The only additional cost relative to standard RM training is the averaging step itself (which is essentially free — it's just element-wise addition and division), and there is no additional inference cost since the result is a single model. This makes WARM extremely practical to adopt.

Design Choices: Why This Approach Over Alternatives

Why weight averaging over prediction ensembling?

The paper's answer is threefold. First, efficiency: WA requires one model at inference time (same memory and compute as a single RM), while ENS requires $M$ models. Second, robustness to label corruption: WA selectively suppresses spurious features (due to the $p_j^2$ effect), while ENS preserves all features in proportion to their prevalence — including memorized corrupted labels. Third, theoretical clarity: WA connects ensembling to invariance regularization, providing a unified perspective on why combining models helps under distribution shifts.

Why independent fine-tunings over moving average?

The moving average approach (averaging checkpoints from a single training run) was tested but found inferior (Remark 2). The diagnosis: early checkpoints (which provide diversity) have poor accuracy and degrade the average, while later checkpoints (with good accuracy) are too similar to provide variance reduction. Independent fine-tunings solve this by allowing all models to achieve high individual accuracy while still being diverse, because they explore different optimization trajectories from different initializations or hyperparameters.

Why Baklava initialization over standard model soups?

Standard model soups require all fine-tunings to start from the exact same SFT checkpoint. Baklava relaxes this to "same SFT trajectory" — allowing different checkpoints along the same training run to serve as initializations. This increases diversity (the featurizers start from different points in weight space) while maintaining LMC (all checkpoints share the same pre-training). Baklava is more efficient than model ratatouille, which requires training models on multiple different auxiliary tasks, because it only requires a single SFT run.

Why linear probing for the classification head?

The paper follows Kumar et al. (2022) in using linear probing: the classification head $\omega$ is pre-trained on the RM task before the featurizer is updated. This prevents the featurizer from distorting its pre-trained representations to accommodate a randomly initialized head — a phenomenon called "feature distortion" that can break LMC. When all RMs share the same linear probe and only the featurizers differ, the weights remain in the same basin of the loss landscape and remain interpolable.

Why not combine RMs from different architectures or pre-trainings?

The paper explicitly acknowledges this as a limitation of WA compared to ENS (Section 6): prediction ensembling can combine RMs with different architectures or different pre-training, while WA requires a shared pre-training and (ideally) the same architecture. So ENS has an advantage in terms of what diversity it can leverage. However, the paper argues that the practical benefits of WA (efficiency and corruption robustness) outweigh this limitation for most deployment scenarios, and that substantial diversity can be achieved even within the constraints of shared pre-training through hyperparameter variation and Baklava initialization.

4. Key Insights and Innovations

Innovation 1: Weight Averaging as a Mechanism for Separating Robust from Spurious Features in Reward Models

The paper's deepest conceptual contribution is not WARM itself but the theoretical insight that distinguishes weight averaging from prediction ensembling — and, by extension, that reveals why simply averaging predictions is insufficient for robust reward modeling. Before this work, the dominant understanding of weight averaging in the fine-tuning literature (model soups, DiWA, SWAD) treated it as an efficient approximation of prediction ensembling: average the weights and you get approximately the same result as averaging the predictions, but with zero inference overhead. This was the "first-order" story, and it was well-supported empirically — WA and ENS curves track each other closely under standard evaluation.

What this paper discovers is that this approximate equivalence breaks down precisely where it matters most for RLHF — when the training data contains corrupted labels. The diagnostic experiment in Section 4.2 is the key intellectual move: by deliberately injecting 25% label noise and then separately evaluating on corrupted training data, clean training data, ID validation data, and OOD test data, the paper reveals a systematic pattern that the first-order story cannot explain. WA performs worse than ENS on the corrupted training samples (it fails to memorize the flipped labels), slightly worse on clean training samples, and better on OOD test data. The accuracy advantage of WA over ENS grows monotonically as data moves away from the training distribution.

This is not an incremental empirical observation — it is a diagnostic finding that reframes what weight averaging actually does. The paper's theoretical analysis formalizes this as the difference between weighting features by p_j (ENS) versus p_j^2 (WA), where p_j is the probability that an individual model uses a given feature. Features that are consistently learned across independent training runs (high p_j, corresponding to genuine predictive patterns) are largely preserved by both methods. But features that are only occasionally learned — the run-specific patterns that often correspond to memorization of corrupted labels or brittle heuristics — are quadratically suppressed under WA compared to ENS. A feature used by 10% of models contributes at 10% strength in ENS but only 1% strength in WA.

The significance of this goes well beyond WARM's empirical performance. The paper effectively demonstrates that weight averaging implements a form of invariance regularization without any explicit invariance loss term. The connection to the domain generalization literature (Arjovsky et al., 2019; Muandet et al., 2013) is not merely analogical — the mathematical mechanism (suppressing features that vary across training runs while preserving invariant ones) is the same principle that methods like Invariant Risk Minimization encode through specialized training objectives. WARM achieves this through a simple post-hoc operation that requires no modification to the training pipeline. This unifies two previously separate paradigms for robustness — ensembling (variance reduction) and invariance (causal feature selection) — under a single mathematical framework, showing that weight averaging benefits from both simultaneously. The p_j^2 insight is the paper's most transferable theoretical contribution: it predicts that the advantage of WA over ENS should grow with the noise level in the training data and with the distance of the test distribution from training, and it provides a principled explanation for why WA has consistently matched or exceeded ENS in the OOD generalization literature.

Innovation 2: The Diagnostic Framing of Reward Hacking as a Two-Challenge Problem

The paper reframes the reward hacking problem in a way that makes the design space for solutions legible. Prior work on reward hacking — including important contributions by Gao et al. (2023) on scaling laws for overoptimization and by Eisenstein et al. (2023) on RM ensembles — treated it largely as a monolithic phenomenon: the policy exploits the RM, and defenses (ensembling, KL regularization, iterative data collection) make exploitation harder. The field lacked a decomposition of why RMs are exploitable into distinct causal mechanisms that might call for different interventions.

This paper proposes exactly such a decomposition by identifying two independent challenges: distribution shifts (the RM must score generations from distributions it was not trained on, and the mismatch grows as the policy drifts during RL) and preference inconsistencies (the binary labels in the training data are noisy, with low inter-annotator agreement and systematic reliance on superficial heuristics like length). These are not merely two contributing factors — they are qualitatively different failure modes that interact in a specific way. The paper articulates how they form a negative feedback loop: distribution shifts cause the RM to become miscalibrated on OOD generations (producing more extreme and less reliable reward values), these unreliable rewards exacerbate policy drift, the larger drift creates larger distribution shifts, and the cycle intensifies until the policy has fully exploited the RM.

What makes this framing intellectually distinctive is that it structures the evaluation of any proposed defense. A method that only addresses distribution shifts (like prediction ensembling, which reduces variance under shift) will still fail if the training labels are noisy — because the RM can learn to rely on spurious features that happen to correlate with preferences in training but don't generalize. A method that only addresses label noise would still fail under distribution shifts because the RM would encounter genuinely novel input patterns it cannot score reliably. The paper's diagnostic experiment (clean vs. 25% corrupted label settings) operationalizes this decomposition: by testing WARM in both regimes, it demonstrates that the method is uniquely positioned to address both challenges simultaneously, while ENS addresses only the first. The interaction between the two challenges — and the negative feedback loop they create — also explains why reward hacking is such a persistent problem despite individual defenses: any residual error in the RM gets amplified through the RL dynamics.

This reframing has implications beyond this paper. It suggests that future work on reward modeling should evaluate methods under both clean and corrupted label conditions, and should report performance on a spectrum from ID to OOD data (not just aggregate metrics). It also clarifies why KL regularization, the most common defense, is fundamentally limited: KL constrains distribution shift but does nothing to address label noise, so it merely delays rather than prevents exploitation. The paper's decomposition thus serves as a conceptual tool for the field, making explicit what was previously implicit in the design of reward modeling approaches.

Innovation 3: Baklava — Diversity Through SFT Trajectory Sampling as an Enabling Technique Without Additional Training Cost

The Baklava initialization strategy is a conceptually elegant solution to a practical tension in weight averaging: WA requires models to be diverse enough that averaging provides statistical benefit, yet similar enough that the weights remain linearly connected and interpolation is meaningful. Prior approaches to this tension fell into two categories. Model soups (Wortsman et al., 2022) used identical initialization and relied on stochastic training dynamics (data order, dropout) for diversity — safe for LMC but limited in the diversity it could achieve. Model ratatouille (Ramé et al., 2023) increased diversity by training models on different auxiliary tasks before fine-tuning — more diversity but at the cost of training multiple auxiliary models, a significant computational overhead.

Baklava carves a third path that is both effective and efficient: use different checkpoints collected along a single SFT trajectory as the diverse initializations for subsequent RM fine-tunings. The key insight is that SFT checkpoints at different training steps — while not identical — share the same pre-training and have not diverged so far that LMC breaks, yet they encode meaningfully different featurizer states because they sit at different points along the optimization path. The paper collects checkpoints at steps 8k, 10k, and 12k of a 12k-step SFT run — a range where the featurizer is already competent (too-early checkpoints would have poor accuracy) but still exhibits meaningful variation.

What makes this contribution distinctive is not the empirical gain it provides (though Figure 3(d) shows it yields the largest WA accuracy gains), but the conceptual move it represents: Baklava demonstrates that diversity for weight averaging can be extracted from the natural variation present in standard training trajectories, without requiring additional training, auxiliary tasks, or specialized diversity-promoting objectives. It is essentially a free lunch — the SFT run already exists (SFT is a standard step in the LLM training pipeline), and checkpointing it at multiple points costs nothing. Yet the diversity gains are sufficient to materially improve WARM's performance. This has implications beyond reward modeling: any setting where fine-tuning from a pre-trained model is followed by an averaging step (which is increasingly common, per the WA literature) could potentially benefit from Baklava-style initialization diversity at zero additional cost.

The paper also implicitly makes a methodological point through the Baklava name and the associated Figure 2, which uses a geometric diamond shape to illustrate the procedure: pre-training at the top, SFT checkpoints fanning out in the middle, RM fine-tunings converging at the bottom, and the final weight average as the lowest point. This visual metaphor captures the intuition that Baklava creates a controlled expansion of diversity (fan-out during SFT) followed by a re-convergence (fine-tunings on the same task bring the models back toward a common basin), with the final average sitting at the intersection. It's a rare case where the diagram genuinely aids conceptual understanding rather than merely illustrating a pipeline.

5. Experimental Analysis

Evaluation Methodology

Dataset. The primary dataset is the Reddit TL;DR summarization benchmark (Völske et al., 2017; Stiennon et al., 2020), consisting of posts from Reddit filtered for quality. The training set contains 123k posts with candidate summaries generated by GPT-3 variants (Brown et al., 2020), and approximately 5% of this data is held out as an in-distribution (ID) validation set. An out-of-distribution (OOD) test dataset D_ood is constructed with 92k pairwise comparisons where the summaries are generated by multiple PaLM-XS policies (pre-trained only, SFT-ed, and RLHF-ed) at high temperature to ensure diverse, potentially challenging evaluation samples.
Base model(s). All reward models are PaLM-XXS models (Anil et al., 2023), trained from a shared pre-training checkpoint and then supervised fine-tuned on the preferred summaries from D_train for 12k steps with a batch size of 128 and the Adafactor optimizer (Shazeer and Stern, 2018) with learning rate 1 × 10^{-5}. For BoN experiments, candidate summaries are generated from two SFT policies: one based on PaLM architecture (generating N = 8 summaries per prompt for D = 15,000 prompts) and one based on T5 architecture (Raffel et al., 2020) (generating N = 1,000 summaries per prompt for D = 1,000 prompts). For RL experiments, both policy and value models are PaLM-XS, initialized from the same SFT model. A separate PaLM-XS RM (achieving 80.1% accuracy on D_ood) serves as the pointwise control reward to detect reward hacking; this control RM benefits from a larger architecture and disjoint pretraining relative to the PaLM-XXS RMs under evaluation.
Metrics. The paper uses three evaluation metrics. (1) Pairwise accuracy on D_ood: the fraction of test pairs where the reward model correctly assigns a higher score to the preferred summary, defined as Acc(r_phi, D) = E_{(x, y^+, y^-) in D}[1_{r_phi(x, y^+) >= r_phi(x, y^-)}]. This serves as the primary metric for ranking individual RMs and selecting which weights to average. (2) Pointwise control reward: a separate, larger PaLM-XS RM evaluates the absolute quality of summaries selected by BoN or produced by the RL policy. Increases in control reward indicate genuine quality improvement, while collapses in control reward (despite the proxy RM reporting high rewards) indicate reward hacking. (3) Pairwise oracle preference win rate: following the RLAIF procedure from Lee et al. (2023), an instruct fine-tuned PaLM-L model is prompted with chain-of-thought (zero-shot, "Detailed + CoT" strategy, temperature T = 0.0, maximum decoding length 512 tokens) to compare summaries and produce preference judgments. To avoid position bias, the AI labeler is run in both possible orderings. This metric captures whether one selection/training strategy produces summaries that are preferred over another according to an independent evaluator that mimics human preferences (this strategy was shown to perform similarly to human labelers with similar inter-agreement in Lee et al., 2023).
Baselines. The paper compares WARM against four categories of baselines. (1) Individual RMs: phi_1 and phi_2, the two best individual RMs ranked by decreasing accuracy on D_ood. These represent the standard approach of training and deploying a single RM. (2) Prediction ensembling (ENS): averaging the scalar reward predictions from multiple RMs: r_ENS(x, y) = (1/M) * sum_{i=1}^M r_{phi_i}(x, y). The paper uses only M = 2 for ENS due to computational constraints, noting the memory and inference costs grow linearly with M. (3) Majority voting / random selection: for BoN experiments with N = 1, this corresponds to randomly selecting a summary from the SFT policy — the baseline against which all selection strategies are compared. (4) Different values of M for WARM: WARM with M = 2, M = 6, and M = 10 (for RL experiments) or M = 6 (for BoN experiments) are compared to assess how the number of averaged weights affects performance.
Generation budget / compute accounting. The paper measures inference cost in terms of the number of forward passes required by the reward model. For BoN with budget N: a single RM requires N reward evaluations, ENS with M = 2 requires 2N evaluations, and WARM requires N evaluations (the same as a single RM). For RL fine-tuning: at each training step, the policy generates a batch of summaries (batch size 128), and the RM must score all of them; ENS doubles this cost while WARM leaves it unchanged. The paper does not account for the one-time training cost of training M RMs in parallel, arguing that this is a fixed upfront cost amortized over many inference queries, and that in practice grid search over hyperparameters already produces multiple candidate RMs.
Cross-validation / statistical protocol. For the weight selection procedure, the 10 trained RMs are sorted by their individual accuracy on D_ood, and only the top M are averaged (where M is the desired WARM size). This ensures that poorly performing RMs do not degrade the average. The paper validates this procedure by comparing best-to-worst, random, and worst-to-best selection in Figure 10 (Appendix), confirming that selecting from best to worst serves as a reliable heuristic. For the label corruption experiments, 25% of the binary preference labels in the training set are randomly swapped before RM training, and results are compared between the clean and corrupt settings to assess robustness.

Main Quantitative Results

Best-of-N Experiments: WARM as a Reward Selector

Headline finding. In best-of-N sampling, where N candidate summaries are generated from an SFT policy and the one with the highest proxy reward is selected, WARM consistently outperforms both individual RMs and prediction ensembling across two SFT architectures (PaLM and T5), in both clean and 25% label corruption settings. Moreover, the performance gap between WARM and alternatives grows as M — the number of averaged weights — increases.

Control reward results (Figure 6). The paper reports control reward gains relative to phi_1, the best individual RM on D_ood. For PaLM-generated summaries in the clean setting (Figure 6a), WARM with M = 6 achieves approximately 0.10–0.12 higher control reward than phi_1 at intermediate KL values (around KL ≈ 0.4–0.8), while ENS with M = 2 shows gains of only ~0.04–0.06. The individual RM phi_2 performs worse than phi_1 despite being the second-best on D_ood, and in the T5 setting (Figure 6c) actually outperforms phi_1 in BoN — the paper notes this highlights that "selecting the appropriate individual RM is not trivial," as ID validation metrics may poorly correlate with BoN performance. In the corrupt setting (Figures 6b and 6d), the advantage of WARM over ENS becomes more pronounced: with PaLM generations and 25% label corruption (Figure 6b), WARM M = 6 maintains control reward gains of ~0.15–0.20 over phi_1 at the highest KL values, while ENS M = 2 shows gains of only ~0.05–0.10. The T5 corrupt setting (Figure 6d) shows the largest absolute gains for WARM M = 6, reaching approximately +0.5 control reward over phi_1 at the maximum N (KL ≈ 6), while ENS M = 2 reaches approximately +0.3. Figure 15 in the appendix reports the absolute control rewards rather than relative gains; the absolute values range roughly between 3 and 7, with WARM consistently achieving the highest values.

Oracle preference results (Figure 7). Using the pairwise oracle preference metric on T5-generated summaries, Figure 7a shows the win rate against the SFT baseline (random selection, or equivalently BoN with N = 1). WARM M = 6 achieves a win rate of 92.5% at the highest N (KL ≈ 6), compared to ~90% for ENS M = 2 and ~88% for the best individual RM phi_2. The performance ordering is consistent across the full range of N: WARM M = 6 > WARM M = 2 > ENS M = 2 > best individual RMs. In the corrupt setting (Figure 7b), the same ordering holds, though all methods achieve slightly lower absolute win rates (WARM M = 6 reaches ~90% at the highest N). Figures 7c and 7d flip the reference: when WARM M = 6 is used as the reference policy, all other methods have win rates below 50% (the dashed line), meaning none can beat WARM M = 6. Specifically, in the clean setting (Figure 7c), ENS M = 2 has a win rate of approximately 46–48% against WARM M = 6 at the highest N, while the best individual RM phi_2 achieves only ~42%. The gap is larger in the corrupt setting (Figure 7d), where ENS M = 2 achieves ~44% and phi_2 achieves ~36% against WARM M = 6. Figure 17 in the appendix confirms these patterns on PaLM generations and further validates with a T5 experiment where WARM M = 6 always uses N = 1,000 while other methods vary N from 1 to 1,000 — WARM with fixed large N outperforms all other methods at all their N settings, confirming the BoN oracle preference advantage.

Interpretation of BoN results. The BoN experiments demonstrate that WARM produces a reward model that better identifies genuinely high-quality summaries, as measured by an independent AI evaluator. The fact that the gap between WARM and ENS widens under label corruption (Figures 7b and 7d vs. 7a and 7c) directly supports the paper's theoretical claim that weight averaging selectively suppresses features that memorize corrupted labels while preserving generalizable features — ENS, by retaining all features including corrupted ones, makes poorer selections when the training data is noisy. The increasing benefit with larger M (comparing WARM M = 2 vs. M = 6) is consistent with the variance-reduction perspective: more ensemble members provide greater variance reduction when individual errors are uncorrelated.

RL Experiments: WARM as a Proxy Reward for Policy Optimization

Headline finding. When used as the proxy RM for REINFORCE-based policy optimization, WARM significantly mitigates reward hacking compared to both individual RMs and prediction ensembling. A policy trained with WARM (M = 6) achieves a 79.4% win rate against a policy trained with the best individual RM (phi_1), and increasing M pushes the Pareto front of solutions toward higher control rewards at higher KL divergences. Importantly, WARM also changes the optimal KL regularization strength — the best policies trained with WARM are obtained at larger KL values than with individual RMs, because WARM's robustness allows the policy to explore further before encountering reward hacking.

Control reward over training (Figures 1b, 8, 18). Figure 8a plots control reward versus KL divergence during RL training in the clean setting. All methods initially show increasing control reward as the policy improves, but eventually exhibit a collapse — the signature of reward hacking. The key difference is when and how severely this collapse occurs. With the best individual RM phi_1 (yellow, lighter shade), the control reward peaks around KL ≈ 200–300 at a value of approximately 9.5–10.0, then declines to ~7.5 by KL ≈ 500. With ENS M = 2 (red), the peak is slightly higher (~10.0) and the decline is somewhat delayed, but the collapse is still substantial (falling to ~8.0 by KL ≈ 600). With WARM M = 2 (blue, lightest shade), the peak reaches ~10.5 and the decline is more gradual, reaching ~9.0 by KL ≈ 800. With WARM M = 6 (blue, medium shade), the peak reaches ~10.5 and remains above 9.5 even at KL ≈ 1,000, with a much slower and later decline. Figure 1b plots the same data against training steps rather than KL (and Figure 18 shows the corrupt-setting version of the steps plot); the pattern is consistent — WARM delays the collapse and maintains higher absolute control rewards throughout training. Figure 8b shows the corrupt setting (25% label corruption). Here, the advantage of WARM over ENS is even starker: WARM M = 6 (blue, darkest) achieves a peak control reward of ~6.5 at KL ≈ 25–50 and maintains values above 5.5 through KL ≈ 175, while ENS M = 2 (red) peaks at ~5.8 and rapidly collapses below the individual RM phi_1 by KL ≈ 75. The individual RMs (yellow) show immediate decline after very small KL values.

Ablation of KL regularization strength (Figure 8c). The paper ablates the KL regularization coefficient alpha for policies trained with WARM M = 6 versus the individual RM phi_1, testing alpha values in {0.01, 0.003, 0.001} (where lower alpha means weaker KL penalty and thus the policy is allowed to drift further). The results reveal an important interaction: with phi_1, the optimal alpha is 0.003 — lower alpha (0.001, allowing more drift) causes a sharp collapse in control reward to ~3–4 by KL ≈ 1,500, while higher alpha (0.01, restricting drift) limits the peak control reward to ~6. With WARM M = 6, the same qualitative pattern holds but with two key differences: (1) the collapse under low alpha is much less severe, with control reward remaining above 7 even at KL ≈ 1,500 for alpha = 0.001, and (2) the optimal alpha shifts to 0.001 — WARM can safely use weaker KL regularization, allowing the policy to explore more and achieve higher peak control rewards (~10 vs. ~6 for phi_1 with alpha = 0.01). The paper interprets this as WARM's robustness enabling larger KL budgets: "from a practical perspective, this highlights that the optimal value of alpha for WARM is lower than for a single RM; this is because WARM can mitigate reward hacking, and thus the optimal policies are obtained for larger values of KL."

Oracle preference over training (Figure 9). The pairwise oracle preference metric provides the most direct evidence that WARM-trained policies produce better summaries. Figure 9a uses the SFT policy as the reference (random selection, analogous to win rate over the initialization). All RL-fine-tuned policies substantially outperform SFT, but WARM M = 6 achieves the highest win rate: 99.8% after 3,500 training steps, compared to ~97% for phi_1 and ~95% for ENS M = 2 at their respective peaks. Figure 9b uses the WARM M = 6 policy at step 3,500 (the peak-performing checkpoint) as the reference. No other policy exceeds a 50% win rate against this reference — meaning the WARM-trained policy is preferred over all alternatives. Specifically, WARM M = 10 achieves a win rate of ~45–48% against WARM M = 6 (still below 50%), WARM M = 2 achieves ~38–42%, ENS M = 2 achieves ~35–40%, and the best individual RM phi_1 achieves only ~25–30%. Figure 9c uses the policy trained with phi_1 at step 3,000 as the reference. The WARM M = 6 policy achieves a 79.4% win rate against this reference — cited in the paper's abstract as the headline RL result. Moreover, there exists a large region of training steps (roughly steps 2,500–4,500) where all WARM variants (even M = 2) exceed 50% win rate against the best individual RM's policy, indicating robust improvement across the entire training trajectory, not just at a single cherry-picked checkpoint. Figure 19a in the appendix shows an alternative view: comparing all policies at fixed training steps rather than using a single reference checkpoint. At each step, WARM M = 6 achieves higher win rates against the other methods than vice versa, with the gap widening in mid-training (steps 2,500–4,500) where reward hacking begins to affect the baseline methods.

Interpretation of RL results. The RL experiments validate the paper's central empirical claim: weight averaging produces a proxy RM that is sufficiently reliable and robust to guide policy optimization without the policy exploiting the RM's vulnerabilities. The control reward curves directly demonstrate mitigation of reward hacking — WARM not only delays the inevitable collapse but also raises the absolute performance ceiling. The oracle preference results establish that this translates to genuinely better policies: summaries from WARM-trained policies are preferred by an independent AI evaluator over summaries from policies trained with individual RMs or prediction ensembles. The finding that WARM permits lower alpha values (weaker KL regularization) is practically significant: it means practitioners can get more improvement from RL fine-tuning without as much risk of reward hacking, effectively expanding the usable KL budget.

A note on M = 10 vs. M = 6. The paper observes that increasing from M = 6 to M = 10 does not improve peak performance, though it may delay reward hacking (Figure 9b shows WARM M = 10 at ~45–48% win rate vs. WARM M = 6, still below 50%). The authors speculate this is related to the weight selection procedure: the RMs phi_7 through phi_10 have lower individual accuracy on D_ood than phi_1 through phi_6, and including them may add more noise than signal despite the variance-reduction benefit. This suggests diminishing returns — and eventually negative returns — from including too many additional RMs, particularly when later RMs are substantially worse than the best ones.

Supplementary Experiments: Baklava and Distillation Setup

Baklava in BoN (Figure 16, Appendix). The paper evaluates Baklava-based WARM (M = 2) where the two RMs have featurizer initializations from different SFT checkpoints (step 12k and step 8k). In the PaLM BoN setting (Figure 16a), Baklava WARM M = 2 achieves control reward gains of ~0.04–0.06 over the best individual RM (phi_1, initialized from step 12k) at intermediate KL values, while ENS Baklava M = 2 performs comparably to phi_1. In the T5 BoN setting (Figure 16b), the Baklava WARM advantage is larger: ~0.10–0.15 over phi_1 at higher KL values, with ENS Baklava showing gains of only ~0.05. These results confirm that Baklava provides useful diversity for weight averaging — the WA of two RMs with different SFT initializations outperforms both the individual RMs and their prediction ensemble.

Distillation setup (Figure 22, Appendix). The paper reproduces the distillation setup from Gao et al. (2023), where the control PaLM-XS RM generates the labels to train PaLM-XXS RMs (i.e., the proxy RM is being trained to imitate a larger, more capable RM). In the T5 BoN setting, WARM M = 6 achieves control reward gains of ~0.05–0.10 over phi_1 at intermediate KL, substantially outperforming ENS M = 2 (~0.02–0.05 gain) and phi_2 (which underperforms phi_1). However, the paper notes that distillation changes the diversity across fine-tuned RMs, potentially altering the significance of this setup, which motivated the use of the more realistic RLAIF labeling approach in the main experiments.

Ablation Studies and Robustness Checks

Number of averaged weights (M). Across nearly all experiments, increasing M improves performance. In BoN with PaLM clean (Figure 6a), WARM M = 6 achieves larger control reward gains than M = 2. In RL clean (Figure 8a), WARM M = 6 achieves higher peak control reward and slower decline than M = 2, and M = 10 further delays the collapse (though without improving peak; see Figures 1b and 9b). In oracle preference against the best individual RM's policy (Figure 9c), WARM M = 2, M = 6, and M = 10 all exceed 50% win rate across a wide range of steps, but M = 6 achieves the highest peak. The diminishing returns from M = 6 to M = 10 are attributed to the lower individual quality of the additional RMs rather than a fundamental limitation of averaging more weights.

Weight selection procedure (Figure 10). When 10 RMs are available, the order in which they are added to the average matters substantially. Averaging from best to worst (sorted by individual accuracy on D_ood) consistently achieves the highest accuracy for any M. Random permutations produce intermediate results, with accuracy varying by ~0.005–0.010 depending on the specific permutation. Averaging from worst to best produces the lowest accuracy, substantially below even the best individual RM for small M. This validates the paper's heuristic of pre-filtering RMs by OOD accuracy before averaging.

Corruption level (clean vs. 25% corrupt). The paper compares clean and corrupt settings throughout all experiments: BoN control reward (Figure 6, clean in 6a/6c vs. corrupt in 6b/6d), BoN oracle preference (Figure 7, clean in 7a/7c vs. corrupt in 7b/7d), and RL control reward (Figure 8, clean in 8a vs. corrupt in 8b). In every comparison, the absolute performance of all methods drops under corruption, but the relative advantage of WARM over ENS and individual RMs grows. For example, in the BoN oracle preference against WARM M = 6 (Figures 7c vs. 7d), the win rate of ENS M = 2 drops from ~48% (clean) to ~44% (corrupt), while the win rate of phi_2 drops from ~42% to ~36%. This systematic pattern supports the paper's claim that WARM provides robustness to label corruption that ENS lacks.

KL regularization strength (alpha) for RL (Figure 8c). For both WARM M = 6 and the individual RM phi_1, the paper sweeps alpha values in {0.01, 0.003, 0.001}. The primary finding is that WARM is less sensitive to alpha and can tolerate lower values (weaker regularization, more policy drift). A secondary finding is that the optimal alpha shifts: for phi_1, alpha = 0.003 is best; for WARM, alpha = 0.001 is best (or at least as good as 0.003). The paper does not sweep additional alpha values or combine alpha optimization with M selection, but the finding is robust enough to support the practical recommendation that WARM permits more aggressive RL exploration.

Sources of diversity (Figure 3). The paper ablates four sources of diversity between pairs of RMs and evaluates the resulting WA and ENS accuracy on D_ood. Different training steps along a single RM run (Figure 3a) yields minimal diversity — the WA curve is nearly identical to the diagonal (individual accuracies interpolated), with almost no accuracy gain. Independent fine-tunings with shared configuration but different data orders (Figure 3b) provide moderate diversity, with a small but visible "bump" above the diagonal. Different learning rates (Figure 3c) provide higher diversity with a more pronounced accuracy gain. Baklava initialization (Figure 3d) provides the largest diversity and the highest accuracy gains from averaging, with WA achieving approximately 0.764 accuracy at lambda = 0.5 versus approximately 0.762 for the better individual model. In all four cases, WA and ENS curves closely track each other, validating Observation 2 (first-order equivalence) and Observation 1 (LMC, since WA is always at or above the diagonal).

Baklava vs. standard initialization (Figure 16 vs. main experiments). Comparing Figure 16 (Baklava M = 2) with Figure 6 (standard WARM M = 2 with identical SFT initialization), the Baklava version achieves comparable or slightly better control reward gains in BoN. This confirms that Baklava is a viable alternative to standard initialization that provides additional diversity without breaking LMC. The paper does not report Baklava results for RL experiments or for M > 2.

Moving average vs. independent fine-tunings (Remark 2, Section 3.3). The paper reports a negative result: averaging checkpoints collected along a single RM fine-tuning trajectory (stochastic weight averaging / moving average) yields inferior performance compared to averaging weights from independent fine-tunings. The diagnosis is that early checkpoints (providing diversity) have poor accuracy, while later checkpoints (with good accuracy) lack diversity — the accuracy-diversity trade-off is not favorable for single-trajectory averaging. This motivates the use of multiple independent fine-tuning runs despite the higher training cost.

Distillation vs. RLAIF labeling (Figure 22 vs. main experiments). The distillation setup (training proxy RMs to imitate a larger control RM) produces qualitatively different diversity patterns across fine-tuned RMs compared to the RLAIF setup (training on AI-labeled preferences). The paper notes this explicitly as motivation for focusing on the RLAIF setup, but does not provide a detailed analysis of how or why the diversity patterns differ. In the distillation BoN experiment (Figure 22), WARM M = 6 still outperforms ENS M = 2 and individual RMs, suggesting that WARM's benefits are not specific to the RLAIF labeling procedure.

Oracle preference metric validation (Section 5 and Appendix B.2). The AI labeling procedure for the oracle preference metric is validated by reference to Lee et al. (2023), who showed that the "Detailed + CoT 0-shot" prompting strategy with PaLM-L performs similarly to human labelers with similar inter-agreement. The paper uses greedy decoding (T = 0.0) for consistency and runs the labeler in both possible orderings to eliminate position bias. This metric is used as the ground-truth evaluator for all pairwise preference claims (win rates in Figures 7 and 9).

Model architecture and scale consistency. All RMs are PaLM-XXS; the control RM is PaLM-XS; the oracle preference labeler is PaLM-L. This represents a 3-tier scale hierarchy: the models being evaluated and combined (XXS), the model used to detect reward hacking (XS), and the model used to provide ground-truth preferences (L). The paper does not experiment with varying the RM scale or the gap between RM and policy scale.

Critical Assessment

The experiments provide substantial evidence for the paper's main claims, but several important gaps and limitations must be noted.

Claim: WARM mitigates reward hacking. (Strongly supported, with boundary conditions.) The RL experiments in Figures 8 and 9 provide clear, multi-metric evidence: WARM-trained policies achieve higher control rewards, maintain them longer (delayed collapse), and are strongly preferred by an independent oracle evaluator (79.4% win rate vs. best individual RM's policy). The BoN experiments in Figures 6 and 7 show analogous selection advantages. However, "mitigates" is the operative word — WARM does not eliminate reward hacking. Even with WARM M = 6, the control reward eventually declines in Figure 8a (the collapse is delayed but not prevented). The paper claims WARM "mitigates reward hacking" and the experiments support precisely this: the severity is reduced and the onset is delayed, but the fundamental vulnerability to proxy RM exploitation remains, just at a higher threshold.

Claim: WARM is more efficient than prediction ensembling. (Trivially supported, but the comparison is asymmetric.) WARM requires one forward pass per reward evaluation; ENS requires M forward passes. This is an architectural fact, not an experimental finding. The experiments consistently show WARM ≈ ENS (or WARM > ENS) in accuracy while using M times fewer evaluations. However, the paper only compares against ENS with M = 2 (stated reason: computational constraints). This is a significant limitation: can ENS with M = 6 or M = 10 match or exceed WARM with M = 6? If ENS with enough members could achieve comparable performance, the "efficiency at same performance" claim holds, but the "better than ENS at any budget" claim would not. The paper acknowledges ENS scales poorly with M, but doesn't empirically demonstrate where the crossover point lies (if any) where ENS with many members catches up to WARM. In fairness, WARM with M = 2 already outperforms ENS M = 2 in most corrupt-setting experiments (e.g., Figure 8b), which cannot be explained by an asymmetric M comparison and supports the claim that WA provides benefits beyond mere variance reduction.

Claim: WARM improves robustness to label corruption. (Supported, with a diagnostic limitation.) The 25% corruption experiments are the paper's strongest empirical contribution, demonstrating a systematic pattern where the WA > ENS gap grows as data moves from ID to OOD (Figures 4, 5). This is replicated across BoN and RL settings, and the pattern is consistent. The limitation is that only a single corruption level (25%) and a single corruption type (random label flipping) are tested. Would the advantage persist at 10% corruption? 50%? What about structured noise (e.g., systematically preferring longer summaries) rather than random noise? The theoretical analysis predicts the WA advantage should scale with the prevalence of low-probability (spurious) features, but this prediction is not empirically tested across different noise levels or types.

Claim: The p_j^2 mechanism explains WA's advantage over ENS. (Theoretically motivated but not directly tested.) The theoretical analysis in Section 4.3 provides an elegant explanation for the empirical pattern in Figure 4, but the paper does not perform experiments that directly test whether the p_j^2 mechanism — as opposed to some other difference between WA and ENS — drives the observed effects. For instance, one could measure the probability p_j that individual RMs rely on specific features (e.g., length heuristics) and verify that WA suppresses low-p_j features relative to ENS. Without such direct tests, the p_j^2 explanation remains a plausible and well-motivated hypothesis rather than a verified mechanism.

Claim: WARM connects ensembling to invariance. (Conceptually supported, but invariance is not directly measured.) The paper argues that WA "only keeps the invariant predictive mechanisms across runs" and connects this to the invariance literature. However, "invariance" is operationalized only indirectly — through the suppression of low-p_j features — and not through standard invariance diagnostics (e.g., evaluating whether the model's predictions are stable across different environments or perturbations). The claim is conceptually appealing but would be strengthened by experiments showing that WARM-trained policies generalize better across different types of distribution shift (not just the specific OOD test set used here) or that WARM RMs are less sensitive to spurious correlations.

Missing baselines and ablations. Several experiments would have strengthened the paper: (1) Comparison against ENS with M > 2, at least for a subset of experiments, to test whether the WA > ENS finding holds at matched ensemble sizes. (2) Varying the corruption rate (10%, 50%) to test the predicted relationship between noise level and the WA-ENS gap. (3) Testing whether the benefits of Baklava initialization compound with hyperparameter diversity, or whether they are redundant — the paper presents them as separate sources of diversity but never ablates their interaction. (4) Evaluating on human-labeled preference data rather than exclusively AI-labeled data; the paper argues WARM should extend to human preferences, but human label noise may have different structure (e.g., systematic biases rather than random errors) that the p_j framework may not handle identically. (5) Testing with different RM and policy architectures; all experiments use PaLM variants, and the paper acknowledges that WA requires shared pre-training, limiting the diversity that can be leveraged compared to ENS which can combine RMs from different model families.

Limited scale of experiments. The RM is PaLM-XXS, the policy is PaLM-XS, and the dataset is TL;DR summarization (123k training examples). This is a substantial experimental setup, but it is a single domain (summarization) with a single preference type (AI-labeled). The paper does not test on other RLHF domains (dialogue, instruction following, harmlessness), other dataset scales, or other model scales. The scaling trend with model size — whether WARM's benefits are larger, smaller, or unchanged for larger RMs — is unexplored.

Weight selection circularity concern. The weight selection procedure sorts RMs by their individual accuracy on D_ood — the OOD test set. This introduces a potential circularity: the test set is used to select which weights to average, and then the averaged model is evaluated on (subsets of) the same test set. The paper's cross-validation consists only of validating the best-to-worst selection heuristic (Figure 10) by comparing against random permutations, but does not hold out a separate RM selection set distinct from the evaluation set. This means the reported results may be optimistic relative to a setting where OOD performance is truly unknown at RM selection time. The paper notes this implicitly by stating that the weight selection is based on OOD accuracy, but doesn't discuss the implications for deployment where OOD labels are unavailable.

Robustness of the "WARM allows lower alpha" finding. Figure 8c shows that alpha = 0.001 works well for WARM M = 6 but poorly for phi_1. This is a practically important finding, but it is based on only three alpha values and a single WARM configuration (M = 6). Does the optimal alpha decrease monotonically with M? Would alpha = 0.0005 be even better for WARM M = 10? The paper does not explore this interaction systematically, so the practical recommendation ("WARM permits lower alpha") is directional rather than quantitative.

Interpretability of the control reward metric. The control RM is a larger PaLM-XS model achieving 80.1% accuracy on D_ood. While it is "disjoint" from the PaLM-XXS RMs being evaluated, 80.1% accuracy means it disagrees with the oracle preference on ~20% of pairs. The paper uses control reward as the primary metric for detecting reward hacking (Figures 6 and 8), but doesn't analyze how the 20% error rate in the control RM affects the reliability of the reward hacking diagnosis. If the control RM has systematic biases (e.g., also preferring longer summaries), its judgments may not perfectly reflect true quality degradation. The oracle preference metric (Figures 7 and 9) addresses this limitation by providing an independent evaluation using a different model (PaLM-L) and a different procedure (pairwise preference rather than pointwise scoring), and the consistent results across both metrics mitigate this concern.

Despite these limitations, the experiments collectively make a convincing case for WARM's practical value: across two downstream tasks (BoN selection and RL training), two policy architectures (PaLM and T5), two data conditions (clean and corrupted), and multiple metrics (control reward, oracle preference), WARM consistently matches or exceeds ENS and the best individual RM, with the advantage growing under label corruption — precisely where the theory predicts it should. The experimental design's strength lies in this multi-faceted triangulation rather than in any single ablation or comparison.

6. Limitations and Trade-offs

The Weight Selection Procedure Depends on OOD Performance Labels That Are Unavailable at Deployment

The assumption or constraint. WARM's recipe requires selecting which of the trained RMs to include in the average. The paper's procedure (Section 3.1, Appendix B.3) first evaluates all M individual RMs on the OOD test set D_ood, sorts them by decreasing accuracy, and averages only the top K. The paper explicitly documents this: "when ranked in decreasing accuracy on D_ood, we denote them {phi_i}_{i=1}^{10} … WARM with M weights is the WA of the M best weights {phi_i}_{i=1}^M" (Appendix B.3). The validation of this procedure (Figure 10) tests different permutations but does not escape the core circularity: the same OOD distribution used for evaluation is used for RM selection.

The consequence. In a real deployment, D_ood does not exist — if you had labeled OOD preference data, you would train on it directly rather than using it for model selection. A practitioner deploying WARM must choose which RMs to average using only in-distribution validation data (or some other proxy). The paper's own results demonstrate why this is problematic: Figure 6(c) shows that the individual RM phi_2 — ranked second-best on D_ood — actually outperforms phi_1 in BoN. The paper notes this explicitly: "the individual RM phi_2 performs better in BoN in Figure 6(c) than phi_1 though phi_1 was better than phi_2 on D_ood, highlighting that selecting the appropriate individual RM is not trivial." If ID validation metrics can mis-rank RMs relative to their downstream BoN performance, then a WARM constructed using only ID validation may underperform the oracle-selected WARM reported in the paper. The headline 79.4% win rate depends on an RM selection procedure that is not replicable in practice.

What evidence exists in the paper. The BoN inversion (Figure 6c, where phi_2 beats phi_1 despite lower OOD accuracy) is the only direct evidence of this selection difficulty. Figure 10 validates the best-to-worst ordering but does not test whether ID validation accuracy (which would be available) can substitute for OOD accuracy in the selection procedure. The paper does not report how WARM performs when RMs are selected using ID validation metrics.

Mitigation status. Not addressed. The paper acknowledges the selection difficulty in passing ("selecting the appropriate individual RM is not trivial") but does not propose or evaluate a deployment-feasible selection procedure. Section 8 (Discussion) does not list this as a limitation or suggest future work on model selection. A practitioner reading the paper would not know how to operationalize the weight selection step without access to OOD labels.

All Experiments Use a Single Benchmark and a Single Model Family, Leaving Domain and Architecture Generalization Unverified

The assumption or constraint. Every experiment in the paper — BoN selection, RL fine-tuning, corruption analysis, Baklava diversity — uses the Reddit TL;DR summarization benchmark with PaLM architectures (PaLM-XXS for RMs, PaLM-XS for policies, PaLM-L for the oracle labeler). The paper does not evaluate on any other domain (dialogue, instruction following, harmlessness, coding) or any other model family. The only architectural variation is an auxiliary BoN experiment using T5 policies for summary generation (Figures 6c, 6d, 7), but even there the RMs themselves remain PaLM-XXS.

The consequence. Three aspects of the findings could be benchmark-specific. First, the importance of label noise: summarization preferences may have different noise characteristics than other domains. In dialogue, for instance, preferences are often more subjective and inter-annotator agreement may be even lower than the 72.6% cited for InstructGPT, potentially changing the relative advantage of WA over ENS (since the p_j^2 suppression mechanism targets features correlated with label noise). In code generation, preferences are often tied to objective correctness (does the code pass unit tests?), which may involve less label noise and reduce the corruption-robustness benefit. Second, the nature of distribution shifts: the OOD test set D_ood is constructed from PaLM-XS policies at high temperature — the shifts are within the same model family. Distribution shifts in deployment may be larger (different architectures, different data distributions, different tasks) and could stress linear mode connectivity differently. Third, the linear mode connectivity property itself: the paper's entire method depends on LMC holding for fine-tuned RMs, which is validated only for PaLM-XXS. Different architectures, pre-training procedures, or fine-tuning dynamics could produce different LMC characteristics.

What evidence exists in the paper. The only cross-architecture evidence is the T5 BoN experiments (Figures 6c, 6d, 7), which show WARM maintaining its advantage when selecting among T5-generated summaries. However, this tests only the selection stage — the RMs themselves are still PaLM-XXS, and the T5 experiments do not test RL fine-tuning where the interaction between RM and policy dynamics is more complex. No evidence exists for non-summarization domains.

Mitigation status. The paper does not claim broader generalization and does not discuss this as a limitation. The related work (Appendix A) describes LMC as a general property of fine-tuned models sharing pre-training, which would imply the approach should transfer, but this is an extrapolation rather than an empirical finding. A single note in Section 4.1 mentions that the experimental setup "is a standard in reward modeling for LLMs" — true but does not establish domain robustness.

Weight Averaging Cannot Exploit Diversity Across Architectures or Pre-Trainings, Unlike Prediction Ensembling

The assumption or constraint. WARM requires all averaged RMs to share the same pre-training and — practically — the same architecture, because weight averaging only makes sense when corresponding parameters represent compatible features. The paper explicitly acknowledges this as a fundamental limitation in Section 6: "prediction ensembling can benefit from the diversity brought by combining RMs from various architectures and pre-trainings." In contrast, the diversity sources available to WARM — different data orders, different hyperparameters, Baklava initialization (different SFT checkpoints) — all operate within a single pre-training and architecture.

The consequence. There is a category of diversity that ENS can exploit but WARM cannot: RMs trained on different model families, with different pre-training data, or with different architectures may learn qualitatively different kinds of features. An ensemble combining a PaLM RM, a T5 RM, and a GPT RM could leverage their complementary inductive biases. WARM is structurally incapable of this — you can only average weights from models with identical architecture and shared pre-training. If cross-architecture diversity is substantial, ENS with heterogeneous members could potentially outperform a homogeneous WARM, even accounting for WA's corruption-robustness advantage. The paper provides no evidence on how the ENS-vs-WARM comparison changes when ENS is allowed heterogeneous members while WARM is restricted to homogeneous ones — in the experiments, both methods use the same pool of PaLM-XXS RMs, so this limitation is hidden.

What evidence exists in the paper. None. The paper does not compare WARM against heterogeneous ENS or discuss the magnitude of cross-architecture diversity benefits. The only ENS experiments use PaLM-XXS RMs from the same pre-training — identical to the pool available for WARM.

Mitigation status. The paper acknowledges this limitation (Section 6) but does not explore it empirically or propose workarounds. The acknowledgment is brief and does not discuss whether the corruption-robustness benefits of WA might compensate for the loss of cross-architecture diversity in practice, or whether hybrid approaches (WA within architecture families, then ENS across families) could capture both benefits.

The Training Cost of Producing M Diverse RMs Is Not Accounted for in the Headline Efficiency Claims

The assumption or constraint. WARM's headline efficiency advantage over ENS is that inference requires only one forward pass rather than M. The paper states this as: "by requiring a single model at inference time, it provides a scalable approximation to the traditional, costlier ensembling of predictions, without its memory and inference burdens" (Section 1). This comparison treats the cost of training the M individual RMs as a sunk cost, amortized over many inference queries. The paper argues that "in practice grid search over hyperparameters already produces multiple candidate RMs" (implicitly, the training cost would be paid anyway).

The consequence. There are two issues with this framing. First, the number of RMs needed for strong WARM performance is non-trivial: the best RL results use M = 6 or M = 10 (Figures 8, 9). Training 6–10 RMs independently represents a 6–10× increase in one-time training compute compared to training a single RM. While inference cost dominates in high-volume deployments, for teams that train RMs infrequently or deploy to low-volume applications, this training overhead may be significant. Second, ENS can also be used with fewer models: the paper only evaluates ENS with M = 2 (stated reason: computational constraints for inference). If 2 RMs are sufficient for ENS to achieve comparable performance to WARM with 6 RMs, then WARM's training cost disadvantage (6 vs. 2) may offset its inference advantage, depending on the ratio of training to inference compute. The paper does not provide the data needed to evaluate this tradeoff.

What evidence exists in the paper. The paper does not report total training FLOPs or wall-clock training time for any experiment. The computational cost comparison is exclusively at inference time. The M = 2 vs. M = 6 WARM comparison in RL (Figure 8a) shows that M = 6 substantially outperforms M = 2, suggesting that the training cost scales with M in a way that matters for performance. ENS M = 6 is never evaluated, so we cannot determine whether 6 ENS RMs would match 6 WARM RMs at higher inference cost.

Mitigation status. The paper acknowledges this implicitly by noting that "WARM … incurs no memory or inference overhead" — carefully scoping the efficiency claim to inference only. Section 6 does not flag training cost as a limitation. The paper argues that grid search amortizes the cost, but this argument weakens as M grows: a typical grid search might train 10–20 models, but WARM uses only the top 6–10, and the remaining models are discarded — their training cost contributes nothing to the final model.

WARM Mitigates but Does Not Eliminate Reward Hacking, and the Residual Vulnerability Is Not Characterized

The assumption or constraint. The paper's central empirical claim is that WARM "mitigates reward hacking" (Section 5, Figure 1b caption). The experiments show that WARM delays the collapse in control reward (Figure 8a: WARM M = 6 maintains high control reward to KL ≈ 1,000 vs. KL ≈ 300 for individual RMs) and produces policies strongly preferred by the oracle (Figure 9b: 79.4% win rate vs. best individual RM's policy). However, reward hacking still occurs — the control reward for WARM M = 6 in Figure 8a eventually declines from its peak of ~10.5 to ~9.0 by KL ≈ 1,000. The paper does not characterize at what point this residual hacking begins, how severe it is compared to the individual RM baseline, or whether it matters for downstream policy quality.

The consequence. A practitioner deploying WARM still faces the checkpoint selection problem that the paper identifies as a core challenge of RLHF (Section 1): "it complicates checkpoint selection due to the unreliability of the proxy RM." Even with WARM, the proxy reward continues to rise while the control reward declines — the proxy and the true objective eventually diverge. The practitioner must still decide when to stop training, and WARM provides no mechanism for detecting this divergence. The paper's oracle preference evaluations (Figure 9) identify the peak-performing checkpoint after the fact, but this requires an oracle that is unavailable in deployment. Without such an oracle, the practitioner risks training past the point of genuine improvement and into the hacking regime — WARM expands the safe training window but does not make it infinite or self-monitoring.

What evidence exists in the paper. Figure 8a shows the control reward decline for all methods, including WARM M = 6. The decline is slower and shallower for WARM, but it is clearly present. Figure 1b (and the same data plotted against steps in Figure 18) shows the decline in absolute terms. Figure 8b (corrupt setting) shows the same pattern in a more compressed KL range. The paper does not analyze the characteristics of the residual hacking — whether it produces the same kinds of degenerate outputs as with individual RMs, or novel failure modes specific to WARM.

Mitigation status. The paper does not address this residual hacking or propose methods for detecting it. Section 6 discusses limitations but focuses on WARM's inability to combine heterogeneous architectures, not on the fundamental ceiling that remains even within its scope. The finding that WARM permits lower KL regularization coefficients (Figure 8c, alpha = 0.001 works for WARM but not for phi_1) suggests a partial practical mitigation: tune alpha to stop training before the hacking regime. But this tuning is still manual and requires monitoring the control reward — which, again, is not available in deployment.

Difficulty Estimation and Allocation Are Not Addressed — WARM Applies Uniformly Regardless of Prompt or Generation Difficulty

The assumption or constraint. WARM produces a single, static reward model that evaluates all prompt-summary pairs using the same weight-averaged parameters. The method does not condition on problem difficulty, distribution shift severity, or any measure of how reliable the RM's judgment is for a particular input. Every input receives the same reward model, regardless of whether it lies close to the training distribution (where the RM is well-calibrated) or far from it (where even the averaged RM may be unreliable). This contrasts with adaptive test-time compute strategies (e.g., the "WARM: On the Benefits of Weight Averaged Reward Models" paper analyzed above does not address this, but the broader RLHF literature increasingly recognizes difficulty-dependent allocation as important).

The consequence. WARM's benefits — variance reduction and invariance regularization — are average-case improvements. They reduce expected error across the distribution of inputs the policy generates, but they may still produce systematically unreliable rewards for certain categories of inputs (e.g., summaries on very long posts, summaries with unusual structure, or summaries near the edge of the policy's exploration frontier). The RL policy, which is actively optimizing to find high-reward regions, may discover and exploit these localized failure regions even if the average RM accuracy is improved. WARM provides no mechanism for the RM to express uncertainty or for the policy to be more conservative in regions where the RM is less reliable. The paper's theoretical analysis (Section 4.3) shows that WA suppresses features with low probability p_j of being learned — but this suppression is global, not input-dependent. A feature that is robust for most inputs but spurious for a specific subpopulation would be preserved (because its overall p_j is high) and could still be exploited.

What evidence exists in the paper. The paper evaluates WARM on aggregate metrics (control reward averaged over test prompts, oracle win rate aggregated across all comparisons). It does not perform difficulty-stratified analysis — for instance, showing whether WARM's advantage over ENS is larger or smaller for prompts where the individual RMs disagree more (which might indicate higher uncertainty), or for summaries that are far from the SFT distribution in embedding space. The OOD test set D_ood is treated as a monolithic entity rather than stratified by distance from the training distribution. The theoretical analysis in Section 4.3 is aggregate (as M goes to infinity) and does not characterize input-dependent behavior.

Mitigation status. Not addressed. The paper does not discuss difficulty estimation, uncertainty quantification, or adaptive reward modeling. Section 6 mentions that "prediction ensembling can incorporate prediction disagreement into the reward to provide uncertainty estimation" and notes this as an advantage of ENS over WA, but dismisses it by citing Eisenstein et al. (2023) who found that "simple averaging of logits often performs comparably to more complex prediction aggregation functions that include uncertainty elements." Whether uncertainty estimation would help mitigate WARM's uniform application limitation is not explored.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper makes a methodological pivot in how the field should think about combining reward models for RLHF. The dominant prior approach — prediction ensembling — treats model combination as a purely statistical operation: average the outputs, reduce the variance, get a better reward signal. WARM demonstrates that how you combine models is as important as that you combine them, and that weight-space averaging provides qualitatively different regularization than output-space averaging. This is not merely an efficiency trick (though the zero-inference-overhead property is practically valuable). It is a finding that weight averaging selectively suppresses run-specific features while preserving run-invariant ones, making it a form of implicit invariance regularization that requires no specialized loss function, no auxiliary training objectives, and no architectural modifications.

The most significant conceptual shift is the reframing of what makes reward models exploitable. Prior work on reward hacking — including the influential scaling laws analysis by Gao et al. (2023) and the RM ensemble studies by Eisenstein et al. (2023) and Coste et al. (2023) — focused primarily on distribution shift as the vulnerability: the RM sees unfamiliar inputs during RL and makes errors that the policy can exploit. Defenses accordingly focused on reducing shift (KL regularization) or reducing error variance (ensembling). The corruption experiments in this paper (Section 4.2) reveal a second, independent vulnerability: label noise in the preference data itself creates spurious features that survive prediction ensembling but are suppressed by weight averaging. This is not a subtle refinement — it is a distinct failure mode that requires a distinct defense. The diagnostic pattern in Figure 4, where WA's accuracy advantage over ENS grows monotonically as data moves from corrupted training to OOD test, would not be predicted by a purely variance-reduction account of why combining models helps.

This insight has an important reconciliatory function. The model soups and DiWA literature had repeatedly observed that weight averaging matches or slightly exceeds prediction ensembling under distribution shifts, but the explanation remained at the level of "WA is an efficient approximation of ENS" — the first-order Taylor expansion story. This paper shows that the first-order story is incomplete: WA and ENS behave systematically differently when the training data contains noise. The p_j^2 analysis (Section 4.3) provides the missing mechanism. A feature that is learned by only 10% of individual models contributes at 10% strength in ENS but only 1% strength in WA; features that are consistently learned across all runs are preserved by both. This connects two previously separate research paradigms — ensembling for variance reduction and invariance for causal feature selection — under a single mathematical framework, showing that weight averaging benefits from both simultaneously without requiring either to be explicitly engineered.

The paper also reshapes the practical economics of reward model deployment. The finding that ENS M = 2 and WARM M = 2 perform comparably under clean labels (Figure 6a, Figure 8a) but that WARM M = 2 substantially outperforms ENS M = 2 under corruption (Figure 6b, Figure 8b) means that the choice between WA and ENS is not just about inference cost — it is about robustness to the noise that is universally present in human and AI preference labels. Given that inter-annotator agreement in real preference datasets hovers around 72–78% (the paper cites 72.6% for InstructGPT), the corruption setting is arguably more realistic than the clean setting for most deployed RLHF systems. If this finding transfers to human-labeled data (an important open question, since the paper uses only AI labels), then WA should be the default combination method even when inference cost is not the primary constraint — not because it is cheaper, but because it is more robust to the label noise that every practical preference dataset contains.

The Baklava initialization strategy (Figure 2) introduces a zero-cost method for increasing diversity across fine-tunings, which matters because it lowers the barrier to achieving the diversity that makes WA effective. Previous approaches required either identical initialization (model soups, limited diversity) or training on multiple auxiliary tasks (model ratatouille, expensive). Baklava shows that meaningful diversity can be extracted from the natural variation present in a single SFT trajectory — a resource that already exists in every standard LLM training pipeline. This is a methodological contribution with implications beyond reward modeling: any weight averaging application that follows an SFT step (instruction tuning, task-specific fine-tuning, domain adaptation) can potentially benefit from Baklava-style checkpoint sampling at zero additional training cost.

Finally, the paper makes an indirect but important methodological point about how to evaluate defenses against reward hacking. The standard approach — report control reward or oracle win rate on clean data — misses the interaction between label noise and model combination strategy that the corruption experiments reveal. The paper's practice of evaluating under both clean and corrupted label conditions, and separately reporting performance on corrupted training data, clean training data, ID validation data, and OOD test data (Figures 4, 5), provides a diagnostic template that future work on robust reward modeling would benefit from adopting. A defense that looks strong on clean data but degrades sharply under realistic label noise is not a robust defense; the corruption experiment exposes this in a way that single-condition evaluation cannot.

Follow-Up Research This Work Enables

Measuring the p_j probabilities directly to validate the invariance mechanism. The paper's theoretical analysis (Section 4.3) predicts that WA suppresses features in proportion to p_j^2 rather than p_j, where p_j is the probability an individual RM relies on a given feature. This is the paper's most important theoretical claim, but it is not directly tested — the evidence for it is the pattern of accuracy differences in Figure 4, which is consistent with the theory but could have alternative explanations. A direct test would measure p_j for specific features in the TL;DR setting. For example, summary length is a known spurious feature in RLHF — RMs often learn to prefer longer summaries even when length does not correlate with quality. One could measure, across M independently trained PaLM-XXS RMs, the correlation between each RM's predicted reward and summary length on a held-out set where length is decorrelated from quality. This gives an empirical p_j for the length feature. One would then predict that the ENS-combined reward retains this correlation at approximately p_j strength, while the WA-combined reward retains it at approximately p_j^2 strength. The same analysis could be repeated for other candidate features: number of bullet points, presence of specific phrases, factual consistency with the source post. Confirming the p_j^2 relationship would validate the paper's core mechanistic claim and provide a quantitative framework for predicting when WA will outperform ENS (when spurious features have moderate p_j — high enough to affect ENS but low enough to be suppressed by WA) versus when they will perform similarly (when all features have p_j near 0 or near 1).

Testing WARM with human-labeled preferences across multiple domains. The paper's experiments use exclusively AI-generated preference labels via the RLAIF procedure (PaLM-L prompted with chain-of-thought). While the paper cites evidence that AI labels correlate with human preferences (Lee et al., 2023), the noise structure in AI labels may differ from human label noise. Human labelers exhibit systematic biases — preferences for length, politeness, or specific formatting — that are correlated across examples, whereas the paper's corruption experiment uses random independent label flipping. Random noise produces features with low p_j (each corrupted label is a unique pattern unlikely to be consistently learned across runs), which WA's p_j^2 mechanism is well-suited to suppress. Systematic biases — e.g., human labelers consistently preferring longer summaries across many examples — would produce features with high p_j, which WA would preserve alongside genuine quality features. A critical follow-up would train WARM on a human-labeled preference dataset (the TL;DR human preference data from Stiennon et al., 2020, or the InstructGPT comparisons from Ouyang et al., 2022) and evaluate whether the WA > ENS gap observed under AI label corruption persists under human label noise. Additional domains — dialogue (e.g., Anthropic's helpfulness/harmlessness data), instruction following (e.g., AlpacaFarm), or code generation (where correctness provides an objective ground truth) — would test whether the benefit depends on the type and structure of label noise. If WARM's advantage disappears under systematic human biases, this would refine the theory: WA helps against unsystematic noise (random errors, low inter-annotator agreement) but not against systematic biases (shared preferences for spurious features). If the advantage persists, it strengthens the case for WA as a general-purpose robustness intervention.

Scaling ENS to match WARM's effective ensemble size. The paper's ENS baseline uses only M = 2 (stated reason: computational constraints for inference). This leaves open the question: can ENS with enough members match or exceed WARM? If ENS M = 6 or M = 10 achieves comparable control reward trajectories to WARM M = 6 in the RL experiments, then WARM's advantage is primarily one of efficiency (same performance at 6× lower inference cost) rather than a fundamental qualitative difference. If ENS with arbitrarily many members cannot match WARM in the corruption setting — because ENS's p_j weighting fundamentally preserves spurious features that WA's p_j^2 weighting suppresses — then the advantage is structural and cannot be overcome by scaling ENS. A clean experiment would train 10 PaLM-XXS RMs using the paper's procedure, construct both ENS and WA versions at M = 2, M = 4, M = 6, M = 8, M = 10 (using the same top-K RMs for both), and evaluate both on the BoN and RL benchmarks in the corrupt setting. If the ENS curve plateaus below the WA curve regardless of M, this confirms the structural advantage. If the ENS curve eventually catches up at some M, this quantifies the "ensemble size premium" that ENS pays to match WA's corruption robustness. A secondary analysis would compare the inference-cost-adjusted performance: at what inference FLOPs budget does ENS M match WARM M', for different values of M'?

Combining WARM with explicitly invariance-regularized RM training. The paper argues that WA provides implicit invariance regularization by selecting mechanisms that are invariant across training runs. But this is a post-hoc operation applied after standard RM training — the individual RMs are trained with the standard Bradley-Terry loss (Equation 1), which has no explicit invariance penalty. A natural extension is to couple WA with training objectives that explicitly encourage the individual RMs to learn invariant features. For example, one could add an Invariant Risk Minimization (IRM) penalty to the RM training loss that penalizes feature representations whose optimal linear classifier varies across different subsets of the training data. Or one could apply CORAL or MMD-based domain alignment to ensure that the RM's internal representations for preferred vs. dispreferred summaries are stable across data partitions. The hypothesis would be that explicit invariance training increases the p_j for genuinely causal features and decreases it for spurious ones, making WA's subsequent p_j^2 suppression more effective (since the gap between high-p_j and low-p_j features widens). Alternatively, explicit invariance training might make individual RMs so similar that WA provides no additional benefit (diversity is reduced), in which case the two approaches are substitutes rather than complements. A well-designed experiment would train RMs with varying invariance regularization strength, apply WA with varying M, and measure the interaction: does invariance regularization reduce the marginal benefit of WA (consistent with the diversity-reduction hypothesis), or does it increase WA's benefit by widening the p_j gap (consistent with complementarity)?

Stress-testing WARM under structured label corruption and adversarial preference attacks. The paper tests only one type of label corruption: random independent flipping of 25% of preference pairs. Real-world preference noise is rarely random and independent. A more realistic stress test would introduce structured corruption: systematically reverse preferences for summaries that are longer than a threshold, or for summaries containing certain keywords, or for summaries generated by policies at a specific KL divergence from SFT. This would create features with intermediate p_j — consistently present in some subset of the data — that sit in the regime where the WA-vs-ENS distinction matters most. An even more adversarial test would be a preference poisoning attack: an adversary with knowledge of the RM training procedure (but not the specific random seeds or data orders) deliberately crafts a subset of preference labels designed to create a spurious feature that individual RMs will learn with moderate probability. The question is whether WARM's p_j^2 suppression is sufficient to neutralize such an attack when the adversary can tune the poisoning to maximize the chance that the spurious feature falls in the "moderate p_j" regime. This connects WARM to the broader literature on robustness to data poisoning and would establish whether WA provides a meaningful security property beyond its statistical benefits.

WARM for direct preference optimization (DPO) and reward-free alignment. The paper briefly notes in Section 6 that "a promising research direction is extending WARM to direct preference optimization (DPO) strategies, where averaging the RMs casts back to averaging the DPO policies." DPO (Rafailov et al., 2023) eliminates the explicit RM by reparameterizing the preference probability in terms of the policy itself, making the policy its own reward model. The connection is that averaging DPO policies in weight space is theoretically equivalent to averaging the implicit RMs that those policies encode. Recent work (e.g., Labonne's NeuralBeagle14, cited in the paper) has already explored averaging DPO-trained policies with promising results. A rigorous follow-up would train M DPO policies on the TL;DR dataset with the same diversity sources used for WARM (different hyperparameters, data orders, Baklava initializations), average their weights, and evaluate the resulting policy against ENS of DPO policies, individual DPO policies, and a WARM-guided RL policy. The key question is whether DPO weight averaging inherits the same p_j^2 robustness property — since DPO policies are not explicitly decomposed into featurizer + linear head in the same way as RMs, the mechanism might differ. This would extend the paper's framework to the rapidly growing DPO literature and test whether the invariance-regularization interpretation of WA is specific to the Bradley-Terry RM architecture or is a more general property of weight averaging applied to preference-learned models.

Practical Applications and Downstream Use Cases

Robust reward modeling for production RLHF pipelines with noisy preference data. The most direct application is in any RLHF training pipeline where preference labels are known or suspected to be noisy. The paper's corruption experiments show that with 25% label noise, WARM M = 6 maintains control rewards of ~6.5 in RL while ENS M = 2 collapses to ~5.8 and individual RMs collapse below 5 (Figure 8b). Given that real human preference datasets typically exhibit 22–28% disagreement (the 72.6% agreement cited for InstructGPT corresponds to ~27% noise, directly comparable to the paper's 25% corruption setting), the experimental conditions approximate real-world label quality. A production RLHF team training a chat model on crowdworker preferences should: (1) run their standard RM hyperparameter grid search (which many teams already do), (2) rank candidate RMs by accuracy on a clean held-out validation set, (3) weight-average the top 4–6 RMs, and (4) use this single WARM model as the proxy reward for PPO/REINFORCE fine-tuning. The implementation requires no changes to the training infrastructure — just an additional averaging script after RM training — and reduces inference cost relative to maintaining an ensemble while improving robustness. The paper's finding that WARM permits lower KL regularization coefficients (Figure 8c, alpha = 0.001 works for WARM but not for individual RMs) means the team can also tune alpha more aggressively, getting more improvement from RL before hitting the reward hacking regime.

Efficient best-of-N decoding with improved selection quality. For applications using best-of-N as a training-free alignment method — generating N candidate responses and selecting the highest-scoring one — WARM provides better selection accuracy than a single RM at the same inference cost, and better efficiency than an ensemble at the same selection quality. The BoN experiments (Figure 6a) show WARM M = 6 achieving ~0.10–0.12 higher control reward than the best individual RM at intermediate N, which translates to a 92.5% oracle win rate against random selection (Figure 7a). For a deployment scenario where a model serves user queries and generates N = 8 candidate responses to be reranked, using WARM instead of a single RM means the selected response will be preferred by users more often, with zero additional latency (one forward pass either way). If the deployment currently uses an ensemble of M = 2 RMs for selection (doubling inference cost), switching to WARM M = 6 reduces inference cost by 2× while improving selection quality — a rare case where cost and quality improve simultaneously. The Baklava strategy makes this particularly practical: the team can generate diverse RM initializations simply by checkpointing their SFT run at multiple steps, without any additional training.

Federated or privacy-preserving reward model training. The paper notes in Section 6 that WARM's decomposition into independent fine-tunings "enables embarrassingly simple parallelization" and "facilitates its use in federated learning scenarios where the data should remain private." In a federated RLHF setting, different organizations or user groups may have preference data that cannot be shared (e.g., due to privacy regulations or proprietary concerns). Each party can train their own RM on their local preference data, starting from a shared pre-trained model (which can be distributed openly), and only the trained weights need to be shared and averaged. The p_j^2 mechanism has a privacy-relevant interpretation: features that are idiosyncratic to a single party's data (potential memorization of private preferences) will have low p_j across the full set of RMs and will be suppressed by weight averaging. This is consistent with the paper's observation that "WA adds a layer of privacy and bias mitigation by reducing the memorization of private preference." A concrete deployment would involve: a central party distributes a pre-trained LLM and SFT checkpoint; each of K participating organizations fine-tunes an RM on their local preference data; the resulting K weight vectors are sent to the central party; the central party averages them (potentially with per-organization weights to encode trust or data quality) to produce a single WARM; and the WARM is distributed back for inference. No raw preference data leaves any organization, and the averaging provides both privacy (through suppression of idiosyncratic features) and robustness (through variance reduction across diverse preference distributions). The paper does not empirically evaluate this scenario, but the architecture directly supports it.

Multi-objective and personalized alignment through non-uniform weight interpolation. The paper alludes in Section 6 to the possibility of "non-uniform interpolation of RMs" to "learn a set of personalized policies." The WARM framework naturally extends to multi-objective alignment: instead of averaging all RMs with equal weight, different linear combinations phi_alpha = sum_i alpha_i * phi_i (with sum_i alpha_i = 1) produce RMs that emphasize different aspects of preference. If the individual RMs are fine-tuned on different preference dimensions (e.g., some on helpfulness, some on harmlessness, some on humor), then sweeping the alpha vector produces a continuous family of RMs encoding different tradeoffs — without retraining any models. Each alpha yields a single RM (one forward pass) that can be used to train or select for a specific preference profile. This connects to the rewarded soups framework (Ramé et al., 2023) but applied to RMs rather than policies. A practical implementation for a customizable assistant would: (1) train RMs on separate preference datasets (helpfulness, safety, conciseness, engagement), (2) expose a small set of user-facing "personality" sliders that map to interpolation weights alpha, (3) at inference time, dynamically construct a user-specific WARM by interpolating the stored RM weights according to the user's slider positions, and (4) use this personalized RM for BoN selection or as a reward signal in continued RL fine-tuning. The key practical advantage over prediction ensembling is that the personalized RM is still a single model — so personalization adds no inference overhead relative to a non-personalized system.

When to Prefer This Method

The paper explicitly positions WARM against prediction ensembling (ENS) as the primary alternative for combining multiple reward models, and articulates specific tradeoffs between the two approaches (Section 6). The decision rule follows naturally from the paper's empirical and theoretical findings:

Prefer WARM over prediction ensembling when all candidate RMs share the same pre-training and architecture (a prerequisite for meaningful weight averaging), the preference data is known or suspected to contain label noise (since the corruption experiments show WA > ENS specifically under label corruption), and inference efficiency matters (since WA requires 1× forward passes vs. M× for ENS). The paper's RL experiments under corruption (Figure 8b) and the p_j^2 theoretical analysis provide the direct evidence for this recommendation.
Prefer prediction ensembling over WARM when the available RMs have heterogeneous architectures or pre-trainings (since WA requires weight-space compatibility that cross-family models lack), or when uncertainty estimation from prediction disagreement is required (e.g., for active learning or anomaly detection). The paper acknowledges in Section 6 that "prediction ensembling can incorporate prediction disagreement into the reward to provide uncertainty estimation," but cites Eisenstein et al. (2023) to argue that simple logit averaging often matches more complex uncertainty-aware aggregation — so the practical value of this ENS advantage is uncertain.
The choice is approximately neutral when the preference data is clean and only a small number of RMs (M = 2) are available — the paper's clean-label experiments (Figures 3, 6a, 8a) show WA and ENS performing very similarly, so either method works. However, since WA costs less at inference time and the paper provides no evidence that ENS outperforms WA even in this regime, WA is the safer default.
Increase M for WARM until diminishing returns set in. The paper shows that increasing from M = 2 to M = 6 provides substantial gains in both BoN (Figure 6a vs. 6b) and RL (Figure 8a) under corruption, but that M = 10 does not improve peak performance over M = 6 (Figures 8a, 9b), likely because the later-added RMs (phi_7 through phi_10) have lower individual accuracy. Practitioners should train a pool of candidate RMs, rank by validation accuracy, and include only those with accuracy above a threshold determined by the best models — not blindly average all candidates.