InfAlign: Inference-aware language model alignment

We thank the reviewer for acknowledging the novelty and sufficiency of the provided evidence in our work. Below we address the reviewer’s concerns. We provide additional experimental results at https://drive.google.com/file/d/1pe4kZeu7JkW0e-o5QtQFwpnvkZ4xVI9s

Q1: However, the evaluations seem somewhat limited. Demonstrating improvements in general instruction-following tasks would likely generate more excitement.

Thank the reviewer for the suggestion. We view demonstrating the method on more RLHF tasks as important next steps.

Q2 Weakness: The paper focuses solely on Best-of-N as the inference-time procedure. While additional experiments may not be necessary, including the formulas for transformed rewards in other inference-time procedures could strengthen

Thanks for pointing this out. To demonstrate the generalizability of our framework, we consider another inference-time strategy called rewind-and-repeat, motivated by the recent work of [Beirami et al 2024, Zhang et al., 2024]. Given a pre-defined threshold $\phi$ on the reward, the procedure keeps generating responses until the response hits the threshold. In Figure 9 of the attached file, we show that our proposed method leads to better win-rate / KL tradeoff in these cases compared to other SOTA methods. We also show that the expected number of generations needed to achieve the predefined threshold is lower with our alignment method.

Additional references (not listed in the paper): [Zhang et al., 2024] Yiming Zhang, Jianfeng Chi, Hailey Nguyen, Kartikeya Upasani, Daniel M. Bikel, Jason Weston, Eric Michael Smith. Backtracking Improves Generation Safety. Arxiv 2024.

We thank the reviewer for acknowledging the theoretical soundness of our work. Below we address the reviewer’s concerns. We provide additional experimental results at https://drive.google.com/file/d/1pe4kZeu7JkW0e-o5QtQFwpnvkZ4xVI9s

Q1: The authors claim that training using calibrated reward improves model's winrate compared to the other reward optimization baselines (Sec. 5.3). However, as depicted in Figure 4, CTRL's performance is quite close to that of BoND

We acknowledge that the tradeoff curve for standard win rate of our method is close to that of BoND, one of the SOTA methods (first row of Figure 4). However, for Anthropic helpfulness and harmlessness datasets, we do see a small improvement in the win-rate compared to BoND.

Moreover, we would like to emphasize that the focus of the paper is on improving the inference-time win-rate, such as best-of-4 win rate or worst-of-4 win rate. We note that for these cases, the improvement we get over BoND is more pronounced, as shown in the second row of Figure 4. We mainly present the results for standard win rate as a side product of our investigation on the effect of reward calibration. We believe the fact that it is on par (and better on two of the three tasks) with SOTA methods is a convincing evidence.

Q2: Unlike most existing work that optimizes for reward and use reward to assess model performance, this paper uses the win-rate over SFT model for training objective and evaluation. However, there is a potential issue of reward hacking when relying solely on win rate over a fixed model. To strengthen the evaluation, the authors could include comparisons of raw rewards and win rates over other models (e.g., models trained with different methods).

Regarding using win-rate as the evaluation metric, see response to Q1 of reviewer QsSA. Further, we have evaluated the raw rewards of the models in Figure 10 of the attached file that show that the raw rewards are correlated with the win-rates shown in Figure 4 (top left).

Q3: The calibrated reward measures the expected winning rate over SFT policy, and thus depends on the SFT model. It's natural that the learned reward function differs from the calibrated reward. The authors should discuss how the mismatch between learned reward and calibrated reward can hurt the performance.

The review is correct that the calibrated reward might differ from the learned reward function. However, the reward models are usually learned from human preference data using a Bradley Terry model to be used as proxies for real human performance. Hence, the expected win rate is a more robust metric compared to the learned reward model. For more discussions on the use of win-rate as the evaluation metric, see response to Q1 for reviewer QsSA.

In this paper, we show that instead of hurting, the calibrated reward consistently improves the standard win rate as shown in the first row of Figure 4, where CTRL with $\Phi(u) = u$ represents PPO with calibrated reward and PPO represents PPO with raw reward.

