Research map / Reinforcement learning / RL from feedback
InfAlign: Inference-aware language model alignment
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes InfAlign, an inference-aware alignment framework that transforms the reward (via reward calibration plus exponential transformation) so RLHF optimizes win rate under inference-time procedures like Best-of-N and Worst-of-N, yielding 3-8% gains on inference-time win rates.
Problem
Standard RLHF optimizes win rate under plain sampling, but at deployment LMs are often decoded with inference-time procedures (Best-of-N, best-of-N jailbreaking), creating a train/test mismatch that makes standard alignment suboptimal.
Approach
Show that the optimal aligned policy for inference-time win rate is the KL-regularized RLHF solution under a transformed reward; implement this as InfAlign-CTRL, which calibrates the reward per prompt against reference-policy samples (empirical CDF) and then applies a procedure-specific transformation (exponential tilting for BoN/WoN) before PPO.
Key results
- x
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use InfAlign-CTRL with identity reward transformation instead of the compared SOTA methods | Standard win rate alignment · Anthropic Helpfulness | win rate-KL trade-off | Helped |
| Use InfAlign-CTRL with identity reward transformation instead of the compared SOTA methods | Standard win rate alignment · Anthropic Harmlessness | win rate-KL trade-off | Helped |
| Use InfAlign-CTRL with identity reward transformation instead of the compared SOTA methods | Standard win rate alignment · Reddit summarization | win rate-KL trade-off | No clear effect |
| Use exponential transformation of calibrated reward instead of the best SOTA method | InfAlign-CTRL · Anthropic Helpfulness, Best-of-4 | inference-time win rate: up to 3% | Helped |
| Use exponential transformation of calibrated reward instead of the best SOTA method | InfAlign-CTRL · Reddit summarization, Best-of-4 | inference-time win rate: up to 8% | Helped |
| Use negative exponential transformation of calibrated reward instead of the best SOTA method | InfAlign-CTRL · Anthropic Harmlessness, Worst-of-N | inference-time win rate: up to 5% | Helped |
| Set exponential transformation to t=10 (ablation) | InfAlign-CTRL · Anthropic Helpfulness, Best-of-4 | inference-time win rate | Helped |
| Set negative exponential transformation to t=-10 (ablation) | InfAlign-CTRL · Anthropic Harmlessness, Worst-of-N | inference-time win rate | Helped |
| Use e^(5u) rather than other exponential transformations (ablation) | InfAlign-CTRL · Best-of-2 | inference-time win rate | Helped |
| Use e^(10u) rather than other exponential transformations (ablation) | InfAlign-CTRL · Best-of-4 | inference-time win rate | Helped |
| Use exponential reward transformation instead of other non-CTRL baselines (ablation) | InfAlign-CTRL · Best-of-N, N=2, 4, 32 | inference-time win rate | Helped |
| Calibrate the reward model before reward optimization (ablation) | Reward model alignment · Poisoned Anthropic Helpfulness and Harmlessness evaluation | accuracy | Helped |
Surprising findings
- Standard win-rate-optimal transformations (identity and log) are suboptimal once inference-time procedures such as Best-of-4 are applied.
Limitations
- Experiments rely on PaLM-2 M reward model as automated judge rather than human ratings; the paper acknowledges gap to human preference.
- Theoretical results assume continuous language models with distinct rewards; practical settings are approximations.
- Results are shown only for BoN and WoN procedures with PPO as the RL solver.
Experimental setup
Datasets: Anthropic helpfulness Anthropic harmlessness Reddit TL;DR
Benchmarks: Anthropic helpfulness Anthropic harmlessness Reddit TL;DR summarization
Metrics: win rate best-of-N win rate worst-of-N win rate KL divergence
Baselines: uncalibrated RLHF (PPO) BoNBoN BoND IPO GRPO
Builds on: RLHF (KL-regularized RL) IPO BoNBoN BoN distillation PPO
Closest papers by meaning
- Theoretical Limits of Language Model Alignment (NeurIPS 2026 · similarity 0.91)
- Learnable Chernoff Baselines for Provable Inference-Time Alignment (NeurIPS 2026 · similarity 0.90)
- Inference-time Alignment in Continuous Space (NeurIPS 2025 · similarity 0.90)
- W2S-AlignTree: Weak-to-Strong Inference-Time Alignment for Large Language Models via Monte Carlo Tree Search (AAAI 2026 · similarity 0.90)
- Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment (ICML 2025 · similarity 0.89)
- Aligner: Efficient Alignment by Learning to Correct (NeurIPS 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.