InfAlign: Inference-aware language model alignment
Authors: Ananth Balashankar
Ziteng Sun
Jonathan Berant
Jacob Eisenstein
Michael Collins
Adrian Hutter
Jong Lee
Chirag Nagpal
Flavien Prost
Aradhana Sinha
Ananda Theertha Suresh
Ahmad Beirami
Venue: ICML 2025 poster (2025)
Keywords: language model
alignment
decoding
inference time procedure
best of n
Primary area: deep_learning->large_language_models
TL;DR: Given an inference-time procedure, such as Best-of-N, standard RLHF suffers from a train/test mismatch between train-time policy and inference-time policy. InfAlign bridges this gap via reward calibration and transformation.

Abstract: Language model alignment is a critical step
in training modern generative language models.
Alignment targets to improve win rate of a sample
from the aligned model against the base model.
Today, we are increasingly using inference-time
algorithms (e.g., Best-of-$N$ , controlled decoding, tree search) to decode from language models
rather than standard sampling. We show that this
train/test mismatch makes standard RLHF framework sub-optimal in view of such inference-time
methods. To this end, we propose a framework for
inference-aware alignment (InfAlign), which
aims to optimize *inference-time win rate* of the
aligned policy against the base model. We prove
that for any inference-time decoding procedure,
the optimal aligned policy is the solution to the
standard RLHF problem with a *transformation*
of the reward. This motivates us to provide the
calibrate-and-transform RL (InfAlign-CTRL)
algorithm to solve this problem, which involves
a reward calibration step and a KL-regularized
reward maximization step with a transformation
of the calibrated reward. For best-of-$N$ sampling
and best-of-$N$ jailbreaking, we propose specific
transformations offering up to 3-8% improvement
on inference-time win rates. Finally, we also show
that our proposed reward calibration method is a
strong baseline for optimizing standard win rate.

InfAlign: Inference-aware language model alignment
Ananth Balashankar"
Ziteng Sun Jonathan Berant' Jacob Eisenstein
Michael Collins'
Adrian Hutter'
Jong Lee'
Chirag Nagpal Ahmad Beirami
Flavien Prost Aradhana Sinha
Ananda Theertha Suresh
Standard RLHF
Abstract
RLHF
Inforeoce- Lime progedure
Inference
InFAlign
Language model alignment a critical step in training modern generative language models. Alignment targets improve win rate of sample from the aligned model against the base model. Today, we are increasingly using inference-tinge algorithms (e.g., Best-of-N, controlled decod- ing, tree search) to decode from language mod- els rather than standard sampling. We show that this train/test mismatch makes standard RL frame- work sub-mmtimal in view of such inference-ting methods. To this end, we propose framework for inference-aware alignment (InfAl gn), which aims to optimize inference-time win rate of the aligned policy against the base model. We prove that for any inference-time decoding procedure, the optimal aligned policy the solution the standard RLHF problem with transformation of the reward. This motivates calibrate-ann and- transffrro RL [InfAlign-CTR) algorithm to solve this problem, which involves a reward calibration step and a KL-regularized reward maximization step with a transformation of the calibrated reward. For best-of-N sampling and best-of-N jailbreaking, we propose specific transformations offering up to 3-8% improvement on inference-ting time win rates. Finally, we also show that our proposed reward calibration method is a strong baseline for optimizing standard win rate.
reward calibrarion
transformation
R
RE
T-+
Figure 1. Given an inference-time procedure such as Best-of-N, standard RLHIF suffers from train/test mismatch between train- time policy æ and inference-time policy Ta. In infAlig bridges the gap by optimizing policy-transfone reward R. yielding policy, x², that is optimized for inference under Ta-.
1. Introduction
Aligning language models (LMs) through reinforcement learning from human feedback (RLHF) is widely adopted finetuning framework to improve a reward (e.g., safety or quality). RLHF generally entails training reward model, and then solving a KL-regularized RL problem (Christiano et al., 2017: Stiennon et al., 2020; Ouyang et al., 2022). Other ways of solving RLHF include variants of direct pref- erence optimization (Rafailov et al., 2023; Azar al., 2023), reward model distillation (Fisch et al., 2024), and Best-of-NV distillation (Gui et al., 2024; Amini et al., 2024; Sessa et al., 2024). The success of RLHF is typically measured through the win rate of the aligned model against the base model, i.e., how often a sample from the aligned model wins against the base model using judge for the task.
provide the
However, rarely is the aligned model used is at inference time; instead an inference-time procedure is typically used to accomplish a task. For example, it is customary to perform one or more of the following procedures at decoding: best- of-N sampling (Nakano et al., 2022; Beirami et al., 2024), best-of-N jailbreaking (Hughes et al., 2024b), chain-of- thought reasoning (Wei et al., 2022: OpenAI, 2024), and self consistency (Wang et al., 2022a; DeepSeek-AI, 2025) (see Appendix A for more comprehensive discussion).
'Equal contribution 'Google DeepMind Google Research. Correspondence
Ananth Balashankar <ananthbshankard googlc.com>>,
Ziteng Sun < <zitengsun @ google.com com.>.
Jonathan Berant <joberant @wwoglggoooooooooo google.com>>,
Jacob Eisenstein <jeisenstein @ google.com>>
Ananda Theertha Suresh <theertha @ google.oo googlc.com>> com
Ahmad Beirami Sahmad.bbeiaan @ gmai..com>>.
In this paper, we address the following question: Can we better align a language model to be served with a known inference time procedure? We propose a family of optimiza- tion objectives, which we call the inference-ave -wwre alignment
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

