PaperlayerView exact revision on arXiv

Source-linked AI summary

InfAlign: Inference-aware language model alignment

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

arXiv:2412.19792v5cs.LGcs.CLcs.ITmath.IT

TL;DR

Aligned models are often used with inference-time decoding procedures rather than directly sampled, while standard RLHF does not target the resulting inference-time win rate. InfAlign transforms the reward so that solving a KL-regularized RLHF problem yields an optimal policy for the chosen inference-time procedure. When optimizing WoN performance across various values of N, InfAlign-CTRL outperforms various SOTA RLHF solvers in inference-time win rate by 3-8%.

  • Problem

    Aligned models are often used with inference-time decoding procedures rather than directly sampled, while standard RLHF does not target the resulting inference-time win rate.

    Source: §1 · p. 1

  • Method

    InfAlign transforms the reward so that solving a KL-regularized RLHF problem yields an optimal policy for the chosen inference-time procedure.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    When optimizing WoN performance across various values of N, InfAlign-CTRL outperforms various SOTA RLHF solvers in inference-time win rate by 3-8%.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Reward calibration with identity transformation is reported as a strong standard-win-rate baseline, with slightly better performance than a variety of SOTA methods.

    Source: §4.1 · p. 5, §6 · p. 9

    The iterative procedure derived from the general characterization is inefficient for general language models because evaluating its reward update can require a large or infinite output space, and whether it reaches the optimum is unclear.

    Source: §3 · p. 4

Abstract

from arXiv · show

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.

1. Introduction

Because models are commonly decoded through inference-time procedures, standard alignment can miss the way they will be used. InfAlign targets inference-time win rate and uses reward transformation to make that objective solvable through RLHF.

1. Introduction

  • InfAlign optimizes inference-time win rate against the base model with KL regularization, using a reward transformation to obtain an optimal solution through RLHF.Direct optimization of inference-time win rate is described as seemingly intractable.Source: §1 · p. 2
  • InfAlign-CTRL calibrates reward-model scores on reference-policy responses, transforms those scores for the decoding procedure, and solves RLHF with the transformed reward.The authors report calibration reduces reward hacking and improves standard win rate; identity transformation is on par with SOTA methods and improves helpfulness and harmlessness win rates.Source: §1 · p. 2
  • The paper derives reward transformations for Best-of-N sampling and Best-of-N jailbreaking, called Worst-of-N from the defender’s perspective.These procedures are studied as popular approaches for inference-time alignment and jailbreak evaluation.Source: §1 · p. 2
  • 3–8% higher inference-time win rates than various SOTA RLHF solvers are reported for InfAlign-CTRL on Best-of-N and Worst-of-N tasks.The experiments cover Anthropic helpfulness and Reddit summarization for Best-of-N, and Anthropic harmlessness for Worst-of-N.Source: §1 · p. 2

2. Problem setup

The setup treats a language model as a policy over responses to prompts and evaluates alignment by reward-based win rate. It extends that comparison to inference-time procedures that transform each policy’s response distribution.

2. Problem setup

  • RLHF maximizes expected response reward while regularizing the aligned policy’s KL divergence from a fixed reference policy.The reward assigns a scalar to each prompt-response pair, and the reference policy can be a supervised-finetuned base policy.Source: §2 · p. 2
  • Standard win rate measures how often a response from one policy beats a response from another under the reward-based comparison.The paper notes that experimental evaluation uses a judge separate from the reward model.Source: §2 · pp. 2–3
  • An inference-time procedure T maps a policy’s response distribution to the distribution from which responses are obtained after decoding.Inference-time win rate compares policies after applying that procedure.Source: §2 · p. 3
  • InfAlign maximizes KL-regularized inference-time win rate; with identity processing, the standard win-rate formulation corresponds to IPO.The paper characterizes the identity case as an extension of the standard formulation to arbitrary inference-time procedures.Source: §2 · p. 3
  • Some theoretical results assume continuous outputs and distinct rewards for responses to a prompt, but the algorithmic developments and experiments do not make these assumptions.The authors note these assumptions can lead to approximations for real-world distributions.Source: §2 · p. 3

