Inference-Time Reward Alignment

• Inference-time reward alignment is a suite of methods that adjusts a pretrained model’s outputs during generation to meet desired reward metrics without modifying model weights.
• It employs calibrated reward transformations and selection procedures like Best-of-N to align outputs with human preferences and mitigate reward hacking.
• The approach is theoretically grounded and empirically validated, balancing exploitative optimization with distributional integrity to optimize performance.

Inference-time reward alignment refers to the suite of methods that adapt a pretrained generative model’s outputs to maximize a target reward function (often a proxy for human preferences or desired behavior) by intervening at the generation stage, not by model weight modification. These techniques have become central in modern language modeling and diffusion model pipelines, especially as inference-time search and selection procedures (e.g., Best-of-N, tree search, resampling) have outpaced naïve output sampling for alignment-sensitive objectives such as win-rate against a base model, preference satisfaction, or constraint satisfaction. The development of general theory, provably optimal algorithms, and practical frameworks for inference-time reward alignment—spanning both discrete (LLMs) and continuous (diffusion models) settings—has enabled practitioners to faithfully target desired reward metrics, mitigate reward hacking, and balance the inherent trade-offs between exploitative optimization and distributional integrity.

1. Motivation and Problem Formulation

In the standard RLHF paradigm and maximum reward-KL alignment, one seeks to identify a policy π maximizing

$$\mathbb{E}_{x,y\sim\pi}[r(x,y)] - \beta\,\mathrm{KL}\bigl(\pi \,\|\, p\bigr),$$

where $p$ is a reference (pretrained) policy, $r$ is a scalar reward model, and $\beta$ trades off reward optimization with KL-regularized deviation from $p$. However, practical decoding rarely corresponds to the training-time sampling assumptions: rather than i.i.d. draws from π, inference-time deployments utilize Best-of-N (BoN) selection, adversarial sampling, tree search, and other strategies. This "train/test mismatch" leads to sub-optimal deployment performance, as standard RLHF does not anticipate the selection-induced shift in output distribution (Balashankar et al., 2024).

Inference-time reward alignment thus formalizes the objective of maximizing inference-time metrics—such as win rate versus a base model under a search procedure $T$—by tuning reward transforms, sample selection, or decoding rules that account for the actual generation and selection workflow. A general goal is:

$$\max_\pi\, W^T_r(\pi \succ p) - \beta\,\mathrm{KL}(\pi\,\|\,p)$$

where $W^T_r(\pi \succ p)$ is the win rate under decoding procedure $T$ and reward $r$.

2. Theoretical Foundations and Reward Transformations

Universal to modern inference-time alignment theory is the demonstration that for any inference-time decoding procedure $p$0, there exists a transformed reward $p$1 such that solving a standard RLHF problem with $p$2 as the reward recovers the optimal policy for $p$3 (Balashankar et al., 2024). The solution is characterized via the coupled stationarity equations: $p$4

$p$5

Standard settings (e.g., identity $p$6) yield the calibrated reward

$p$7

which encodes the average pairwise win against $p$8—the metric most directly linked to human/language-model evaluation and which is robust to reward scale (Balashankar et al., 2024, Sriraman et al., 5 Mar 2026).

For selection procedures such as BoN or "Worst-of-N" (adversarial/jailbreaking), closed-form transformations $p$9 of the calibrated reward deliver optimal or near-optimal alignment. For BoN of size $r$0, this is typically an exponential tilt,

$r$1

yielding distributions that emphasize high-quantile calibration. For adversarial settings, a negative exponential is used, penalizing low-reward quantiles. This framework makes reward transformation and optimal decoding explicit for arbitrary $r$2.

3. Algorithmic Developments: Calibrate-and-Transform RL and Extensions

The practical implementation of inference-time reward alignment is exemplified by the InfAlign-CTRL (Calibrate-and-Transform RL) algorithm (Balashankar et al., 2024):

1. Reward Calibration: Empirically estimate the calibrated reward $r$3 using base model rollouts, ensuring quantile normalization and (optionally) monotonic anchor fitting.
2. Reward Transformation: Apply a $r$4 transformation tailored to the desired $r$5 (e.g., BoN).
3. KL-Regularized Optimization: Solve the standard KL-regularized RL objective using $r$6 as the reward.

