# Safe Inference-Time Alignment via Lagrangian Reward Augmentation

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

## Key facts

- Cited by 1 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 24 of its 35 references that are in the map.
- First posted on arXiv 2 Jul 2026.

## Placed in

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

## Builds on (most cited first, 24 of 24)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [RealToxicityPrompts: Evaluating Neural Toxic Degeneration in Language Models](https://mlsanity.com/paper/2009.11462/) (2020)
- [Scaling Laws for Reward Model Overoptimization](https://mlsanity.com/paper/2210.10760/) (2022)
- [Ethical and social risks of harm from Language Models](https://mlsanity.com/paper/2112.04359/) (2021)
- [BeaverTails: Towards Improved Safety Alignment of LLM via a Human-Preference Dataset](https://mlsanity.com/paper/2307.04657/) (2023)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Red Teaming Language Models to Reduce Harms: Methods, Scaling Behaviors, and Lessons Learned](https://mlsanity.com/paper/2209.07858/) (2022)
- [Safe RLHF: Safe Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2310.12773/) (2023)
- [FUDGE: Controlled Text Generation With Future Discriminators](https://mlsanity.com/paper/2104.05218/) (2021)
- [GatorTron: A Large Clinical Language Model to Unlock Patient Information from Unstructured Electronic Health Records](https://mlsanity.com/paper/2203.03540/) (2022)
- [Scaling Laws for Reward Model Overoptimization in Direct Alignment Algorithms](https://mlsanity.com/paper/2406.02900/) (2024)
- [Controlled Decoding from Language Models](https://mlsanity.com/paper/2310.17022/) (2023)
- [Theoretical guarantees on the best-of-n alignment policy](https://mlsanity.com/paper/2401.01879/) (2024)
- [ARGS: Alignment as Reward-Guided Search](https://mlsanity.com/paper/2402.01694/) (2024)
- [Reward-Augmented Decoding: Efficient Controlled Text Generation With a Unidirectional Reward Model](https://mlsanity.com/paper/2310.09520/) (2023)
- [Value Augmented Sampling for Language Model Alignment and Personalization](https://mlsanity.com/paper/2405.06639/) (2024)
- [SafeDPO: A Simple Approach to Direct Preference Optimization with Enhanced Safety](https://mlsanity.com/paper/2505.20065/) (2025)
- [A Critical Look At Tokenwise Reward-Guided Text Generation](https://mlsanity.com/paper/2406.07780/) (2024)
- [Towards Cost-Effective Reward Guided Text Generation](https://mlsanity.com/paper/2502.04517/) (2025)
- [Safe RLHF Beyond Expectation: Stochastic Dominance for Universal Spectral Risk Control](https://mlsanity.com/paper/2603.10938/) (2026)

## Built on by (most cited first, 1 of 1)

- [Safety Hacking in Constrained Best-of-N Inference-time Scaling](https://mlsanity.com/paper/2608.22915/) (2026)

## Abstract (arXiv)

Inference-time alignment steers a frozen language model during decoding using auxiliary reward signals, avoiding the cost of repeated weight updates. However, existing inference-time alignment methods typically optimize a single scalar score, so explicit safety constraints must either be ignored or encoded through manually tuned penalties. We propose Lagrangian Reward Augmentation (LARA), a general inference-time alignment framework under safety constraints. Starting from a KL-regularized constrained objective with a reward model and a cost model, LARA dualizes the constraint and reduces the optimization problem to a one-dimensional convex problem over a nonnegative dual variable. Estimated on a small calibration set, this dual variable defines an augmented reward that can be used as a drop-in scoring signal within existing inference-time alignment methods. For sequence-level sampling methods, such as Best-of-N reranking, the calibrated dual variable corresponds to the solution of the expected-cost constrained problem. For token-level reward-guided decoding methods, the same construction yields a principled dual-calibrated heuristic rather than an exact constrained-policy guarantee. We evaluate LARA on both sequence-level and token-level inference-time alignment methods, and find that LARA improves the helpfulness-harmlessness tradeoff, with Best-of-N achieving the best performance among inference-time methods, approaching finetuning-based direct alignment baselines.