3. Reinforcement learning with reward transformation

The paper shows that InfAlign can be solved by standard KL-regularized RLHF with a reward transformed using the reward model, reference policy, and inference-time procedure. The general characterization motivates an iterative approach, while the identity case reduces to calibrated reward.

3. Reinforcement learning with reward transformation

  • For any base policy, reward model, inference-time procedure, and positive regularization parameter, a transformed reward makes the RLHF maximizer solve InfAlign.The optimal transformed reward depends on the base policy, inference-time procedure, and reward model.Source: §3 · pp. 3–4
  • Theorem 1 characterizes the optimal transformed reward and policy through coupled equations, assuming the procedure’s derivative with respect to the policy exists.The characterization supports iterative updates of the policy and reward in turn.Source: §3 · p. 4
  • The iterative characterization is inefficient for general language models because evaluating the reward update can require scoring a large or infinite output space, and its convergence to the optimum is unclear.These drawbacks motivate the paper’s practical solver.Source: §3 · p. 4
  • With no inference-time procedure, the transformed reward becomes the reference-policy calibrated reward, and its RLHF maximizer solves the standard win-rate objective.The paper relates this identity-case result to prior observations by Azar et al. and Gui et al.Source: §3 · p. 4

4. InfAlign-CTRL: Calibrate-and-transform reinforcement learning

InfAlign-CTRL calibrates reward scores, transforms them for the chosen inference-time procedure, and then solves RLHF with the transformed reward. The paper establishes calibration properties and shows that different transformations suit Best-of-N and Worst-of-N decoding.

4. InfAlign-CTRL: Calibrate-and-transform reinforcement learning

  • InfAlign-CTRL calibrates the reward, transforms the calibrated scores for the inference-time procedure, then solves RLHF with the transformed reward.The final optimization step can use standard RLHF solvers.Source: §4 · pp. 4–5

4.1. Properties of reward calibration

  • Calibration preserves reward ordering, is invariant to strictly increasing reward transformations, and maps reference-policy reward draws to a uniform distribution over [0, 1].These properties support a calibrated reward that is more robust to reward-model learning.Source: §4.1 · p. 5

4.2. Reward transformation for calibrated inference-time procedure

  • For calibrated inference-time procedures and continuous language models, aligned-policy win rate and KL divergence are independent of the base policy and reward model.This lets transformation functions be evaluated using simple, simulatable continuous language models.Source: §4.2 · p. 5

4.3. Finding reward transformations for BoN and WoN

  • The paper derives fixed-point transformations and considers exponential tilting, where positive t targets higher calibrated-reward quantiles and negative t targets lower quantiles for WoN.The fixed-point transformations are found iteratively from coupled equations.Source: §4.3 · p. 6
  • For Best-of-4, bon fp gives the best tradeoffs and exp(10x) is nearly as good; for Worst-of-4, won fp leads, with exp(−10x) nearly as good.Identity and log transformations are suboptimal for both Best-of-4 and Worst-of-4.Source: §4.3 · pp. 6–7, Figure 2 · p. 6

4.4. Practical implementation of InfAlign-CTRL

  • Practical calibration estimates calibrated rewards by ranking K reference-model samples per training prompt, then transforms those scores and runs PPO on the resulting reward.The paper also gives a probabilistic error bound for empirical calibration.Source: §4.4 · p. 7

5. Experiment results

InfAlign-CTRL improves inference-time win rates with transformations tailored to Best-of-N and Worst-of-N, while its identity-transform version is competitive on standard win rate. The experiments also find reward-model miscalibration and report that offline calibration performs similarly to GRPO on standard win-rate–KL tradeoffs.

5.1. Evaluation setup

  • The evaluation uses Anthropic helpfulness and harmlessness and Reddit summarization tasks, with a larger PaLM-2 M reward model as the automated judge.The authors note a gap between model-generated and human preferences; reported judge accuracy is 77.7%, 77.0%, and 76.4% on the three tasks, respectively.Source: §5.1 · p. 7