Ablation studies demonstrate that calibration alone can outperform classical batchwise preference optimization and that increasing the number of anchor points and base model rollouts improves the trade-off between win rate and KL divergence.

The generality of this approach supports arbitrary inference-time search or filtering algorithms, connecting their optimal policies to explicit reward transformations, and yielding robust empirical improvements in win-rate and robustness to reward hacking (Balashankar et al., 2024, Sriraman et al., 5 Mar 2026).

4. Special Cases: Best-of-N, Soft Best-of-N, and Reward Hacking

Best-of-N (BoN) selection—sample $r$7 outputs, pick the one maximizing the proxy reward—has emerged as the most widely used inference-time alignment method. The mapping between BoN and the exponential-tilted distribution is exact in the calibration limit, and extensive theory has clarified both its efficiency and failure modes (Balashankar et al., 2024, Sriraman et al., 5 Mar 2026, Khalaf et al., 24 Jun 2025).

A key limitation of BoN (and related search-based methods) is reward hacking: as $r$8 increases, the likelihood of selecting a candidate that simultaneously achieves high proxy reward on the learned reward model but low true reward increases due to over-optimization (the "winner’s curse"), leading to an eventual collapse of true reward/accuracy beyond an optimally tuned $r$9 (Khalaf et al., 24 Jun 2025). This behavior is provably inevitable for a broad class of monotone-likelihood-ratio selection mechanisms.

Mitigating reward hacking requires hedging, or the calibration of the selection intensity via either parametric (e.g., Soft-BoN with tunable temperature) or coverage-regularized (e.g., $\beta$0-monotone selection) variants. Efficient algorithms such as HedgeTune can identify the peak of the true reward curve and set the relevant parameter, maximizing global reward without overstepping into failure regions (Khalaf et al., 24 Jun 2025, Sriraman et al., 5 Mar 2026).

5. Empirical and Practical Considerations

Empirical studies demonstrate that inference-time reward alignment via tailored reward transformation or selection consistently outperforms both naive RLHF and off-the-shelf selection approaches on win-rate and robustness metrics (Balashankar et al., 2024, Sriraman et al., 5 Mar 2026). Notably:

• Calibration alone (even without further transformation) yields significant improvements.
• The optimal BoN size $\beta$1 or softmax temperature must be tuned carefully. Over-optimization leads to reward hacking.
• Limitations include dependence on the accuracy of the base reference policy $\beta$2 and reward model $\beta$3, and approximation error in $\beta$4.
• Current frameworks are best suited to rank-based or pairwise evaluation procedures. Extension to complex multi-task or constrained decoding requires further theory and empirical tuning.

The field is moving toward both tighter theoretical guarantees (e.g., total-variation error bounds for selection schemes) and more robust calibration/regularization against adversarial prompt or reward distributions (Balashankar et al., 2024, Sriraman et al., 5 Mar 2026).

6. Limitations and Ongoing Directions

Major challenges and open problems in inference-time reward alignment are concentrated in several areas:

• Generalization beyond rank-based selection: While current theory handles Best-of-N, pessimistic variants, and "jailbreaking," more complex inference-time procedures (chain-of-thought, latent search, hard constraint satisfaction) may require higher-order or structured calibrations (Balashankar et al., 2024).
• Online calibration: As the aligned policy $\beta$5 evolves, recalibration of $\beta$6 may be needed for nonstationary deployments.
• Multi-task and multi-reward alignment: Extending single-reward calibration to joint or constrained objectives remains open.
• Reliance on base model and reward quality: Approximation errors, off-distribution reward hacking, and base policy bias directly impact alignment quality.
• Scaling and compute: For very large $\beta$7, both calibration and online optimization incur significant computational costs.

Nonetheless, inference-time reward alignment provides a comprehensive, unifying framework for principled, practical, and robust alignment under arbitrary post-training decoding, and has established itself as the theoretical backbone of modern large-model alignment post-processing (Balashankar et al., 2024, Sriraman et al., 5 Mar 2026, Khalaf et al., 24 Jun 2025).