Safe Inference-Time Alignment via Lagrangian Reward Augmentation

2026-07-02

AI Summary

Existing inference-time alignment methods struggle to explicitly satisfy safety constraints while optimizing reward, often relying on manually tuned penalty terms. This work proposes LARA, a novel framework that introduces Lagrangian duality into the inference phase for the first time, transforming the constrained optimization problem into a one-dimensional convex problem. By estimating dual variables on a calibration set, LARA constructs an augmented reward signal compatible with both sequence-level and token-level alignment strategies. Without requiring any fine-tuning, LARA principledly balances helpfulness and harmlessness, achieving state-of-the-art performance among inference-time methods when combined with Best-of-N reranking. It substantially improves the helpfulness–harmlessness trade-off, closely approaching the performance of fine-tuning-based direct alignment baselines.

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.

Problem

• inference-time alignment
• safety constraints
• reward augmentation
• language model alignment
• constrained optimization

Innovation

• Lagrangian Reward Augmentation
• inference-time alignment
• safety constraints
• reward-cost tradeoff
• dual calibration