# InfAlign: Inference-aware language model alignment

arXiv: https://arxiv.org/abs/2412.19792 · HTML: https://mlsanity.com/paper/2412.19792/

## Key facts

- Cited by 35 papers in the map (data through 14 Oct 2026).
- Placed in LLM preference optimization methods, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 54 of its 68 references that are in the map.
- First posted on arXiv 27 Dec 2024.

## Placed in

- LLM alignment / RLHF / preference optimisation / [LLM preference optimization methods](https://mlsanity.com/alignment/llm-preference-optimization-methods/)

## Builds on (most cited first, 30 of 54)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2204.05862/) (2022)
- [Self-Refine: Iterative Refinement with Self-Feedback](https://mlsanity.com/paper/2303.17651/) (2023)
- [Tree of Thoughts: Deliberate Problem Solving with Large Language Models](https://mlsanity.com/paper/2305.10601/) (2023)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [Concrete Problems in AI Safety](https://mlsanity.com/paper/1606.06565/) (2016)
- [Jailbreaking Black Box Large Language Models in Twenty Queries](https://mlsanity.com/paper/2310.08419/) (2023)
- [PaLM 2 Technical Report](https://mlsanity.com/paper/2305.10403/) (2023)
- [Scaling Laws for Reward Model Overoptimization](https://mlsanity.com/paper/2210.10760/) (2022)
- [Large Language Monkeys: Scaling Inference Compute with Repeated Sampling](https://mlsanity.com/paper/2407.21787/) (2024)
- [A General Theoretical Paradigm to Understand Learning from Human Preferences](https://mlsanity.com/paper/2310.12036/) (2023)
- [RewardBench: Evaluating Reward Models for Language Modeling](https://mlsanity.com/paper/2403.13787/) (2024)
- [Fine-Grained Human Feedback Gives Better Rewards for Language Model Training](https://mlsanity.com/paper/2306.01693/) (2023)
- [A StrongREJECT for Empty Jailbreaks](https://mlsanity.com/paper/2402.10260/) (2024)
- [A Real-World WebAgent with Planning, Long Context Understanding, and Program Synthesis](https://mlsanity.com/paper/2307.12856/) (2023)
- [The Effects of Reward Misspecification: Mapping and Mitigating Misaligned Models](https://mlsanity.com/paper/2201.03544/) (2022)
- [Reward Model Ensembles Help Mitigate Overoptimization](https://mlsanity.com/paper/2310.02743/) (2023)
- [Helping or Herding? Reward Model Ensembles Mitigate but do not Eliminate Reward Hacking](https://mlsanity.com/paper/2312.09244/) (2023)
- [From Decoding to Meta-Generation: Inference-time Algorithms for Large Language Models](https://mlsanity.com/paper/2406.16838/) (2024)
- [RL with KL penalties is better viewed as Bayesian inference](https://mlsanity.com/paper/2205.11275/) (2022)
- [Controlled Decoding from Language Models](https://mlsanity.com/paper/2310.17022/) (2023)

## Built on by (most cited first, 30 of 35)

- [A Survey of Frontiers in LLM Reasoning: Inference Scaling, Learning to Reason, and Agentic Systems](https://mlsanity.com/paper/2504.09037/) (2025)
- [Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment](https://mlsanity.com/paper/2503.21878/) (2025)
- [Pass@K Policy Optimization: Solving Harder Reinforcement Learning Problems](https://mlsanity.com/paper/2505.15201/) (2025)
- [What Makes a Reward Model a Good Teacher? An Optimization Perspective](https://mlsanity.com/paper/2503.15477/) (2025)
- [Optimizing Language Models for Inference Time Objectives using Reinforcement Learning](https://mlsanity.com/paper/2503.19595/) (2025)
- [Chasing the Tail: Effective Rubric-based Reward Modeling for Large Language Model Post-Training](https://mlsanity.com/paper/2509.21500/) (2025)
- [POPE: Learning to Reason on Hard Problems via Privileged On-Policy Exploration](https://mlsanity.com/paper/2601.18779/) (2026)
- [Inference-Time Reward Hacking in Large Language Models](https://mlsanity.com/paper/2506.19248/) (2025)
- [Single-stream Policy Optimization](https://mlsanity.com/paper/2509.13232/) (2025)
- [Representation-Based Exploration for Language Models: From Test-Time to Post-Training](https://mlsanity.com/paper/2510.11686/) (2025)
- [Reuse your FLOPs: Scaling RL on Hard Problems by Conditioning on Very Off-Policy Prefixes](https://mlsanity.com/paper/2601.18795/) (2026)
- [REINFORCE Adversarial Attacks on Large Language Models: An Adaptive, Distributional, and Semantic Objective](https://mlsanity.com/paper/2502.17254/) (2025)
- [Risk-Sensitive RL for Alleviating Exploration Dilemmas in Large Language Models](https://mlsanity.com/paper/2509.24261/) (2025)
- [KL-Regularized RLHF with Multiple Reference Models: Exact Solutions and Sample Complexity](https://mlsanity.com/paper/2502.01203/) (2025)
- [Best-of-Majority: Minimax-Optimal Strategy for Pass@k Inference Scaling](https://mlsanity.com/paper/2510.03199/) (2025)
- [Why Pass@k Optimization Can Degrade Pass@1: Prompt Interference in LLM Post-training](https://mlsanity.com/paper/2602.21189/) (2026)
- [Vector Policy Optimization: Training for Diversity Improves Test-Time Search](https://mlsanity.com/paper/2605.22817/) (2026)
- [Why is Your Language Model a Poor Implicit Reward Model?](https://mlsanity.com/paper/2507.07981/) (2025)
- [EvoIdeator: Evolving Scientific Ideas through Checklist-Grounded Reinforcement Learning](https://mlsanity.com/paper/2603.21728/) (2026)
- [Power Distribution Bridges Sampling, Self-Reward RL, and Self-Distillation](https://mlsanity.com/paper/2605.04542/) (2026)
- [PICACO: Pluralistic In-Context Value Alignment of LLMs via Total Correlation Optimization](https://mlsanity.com/paper/2507.16679/) (2025)
- [Best-of-Tails: Bridging Optimism and Pessimism in Inference-Time Alignment](https://mlsanity.com/paper/2603.06797/) (2026)
- [CLR-voyance: Reinforcing Open-Ended Reasoning for Inpatient Clinical Decision Support with Outcome-Aware Rubrics](https://mlsanity.com/paper/2605.09584/) (2026)
- [Using Reward Uncertainty to Induce Diverse Behaviour in Reinforcement Learning](https://mlsanity.com/paper/2606.03962/) (2026)
- [Saffron-1: Safety Inference Scaling](https://mlsanity.com/paper/2506.06444/) (2025)
- [When Errors Can Be Beneficial: A Categorization of Imperfect Rewards for Policy Gradient](https://mlsanity.com/paper/2604.25872/) (2026)
- [Calibrating the Evaluator: Does Probability Calibration Mitigate Preference Coupling in LLM Agent Feedback Loops?](https://mlsanity.com/paper/2606.31371/) (2026)
- [Learnable Chernoff Baselines for Inference-Time Alignment](https://mlsanity.com/paper/2602.07738/) (2026)
- [What should post-training optimize? A test-time scaling law perspective](https://mlsanity.com/paper/2605.10716/) (2026)
- [ExpRL: Exploratory RL for LLM Mid-Training](https://mlsanity.com/paper/2606.17024/) (2026)

## Abstract (arXiv)

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.