InfAlign: Inference-aware language model alignment
(InfAl ign) framework. The objective optimizes for the inference-time win rate against the base model with KL regularization, where inference-time win rate entails obtain- ing a response from each model through the inference-ie time procedure and counting which sample wins. While directly optimizing the inference-time win rate seems intractable, we prove that the solution could be obtained by solving RLHF with a specific transformation of the reward (Lemma 1). Therefore, the challenge of optimizing for inference-time win rate can be captured by designing reward transforma- tion that is suited to the specific inference-time procedure. We present a diagram of InfAlia ign in Fig. I.
mance and Anthropic harmlessness for optimizing WoN per- formance with various values of N. We show that solving InfAlign through InfAl i outperforms vari- ous SOTA RLHF solvers at inference-time win rate by 3-8%. We also show In i gn-CTRL (with no reward transfor- mation) is on-par with SOTA methods for optimizing the standard win rate, with improved win rate for Anthropic helpfulness and harmlessness tasks.
Organization. In Section 2, we provide the InfAlign problem setup. In Section 3, we introduce RLHF with reward transformation as the main framework to solve InfAlign. In Section 4, we present the InfAlign- CTRL method, discuss properties of calibration, and present practical implementations. In Section 5, we provide experi- mental results.
We show that the optimal reward transformation satisfies a coupled-taasformeed reward/policy optimization objective which lends itself to iterative optimization for a large class of inference-time procedures (Theorem 1). However, the approach is unfortunately computationally inefficient and infeasible for real-world models.
2. Problem setup
We consider a generative language model that produces a response conditioned on an input prompt. Given prompt xr e.g., x What is a large language model?, a genera- tive language model or a policy, specifies a conditional distribution x-- æ)), from which the response y should be generated. We use X and y to denote the space of possible inputs and outputs, respectively. Throughout the paper, we assume Free is a fixed base policy, e.g., obtained from supervised finetuning. We will often use the notation T ry x)) x f'y) for some y - IR+ to denote the con- ditional distribution r x)) = f(y)/(E, f1)) obtained after normalization.
To enable practical solutions to InfAlign, we pro- pose Calibratc-and-Trrasffom Reinforcement Learning (InfAlign-CTRL) which adopts a three-step approach: (1) calibrate the scores of the reward model with respect to responses sampled per- per-prompt by the reference model; (2) transform the calibrated scores for given inference-time procedure; and (3) solve the RLHF problem with the trans- formed reward. We prove various desirable properties of the reward calibration step and empirically show that it reduces reward hacking and improves standard win rate. Reward transformation allows us further tailor the objective based on the inference-ting time procedure. Different choices of trans- formation lead to popular existing alignment objectives such as IPO (Azar et al., 2023) and best-of-N distillation (Gui et al., 2024; Amini et al., 2024; Sessa et al., 2024).
A Alignment of language models. Letr: x - I be a reward function that assigns a scalar value to any (prompt, response) pair, e.g., a model trained side-by-side on hu- man preferences. The reward determines the goodness of response y in context The goal of language model align- ment is to construct aligned distribution that improves the reward of the response while being close to Tref. We in- troduce the popular KL-regularized reinforcement learning (RL) framework, also known as RLHF
We particularize the study to two simple yet popu- lar inference-time strategies: Best-of-N (BoN) sam- pling (Nakano et al., 2022; Beirami et al., 2024) and Best- of-N jailbreaking (Hughes et al., 2024a), which we call Worst-of-N (WoN) since the defender may assume that the attacker is choosing the worst outcome of N samples. De- spite simplicity, BoN is known to be an effective procedure for inference-time time alignment (Beirami et al., 2024; Gui et al., 2024; Mudgal et al., 2024) and is the dominant approach in scaling inference-time compute (Snell et al., 2024; Brown et al., 2024). Variants of WoN are effective and popular for evaluating safety against jailbreaks (Yohsua et al., 2024; Chao et al., 2023; Souly et al., 2024; Hughes et al., 2024b; Beetham et al., 2024; Mehrabi et al., 2023). We find suit- able reward transformations for these procedures through the In InfAlign ign framework.
Definition 1 (KL-regularizzd RL). Let 3 >> 0 be a regular- ization parameter, the KL-regularized RL (RLHF) problem aims to maximize the expected reward with a KL regularizer below:"
T" |x))= = arg max (Lrrd 7
where
F fr(ar, b)) (1x)))))))))))))()))
When evaluating an aligned policy, common measure to use is the win rate (Stiennon et al., 2020; Hilton & Gao,
Empirically, we apply InfAlign-Cre to the Anthropic helpfulness (Bai et al., 2022), and Reddit summarization dataset (Stiennon et al., 2020) for optimizing BoN perfor-
'The solution to this optimization problem is unique and admits closed-fort form expression (Korbak al., 2022; Rafailov al., 2023; Yang et al., 2024).
2

---

InfAlign: Inference-aware language model alignment
2022) over the base policy Fret- Let
Definition 5 (InfAlign), For given inference-lime pro-
cedure T and 3> 00 InfAli An solves the following KL-
wr ==)x===))==(rr( y) () +)}+ 1 1( y) r }}}
regularized inference-time win rate maxmization problem:
2
max [ (T>>)et(==)--- - 30x (I (-1z))) (-1x))) (3)
be the win random variable under reward r. We define the calibrated reward and win rate below.
The formulation of optimizing standard win-rate vs KL tradeoff has been previous studied for RLHF in Gui et al. (2024); Azar et al. (2023). The above formulation reduces to the IPO objective (Azar et al., 2023, Equation (8)) when T is the identity transformation and extends it otherwise for an arbitrary inference-time procedure. In practice, the win rate is often evaluated by judge different from the training time reward r. In this paper, we stick to when developing the algorithms as in standard RLHF framework (Definition 1). In experiments, we use a separate judge from the reward model.
Definition 2 (Calibrated reward). The calibrated reward" (r, y) under policy iT is defined below
Ca (Z., v) == E (essn(lll)]W-ll
æ)}.
Definition (Standard win rate). For any policy ai and no.
the standard win rate (or win rate) of policy Th over policy 22 given prompt x, as measured by reward is defined as:
W > n2 x)= == Ey x) (-bn)) (w. x)}.
Continuous language models. For simplicity of the pre- sentation and analysis, some of the theoretical results will be based on the assumption that y is a continuous set, and the language model has a density over y. We also assume that r assigns distinct rewards to different y's for a given æ. Note that we don't make any such assumptions when providing our algorithmic developments, and the experimental results.
Moreover, it can be shown that ( (- (Crr (3)
Inference-time procedure (T). As mentioned earlier, in many cases, decoding is done through an inference-time procedure. The obtained sample follows a transformed distribution that depends on the policy and the inference- time procedure. Let Ay be the set of possible distributions over the output set V. In this work, we model an inference- time procedure mapping between distributions over the output space T Ay Ay.
These assumptions have been made in the past implicitly by (Hilton & Gao, 2022) to estimate the KL divergence of Best-of-n policy, and by (Gui et al., 2024) to charac- terize the KL divergence and win rate tradeoffs. While these assumptions lead to approximations when analyzing real-world distributions, the results derived under them are reasonably tight when the actual likelihood of the language model outcomes are small (Beirami et al., 2024).
π(- x)) T. x))
where [: ax) denotes the distribution of responses condi- tioned on æ after the inference-time procedures is applied. In these cases, it is customary to compare the models by considering the following inference-time win rate of the
3. Reinforcement learning with reward transformation
aligned policy x.
Definition& 4 (Inference-time win rate). Under inference- time processing T. the inference-time win rate of policy Tı over n2 is defined as
In this section, we propose a general framework for solv- ing InfAlign (Definition 5). Our approach is based on designing new reward function R based on the reward model r, the inference-time procedure T, and the base pol- icy Tref, such that solving the RLHF problem (Definition 1) with the transformed reward R. leads to an optimal solution to In ign. More precisely, the aligned policy is the maximizer of the following regularized objective:
W. (T] > n2l [+) := ByTTT, (lo)), ae) twr 2|x)}.
Similarly, it can be shown that
W (m > ne x)) (o, (3 (
(2)
Re 7 ((x, y) 3DL (TT(- xx) [mrrr( x)), (4)
where Cr,T., (xx, y) is the calibrated reward under the inference-time policy Tmp (see Proof in Appendix B.2).
where Ry, (xx y) is transformed reward function. In- terestingly, we show that such reward transformation is
With the above definitions at hand, the goal of InfAlign is solve the following KL-regularized inference-time win rate maximization problem.
"One question that arises is the role of the KL divergence regularizer Eq. (3). We argue that the regularizer essentially enables multi- tasking between the SFT task and the RL task, which we formally prove for log- linear models in Appendix C. In other words, the KL divergence regularizer enables to preserve/tistill the core capabilities of the SFT model while acquiring new one through the RL.FF process.
Thh definition is similar to the cumulative density function (CDF) of the reward r(r, y) under except for how ties are decided.
3

---

InfAlign: Inference-aware language model alignment
sufficient solve InfAlign.
generalization of these results with general inference-time procedures. The observation motivates us to consider RLHF with specific family of reward transformations that in- volves a reward calibration step, described next.
Lemma 1. For any base policy Tref reward model r,
inference-time procedure T, and B > O, there exists a
reward function R T such that the maximizer of Eq. (4)
solves the optimization problem in Eq. (3) (Definition 5).
4. InfAlign-CTRL: Calibrate- and-transform reinforcement learning
In general, such optimal reward transformation will depend on the base policy Tref the inference-time procedure T, and the reward model r. In the lemma below, we list the property that the reward transformation and the resulting optimal aligned policy must satisfy.
In this section, we propose the InfAlign-CIRE method, which is our proposed solver for the InfAlign problem. The method consists of: (1) Calibration: Approximate the calibrated reward Cr, (2) Transformation: Apply trans- formation & on top of Crames to obtain Rs = o o Cram (3) Solve RLHF with the transformed reward. We discuss each step in more details below.
Theorem 1 (Characterization of InfAlign solution). As- suming that T is such that (y) 2=)/T(y>> 2)*² exists for all x, yr, y2, then we have the optimal transformed re- ward R and the optimal policy x² in Eq. (3) must satisfy the following coupled equations: Vx,y
Reward calibration. The goal of this step is to obtain the calibrated reward Cr,r,et (Definition 2). We first assume (5) Crarn could be obtained perfectly and discuss practical ways to approximate it in Section 4.4. In Section 4.1, we (6) discuss various properties of Cr, In particular, it can be shown that for continuous language models, the distribu- (7) tion of Cra is independent from the base policy, and the reward model (Lemma 4), which provides unified view of the outputs from a language model through the lens of calibrated reward. This allows us to focus the design of the transformation function T on T in the following step.
x² xy(y(x)) x Tret (y|
x 1)) R(x.))
R(0, θT x))
W (T > Tref x) IT
aTa (z x))
= - (x, 2)
Om x²)
where the last inequality is due to Eq. (2).
Missing proofs are presented in Appendix B. Theorem naturally leads to an iterative EM-style algorithm that (I) updates π with fixed based on Eq. (5) and (II) updates R with π fixed based on Eq. (7) until convergence. However, such algorithm suffers from two drawbacks: first, for gen- eral language models, it is inefficien//ntraattblle to evaluate Eq. (7) since it involves evaluating the policy on a large, or even infinite output space; second, it is unclear whether such an algorithm could lead to the optimal solution.
As mentioned in Corollary 1, using the calibrated reward directly in the RLHF problem leads to an optimal solution to InfAli gn with no inference-time procedure. Note that the resulting solution is the same as the alignment objective of Azar et al. (2023). Moreover, it can also be shown that using log Cra as the reward, the solution to the RLHF problem recovers the popular best-of-N distillation objective, which has been studied in a recent line of works (Gui et al., 2024; Amini et al., 2024; Sessa et al., 2024), and shown to be nearly optimal for standard win rate (Yang et al., 2024; Gui et al., 2024). We note that while these methods lead to simi- lar optimization objectives, InfAll FAlignn-TRL makes this calibration step explicit, resulting in different algorithmic approach. In Section 5, we compare the results with the abovementioned baseline approaches on standard win rate, and show that it leads to improved win rate, We also empiri- cally show that reward calibration is beneficial to mitigate reward hacking (Appendix D).
To find more efficient ways to design reward transforma- tions, we examine the case when no inference-time proce- dure is performed. In this case, Tx x and
i
[x x²) = 1 1f = y}.
x)
Eq. (7) will reduce to R(,,,) Cr. (x., y), the cali- brated reward under Tyof.
Corollary 1. When no inference-time procedure is per- formed, i.e. VT, Tx = π, the maximizer of Eq. (4) with R(x,, y) = Cr (x., y) is the solution to Eq. (3).
The goal is to further trans- form the calibrated reward using a transformation function T : 10 + R. The function & is chosen based on the inference-time procedure T so that & o Cr, is a good reward transformation to use for solving InfAl ign..
Reward transformation.
The above corollary is also observed in Azar et al. (2023); Gui et al. (2024). Hence Theorem can be viewed as a
"To make sure that the partial derivative is well-defined, we assume that T is well-defined for inputs within an infinitesimal ex- pansion of Ay. We note that the assumption holds for procedures like BoN and WON, as shown I Lemma 5.
Ideally, we would like the design of T to only depend on the inference- procedure T so that it is transferable among different and Tref- A natural question for what
4

---

InfAlign: Inference-aware language model alignment
type of T's this is possible. In Section 4.2, show that for calibrated inference-time procedures (Definition 6), and continuous language models, it is sufficient for & to depend on T. And in these cases, different d's could be evaluated efficiently with simple language models that can be easily simulated, which enables the search for good or even opti- mal transformations. We use Best-of-N /V and Worst-of- N as examples of inference-time procedures to demonstrate the effectiveness of such approach in Section 4.3.
term calibrated procedures, and discuss how to design suitable c for this family of transformations. We first define calibrated procedures below.
Definition 6 (Calibrated inference-time procedure). An
inference-time procedure TT is called calibrated proce- dure if there exists mapping function 9T : [0, 1] + R such that for any n, T, and æ, y, we have
Tx x)) x Ty x) (C,, (2,
RLHF with the transformed reward. At this step, we solve the RLHF problem with reward function,
Our next result shows that for calibrated inference-time procedures and continuous language models, the aligned policy obtained from InfAlign-CTRI with any & has a win rate and KL divergence independent of Tyef and r. Theorem (Model-agnostic property of calibrated infer- ence-time procedures, informal version of Theorem 4). If T is a calibrated inference siteenncettieeee procedure, for any contin- uous language model B O and transformation func- tion Φ, we have that both W.T RR > Nyet | x)) and Dxz. ( Ra [all are independent of T and Tyef-
R. (a. y) = (Crr, (y) x)),
to obtain the aligned policy R. Since we only modify the reward function, the step can be performed with stan- dard RLHF solvers. We present a practical version of the In Frlllign-TTRL method in Section 4.4.
4.1. Properties of reward calibration
We present properties of Cr. The first property states that reward calibration preserves the ordering of the reward. Lemma 2 (Calibration is a bounded monotone increasing transformation of reward). We have Cr. (z. 3) E [o, 1]. Furthermore, we have for any y and z
The above theorem allows us to evaluate transformation & by focusing on simple continuous language models that are easy to compute and simulate. the next section, we focus on Best-of-N and Worst-of-NN N, as examples to demonstrate how the theorem enables us efficiently evaluate different transformations in practical scenarios.
r(x, y) 2 r(xx,2) = Crowie (x, v) Craszer (2,2). 2).
Morcover, Cr. is invariant under all monotone increasing transformations of the reward function, stated below. Lemma (Calibration is invariant under monotone increas- ing transformations). Let m R + R be any strictly mono- tone increasing function. Then Cm =G
4.3. Finding reward transformations for BoN and WoN
Best-of-N inference-time procedure (BoN). During infer- ence, N i.i.d. responses from policy æ are generated. The final output is the one with the highest reward, i.e.
This property is useful since as long as the learned reward model r can capture relative human preference between each pair of generations, the calibration of r will remain unchanged, making Cr, more robust to the learning of
YBO arg
max (vI....yu
r(z, (2,
Worst-of-N inference-time procedure (WoN). N i.idd. re- sponses from policy π are generated, and the final output is the one with the lowest reward. i.e.,
The next property shows that the calibration operation al- lows us to transform the distribution of the reward under the base policy to a uniform distribution over [0, 1] regardless of the base policy Feef and the reward model r.
yKON arg
min WE (w NN)
r(æ, y).
Lemma 4. ग. is a continuous language model, let y be sampled from Wres(. ( α)), then we have Vaz
The lemma below presents the distribution of outputs after the inference-time procedure is performed.
Cr, (xx y) Uni 1)).
Lemma 5. For any N and continuous language model n
BoN, (y xx) N < x) (x, y) N-1.
The lemma provides us a unified view of the output from language model through the space of calibrated reward.
WoN, (3 x)) N ( x)) Cr, y)) N-1
4.2. Reward transformation for calibrated inference-time procedure
Note that the results for BoN have already been derived previously (Beirami et al., 2024; Gui et al., 2024; Amini et al., 2024). The lemma shows that these two inference- time procedures are calibrated procedures so that as claimed
We consider a family of inference- procedures that only depend on the calibrated reward of the outputs, which we
5

---