(Calibration in fact mitigates reward hacking) Moreover, we conduct controlled experiments to show that calibration in fact mitigates reward hacking, as discussed in Appendix D. To induce reward hacking, we injected specific phrases into the start of preferred responses of our preference datasets: “Sorry, I can’t help you with that” for Harmlessness and “Sure” for Helpfulness. We then evaluated the model’s accuracy on a poisoned evaluation set where these phrases were inverted (added to the unpreferred responses). A significant drop in accuracy on this poisoned set would indicate reward hacking: a reliance on the spurious correlation. Figure 5 in the appendix shows that calibrated reward models are more robust to these manipulated correlations, maintaining higher accuracy compared to uncalibrated models. Thus, calibration improves the reward model’s robustness to hacking based on training data poisoning.

Q4: The generalization of the proposed framework to other inference-time strategies (e.g., self-consistency, chain-of-thought) is not discussed.

In Figure 9 of the attached file, we consider an additional inference-time strategy called Rewind-and-Repeat inference-time strategy on the Anthropic helpfulness task. We refer to the response to Q2 of Reviewer XbH2.

Q5: How much additional computational overhead does the calibrated reward estimation introduce?

While the calibration step induces extra computation overhead, we remark that it involves only forward pass on the model and only needs to be performed once per prompt before performing the policy optimization algorithm. In our experiments, we take training steps equivalent to 80 epochs for all datasets. Based on the 2:1 FLOPS ratio between back propagation and forward pass, the reward calibration step takes about 29% of the total training FLOPS.

We thank the reviewer for the detailed reading and comments. Below we address the reviewer’s questions in detail. We provide additional experimental results at https://drive.google.com/file/d/1pe4kZeu7JkW0e-o5QtQFwpnvkZ4xVI9s

Q1: I am skeptical of the problem setting. The reward in definition 5 is the win rate against the reference policy. Obviously this is different from the original RLHF problem. And I am not sure it is always reasonable, especially when the reference policy is weak. The experiments also report the win rate against the reference policy.

Rooted in preference learning, reporting the win-rate vs KL tradeoff is a standard practice in the RLHF literature to evaluate the effectiveness of the language model alignment, ranging from the canonical work of Stiennon et al 2020 and the blog of Hilton and Gao 2022, which use win-rate with human preferences, to more recent works of Eisenstein et al., 2024; Mudgal et al., 2024, which uses more powerful models as judges. These win rates are direct reflections of human preferences while alternatives such as expected reward scores are indirect proxies, especially in cases where these reward models are learned from human preference data using a Bradley Terry model and their raw values may not have a physical meaning.

Moreover, Azar et al 2023, Gui et al 2024 formally formulated the optimization for win-rate vs KL tradeoff as the objective of RLHF. We follow these works and generalize the objective to include inference-time win rates, which better suits the modern regime of increasing inference-time compute. Hence we believe considering the win rate as the RLHF objective should not be viewed as a limitation of the work.

Further, we have evaluated the raw rewards of the models in Figure 10 of the attached file, which shows that the raw rewards are correlated with the win-rates shown in Figure 4 (top left).

Additional references (not listed in the paper): [Hilton and Gao 2022] Hilton, J. and Gao, L. Measuring Goodhart’s law, April 2022. URL https://openai.com/research/measuring-goodharts-law. Accessed: 2024-0103.

Q2: What is the accuracy improvement of the larger reward model based on PaLM-2 M against PaLM-2 S? Is it enough to model the generalization error of reward models in practice?

For Anthropic helpfulness dataset, the pairwise preference accuracy increases from 73.0% to 77.7%. There is still a gap between the accuracy of the large model and the true human preference, which is indeed a limitation of our evaluation approach. However, we want to remark here that collecting real human preference data is costly, and prior work [Stiennon et al., 2022, Wang et al., 2024] takes a similar approach where a larger and more accurate reward model is used as a judge to compute win-rate over models aligned with smaller reward models. We will add discussions on this limitation in the revised draft.

Q3: When measuring the BoN win rate in the experiments, do you use true rewards or learned rewards for sampling?

Following the literature (e.g., Eisenstein et al., 2024; Mudgal et al., 2024), we use a separate, more powerful model than the reward model as the judge to measure the win rate. More specifically, we use the PaLM-S fine-tuned reward model during RL training and BoN/WoN selection. We then evaluate the inference-time win-rate w.r.t base reference policy model using the PaLM-M reward model.

We thank the reviewer for acknowledging the novelty of proposed method and theoretical investigations. Below we address the reviewer’s concerns and provide additional experimental results: https://drive.google.com/file/d/1pe4kZeu7JkW0e-o5QtQFwpnvkZ4xVI9s

