Research map / Reinforcement learning / RL from feedback
Inference-time Alignment in Continuous Space
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 Simple Energy Adaptation (SEA), which refines base-model logits via Langevin dynamics on an RLHF energy function at inference time, instead of searching over discrete sampled responses.
Problem
Inference-time alignment methods such as Best-of-N search over a discrete set of base-model responses, so they struggle when the base policy is weak or the candidate set is small.
Approach
Treat the optimal RLHF policy as an energy-based model over continuous logits and run gradient-based Langevin dynamics initialized from base-policy logits, using reward-model gradients (via softmax straight-through estimation) to adapt responses toward high-reward regions before decoding.
Key results
- On AdvBench with LLaMA-3.2-1B-Base, SEA reduces harmful rate to 5.58% versus 65.96% for SFT and 43.85% for BoN-64 (91.54% relative gain over the second-best).
- On MATH with LLaMA-3.2-1B-Instruct, SEA reaches 32.00% accuracy versus 27.50% for SFT and 16.00% for BoN-64.
- On TruthfulQA with LLaMA-3.2-1B-Base, SEA achieves truthful rate 78.0% and informative rate 100.0%, matching or exceeding BoN-64 (78.0%).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use SEA instead of CBS | LLaMA-3.2-1B-Base · AdvBench | Harmful Rate (%): 24.81 -> 5.58 | Helped |
| Use SEA instead of BoN-64 | LLaMA-3.2-1B-Base · TruthfulQA | Truthful Rate (%): 78.0 -> 78.0 | No clear effect |
| Use SEA instead of CBS | LLaMA-3.2-1B-Base · TruthfulQA | Diversity: 0.86 -> 0.90 | Helped |
| Use SEA instead of BoN-64 | LLaMA-3.2-1B-Instruct · GSM8K | Accuracy (%): 57.00 -> 58.00 | Helped |
| Use SEA instead of BoN-8 | LLaMA-3.2-1B-Instruct · MATH | Accuracy (%): 19.50 -> 32.00 | Helped |
| Replace original-response initialization with Gaussian-noise initialization (ablation) | SEA, LLaMA-3.2-1B-Base · AdvBench | Harmful Rate (%): 5.58 -> 4.04 | Helped |
| Remove reward model guidance (ablation) | SEA, LLaMA-3.2-1B-Base · AdvBench | Harmful Rate (%): 5.58 -> 19.62 | Hurt |
| Remove reference-model regularization (ablation) | SEA, LLaMA-3.2-1B-Base · AdvBench | Harmful Rate (%): 5.58 -> 12.69 | Hurt |
| Remove Gaussian noise from Langevin dynamics (ablation) | SEA, LLaMA-3.2-1B-Base · AdvBench | Harmful Rate (%): 5.58 -> 6.73 | Hurt |
| Use one initialization instead of four (ablation) | SEA, LLaMA-3.2-1B-Base · AdvBench | Harmful Rate (%): 5.58 -> 13.65 | Hurt |
| Use one initialization and replace original-response initialization with Gaussian noise (ablation) | SEA, LLaMA-3.2-1B-Base · AdvBench | Harmful Rate (%): 13.65 -> 11.73 | Helped |
| Use one initialization and remove reward model guidance (ablation) | SEA, LLaMA-3.2-1B-Base · AdvBench | Harmful Rate (%): 13.65 -> 31.15 | Hurt |
Surprising findings
- Even without reward guidance (w/o Reward ablation), Langevin noise-driven exploration still improves over SFT on AdvBench.
- Random initialization sometimes outperforms initialization from the original response on AdvBench, because original harmful responses have very low reward.
Limitations
- Higher inference latency than training-time alignment methods.
- Requires a revision prompt to handle unknown output length in non-autoregressive optimization.
Experimental setup
Datasets: AdvBench TruthfulQA GSM8K MATH HH-RLHF
Benchmarks: AdvBench TruthfulQA GSM8K MATH HH-RLHF
Metrics: average reward harmful rate truthful rate informative rate diversity accuracy win rate KL divergence
Baselines: SFT BoN-8 BoN-32 BoN-64 Rejection sampling ARGS CBS Re-Control
Builds on: Best-of-N sampling COLD COLD-Attack ARGS Chunk-level beam search
Closest papers by meaning
- Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment (ICML 2025 · similarity 0.91)
- InfAlign: Inference-aware language model alignment (ICML 2025 · similarity 0.90)
- Learnable Chernoff Baselines for Provable Inference-Time Alignment (NeurIPS 2026 · similarity 0.89)
- Aligning Large Language Models with Representation Editing: A Control Perspective (NeurIPS 2024 · similarity 0.89)
- Bounded Rationality for LLMs: Satisficing Alignment at Inference-Time (ICML 2025 · similarity 0.89)
- W2S-AlignTree: Weak-to-Strong Inference-Time Alignment for Large Language Models via Monte Carlo Tree Search (AAAI 2026 · similarity 0.88)
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.