Test-Time Alignment via Hypothesis Reweighting

Thank you for your constructive feedback. We address each of your concerns below.

Comparisons with alignment baselines

Thank you for this suggestion. We have modified our original evaluation of RewardBench datasets to be directly comparable to the methods on the official leaderboard. Our evaluation of HyRe builds on the GEM-Gemma-2B reward model, which was the state-of-the-art at 2B scale at the time of submission.

The table below shows representative results, comparing with frontier generative models and open-source models trained via alignment methods such as DPO. Notably, HyRe achieves consistent improvements in overall score, with strong gains in the Chat, Safety, and Reasoning splits. We note a slight performance drop on the Chat Hard split, which includes less task ambiguity as the preference datasets used to train these models primarily focus on challenging chat responses.

We note that this is not a direct head-to-head comparison, as HyRe leverages labeled data not utilized by the other methods. The primary purpose of this experiment is to demonstrate the potential performance gains achievable through test-time alignment.

Your motivation, the claim that "the best single model (A) can substantially outperform the ensemble average (B)" does not directly lead to the conclusion that "it is more advantageous to view the ensemble as representing a set of candidate models (C) rather than aiming for a single 'best' function through uniform averaging (D)". The relationship between A, B, C, and D needs clearer justification. Are B and D describing the same approach? This logical connection requires more elaboration to be convincing.

We apologize for any confusion. B and D are indeed the same. Both describe a uniform average of ensemble members. We observe empirically that A (the best single model) can substantially outperform B (the uniform ensemble) in underspecified tasks (see Fig 1). This motivates C: dynamically selecting or weighting the models at test time to better align with the target task. HyRe is an instantiation of C. To improve clarity, we have made a major revision to section 3.1 (most important changes in red).

Your method is based on Strong Assumptions. It relies on training multiple heads as the basis for test-time adaptation, which implies: You assume that test tasks can be represented by a limited set of basis functions, which may not hold true in many real-world applications. You also assume that test tasks are linear combinations of these basis functions, another strong assumption that is often unrealistic.

We acknowledge that modeling target functions as linear combinations of basis functions is a simplifying assumption. However, this is a practical modeling assumption to strike a good bias-variance tradeoff and does not need to hold exactly for HyRe to be effective. As HyRe operates in a low-data regime (5-50 labeled examples), mitigating the risk of overfitting is important, even at the cost of constraining the hypothesis space. Despite the constrained hypothesis space, HyRe demonstrates strong performance across diverse real datasets and scales effectively to many basis functions (100), showing its practicality for real-world problems.

Thank you for your constructive feedback. We address each of your concerns below.

Comparison with reweighting baselines.

Thank you for pointing out this relevant literature. We have added new experiments comparing HyRe against these methods for ensemble reweighting. The table below shows the average accuracy across 15 RewardBench datasets. HyRe demonstrates substantial improvements over prior methods, with most benefits coming from the first few examples.

[1] D. Jimenez, "Dynamically weighted ensemble neural networks for classification," 1998 IEEE International Joint Conference on Neural Networks Proceedings.

[2] Shahhosseini, Mohsen, Guiping Hu, and Hieu Pham. "Optimizing ensemble weights and hyperparameters of machine learning models for regression problems." Machine Learning with Applications 7 (2022).

Isolating the effects of ensemble learning vs. weight learning. The authors’ main argument / hypothesis is that given task underspecification, a single model can outperform a naive ensemble. Can the authors also provide how the performance of a single model fares compared to an ensemble + HyRE — beyond the toy experiment in Figure 3? Essentially, it would be good to quantify if HyRE is able to learn the “optimal” weighting? If the authors could provide some additional results on perhaps the language model or vision (WILDS) experiments.

In the table above, all methods other than GEM and HyRe are based on ensemble learning only, without task-specific weight learning. We see that learning task-specific weights (GEM and HyRe) provides a substantial boost over methods for aggregating the ensemble alone.