Q1: …the assumption that both the win-rate and inference-time procedure are differentiable … seems unrealistic for BoN/WoN …

We acknowledge that Theorem 1 for general inference-time procedures needs the assumption on differentiability of the inference-time win-rate. However, this assumption naturally holds for the considered procedures of BoN/WoN.

To see this, under a fixed prompt $\mathbf{x}$, each policy $\pi(\cdot \mid \mathbf{x})$ can be viewed as a $\mathcal{Y}$-dimensional vector. For a pair of policies $\pi\_1$ and $\pi\_2$, the win-rate as defined in Definition 3 is an expectation of the win random variable with samples from $\pi\_1(\cdot \mid \mathbf{x})$ and $\pi\_2(\cdot \mid \mathbf{x})$. Hence it is a linear function of both policies, making it differentiable with respect to both policies. Correspondingly, the inference-time win rate as defined in Definition 4 is linear with respect to both inference-time policies. By the chain rule, it is sufficient to show that the inference-time policy is differentiable. For BoN/WoN, the inference-time policy has explicit forms stated in Lemma 5. For example, $BoN\_{\pi}(\mathbf{y} \mid \mathbf{x}) = N pi(\mathbf{y} \mid \mathbf{x}) \mathcal{C}\_{r, \pi} (\mathbf{x}, \mathbf{y})^{N\_1}$. Since $\mathcal{C}\_{r, \pi} (\mathbf{x}, \mathbf{y})$ is a linear function of $\pi(\cdot \mid \mathbf{x})$ as defined in Definition 2, for all $\mathbf{y}$, $BoN\_{\pi}(\mathbf{y} \mid \mathbf{x})$ is a polynomial function of the base policy $\pi(\cdot \mid \mathbf{x})$, making it differentiable w.r.t. $\pi(\cdot \mid \mathbf{x})$. These two facts combined leads to differentiability of inference-time win-rate for BoN/WoN. We will make the above discussion clear in the updated version.

Q2: The actual transformations for BoN/WoN are exponential transformations with a hyperparameter t dependent on N. …. requires the hyperparameter search on t and retraining of the aligned model for different N's.

While we use the exponential transformation in the empirical experiments, we did show that they lead to near optimal KL / win-rate tradeoff curves in Section 4.3. For example, in Figure 2 (a), the curve marked bon_fp is obtained by using the transformations obtained from solving the fixed point in Corollary 2. As shown, the tradeoff curve is numerically close compared to the one obtained from $exp(10u)$. Additionally, the exponential transformation is motivated by exponential tilting of loss functions (Li et al., 2021; 2023), which has been shown to help optimize different quantiles of the reward with different values of t, making suitable transformation for the BoN/WoN inference-time procedure.

When choosing the hyperparameter $t$, our method doesn’t need to retrain different models to perform the hyperparameter search. Instead, the search can be done efficiently using analytical tools with closed form expressions on KL and win rate (Theorem 3). To demonstrate that the findings will generalize to practical settings, we present more results in the attached file. For the analytical analysis (middle column of Figure 7 in the submission), $e^{5u}$ works better than $e^{10u}$ for best-of-2 and the reverse holds for best-of-4. In Figure 8 of the attached file, we show that for real models and tasks, we see the same trend, demonstrating the transferability.

Q3: Figure 4 .. most of results are shown with only N=4 and there are no argument on the possibility of overfitting to N.

Regarding the overfitting to $N$, we do find different $t$’s work the best for different $N$s. However, we would like to mention that in many practical settings, the $N$ that is going to be used at inference-time is known during the RLHF phase. And we could find the best exponent $t$ efficiently for the $N$ to be used as mentioned in the response to Q2. In the attached file, we also provide results for N=2 and N=32 to demonstrate the generality of our approach. In cases where $N$ might change due to the deployment resources, we show that while the best $t$ is different for different $N$’s, the gains are generalizable for mismatched cases as well. In Figure 8 of the attached file, we obtain consistent gains for different $N$s with different $t$’s. For example, with $t = 10$, we obtain significant gains for all three cases of $N = 2, 4, 32$, showing that overfitting to $N$ is not a major limitation.

Q4: why the maximization of calibrated rewards leads to improved and robust alignments

Response: In Appendix D, we present experiments to demonstrate that calibrated reward models are less susceptible to reward hacking. See response to Q3 of reviewer dJSD for more discussion on this. We will add more discussions in the future revisions.