InfAlign: Inference-aware language model alignment
in Theorem 2, for the aligned policy, the inference-time win rate and KL divergence deviation from the base policy are independent of the base policy and reward model. Below we present the precise formula for these two procedures.
1.0
0.9
0.9
0.8
0 bon !P
0.8
won_
ecu) - logtu)
0(0) - log(o)
0.7
0.7
Theorem 3 (Properties of BoN and WoN procedures). For
0.6
a(0) exp(1O0)
Wersst
0.6
0(0)
((1 -Xp(--A)
any transformation function Q, the solution T to InfALign-CTRI satisfies the followings: Let Fo.c(u)=
0.5
(t0) exptbut
0.5
0(o) exp(-10u)
0.5
2.0
0.5
2.0
KLiRlnud
Figure 2. Best-of-N (left) and Worst-of-N (right) win rate vs KL
tradcoff curves for N = 4 with different transformation functions.
For any x, Who > Tref xx)
where F(v) So f(u)du is the CDF with density function f. and for T = WoN, is satisfies that
I-N S.' (o. NuN--luu
2 wen (u)(u) N2 [. (1 (1-
F(v))
Forr any 2, WWWO (TRO, > Toot x))
f(u) xe "mogin)
N L (1- (o. (1 --=)x u)
(9)
Hence, we derive transformation function by iteratively finding the fixed point of the coupled equations in Corol- lary 2. We call them bon fp and won fp.
. DT Ro. [TTett
1 So (())e - log (fr ((woyea)) (10)
Exponential tilting for reward transformation. In addition to deriving the optimized transformation, motivated by the exponential tilting of loss functions (Li al., 2021; 2023), we consider the following exponential transformation:
3 So e
The above theorem generalizes the win rate calculation in Gui al. (2024, Lemma 5) to inference-time win rate. By varying 3 in Theorem 3, we obtain an alignment curve plotting the inference-time win rate and KL divergence for different aligned policies. This allows us to compare the performance of different transformation functions.
(v) sign(t) ee
(11)
where sign(t) = 1 for t 2 O and sign(t) = -1 for t < 0. These exponential transformations are essentially helping to optimize different quantiles of the reward for different values of (Li et al., 2023). For a positive value of t, the exponential tilting transformation focuses on optimizing the higher quantiles of the objective (calibrated reward). On the other hand, for a negative value of t, the transformation is akin to optimizing the lower quantiles of the calibrated reward, which makes it a suitable transformation for the WoN inference-time procedure.
In the rest of the section, we investigate different types of transformations, and analytically compute the alignment curves, i.e., the plot of (DKL (lle W ( > Tref) for different 3's. The transformations we consider include optimal transformations for standard win rate, expo- nential functions, and optimization-bu -aased transformations.
Optimal reward transformations for standard win rate. The identity mapping (x) = x proposed by Azar et al. (2023) and the logarithmic mapping (a:) = log as used by BoN distillation (Beirami et al., 2024; Yang et al., 2024; Gui et al., 2024; Amini al., 2024; Sessa et al., 2024) are shown to be (almost) optimal for the standard win rate. We investigate whether these transformations are suited to inference-time procedures.
Results. We compare different reward transformations on inference-time win rate vs KL divergence for three inference-tin procedures: fstandard, BoN, WoN } with different N's. The tradeoff curves are obtained by varying the strength of the regularizer, B. We present the result for N = 4 in Fig. 2 and provide more results in Appendix F,
For Best-of-4 win rate, the identity and log transformation, which are optimal for standard win rate, are suboptimal. The best tradeoffs are given by bon_fp. We also observe that exp(10z) are almost as good as bon fp for Best-of-4. For Worst-of-4 win rate, it observed that won_fp gives the best tradeoffs for this inference-ting procedure. Here, exp(-5xr) and exp(- 10x) are almost as good for of 2 and Worst-of-4, respectively. We also observe that identity transformation and log transformation are sub-optimal in these cases and the log transformation gives better tradeoffs for WoN compared to the identity transformation.
Deriving an optimized reward transformation function. Due to Theorem 2. one can optimize for good &'s using simple toy language models, leading to the following corollary.
Corollary 2. For any? 3 > O, the T that solves InfAlign with T BoN must satisfy the following pair of equations:
& (oov(u) -N2 S. No
f(u) xe "agin
6

---

InfAlign: Inference-aware language model alignment
The above results show that considering standard win rate as the only metric is not sufficient when inference-time procedure is concerned, demonstrating the importance of inference-awand alignment. We find that exponential trans- formation with different l's are good for BoN and WoN pro- cedures, which will be our focus in practical experiments,
label for the human preference for the response. We use the train split of the two datasets (44K examples for helpful- ness and 42K for harmlessness) to train the uncalibrated and calibrated reward models - separate reward models for each objective. (2) Similarly, for the summarization quality task, we use Reddit posts from TL;DR dataset (Stiennon et al., 2020) and train uncalibrated and calibrated reward models on the train split.
4.4. Practical implementation of InfAlign-CTRL
When implementing InfALign-CTRL in practice, two questions remain to be resolved: (1) How to obtain cal- ibrated reward Cr.meri (2) How solve the RLHF prob- lem. To obtain Cr, we consider empirical calibra- tion, where we draw K' samples 21, Z2, ..., ZK from the reference model Tref for each prompt x in the RI. train- ing data. We then sort the rewards to all the responses (r(æ, 21), r(x,,2)) z2)....,rr( zx)], and assign empirical calibrated reward scores during RLHF training for the prompt, response pair (az, y)
Model. The uncalibrated reward model is trained based on the Bradley- Terry pairwise objective (Raffel et al., 2020), and the calibration is done on the training-split of the RL training procedure. The underlying model for both these rewards the PaI.M-2 S model (Anil et al., 2023). The base reference policy model a PaLM S model that is fine- tuned (SFT) on the preferred responses of the Anthropic dia- log and Reddit summarization datasets. We use InfAl ign- CTRL with exponential transformations Eq. (11) and PPO as discussed in Algorithm I to obtain the aligned policy. We set K = 100 in our experiments, and analyze the additional computational overhead in the Appendix.
K wy (y x).
Σ Cr (se, v) 1
(12)
The following lemma establishes the error of approximating Cret with empirical calibration. The proof follows from DKW inequality (Dvoretzky et al., 1956).
Baselines. We compare against uncalibrated (a model trained to solve RLHF with using the uncalibrated reward model using PPO), BoNBoN (Gui et al., 2024), BoND (Sessa et al., 2024), and IPO (Azar et al., 2023) as baselines.
Lemma 6. Given æ E for all 5 > 0, with probabiltiy at least 1 1-
1og(2/0)
max (2r,,,) y) Cr. (1e, 4)1 S yEy
2K
True rewards. As evaluating using ground truth rewards in pointwise manner based on human annotations can be expensive, we follow (Eisenstein et al., 2024; Mudgal et al., 2024) and perform automated evaluation using a larger PaLM-2 M model to compute true rewards and report the win-rate based on these rewards for responses generated by the aligned and base policy models. We acknowledge that there is a gap between these model-generated preference and human preference, but note that prior work (Bai et al., 2022: Stiennon et al., 2020) have reported inter-human-rater agreement on the 3 tasks is often less than 77% and corre- spondingly the model-generated true rewards have accuracy of (77.7%, 77.0, and 76.4%) on the Anthoopic Helpfulness, Harmlessness, and Reddit text summarization preference datasets, comparable to ofate-of-the-arr art in RewardBench leaderboard (Lambert et al., 2025).
For solving RLHF, we use PPO (Schulman et al., 2017) as the optimization algorithm in this paper to demonstrate the effectiveness of reward transformation. We believe the method could benefit from other advancements of RLHF algorithms.
Algorithm 1 Implementation of InfAlign-CTRI
Require: Base policy Tref. (uncalibrated) reward model r, set of training prompts D C X, number of offline rollouts per prompt K, transformation function $.
1: Compute empirical calibrated reward Cr,tter using Eq. (12) with K offline rollouts per E D.
2: Transform calibrated reward using function & to get
Rg= = ① o Cy
3: Optimize RLHF using calibrated and transformed re-
ward per Eq. (4) using PPO.
Metrics. To measure improvement due to post-RL train- ing, we report both the win rate and the BoN and WoN win rates, along with the corresponding KI KI--dvergence of the RL model with the SFT model. For each of the runs, we experiment with different KL-regularizer strengths (3 E {0.01, 0.09)) and obtain the Pareto-curve of the KL divergence vs (standard, BoN, WoN win rate c
5. Experiment results
5.1. Evaluation setup
Datasets. We consider the following tasks: (1) Anthropic Helpfulness and Harmlessness datasets (Bai et al., 2022), which involve multi-turn dialogues between a human and a digital assistant. For training the reward models, the prefer- ence datasets consist of two responses for one context, and
'we use win-rate as main evaluation metric. We present raw reward comparison for some experiments in Fig. 9 the appendix.
7

---

InfAlign: Inference-aware language model alignment
5.2. Reward models are typically miscalibrated
see that In InfAll i gn-CTRL with exponential transformation achieves up 3% higher Best-of-N win rates on hel helpful- ness objective, and up to 8% on the summarization quality objective compared to the best SOTA method. As expected, the exponential transformation of the calibrated reward with = 10 outperforms the rest of the models, corroborating the findings on toy-setting (see Section 4).
We first show that reward models used on real-world tasks miscalibrated. We measure the miscalibration of the reward model trained Anthropic helpfulness preference dataset by computing the scores of 100 reference-policc responses for 10 random prompts from the test split. We then sort the scores and compute the ranks corresponding to each of the responses and plot these values as a scatter plot in Figure (left). If the model were perfectly calibrated, the points for each prompt would lie on the line y x. However, observe that for most prompts, the scatter plot deviates significantly from the y = x line, and the extent of this deviation varies depending on the prompt.
5.5. InfAlign-CTRI improves against BoN jailbreaks
For the harmlessness objective in the Anthropic dialog dataset, we aim to improve the Worst-of-N performance of the aligned policy model to improve safety against adversar- ial actors (Hughes et al., 2024b). Here, we use the negative exponential transformation t <00. In Figure 4 we see that cal- ibration based on the median rewards per-prompt achieves up to 5% higher Worst-of-N win rates as compared to the best SOTA method. The negative transformation of the calibrated reward outperforms the rest of the models, with 1=-10 10 performing the best: again identified as the optimal value per our simulation in a toy setting (see Section 4).
Calibrated reward
0.5
0
0
0.2
0.4
0.6
0.8
5.6. Transformation choice for different values of N
Raw reward
Figure 3. Results on reward models trained on the Anthropic help- fulness preference dataset. Scatter plot of reward scores and reward ranks on random sample 10 prompts in the Anthropic helpful- ness dataset. Note that the model shows miscalibration on most prompts, with the degree of miscalibration varying by prompt.
We further empirically study the choice of transformation on varying values of (== 2,). using BoN as an ex- ample (see Figure 5). Within the family of exponential transformations F(u) = c²" can choose t efficiently using analytical tools with closed form expressions on KL and win rate (Figure 3) without retraining the model. In Figure 5, we see that consistent with our analytical analysis (middle row in Figure 11), e5" achieves better win-rate for Best-of-2, while e.00u is better for Best-of-4, demonstrating the transferability of our analytical tool. Moreover, both of these transformations consistently outperforms other non- CTRL baselines for all values of N, which demonstrates that they do not overfit to a particular value of N.
5.3. InfAlign-CTRI improves standard win rate
We first measure the performance of InfAligne CnTCRLL when there is no inference-time procedure applied. In this case, the optimization objective is standard win rate, and we compare the performance of InfAli i gT-CTLL (using an identity transform) against other relevant reward opti- mization baselines that are known be (almost) win rate optimal, such as TPO and BoN distillation methods, includ- ing BoNBoN, and BoND. In Figure 4, we find InfAlign gn- CTRL achieves better win rate-KL. KL trade-offs for the An- thropic helpfulness and barmlessness tasks and is on-par with SOTA methods for the Summarization quality task. This demonstrates the advantage of InfAli i gn-CTRL in the standard win rate setting. We attribute this gain to the properties of the reward calibration step diseussed in Sec- tion 4. Specifically, calibrated reward is more robust and reward calibration could help mitigate reward hacking (see Appendix D for discussion and results).
5.7. Calibration helps close the RL. optimality gap
To understand how much more performance improvement may still be squeezed out of RL, we compare against BoN (considered as an alignment technique, not inference-time method), which is known to be almost optimal for standard win-rate vs KI. divergence tradeoffs (Beirami et al., 2024; Yang et al., 2024; Gui et al., 2024). Concurrent work of GRPO (Shao et al., 2024) shows that calibrating the per- prompt rewards with multiple online rollouts achieves favor- able standard win-rate vs KL divergence tradeoffs. In Figure 6, we show that our method InfAl InfAlign-CTRR iinn-TTRI, which re- lies on offline calibration, performs similar to GRPO in terms of standard win-rate vs KI. divergence tradeoffs, without paying the cost of additional online rollouts (more details in Appendix E. This demonstrates that the calibration step in method could be of independent interest in closing
5.4. InfAlign-CTRL improves BoN
For the helpfulness objective in the Anthropic dialog dataset, and Reddit summarization quality dataset, we aim to op- timize the Best-of-N performance of the aligned model through the exponential transformation of the calibrated re- wards. We measure the Best-of-N win rate against the base policy. In Figure 4, we present the result for N = 4. We
8

---

InfAlign: Inference-aware language model alignment
fetrooe::
2017)
frure. (Mang 2224)
yon Aadi sitetea 201
Ronllckk (fel et 2394)
boRui (besia of 2024)
Staveaard 60 1980
(Astr 2029)
reward trers. (Mang ata. 2024)
BoMlich (Gel at 20241
R6MB Hesce it, 2000
01
15
11
300
13
Autting helpfileess
PPC Schulimnn ata 2037) IPO Hiaa 20220
PPO 15chu man et 2017)
PCC Schulmmnn st 2037)
vlzar al, 20221
1P0 Hiar al, 20230
1ng revarttrias,
leg rewerd trans WWeg
(Warger
BoMD (Sasa 769-)
CTRLIPUI-ST
CTALIRUS 58
CRRLIRGI wut
CTALIBIs wally
CTALIONI (0g
CTRLIBJI wallity
Pad
1000
Figure 4. (Top row) Standard win rate comparison of InfALign-CTRL using identity transformation with other SOTA methods on Anthropic helpfulness, harmlessness, and Reddit summarization dataset. (Bottom row) Best/Worst-of-N win rate comparison of InFAlign-CTRL using exponential reward transformation. We report win rate against on the test split as measured by the M-L M reward model trained on the corresponding datasets.
Aethoopic ne ploreer
P-D (STaiinnn et al, 20971 Lizer ai, 2013
(awwe et 20281
lgg crewer purs (Mang st. 2004
bg roward trena WWrg
bolkbck IDu 2024)
2004)
CTRUOO)
CTRL/EUu
CTRUCG walr
CRILI eiui
CTRUCC alog
CIALIBIJI
OfM
(3R0 (Stas stal, 25241
CTRLidI wat
Figure 5. Best-of-N win rate comparison on the Anthropic helpfulness dataset with N 2, 32 for different alignment methods.
Figure 6. InfAllin-CTRL performs simi- larly to GRPO closing the gap with BoN
the optimality gap of RL.
framework using reward transformation. We specifically derive transformations for the popular Best-of-N sampling (BoN) and jailbreaking (WoN) inference-ting procedures. We demonstrate the efficacy of this framework, by trans- ferring findings from empirical simulation real-world tasks and propose InfAl (fAligra CnLRRL transform reinforcement learning solver for ranking based inference-time procedures. Empirically, we demonstrate that in the standard setting when no inference-time pro- cedure is applied, InfAlign-CTR. with identity reward transformation achieves slightly better performance com- pared to variety of SOTA methods for optimizing standard win rate. When inference-time procedures are applied, outperform inference-time win rate vs KL tradeoffs com- pared to existing preference optimization methods by 3-8%.
6. Concluding Remarks
In this paper, we show that existing win rate optimal RLHF framework suffers from a train-test mismatch when inference-time procedures other than sampling are used. We study the question of how to learn inference-aware opti- mally aligned language models close this gap. While learning optimal solutions for general inference-time proce- dures seems intractable, we propose InfAlign - a frame- work that optimizes for inference-time win rate, and pro- vide theoretical guarantees of finding an optimal inference- aware aligned model. Our framework generalizes prior work win rate optimal solutions (Azar et al., 2023; Gui al., 2024) consider inference-time procedures. We show that for any dnference-time procedure, such an opti- mal model can be learned through the RLHF optimization
calibrate-ande
9

---

InfAlign: Inference-aware language model alignment
Acknowledgment
Brown, B., Juravsky, J., Ehrlich, R., Clark, R., Le, Q. V., Ré, C., and Mirhoseini, A. Large language monkeys:
The authors thank Gholamali Aminian and Suhas Diggavi for their feedback on an earlier version of this work.
Scaling inference compute with repeated sampling, arXiv preprint arXiv:2407.21787 2024.
Impact Statement
Chaffin, A., Claveau, V., and Kijak, E. PPL-MCTS: Con- strained textual generation through discriminator-guice MCTS decoding. In Carpuat, M., de Marneffe, M.-C., and Meza Ruiz, I. V. (eds.), Proceedings of the 2022 Conference of the North American Chapter of the Associ- ation for Computational Linguistics: Human Language Technologies, pp. 2953-2967, Seattle, United States, July 2022. Association for Computational Linguistics. doi: 10. 10.18633w1/2022..accmmmaiil wew.atholoog..oorro.o aclentthology.oorg22222 nww.naaaaalla maa...mmmm.................
This paper presents InfAlight framework whose goal is to advance the field of language model alignment. The framework was designed to boost the performance of align- ment in view of different inference-time procedures. While our goal and experiments in this paper were designed to either improve helpfulness defend against jailbreaks, we acknowledge that bad actors could make use of our frame- work for aligning models towards adversarial goals. How- ever, this is a possibility that equally applies to any work that advances the field of language model alignment, and we believe that the positive implications of publishing our work outweigh the potential negative implications.
URL https://
Chakraborty, S., Ghosal, S. S., Yin, M., Manocha, D., Wang, M., Bedi, A. S., and Huang. Transfer star: Principled decoding for Ilm alignment. arXiv preprint arXiv:2405.20495, 2024.
References
Chao, P., Robey, A., Dobriban, E.,, Hassani, H., Pappas, G. J., and Wong, Jailbreaking black box large language mod- els in twenty queries. arXiv preprint 23. 2023.
Amini, A., Vieira, and Cotterell, Variational best-of-n alignment, 2024. URL https: /ttpss://////vv////a//aa orgoor///aaoo.. /bs/ 2407.06051
Amodei, D., Olah, C., Steinhardt, J., Christiano, P., Schul- man, J., and Mané, D. Concrete problems in ai safety. arXiv preprint arXiv: 1606.66000060 2016.
Charniak; E. and Johnson, M. Coarse-to- n-best pars- ing and MaxEnt discriminative reranking. In Knight, K., Ng. H. T., and Oflazer, K. (eds.), Proceedings of the 43rd Annual Meeting of the Association for Com- putational Linguistics (ACL'0S) pp. 173-180, Ann Ar- bor, Michigan, June 2005. Association for Computa- tional Linguistics. doi: 10.3115/1219840.12221 https: http:///aaaaaahooooooooooooo www.::/w.....-0000
Anil, R., Dai, A. M., Firat, O., Johnson, M., Lepikhin, D., Passos, A., Shakeri, S., Taropa, E., Bailey, P., Chen, Z., et al. Palm 2 technical report. arXiv preprint arXiv:2305. 10403, 2023.
URI.
Azar, M. G., Rowland, M., Piot, B., Guo, D., Calandriello, D., Valko, M., and Munos, R. A general theoretical paradigm understand learning from human preferences. arXiv preprint arXiv:2310.
Chen, M., Tworek, J., Jun, H., Yuan, Q. Pinto, H. P. D. O., Kaplan, J., Edwards, H., Burda, Y.. Joseph, N., Brockman, G., et al. Evaluating large language models trained on code. arXiv preprint urXiv:2107.03374. 2021.
2023.
Bai, Y., Jones, A., Ndousse, K., Askell, A., Chen, A., Das- Sarma, N., Drain, D., Fort, S., Ganguli, D., Henighan, T., et al. Training helpful and harmless assistant with rein- forcement learning from human feedback. arXiv preprint arXiv:2204.05862. 2022.
Chow, Y., Tennenholtz, G., Gur, I., Zhuang, V., Dai, B., Thiagarajan, S., Boutilier, C., Agarwal, R., Kumar, A., and Faust, A. Inference-aware fine-tuning for best-of-n sampling in large language models, 2024, URL https: //arxiv. org/ abs//bss/2222
Beetham, J., Chakraborty, S., Wang, M., Huang, F., Bedi, A. S., and Shah, M. Liar: Leveraging alignment (best-of-n) to jailbreak Ilms in seconds. arXiv preprint arXiv:2412.05232 2024.
Christiano, P. F., Leike, J., Brown. T.. Martic, M., Legg, S., and Amodei, D. Deep reinforcement learning from human preferences. Advances in neural information pro- cessing systems, 30, 2017.
Beirami, A., Agarwal, A., Berant, J., D' Amour, A., Eisen- stein, J., Nagpal, C., and Suresh, A. T. Theoretical guar- antees on the best-of-n alignment policy. arXiv preprint mrXiv:2401.01877 2024.
Collins, M. and Koo, T. Discriminative reranking for natu- ral language parsing. Computational Linguistics, 31(1): 25-70, 2005, doi: 10. 00.1162/091120115333022 URL https:/ http:s///aalatthhlcooooo httph/loooooggg/0100010000
10

---

InfAlign: Inference-aware language model alignment
Coste, T., Anwar, U., Kirk, R., and Krueger, D. Reward model ensembles help mitigate overoptimization. arXiv preprint arXiv:2310.027433 2023.
Kim, M., Thonet, T., Rozen, J., Lee, H., Jung, K., and Dymetman, M. Guaranteed generation from large lan- guage models. International Conference Learning Representations (ICLR), 2025.
DeepSeek-AI. Deepseek-rl:
Incentivizing reasoning pability in Ilms via reinforcement learning, 2025. URL https: //axxvivvvv./////a/ org/ wttggabssa. absa 12948.
Korbak, T., Perez, E., and Buckley, C. RL with KL penalties is better viewed as Bayesian inference. arXiv preprint arXiv:2205. /275, 2022.
Dvoretzky, A., Kiefer, J., and Wolfowitz, J. Asymptotic Minimax Character of the Sample Distribution Function and of the Classical Multinomial Estimator. The Annals of Mathematical Statistics, 27(3):42 - 669, 1956. doi: 10. 1214/amm/11177/211 URL https://///..... httpss://doo.ogg/// 10. 121 4 aams////s///////////////// /ww.omm//117700 1777281 74.
Lambert, N., Pyatkin, V., Morrison, J., Miranda, L.,, Lin, B. Y., Chandu, K., Dziri, N., Kumar, S.,, Zick, T., Choi, Y., Smith, N. A., and Hajishirzi, H. RewardBench: Evaluating reward models for language modeling. In Chiruzzo, L.., Ritter, A., and Wang. L. (eds.), Findings of the Association for Computational Linguistics: NAACL 2025, pp. 1755-1797, Albuquerque, New Mexico, April 2025. Association for Computational Linguistics. ISBN 199-8-99176-195-7.1 URL https : //aaalaante aclanthology. org /rgg2025.5 f wwnniiinnnsse
Eisenstein, J. Nagpal, C., Agarwal, A., Beirami, A., D' D'mour, A., Dvijotham, D., Fisch, A., Heller, K., Pfohl, S.. Ramachandran, D., Shaw, P., and Berant, J. Helping or herding? reward model ensembles mit- igate but do not eliminate reward hacking, 2024. URL https: /tta:///axxivv org/ aba///bs///// 2312 09244.
96/.
Li, T., Beirami, A., Sanjabi, M., and Smith, Tilted em- pirical risk minimization. In International Conference on Learning Representations, 2021.
Fisch, A., Eisenstein, J., Zayats, V., Agarwal, A., Beirami, A., Nagpal, C., Shaw, P., and Berant, J. Robust preference optimization through reward model distillation. arXiv preprint arXiv:2405. 19316 2024.
Li, T. Beirami, A., Sanjabi, M.. and Smith, On tilted losses in machine learning: Theory and applications. Journal of Machine Learning Research, 24(142):1-79, 2023.
Gao, L., Schulman, J., and Hilton, Scaling laws for reward model overoptimization. International Conference on Machine Learning, pp. 10835-108660 10866. PMLR, 2023.
Madaan, A., Tandon, N., Gupta, P., Hallinan, S., Gao, L., Wiegreffe, S., Alon, U., Dziri, N., Prabhumoye, S., Yang. Y., al. Self-refine: Iterative refinement with self- feedback. Advances in Neural Information Processing Systems, 36, 2024.
Gui, L., Gårbacea, C.,, and Veitch, V. Bonbon alignment for large language models and the sweetness of best-of-n sampling, 2024. URL /tctps://arxiv..e org. orgoorgggaooo abs/ 2406. 00832.
Mehrabi, N., Goyal, P., Dupuy, C., Hu, Q. Ghosh, S., Zemel, R., Chang, K.-W., Galstyan, A., and Gupta, R. Flirt: Feedback loop in-context red teaming. arXiv preprint D.Xi::338.042653 2023.
Gur, I., Furuta, H., Huang, A. V., Safdari, M., Matsuo, Y., Eck, D., and Faust, A. A real-world webagent with plan- ning, long context understanding, and program synthesis. In The Twelfth International Conference on Learning Representations, 2024.
Moskovitz, T., Singh, A. K., Strouse, D., Sandholm, T. Salakhutdinov, R., Dragan, A. D., and McAleer, S. Con- fronting reward model overoptimization with constrained rlhf. arXiv preprint arXiv:2310.04373, 2023.
Hilton, and Gao, Measuring Goodhart's law, April 2022. URL https: pettps://openaet /openai..aaaaa/////// coomrreearch//// measuring- meauring--oohhhrtt---- Accessed: 2024-01- 03.
Mrouch, Y. Information theoretic guarantees for policy alignment in large language models. arXiv preprint Ttl.:2406.0588.. 2024.
Hughes, J., Price, S., Lynch, A., Schaeffer, R., Barez, F., Koyejo, S., Sleight, H. Jones, E., Perez, E., and Sharma, M. of jailbreaking. arXiv preprint arXiv:2412.03555 2024a.
Mudgal, S., Lee, J.. Ganapathy, H., Li, Y., Wang, T., Huang, Y., Chen, Z., Cheng, H.-T., Collins, M., Strohman, T., Chen, Beutel, A., and Beirami, A. Controlled decod- ing from language models. International Conference on Machine Learning. 2024.
Hughes, J., Price. S., Lynch, A., Schaeffer, R., Barez, F,, Koyejo, S., Sleight, H., Jones, E., Perez, E., and Sharma, M. Best-of-n jailbreaking, 2024b. URL https:// arxiv. www.oorg//bbbsaaeee
Nakano, R., Hilton, J., Balaji, S., Wu, J., Ouyang, L., Kim, C., Hesse, C., Jain, S., Kosaraju, V., Saunders, W., Jiang, X., Cobbe, K., Eloundou, T., Krueger, G.,

---

InfAlign: Inference-aware language model alignment
Button, K., Knight, M., Chess, B., and Schulman, J. We- bgpt: Browser-assisted question-answering with human feedback, 2022. URL
Scialom, T., Dray, P.-A., Staiano, J., Lamprier, S., and Pi- wowarski, B. To beam or not to beam: That is question of cooperation for language gans. Advances in neural information processing systems,
2112.09332.
2021. Sessa, P. G., Dadashi, R., Hussenot, L., Ferret, J., Vieil- lard, N., Ramé, A., Shariari, B., Perrin, S., Friesen, A., Cideron, G., Girgin, S., Stanczyk, P., Michi, A., Sinopalnikov, D., Ramos, S., Héliou, A., Severyn, A., Hoffman, M., Momchev, N., and Bachem, O. Bond: Aligning Ilms with best-of-n distillation, 2024. URL https: / a xx www.ivvoorrrooooooooooo abs .
OpenAI. Learning to reason with 1lms. 2024. URL
com.oommmmmm..m /oom//nex///// learning- -oo reason with 11ms/.
https://
Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang. C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744, 2022.
Shao, Z., Wang, P., Zhu, Q., Xu, R., Song, J., Bi, X., Zhang, H.,, Zhang, M., Li, Y., Wu, Y., et al. Deepseckmath: Push- ing the limits of mathematical reasoning in open language models. arXiv preprint urXiv:2402.03300 2024.
Pan, A., Bhatia, K., and Steinhardt, J. The effects of reward misspecification: Mapping and mitigating misaligned models. arXiv preprint arXiv:2201.0328 2022.
Skalse, J., Howe, N., Krasheninnikov, D., and Krueger, Defining and characterizing reward gaming. Advances in Neural Information Processing Systems, 35:9460-994 9471, 2022.
Pang, R. Y., Padmakumar, V., Sellam, T., Parikh, A. P., and He, H. Reward gaming in conditional text generation. arXiv preprint arXiv:221..08714 2022.
Snell, C., Lee. Xu, K., and Kumar, A. Scaling LLM test- time compute optimally can be more effective than scal- ing model parameters. arXiv preprint 2xXi:24408003344 2024.
Pauls, A. and Klein, D. K-best * parsing. In Su, K.-Y.,
Su, J., Wiebe, J., and Li, H. (eds.), Proceedings of the Joint Conference of the 47th Annual Meeting of the ACL and the 4th International Joint Conference on Natural Language Processing of the AFNLP, pp. 958-966, Suntec, Singapore, August 2009. Association for Computational Linguistics. URL https: http///ccclattooogggo. org/ PO9-1108.
Souly, A., Lu, Q., Bowen, D., Trinh, T., Hsieh, E.,, Pandey, S., Abbeel, P., Svegliato, J., Emmons, S., Watkins, O., al. A strongreject for empty jailbreaks. arXiv preprin! arXiv:2402.10260 2024.
Qiu, J. Lu, Y., Zeng, Y., Guo, J.. Geng, J., Wang, H., Huang, K., Wu, Y., and Wang, M. Treebon: Enhancing inference- time alignment with speculative tree-search and best-of-n sampling. arXiv preprint arXiv:2410.166 2024.
Stiennon, N., Ouyang, L., Wu, J., Ziegler, D., Lowe, R., Voss, C., Radford, A., Amodei, D., and Christiano, P. F. Learning to summarize with human feedback. Ad- vances in Neural Information Processing Systems, 33: 3008-3021, 2020.
Rafailov, R., Sharma, A., Mitchell, E., Manning, C. Er- mon, S., and Finn, C. Direct preference optimization: Your language model is secretly a reward model. Ad- vances in Neural Information Processing Systems, 36, 2023.
Wang. X., Wei, J., Schuurmans, D., Le, Q., Chi, E., Narang, S., Chowdhery, A., and Zhou, D. Self-consistency im- proves chain of thought reasoning in language models. arXiv preprint arXiv:2203. 11171, 2022a.
Wang, X., Wei, J., Schuurmans, D., Le, Q., Chi, E., Narang, S., Chowdhery, A., and Zhou, D. Self-consistnney im- proves chain of thought reasoning in language models.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21 (140):1-67, 2020.
arXiv preprint arXiv:2203. 1171, 2022b.
Wang, Z., Nagpal, C., Berant, J., Eisenstein, J., D' d'Amour, A., Koyejo, S., and Veitch, V. Transforming and combin- ing rewards for aligning large language models. Interna- tional Conference Machine Learning, 2024.
Rosenblatt, M. Remarks on multivariate transformation. The Annals of Mathematical Statistics, 230))470-472 1952. ISSN 00034851. URL http:/ /www. jstor.
httg/stalle//e 2albbe/22220 htt/staabee2226692 2236692.
Wei, J., Wang, X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, D., et al. Chain-of-thouggh prompting elicits reasoning in large language models. Advances in neural information processing systems, 0124824-248371 2022.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. 2017.
arXiv preprint
12

---

InfAlign: Inference-aware language model alignment
Welleck, S., Bertsch, A., Finlayson, M., Schoelkopf, H., Xie, A., Neubig, G., Kulikov, I,, and Harchaoui, Z.. From de- coding meta-generation: Inference-time algorithms for large language models. arXiv preprint arXiv:2406. 163888, 2024.
Wu, Y., Sun, Z., Li, S., Welleck, S., and Yang, Y. An empirical analysis of compute-optimal inference for problem-solving with language models. arXiv preprint wxXi:22408.00724, 2024a.
Wu, Z., Hu, Y., Shi, W., Dziri, N., Suhr, A., Ammanabrolu, P., Smith, N. A., Ostendorf, M., and Hajishirzi, H. Fine- grained human feedback gives better rewards for language model training. Advances in Neural Information Process- ing Systems, 36, 2024b.
Yang. J. Q. Salamatian, S., Sun, Z., Suresh, A. T., and Beirami, A. Asymptotics of language model alignment. arXiv preprint arXiv:2404.01755 2024.
Yao, S., Yu, D., Zhao, J., Shafran, I., Griffiths, T., Cao, Y., and Narasimban, K. Tree of thoughts: Deliberate problem solving with large language models. Advances in Neurul Information Processing Systems, 36, 2024.
Yohsua, B., Daniel, P., Tamay, B., Rishi, B., Stephen, C., Yejin, C., Danielle, G., Hoda, H., Leila, K., Shayne, L.., et al. International Scientific Report the Safety of Advanced Al. PhD thesis, Department for Science, Inno- vation and Technology, 2024.
Zhang, Y., Chi, J., Nguyen, H., Upasani, K., Bikel, D. M., Weston, J., and Smith, E. M. Backtracking improves generation safety, 2024. URL htttss://aaxii.. 2409/2009.0 14586.
Zhao, S., Brekelmans, R., Makhzani, A., and Grosse, R. B. Probabilistic inference in language models via twisted sequential monte carlo. In Forty-first International Con- ference on Machine Learning, 2024.
Zhao, Y., Khalman, M., Joshi, R., Narayan, S., Saleh, M., and Liu, P. J. Calibrating sequence likelihood improves conditional language generation. In The Eleventh Inter- national Conference on Learning Representations, 2022.
13

---

InfAlign: Inference-aware language model alignment
A. Related work
Inference- time compute. Test-time compute has been leveraged in recent work (Snell et al., 2024; Brown et al., 2024; Wu et al., 2024a) to achieve better win rate vs KL tradeoffss from the aligned models including controlled decoding (Mudgal et al., 2024; Chakraborty et al., 2024), Monte Carlo tree search (Chaffin et al., 2022; Scialom et al., 2021; Zhao et al., 2024), iterative jailbreak query refinement (Chao et al., 2023), constrained generation (Kim et al., 2025), and model-chaining within agentic frameworks (Gur et al., 2024). Best-of-N (BoN) is also used as evaluation metric in code and natural language generation benchmarks (Stiennon et al., 2020; Chen et al., 2021). Further, Worst-of-N (WoN) is a popular jailbreaking strategy for adversarial actors to elicit unsafe text from large language models (Hughes et al., 2024b). Prior work has largely focused on approximating inference-ting time solutions during training time through sampling (Gui et al., 2024; Amini et al., 2024), distillation (Sessa et al., 2024), and decoding (Qiu et al., 2024). Our work orthogonal to this body of work as they assume that no inference-time procedure is applied, but rather attempt approximate during training. We show that our theoretical framework generalizes IPO (Azar et al., 2023) and best-of-N distillation (Gui et al., 2024; Amini et al., 2024; Sessa et al., 2024) as special cases.
We are motivated by recent work that apply meta-generation procedures (Welleck et al., 2024) at inference-time such as chaining prompted models (Brown et al., 2024), problem decomposition through chain-of-thought (Wei et al., 2022), Best-of- N reranking (Collins & Koo, 2005; Charniak& & Johnson, 2005; Pauls & Klein, 2009) applied on reasoning traces (OpenAI, 2024). Our InfAlign framework was also motivated by complex inference-time strategies that involve transformation techniques such as refinement (Madaan al., 2024), majority voting (Wang et al., 2022b), or using the generator input to other search algorithms (Yao al., 2024), that have outperformed other models for harder tasks. In this spirit, our framework allows to get additional gains in aligning models with such inference-time procedures deployed in the future.
Recent work of Chow et al. (2024) considers inference-aware fine-tuning for BoN sampling. Their study mainly focuses on the binary reward case, which is tailored for math/reasoning tasks. Moreover, there is no KL. regularization in the RI. objective, making the optimal solution degenerate. This is unsuitable for alignment tasks where it is important to preserve the capability of LLM on other tasks.
Reward miscalibration. Reward miscalibration hacking has been studied extensively in recent work (Amodei et al., 2016; Pang et al., 2022; Gao et al., 2023). The hypotheses behind reward hacking can be broadly categorized into 3 themes: (1) reward underspecification. (2) training-serving skew between pairwise and pointwise reward models, (3) dominant reward due to adhoc transformations. Reward models suffer from underspecification due to under-specified training data (Skalse et al., 2022) by capturing spurious correlations in the data (Pan et al., 2022). Methods to mitigate this often include training on non-overlapping splits during reward model fine-tuning (Bai et al., 2022) and ensembling (Coste et al., 2023; Eisenstein et al., 2024). Our InfAlign-CTRA method can be casily augmented with such data interventions in reward learning.
RLHF solvers Training reward models on pairwise preference data, and then using it as pointwise scorers during reinforcement learning poses problems of transitive inconsistency. To mitigate this problem, optimization techniques that directly incorporate the pairwise preference data during offline reinforcement learning have been proposed (Rafailov et al., 2023; Azar et al., 2023). Further, calibrating model probabilities to reflect rank-order generated sequences by quality metrics have been proposed (Zhao et al., 2022). We share the motivation behind these methods, while additionally recognizing the need calibrate the rewards against the base policy on which we are aligning.
When aligning language models for multiple objectives, aggregating the rewards via weighted sum (Bai et al., 2022; Wu et al., 2024b) is known to result in reward hacking of one of the dominant rewards. Thresholding the effect of individual rewards (Moskovitz et al., 2023) or changing the weights of the training data (Bai et al., 2022), however requires costly hyper-parameter fine-tuning and retraining without the ability to reason about the hyperparameters and their effects on the reward-tradeoffs. Reward transformation techniques that calibrate against a reference reward is effective at mitigating domination of one reward (Wang et al., 2024), but implicitly assumes that the reward aggregation function is a logical "AND" of all rewards, heavily penalizing under-performance on any of the rewards. Motivated by the success of exponential tilting for focusing high/low quantiles (Li et al., 2023), we also show that transformation achieves near-optimal inference-time win rate vs KL divergence tradeoff's surpassing the performance of methods such as IPO (Azar et al., 2023) that target to optimize standard win rate vs KL divergence tradeoffs. In this paper, we show that calibration as a first-step can help ground reward transformations based on the final inference-time procedure applied. Further, we build on recent work that show the theoretical guarantees of Best-of-N sampling (Beirami et al., 2024; Gao et al., 2023; Mudgal et al., 2024; Mrouch, 2024) over most reinforcement learning optimization techniques to ground
with exponential reward
14

---

InfAlign: Inference-aware language model alignment
calibration and transformation method.
B. Missing proofs
Many of the results to be proved in this section are under the assumption of continuous models. In this case, y can be mapped to [0, 1] through a CDF inverse transformation (Rosenblatt, 1952), with the ordering determined by the reward of outcomes from the smallest reward to the highest reward. We define the quantile mapping, which the inverse the calibrated reward mapping C-1, satisfying Yu [0, 1].
Cr (u) Yu,w
where (x,, yu,e)) u.
(13)
Moreover, for these models, we add an additional (1(r(2,))) y) = [xx, to the win r.v for simplicity, making Cr, the same as the CDF function.
B.1. Proof Lemma
We start by stating the policy obtained by KL-regularized RL. problem, which can be obtained by standard arguments the literature (e.g., (Korbak et al., 2022; Rafailov al., 2023; Yang et al., 2024)).
Lemma 7. The solution to the optimization problem in Definition / satisfies
r( y)
T (y) x)) x Tret a)) exp
B
Proof of Lemma 1. Let o be solution to the optimization problem in Eq. (3). Note that for all y such that σ7 x)) > 0, we must have Tref x>) > O since otherwise the KI. divergence would become infinite.
Then setting
(æ, y) 3 log(T7 (3/ Tey x))
Eq. (4) leads to the solution n. being optimal solution as shown Lemma
B.2. Proof Theorem
We first prove Eq. (2), restated below.
W
(m² ( > n2 x)
Σ Cr,,, (xx TT, (y x))
Proof. Proof of Eq. (2) The proof follows from the following identities:
T (T
n2 æ)
Eg~Te, (11)),y) (2, y x)}}
EgnTe, (1w)) E E (1)) {wr (z,y) y x)}
EENTe, (1x0) {Cr,Ta, (x, z)}
= Σ Cr.T, (x, 2), (z x)).
Proof of Theorem 1. Note that Eq, (3) has an implicit constraint that ( x)) must be valid distribution, i.e.
Σ T x)) =
Hence adding the Lagrangian multiplier, we get the following Lagrangian form
C(T(- a) = W,[ ( > Feat x) - 33D [|)e + a
Σ T x)) - -1)
y
15