As an upper bound for the achievable performance of ensemble reweighting, we also evaluate GEM [2] with access to the entire test set (no held-out examples). We note that this is not a fair point of comparison as it directly uses the test set, and should rather be seen as the highest achievable accuracy from “overfitting” during the weight learning stage.

Also, what is also the entropy of the learned weights over the members of the ensembles? Do the authors observe that it collapses onto a single model, if the hypothesis is indeed true that a single model can outperform all others in task underspecification settings?

Following your suggestion, we measured two metrics over all 15 RewardBench datasets after 40 examples each: Normalized entropy of the learned weights (scaled to [0, 1]). For a uniform ensemble, this would be 1.0 Maximum weight assigned to any single model. For a uniform ensemble of 100 members, this would be 0.01. The average normalized entropy was 0.4466, and the average maximum weight was 0.4395, indicating that the weights are far from uniform. The learned ensemble does not completely collapse onto a single model with finite data, reflecting the appropriate behavior to avoid overfitting. With infinite data, the method would necessarily converge to the best-performing model(s).

The reviewer acknowledges that running scaling experiments may be difficult given the time constraint, but an obvious argument against this ensemble approach would be that it then takes up to K times the amount of training compute to train K models. How does this fare against training 1 larger model in the paper’s experiments?

We clarify that HyRe does not require training K separate models. We use a single pre-trained backbone and K prediction heads, i.e., K small MLPs that take backbone embeddings as input. The computational overhead is negligible; for example, in our reward model experiments, the 100 ensemble heads (5.5e5) add less than 0.03% to the parameter count of the gemma-2b backbone (2.0e9).

Thank you for your constructive feedback. We address each of your concerns below.

HYRE is compared against fine-tuning and other ensemble-based models but lacks direct comparisons with recent advances in task alignment or ensemble calibration methods. The results show that HYRE outperforms conventional ensemble approaches. Could the authors clarify how HYRE compares with models explicitly trained for task alignment, particularly in settings where task ambiguity is less pronounced?

Thank you for this suggestion. We have modified our original evaluation of RewardBench datasets to be directly comparable to the methods on the official leaderboard. Our evaluation of HyRe builds on the GEM-Gemma-2B reward model, which was the state-of-the-art at 2B scale at the time of submission.

The table below shows representative results, comparing with frontier generative models and open-source models trained via alignment methods such as DPO. Notably, HyRe achieves consistent improvements in overall score, with strong gains in the Chat, Safety, and Reasoning splits. We note a slight performance drop on the Chat Hard split, which includes less task ambiguity as the preference datasets used to train these models primarily focus on challenging chat responses.

We note that this is not a direct head-to-head comparison, as HyRe leverages labeled data not utilized by the other methods. The primary purpose of this experiment is to demonstrate the potential performance gains achievable through test-time alignment.

Empirical studies are concentrated on well-known datasets, but the paper could benefit from evaluating HYRE on additional real-world datasets, especially those with more nuanced or complex underspecification.

Thank you for the suggestion. We have additionally evaluated our ensemble reward model on the PERSONA dataset [1], which emphasizes underspecification by curating inputs designed to provoke disagreement. We test across ten personas with 200 preference pairs each. Details on each persona are in our appendix. As shown below, HyRe shows a significant improvement, achieving 83.0% accuracy with N=40 examples per persona compared to 14.8% for the base model.

[1] Castricato, Louis, et al. "PERSONA: A Reproducible Testbed for Pluralistic Alignment." arXiv preprint arXiv:2407.17387 (2024).

The paper lacks detailed analysis on how different active learning criteria (entropy, BALD, variance) affect performance across tasks.

We have added a new experiment comparing the effect of different active learning criteria for HyRe in our RewardBench preference alignment tasks. We consider random sampling, BALD, and entropy, measuring their performance over 0 to 40 target examples. Across the acquisition of 40 examples, active learning methods (BALD and entropy) demonstrated slightly better performance compared to random sampling. We note that even random sampling consistently improved performance, suggesting that our reweighting process is robust to datapoint selection strategy.