5.2. Reward models are typically miscalibrated

  • On most sampled helpfulness prompts, reward-score ranks deviate substantially from the perfect-calibration line, with the degree of deviation varying by prompt.The analysis uses 100 reference-policy responses for each of 10 random test prompts.Source: §5.2 · p. 8, Figure 3 · p. 8

5.3. InfAlign-CTRL improves standard win rate

  • InfAlign-CTRL with identity transformation achieves better win-rate–KL tradeoffs on helpfulness and harmlessness and is on par with SOTA on summarization.The authors attribute the gains to reward calibration, which they report can help mitigate reward hacking.Source: §5.3 · p. 8, Figure 4 · p. 9

5.4. InfAlign-CTRL improves BoN

  • BoN win-rate comparisons cover Anthropic helpfulness and Reddit summarization, while WoN comparisons cover Anthropic harmlessness.The study evaluates multiple values of N and varies KL-regularizer strength to obtain win-rate–KL tradeoff curves.Source: §5.1 · p. 7, §5.4 · p. 8, §5.5 · p. 8
  • Up to 3% higher Best-of-N win rates on helpfulness and up to 8% on summarization, while median-calibrated negative-exponential rewards yield up to 5% higher Worst-of-N win rates on harmlessness.These comparisons are against the best SOTA method; Figure 4 reports test-split win rates measured by the PaLM-2 M reward model.Source: §5.4 · p. 8, §5.5 · p. 8, Figure 4 · p. 9

5.6. Transformation choice for different values of N

  • For Best-of-N, exp(5u) performs better at N=2 and exp(10u) at N=4; both outperform non-CTRL baselines across the tested N values.The authors report this pattern as evidence that the transformations do not overfit to a particular N.Source: §5.6 · p. 8, Figure 5 · p. 9

5.7. Calibration helps close the RL optimality gap

  • Offline-calibrated InfAlign-CTRL performs similarly to GRPO on standard win-rate–KL tradeoffs without the cost of additional online rollouts.The authors suggest the calibration step may be independently useful for closing RL's optimality gap.Source: §5.7 · pp. 8–9

6. Concluding Remarks

InfAlign addresses train–test mismatch by optimizing inference-time win rate and proving that an optimal policy can be learned through RLHF with a reward transformation. InfAlign-CTRL applies this approach, with reported gains for inference-time procedures.

6. Concluding Remarks

  • For any inference-time procedure, an optimal aligned model can be learned by solving RLHF with a reward transformation.The paper frames this as a way to address train–test mismatch when procedures other than sampling are used.Source: §6 · p. 9
  • InfAlign-CTRL performs similarly to GRPO in closing the gap with Best-of-N.Source: Figure 6 · p. 9
  • 3–8% better inference-time win rate-versus-KL trade-offs than existing preference-optimization methods are reported when inference-time procedures are applied.The conclusion specifies Best-of-N sampling and jailbreaking transformations.Source: §6 · p. 9

Impact Statement

The section discusses InfAlign’s potential for both beneficial alignment and adversarial goals, relates it to inference-time methods, and includes proofs and mappings for continuous models.

Impact Statement

  • The framework is intended to improve alignment with inference-time procedures, but could also be used to align models toward adversarial goals.The authors say this possibility also applies to other work advancing language-model alignment.Source: Impact Statement · p. 10

A. Related work

  • Prior work approximates inference-time solutions during training, whereas InfAlign addresses procedures applied at inference and generalizes IPO and Best-of-N distillation as special cases.Source: A. Related work · p. 14
  • InfAlign is motivated by inference-time strategies such as chaining, chain-of-thought, reranking, refinement, and majority voting.The authors say the framework may provide additional alignment gains with such procedures.Source: A. Related work · p. 14

B. Missing proofs

  • For continuous models, the appendix maps outcomes to [0, 1] by inverse CDF, ordered by reward, and defines the inverse calibrated-reward mapping.Source: B. Missing proofs · p. 15

B.6. Proof of Lemma 5

  • For Best-of-N and Worst-of-N, the appendix gives inference-time reward probabilities as powers of the calibrated reward mapping and its complement.Source: §B.6 · p. 17