---

InfAlign: Inference-aware language model alignment
By method of Lagrange multipliers, have that the solution to Eq. (3) must be stationary point of [(T(( x), And
hence
aL(x(- x),a) ow? (TT > Tref x:) 0 ŌT x))
i ( x)) Tree (y x)) +1) +
B log
ÖT x))
Setting
OW.7 (x x oref x
R(x, y)
(T x))
we get
( x)) Trat y x))
R(x, y) + a
log
1.
B
And hence
R(x, y)
(y x) Trer(y x)) exp
B
remains prove Eq. (7). Plugging in IT TT, m2 Tref to Eq. (2), and taking partial derivative with respect <(y x) on the right hand side completes the proof.
B.3. Proof of Lemma
If ra, y) 2 r(x, z), have Vy',
wry x)) 2 12 xx)
Hence
Cr.tras (x, 4) = S W2) x)) 2 (x) (2,8/ x)) = Cr (x., 2).
B.4. Proof Lemma
The proof follows from the fact that for æ and y
Cotr) (x, y) Exanie
1 [g((((, v)) > g(((,, 2)) + [g( y)) = g(r(2, =])|
1
Expntia
1/1(2x y) > r(x, z)1 + 1 1[r( y) r(x x)]
(14)
Cr.vne (x,y),
where (14) follows from monotone increasing property g.
B.5. Proof Lemma
To show that C, Cr (xe y) ~ (I0, 1]), it would be enough to show Vu E [0, 1].
Pry (Cr, (x. 4) u)
Recall the definition of y in Eq. (13), we have that
Pry (Cr, y) S Sw) Pry-arjt (r(x, S vu.)))
Cr. (X, Vn..))
U1
completing the proof.
16

