Safe Inference-Time Alignment via Lagrangian Reward Augmentation
Yaswanth Chittepu, Ativ Joshi, Sohini Chintala, Scott Niekum
Why It Matters
What makes this one worth your time
This work is relevant for AI researchers and engineers interested in improving language model alignment without the computational cost of retraining, particularly in safety-critical applications.
LARA enhances inference-time alignment with safety constraints using a novel dualized optimization approach.
Summary
The paper introduces Lagrangian Reward Augmentation (LARA), a framework for inference-time alignment of language models under safety constraints, utilizing a dualized optimization approach to enhance reward signals without modifying model weights.
Key contributions
- Proposes Lagrangian Reward Augmentation (LARA) for inference-time alignment under safety constraints.
- Introduces a dualized optimization approach to reduce the problem to a one-dimensional convex problem.
- Demonstrates improved helpfulness-harmlessness tradeoff in language models using LARA.
Notable insights
- The use of a dualized optimization problem to handle safety constraints in inference-time alignment is a novel approach.
- LARA's framework allows for the integration of safety constraints into existing inference-time methods without altering model weights.
Possible limitations
- Not stated in the abstract
Abstract
arXiv:2607.02781v1 Announce Type: cross Abstract: 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.