B.7. Proof of Theorem 2

The appendix states conditions for a theorem on calibrated inference-time procedures, relates specific decoding strategies to its mapping, and discusses KL regularization’s role in multitask alignment.

B.7. Proof of Theorem 2

  • Theorem 4’s stated result applies to calibrated inference-time procedures, continuous language models, β > 0, and a reward transformation Φ.Source: §B.7 · p. 17

B.8. Proof of Theorem 3

  • The win rate is obtained by plugging gBoN(u) = u^(N−1) and gWoN(u) = (1 − u)^(N−1), as shown in Lemma 5, into Theorem 4.Source: §B.8 · p. 18

B.9. Proof of Corollary 2

  • Corollary 2 is presented as a special case of Theorem 1, while Theorem 2 generalizes Φ to arbitrary continuous language models.Source: §B.9 · p. 18

C. The role of KL divergence in model alignment

D. Calibration Reduces Reward Hacking

The section tests whether reward calibration reduces reward hacking by flipping spurious phrase correlations between training and evaluation. Calibrated reward models maintain higher accuracy than uncalibrated ones on the inverted test set.

D. Calibration Reduces Reward Hacking

E. Additional experimental results

Additional experiments report InfAlign-CTRL’s Best/Worst-of-32 win-rate gains, higher expected raw rewards at a given KL tradeoff, and comparisons with GRPO.

E. Additional experimental results

E.1. Comparison with GRPO

  • With a fixed base policy and matched reward-rollout count, InfAlign-CTRL outperforms GRPO in Best-of-4 and is comparable in standard win rate.The GRPO comparison uses two epochs and K=G=100 samples per prompt.Source: §E.1 · p. 23, Figure 10 · p. 24

F. Analytical comparison of different transformations

The analytical comparisons show that reward transformations suited to standard win rate differ from those suited to Best-of-N and Worst-of-N. Exponential transformations also perform well for rewind-and-repeat, while calibration adds training compute.

F. Analytical comparison of different transformations

  • For Best-of-2 and Best-of-4, bon fp gives the best tradeoffs; exp(5x) and exp(10x) are almost as good, respectively.Identity and logarithmic transformations are sub-optimal for these inference-time win rates.Source: Figure 11 · p. 25, F. Analytical comparison of different… · pp. 25–26
  • For Worst-of-2 and Worst-of-4, won fp gives the best tradeoffs; exp(−5x) and exp(−10x) are almost as good, respectively.Identity and logarithmic transformations are sub-optimal in these cases.Source: Figure 11 · p. 25, F. Analytical comparison of different… · p. 26
  • For standard win rate, IPO is optimal, the logarithmic transformation is nearly optimal, and all other tested transformations are sub-optimal.The logarithmic transformation corresponds to best-of-N distillation.Source: F. Analytical comparison of different… · p. 25
  • The results support choosing the reward transformation for the inference-time procedure rather than optimizing standard win rate alone.Different exponential transformation parameters perform well for different procedures.Source: Figure 11 · p. 25, F. Analytical comparison of different… · p. 26
  • For K=100 rollouts, reward calibration takes about 29% of total training time, and its performance-overhead tradeoff may be task-specific.Calibration uses a forward pass and is performed once per prompt before policy optimization.Source: F. Analytical comparison of different… · p. 26
  • For rewind-and-repeat with an 85th-percentile reward threshold and at most 32 repeats, exponential transformations outperform baselines.The figure also reports repeat trials needed to exceed the threshold.Source: Figure 12 · p. 27

G. Analysis on Rewind-and-Repeat

InfAlign-CTRL also improves performance under the adaptive-compute rewind-and-repeat inference procedure.

G. Analysis on Rewind-and-Repeat

  • Rewind-and-repeat repeatedly samples until reward meets a threshold or the maximum of 32 generations is reached; InfAlign-CTRL improves performance in this setting.The experiment uses an 85th-percentile reward threshold.Source: G. Analysis on Rewind-and-Repeat · p. 27

Related papers

Loading 2412.19792v5…