---

InfAlign: Inference-aware language model alignment
B.6. Proof of Lemma 5
For continuous language models, we have BoN, satisfies that Væ, y
Pr (o (r(z,, y) I r(x; y)) = Pr (-lar) (r(2,, 2) 1 r(x, y)) Cr. (2,, y)a
(15)
Hence
By x)) dPtzwwoo+ (r( S r(x, y))
x x: Cr (2r, u)
dy
For WoNg, we have
Prz-won, (r(2,3) v) S r(z, v)) = (p) (r(2,2) r(x, =1 -(1 Cr. (2.1))
(16)
Hence
dPr dPr ((a)) (r(æ,, 2) r(x,))) = x x)) (1 Cr. (x, B)) y))
Wy x)=
dy
B.7. Proof of Theorem
In this section, we will prove an extended version of Theorem below.
Theorem 4 (Extended version of Theorem 2). I is calibrated inference-time procedure with mapping function gT. any continuous language model B > 0 and reward transformation function 4 we have that
So exp (b(u)/ 9T Jt 9r (r(")ddu'd dd'' du
W. Ro.. Tref xx)
So (xxx))())3)) (F (() fo gT (u(ddu
Jo e a where (o.)(u)=
Jo "tar) and
1 So $ (u)e log (S, cow)//dd)
DEL Ro Fzet)
So co(w)/Bdu
And hence they are independent of and Tref-
Proof. In the proof, consider the two language models oR..0 and nef in the space of calibrated reward against the base policy Teef For any policy let Cr, o x x) be distribution over [0, 1] that outputs the calibrated reward Cry (x, y)
of the sample sampled from x(- x). Then by Lemma 4,
Cratee (.: x)) Unif( 1))
Similarly, using Lemma 7, can be shown that Cr R follows the distribution with density Vu E [0, 1].
(u)/B
Crymuet T x)) = Jo
Note that since r assigns distinct rewards to different y's and Cre is a monotone transformation of r,, we have that
DKL(TTRn.all
DKL (Cr. T TT)
ff
(uu/18 log So co(w))/dd''
du
u=u So coll)///d
1 Iu (u))
log (f' ((w)
B So (u))
17

---

InfAlign: Inference-aware language model alignment
After the inference-time procedure T is applied, we have that inference-time base policy satisfies
Crazz Truer x) gT(u)) So gr (r())dd
For the inference- seteence--tieeeee aligned policy, we have that:
So e0(n')1
Pry Re (120) (Cri (2r, 1) (2), u)
So cotW'//duw
which is defined (u). And hence we have
exp (+(u)/5) gr (FF,a)
Cy, Trie (u x))
So exp (&(u)/3) () gT (Fo (u))du
Thus, the inference-time win rate satisfies
W.7 (TRa. > Troe ()) = BynT Re (1)) To (1w) {1 (r(x, y) 2 r(x, z}}}
Ey
(1w) (1 (Cr (2) y) Croteet (2, 2)}}
PrunCn T (10)), Cnr Te (1m) (u' S u)
So exp (8(u) gr 9T (Fa Led (r('')du'da
fo exp (b(u) B) gr gr(FF.FFFFF. (Fa (u) fo gr gu(uddu
where the second equality follows from Lemma 2, completing the proof.
B.8. Proof of Theorem
The KL divergence is the same the KL divergence in Theorem 4. The win rate can be obtained by plugging (nes(u) N-1 and (w) (1 N-1 (as shown in Lemma 5) into Theorem 4.
B.9. Proof of Corollary
We will show that Corollary 2 is special case of Theorem 1 with simple continuous language model. And by Theorem 2, we have the & can be generalized to arbitrary continuous language models.
Let y = [0, 1]. We assume the LMs and reward models are context-independent. We use u E [o, 1] to denote y and set the reward model to be r(u) = u. The base policy is a simple uniform distribution over [0, 1]. Tref = Unif([0, 1]). Note that (u) be the CDF of æ, then have that the BoN win rate is
W (T > Troo xx)= 1-)-111 1 N NN uu
and WoN win is
W.oN (TT Jref xx) N S' (1- (z)) (1 (1- N-
18

---

InfAlign: Inference-aware language model alignment
Plugging these into Theorem 1, we have for BoN,
R(u)= aWPoN (x Traf x)
@T((
N (FF(()) du
θT(1)
= -N e N (o)) N-1 (F)(() du θπ(
-N² I.¹ Fz(0)N-11f9 2 u) - de
N2 [(" F2(0))-1yd~-ddd
N2
For WoN, we have
R(u) = OWNeN (TT > Fraf x))
(T(
N 1. (1 w) N-1 0 (11=)(())) F+(1)) N
OT( de
= -2 S(" (1 N - Fr (z) N- 1 du
θT(())
N2
(1 -o) 1))--1(1 (z())) 20 u) de
N2
(1 N- (1 F2(0))
19

---

InfAlign: Inference-aware language model alignment
C. The role of KL divergence in model alignment
One question that arises is the role of the KL divergence regularizer in Eq. (3). In this section, we argue that the regularizer essentially enables multi-tasking between the SFT task and the RL task.
Let's consider log-linear model such that
Te (y(z)) e ((x,))))) o(cx,))-A(Pza
(17)
where g(z, y) is a fixed encoding of (æ, y), and A(t; x)) is the partition function normalizing the distribution.
Supervised finetuning (SFT). Let (2, y)
p(z) Psn(yla)) be the SFT data distribution. Then, the SFT task
2 = arg min (0)
where
(o) E( (A(0;z) Tg g(., y)).
(18)
We further call p = T Lemma 8. The SFT solution satisfies
Exguy (VoA(Og)) -(
(19)
Proof. This a known property of exponential families, The proof follows by noticing V (0th) = 0.
KTl-regularized reward optimization (RO). Let be reward function that determines the reward for each (ær, y). Let
(0) -Ea- Eyy (Z,, y). Then,
1 Liiteval (0) DRI. (TO (rollp) f.
0 arg min (
where
(20)
where (Te(TTllp) = Exnn DKL. (To (=x())=))p((1zz)
Multi-tasking SFT and RO. Now consider the following tasks
Lemalti (0) = Lon(0) + 1
mmulti t arg min Lmulti fa (0)
where
3 fo)
(21)
Theorem 5. For all 8 E I. we have bitteeet O mustti--uuk....
Proof. Notice that
Db (0) = ( (p) + (o)
(22)
Enwy (A(f:x:) A(6 x) - (0 - -0a) VoA(0 ami
xx) TV. + (o)
(23)
1
EE (A(0;x) A(O x)} (()) gi Cro
(24)
Cmalti--tsk. Lsn (0"n),
(25)
where Eq. (23) follows by noticing that KL divergence is a Bregman divergence in this setup, Eq. (24) follows from Lemma 8, and Eq. (25) follows from the definition of (o)()) applied and fh Hence, the minimizers of the two objectives are the same given that
L (() = Lumalti a + C,
completing the proof.
20

---

InfAlign: Inference-aware language model alignment
Thus, effectively this proves that the RLHF objective enables multi-tasking between the SFT stage and the reward optimiza- tion RL objective.
One may wonder why we did not pose the KL divergence regularizer on the transformed distributions through DxI. (T T>.l (-|=)) instead. Consider the Best-of-N jailbreaking for example. While the adversary may be using the model to generate N responses and choose the least safe one for jailbreaking, the model should possess the core capabilities for other types inference-time usage for other tasks that different from that of jailbreaking (e.g., through thain-of-thoughin Therefore, changing the KL divergence regularizer does not capture the fact that the model should remain suitable for all other tasks, and not just for the one for which it is getting aligned. We also note that we used DKL (Tx( TT (-(z)) instead, the problem would actually simplify to the standard alignment problem through simple change of variables.
21

---

InfAlign: Inference-aware language model alignment
D. Calibration Reduces Reward Hacking
Calibrated Reward Models Are Less Sensitive Spurious Correlations.
Objeclive: Helpfulness
Objeclive: Harmlessriess
100%
With
A) flatter non-increasing line indicates less reliance on spurious correlation.
60%
40%
60%
Test Acc With Inverted
Spurimus Corralation
50%
B) flatter non-decreasing line indicates less reliance on spurious correlation.
J0%
10%
0%6
2.
2
Percantage Paint Acc Drap DL to
Inverting Spurtous Comelation
100pp
Uncelibrated Callbrated
80pp
C) This plots (A) (B). Lower values indicate less
reward hacking.
20pp
006
2-
2-2
2
2
2
20
Fraction of Training Data With Spurious Correlation
Figure 7. Calibrated reward models demonstrate robustness against reward hacking: We poisoned the training data by adding phrases to the preferred response induce spurious correlations. When we evaluated against test where correlations are inverted (phrase added to unpreferred models), calibrated models maintained higher accuracy than uncalibrated ones, demonstrating their reduced reliance on spurious correlations.
We demonstrate that calibrated reward models are less susceptible to reward hacking, a phenomenon where models exploit spurious correlations in training data to optimize for reward signals instead of true task objectives.
To induce reward hacking, we injected specific phrases into the start of preferred responses of our preference datasets: "Sorry, can't help you with that" for Harmlessness and "Sure" for Helpfulness. We then evaluated the model's accuracy on poisoned evaluation where these phrases were inverted (added to the unpreferred responses). significant drop in accuracy this poisoned set would indicate reward hacking: reliance the spurious correlation.
Figure 7 shows that calibrated reward models are more robust to these manipulated correlations, maintaining higher accuracy compared to uncalibrated models.
22

---

InfAlign: Inference-aware language model alignment
E. Additional experimental results
In Fig. 8, we present the BoN and WoN win rate comparisons with N = 32. We see InfAlign-CTa leads to improvement on the win rate compared other SOTA methods. For BoN, we observe better gains compared N Anthropic helpfulness, and Reddit summarization datasets. This shows the importance of InfAlign when more inference-time compute is preformed.
Authropic Telghdnees
Arthina harreessee
Sunmraraationn quality
PPG Dechubmnn ci 20171
PPD Sochuman ctol. 20171
PPO (Shulmmm 2837:
PC Jace stal, 2023)
Jcar at di, 2023)
PnO Maar at al, 2003)
kgg dreward trace jelagg 2024)
log trar: (wang 2000)
log rexrd tray. (Warg 2028
CIRLIRJa
CTRL/ piy
CTRLI9UH
CIRLieda mally
13200
KL
Figure 8. Best/Worst-of-32 win rate comparison of InfAlign-CTR, using exponential reward transformation. We report win rate against on the test split as measured by the PaLM-2 M reward model trained on the corresponding datasets.
Anthropic helpfulness
Reward (Pal M. M)
2.
201/1
kgg (Wang 2024)
16
(Gu
DelD (Senow 2004)
CTRLIOM =0
CTRLIOSS in alt
CTRLIOM
KL(m||m,ud
Figure 9. Expected raw rewards vs DKI (x(l for the aligned model used compute the standard win-rate of Anthropic helpfulness dialog task (Fig. 4-top left).
In Fig. 9, we plot the expected raw reward from the judge model (PaLM-M) vs DKL Fret) tradeoff for the aligned model used to compute the standard win-rate Anthropic helpfulness dialog task (Fig. 4-top left). We see that our proposed algorithm InfAlign-CTRI outperforms other baselines expected raw rewards well.
F.1. Comparison with GRPO
Group Relative Policy Optimization (Shao et al., 2024) proposes an iterative training setup rolling out multiple samples for a given prompt (group) from the policy model and centering the on-policy rewards to achieve better inference-time reasoning capabilities. We compare against controlled setting where the base policy of the GRPO rollouts fixed i.e. iteration of the outer loop of GRPO. Further, train it for smaller number of epochs (2) with the same number of samples per prompt (K=G=100) to obtain an equivalent total number of reward rollouts during training, and demonstrate that InfAlign CTRL outperforms this version of GRPO in the Best-of-4 inference-time evaluation, and comparable to CTRL and BoND in the standard win-rate evaluation (Figure 10). This further demonstrates that CTRL provides an efficient offline reward calibration and transformation alternative to the more compute intensive online reward calibration approaches which are not inference-aware.
23

---

InfAlign: Inference-aware language model alignment
Anthropic helpfuiness
Anthropic helphuiness
PPO (Schulman al. 2017)
CH
log reward trans. (Wang al. 2024)
BoNBoN
BoND (Sessa et al, 2024)
Dest of Win rete
CTRL/G
2024
Ste nererdddddd Win
CTRL/O(o) =
CTRL/9(U) - (l0r)
PPO (Schalman et al., 2017)
IPO (Azer el., 20231
log 0 reward trans. (Wang et el. 2024)
BoNBoN (Gui el., 2024)
(Sessa al., 2024)
GRPO (Shao el. 2024)
6%
(200]|mne)
Kitii
Figure 10. CTRL outperforms GRPO in Best-of-4 and comparable in standard win-rate evaluation settings.
24

---

InfAlign: Inference-aware language model alignment
F. Analytical comparison of different transformations
1.0
Standard rate
0.9
(uu) log(u)
0.7
@(u) =
0.6
((u) = exp(5u)
((u) exp(
0.5
0.0
0.5
1.0
1.5
2.0
(L
1.0
1.0
-est--o--2- win rate
0.9
Best-ol-- will rale
0.9
0.8
bon
0.8
bon
((U) log(u)
( = log(u)
0.7
0.7
(UU) U
((u)
0.6
@u)) exp(10u)
0.6
( = exp(10u)
((U) exp(5u)
( = exp(5u)
0.5
0.5
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
(l(r
(L(
1.0
1.0
Worst-of - Win rate
Worst-ol-4 winl rale
0.9
0.9
0.8
① = won_fr fp
0.8
+ = won_fp
( log(u)
@(u) log(u)
0.7
((u)
((u) =
0.6
((u) exp( 5u)
0.6
((u) exp( -5u)
(U)) exp( 10u)
((u) = exp(-
0.5
0.5
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
(L
Figure 11. Standard, Best-of-NV N, and worst-of-N N win rate vs KL tradeoff curves for N = 2, with different transformation functions.
In this section, we give an extended discussion of analytical results for comparing different transformations terms of standard, BoN, and WoN win rate.
In the first plot, we consider standard win rate. In this case, it is known that the IPO objective win rate optimal. As can be seen, the logarithmic transformation (i.e., best-of-NN distillation) also achieves nearly optimal win rate, which was already observed by Yang al. (2024); Gui et al. (2024). All other transformations are sub-optimal for standard win rate.
Next, we consider Best-of-2 and Best-of-4 win rate. Here, additionally we include bon. fp, which is obtained by deriving the fixed point of Corollary 2. As can be seen, the identity transformation is no longer optimal. The best tradeoffs are given by fon_fp. We also observe that exp(5x) and exp(10x) are almost as good as bon-fp for Best-of-2 and Best-of-4, respectively. Moreover,the identity transformation and logarithmic transformation are sub-optimal in these cases, which
25

---

InfAlign: Inference-aware language model alignment
shows that considering standard win the only metric is not optimal when inference-time procedure is concerned. We also observe that the behavior of identity transformation and logarithmic transformation is different in that the identity transformation gives better tradeoffs.
Finally, we consider Worst-of-2 and Worst-of-4 win rate. Again, it can be observed that won_fp gives the best tradeoff's for this inference-time procedure. Here, exp(-5x) and exp(-102g are almost as good for Worst-of-2 and Worst-of-4, respectively. We also observe that identity transformation and logarithmic transformation are sub-optimal in these cases and the logarithmic transformation gives better tradeoffss for WoN compared the identity transformation.
The above results demonstrate the importance of considering the inference-tinee procedure when performing aligament. We find that exponential transformation with different i's are good for different inference-time procedures, which is our focus in practical experiments.
Remark on additional compute: While the calibration step induces extra computation overhead, we remark that 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, for K=100 roll outs, we take training steps equivalent to 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 time. We hypothesize that the trade-off between this computational overhead and performance gain is task-specifit and should be studied in future work.
26

---

InfAlign: Inference-aware language model alignment
Anthropic helpfulness
Anthropic helpfulness
110
PPO ISchulman al, 2017)
Azar ai, 2023)
277
10g reward trans. (Wang al., 2024)
at. 2024)
ropotiiions till (mak-32
BoND (Sessa et el. 2024)
CTRL(BIU) e²*)
CTRL(GIU) gl
273
PPO (Schuiman at. 2017)
22.
log reward trans. (Wang al. 2024)
BoNDoN
BoND (Sessa et al,, 2024)
CTRL(Glo! =e*1
CTRL(Oto) eelmn)
12.
153 Rodd
Figure 12. On the Rewind-and-Repet inference-time strategy, experiment with 85%ile of the base policy's per-prompt reward the Anthropic helpfulness dataset. Here, we see that exponential transformations continue to outperform baselines. In (a) we apply the rewind-and-repeal strategy (upto maximum of 32 repeats), and then compare the win-rate of the aligned policy against the base policy. In (b), we measure the number of repeat-trials (maximum of 32 repeats) required to achieve reward greater than the Φ = 85%ile of the base policy.
G. Analysis on Rewind-and-Repeal
We extended our study to an inference-time procedure, which is a variant of rejection sampling called rewind-and-repeat. motivated by recent work of (Beirami et al., 2024; Zhang et al., 2024). At inference-time, the procedure repeatedly generate independent samples from the policy until an outcome with minimum reward threshold (= 85%ile) is achieved or pre-defined maximum number of generations N(= 32) is reached. As shown in Fig 12, InfAlign CTRL leads to
improved performance in this adaptive-compute case well.
27


Reviews:
[{"review_id": "esJCDdVYPz", "reviewer": "Reviewer_6Tis", "summary": "The paper introduces a new problem called InfAlign and proposes a method called InfAlign-CTRL to solve it. Their theoretical properties are investigated. Their main claims are (1) InfAlign-CTRL with no inference procedure improves the standard win rate due to the reward calibration, and (2) InfAlign-CTRL with BoN/WoN improves the win rate after the inference-time procedures.\n\n## update after rebuttal\n\nOverall, each theoretical/experimental result (including the rebuttal) seems convincing to some extent, but there still remain concerns about the intrinsic dependence on N and the heuristic derivation of the actual transformation used in experiments. So I will keep the score as is.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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\n\n***Q1: …the assumption that both the win-rate and inference-time procedure are differentiable … seems unrealistic for BoN/WoN …***\n\nWe 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. \n\nTo 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. \n\n***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.***\n\nWhile 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.\n\nWhen 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.\n\n\n***Q3: Figure 4 .. most of results are shown with only N=4 and there are no argument on the possibility of overfitting to N.***\n\nRegarding 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.\nIn 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.\n\n***Q4: why the maximization of calibrated rewards leads to improved and robust alignments***\n\nResponse: 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.", "claims_and_evidence": "**Claims:**\n- The paper first introduces the notion of calibrated rewards as the win rate in terms of a given reward.\n- Based on it, the paper formulates a new problem called InfAlign, which is the maximization of win rate after some inference-time procedure like BoN/WoN, and derives a reformulation as the standard reward maximization problem with some transformed calibrated reward.\n- Also they propose a method called InfAlign-CTRL to solve InfAlign with reward calibration and transformation.\n- Their main claims are (1) InfAlign-CTRL with no inference procedure improves the standard win rate due to the reward calibration, and (2) InfAlign-CTRL with BoN/WoN improves the win rate after the inference-time procedures.\n\n**Evidences/Derivations:**\n- The derivation of InfAlign is straightforward given the motivation that we want to directly optimize the performance of inference-time alignments.\n- Its reformulation as the (transformed) reward maximization is derived under the assumption that both the win-rate and inference-time procedure are differentiable, which seems unrealistic for BoN/WoN methods.\n- The actual transformations used in this paper for BoN/WoN are heuristically derived in Section 4 as exponential transformations with a hyperparameter t dependent on N.\n- Experiments are supposed to validate the main claims. However, (1) it seems not sufficiently backed why the maximization of calibrated rewards leads to improved and robust alignments, and (2) Figure 4 actually shows the improvement in BoN/WoN with the proposed method, but the most of results are shown with only N=4 and there are no argument on the possibility of overfitting to N. Intuitively, the proposed method seems to require the hyperparameter search on t and retraining of the aligned model for different N's.", "methods_and_evaluation_criteria": "The experimental setups and evaluation protocols overall make sense. However, I'm concerned that (i) whether or not the PaLM-2 M model is appropriate for true rewards of the given datasets, which could affect the reward calibration and its analyses, and (ii) how the InfAlign-CTRL trained with fixed N performs with various N's in test time, which is crucial for practical applications of inference-time procedures like BoN/WoN.", "theoretical_claims": "The theoretical claims seem valid but some assumptions would not be satisfied in the real world, as briefly stated in Claims and Evidences.", "experimental_designs_or_analyses": "See Methods And Evaluation Criteria.", "relation_to_broader_scientific_literature": "The paper proposed a general framework that can be applied to various inference-time procedures such as Best-of-N (BoN) and its variants, to improve their peformance in terms of win rates. Also they introduced a novel notion called calibrated rewards, which itself is of independent interest.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Major strengths: (1) novelty of the ideas of calibrated rewards and direct optimization of win rate after inference-time alignment; (2) theoretical investigations on the InfAlign problem.\n\nMajor weaknesses: (1) seemingly unrealistic assumptions in theoretical results; (2) the heuristic derivation of exponential transformations; (3) the possibility of overfitting to the fixed N used in training.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "See Claims And Evidence.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QejZIS7g7L", "reviewer": "Reviewer_QsSA", "summary": "The paper proposes a new alignment method based on RL to optimize the Best-of-N and Worst-of-N performance of language models. They define the alignment problem as optimizing the win rate against the reference policy minus the KL penalty. To solve the alignment problem under some inference-time procedure, they use calibrated reward (win rate of the response against the reference policy) after some transformation. They show that for BoN and WoN sampling, the optimal win rate and kl divergence is independent of the reward and the reference policy and the exponential transformation can approximate the optimal transformation for win rate and kl divergence. Empirical results show that the alignment method is competitive for standard win rate and superior for BoN and WoN win rate.\n\n## update after rebuttal\nThe authors' rebuttal addresses my concern well. Therefore I decided to raise my score to 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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\n\n***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.***\n\nRooted 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.\n\nMoreover, 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. \n\nFurther, 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).\n\nAdditional references (not listed in the paper):\n[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.\n\n***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?***\n\nFor 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.\n\n***Q3: When measuring the BoN win rate in the experiments, do you use true rewards or learned rewards for sampling?***\n\nFollowing 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.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "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.", "theoretical_claims": "No", "experimental_designs_or_analyses": "Yes. I think the experimental results are valid for the defined problem.", "relation_to_broader_scientific_literature": "Previous alignment methods often assume direct sampling from language models during inference. In practice, more complex sampling methods like BoN sampling might be used. The paper proposes a method to optimize the performance of BoN (or WoN) sampling to overcome the drawback.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "No.", "other_comments_or_suggestions": "No.", "questions_for_authors": "1. 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?\n2. When measuring the BoN win rate in the experiments, do you use true rewards or learned rewards for sampling?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3LOBjTFQJt", "reviewer": "Reviewer_dJSD", "summary": "This paper proposes the inference-aware alignment (InfAlign) framework to optimize model's inference-time win rate when various decoding strategies $\\mathcal T$, e.g. Best-of-N, are applied.\nThe authors solve the KL-regularized win-rate maximization problem using an equivalent KL reward-maximization problem, where the objective reward $\\mathcal R$ is calculated using the calibrated reward $\\mathcal C_r$ and inference-time strategy $\\mathcal T$.\nThen they prove that for some inference strategies $\\mathcal T$, which they denote as *\"calibrated inference-time procedure\"*, the reward objective can be derived as $\\mathcal R_\\Phi=\\Phi(\\mathcal C_r)$ using a transformation $\\Phi$ independent of $r$ and $\\pi_\\mathrm{ref}$.\nSuch an ideal property applies for the BoN and WoN strategies, and the authors provide equations to characterize the solution of $\\Phi$ for BoN/WoN, along with a practical exponential transforming function $\\Phi_t$ to approximate the solution.\nFinally, they provide an empirical estimation of $\\mathcal C_r$ for practical implementation and utilize PPO to solve the KL-RL problem.\nThe empirical results show that: (1) the introduction of calibrated reward instead of original reward function improves model winrate; (2) the proposed method improves BoN decoding model's winrate.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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 \n\n***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***\n\nWe 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. \n\nMoreover, 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.\n\n\n***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).***\n\nRegarding 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).\n\n***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.***\n\nThe 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.\n\nIn 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.\n\n*(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.\n\n\n***Q4: The generalization of the proposed framework to other inference-time strategies (e.g., self-consistency, chain-of-thought) is not discussed.***\n\nIn 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.\n\n***Q5: How much additional computational overhead does the calibrated reward estimation introduce?***\n\nWhile 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.", "claims_and_evidence": "The authors claim that training using calibrated reward improves model's winrate compared to the other reward optimization baselines (Sec. 5.3).\nHowever, as depicted in Figure 4, CTRL's performance is quite close to that of BoND.", "methods_and_evaluation_criteria": "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.\nHowever, there is a potential issue of reward hacking when relying solely on win rate over a fixed model. \nTo strengthen the evaluation, the authors could include comparisons of raw rewards and win rates over other models (e.g., models trained with different methods).", "theoretical_claims": "While I did not verify the detailed proofs, the derivations and arguments presented are clear and logically consistent.", "experimental_designs_or_analyses": "The experiment setting in Sec. 5.2 is strange to me.\nThe calibrated reward measures the expected winning rate over SFT policy, and thus depends on the SFT model.\nIt's natural that the learned reward function differs from the calibrated reward.\nThe authors should discuss how the mismatch between learned reward and calibrated reward can hurt the performance.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Other Strengths:\n- Theoretical soundness: this paper provides a clear and rigorous theoretical framework for inference-aware alignment, with well-justified derivations and proofs.\n\nOther Weaknesses:\n- The paper focuses primarily on BoN and WoN strategies. The generalization of the proposed framework to other inference-time strategies (e.g., self-consistency, chain-of-thought) is not discussed.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "How much additional computational overhead does the calibrated reward estimation introduce?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uVzwZEcdHN", "reviewer": "Reviewer_XbH2", "summary": "This paper explores a novel problem in LLM alignment **considering inference-time procedures**. More specifically, it aims to maximize the reward given a fixed LLM and an inference-time procedure, using reinforcement learning (RL). The focus is on Best-of-N as the inference-time procedure, while also providing general mathematical guidance for other potential inference-time procedures.\n\n## update after rebuttal\nI decide to keep my high score for this good paper. This paper should be accepted IMO!", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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\n\n***Q1: However, the evaluations seem somewhat limited. Demonstrating improvements in general instruction-following tasks would likely generate more excitement.***\n\nThank the reviewer for the suggestion. We view demonstrating the method on more RLHF tasks as important next steps.\n\n\n***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***\n\nThanks 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. \n\nAdditional references (not listed in the paper):\n[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.", "claims_and_evidence": "Overall, the claims made in the submission are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The proposed method and evaluations are reasonable.\n\nHowever, the evaluations seem somewhat limited. Demonstrating improvements in general instruction-following tasks would likely generate more excitement.", "theoretical_claims": "I didn't check the theoretical claims, but there do exist a lot of proofs in the Appendix.", "experimental_designs_or_analyses": "The soundness/validity of experimental designs and analyses makes sense.", "relation_to_broader_scientific_literature": "A key contribution of this paper is proposing a new problem setup, which is valuable. Studying how to improve a model during training while considering the inference-time procedures used for deployment is a meaningful direction.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "**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", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces InfAlign, a novel RLHF framework designed to optimize language model performance for specific inference-time procedures like Best-of-N (BoN). It addresses the train/test mismatch in standard RLHF by maximizing the inference-time win rate against a base model. This is achieved by optimizing a KL-regularized objective using a transformed version of a "calibrated reward" (essentially the win rate). The authors show theoretically and empirically that this approach, particularly with a near-optimal exponential transformation for BoN/WoN, improves inference-time win rates (3-8% reported for BoN). Reward calibration alone is also presented as a strong baseline improving standard win rates and robustness.

Reviewers acknowledged the paper's novelty in addressing inference-aware alignment and introducing calibrated rewards. The theoretical framework was generally seen as sound. Key strengths included the novel problem formulation, theoretical grounding, and empirical gains in inference-time win rates. Despite minor concerns about the practical implications of the heuristic transformation and the method's dependence on a pre-defined N temper enthusiasm, the authors’ rebuttal addressed these concerns. Based on the solid contribution, I recommend to accept this paper.
Decision: Accept (poster)