💬 NLP

Safe Inference-Time Alignment via Lagrangian Reward Augmentation

The paper proposes Lagrangian Reward Augmentation (LARA), a general inference-time alignment framework that dynamically calibrates an augmented reward signal via a dual variable to enforce safety constraints without retraining, thereby improving the helpfulness-harmlessness tradeoff and approaching the performance of finetuning-based methods.

Original authors: Yaswanth Chittepu, Ativ Joshi, Sohini Chintala, Scott Niekum

Published 2026-07-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yaswanth Chittepu, Ativ Joshi, Sohini Chintala, Scott Niekum

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Helpful but Dangerous" Robot

Imagine you have a very smart robot assistant (a Large Language Model). You want it to be helpful (give good answers) but also harmless (not give dangerous advice, like how to build a bomb or give bad medical advice).

Usually, to make a robot safe, you have to go back to the drawing board and retrain it from scratch. This is like hiring a new teacher to re-educate a student every time the rules change. It's expensive, slow, and requires a lot of data.

Some newer methods try to guide the robot while it's speaking (during "inference") without retraining. But these methods have a flaw: they usually ask the human operator to guess a "magic number."

Example: "Okay, let's add a penalty of 5 points for bad words."
The Problem: If you pick 5, the robot might still be dangerous. If you pick 10, the robot might become too scared to say anything useful. It's a game of "guess and check" that relies on human intuition rather than math.

The Solution: LARA (The "Smart Budget" System)

The authors propose a new framework called LARA (Lagrangian Reward Augmentation). Instead of guessing a magic number, LARA uses a mathematical trick to calculate the exact right balance between being helpful and being safe.

Think of it like a family road trip with a strict gas budget.

1. The Setup: The Car and the Map

The Car: The frozen language model (the robot). It's already built and we aren't changing its engine (weights).
The Map: The "Reward Model." This tells the car how fast it can go to get to the destination (Helpfulness).
The Gas Gauge: The "Cost Model." This measures how much "danger" or "harm" is being used up (Harmlessness).
The Budget: You have a strict limit on how much "harm" you can spend (e.g., "We can only spend 10 units of danger").

2. The Old Way vs. The LARA Way

The Old Way (Manual Tuning): You tell the driver, "Drive fast, but if you think you're using too much gas, slow down a little." The driver has to guess how much to slow down. Sometimes they drive too fast and run out of gas (unsafe). Sometimes they drive too slow and never get there (unhelpful).
The LARA Way (The Dual Variable): LARA acts like a smart GPS calculator. Before the trip starts, it looks at a small sample of the road (a calibration set) and calculates the exact price of gas needed to stay within the budget while getting the best speed possible.
- It finds a single number (called $\lambda$ or "lambda").
- This number represents the "shadow price" of safety. It tells the car: "For every unit of danger you take, you lose X amount of helpfulness."
- The car then drives using a new rule: Helpfulness minus (Lambda $\times$ Danger).

3. How It Works (The "Calibration" Step)

LARA doesn't need to retrain the whole robot. It just needs a small "test drive" (a calibration set).

It generates a few sample answers.
It checks how "helpful" and how "risky" they are.
It runs a quick math search (like finding the perfect temperature on a thermostat) to find the one specific number ( $\lambda$ ) that keeps the total risk under the limit while maximizing the helpfulness.
Once it finds this number, it plugs it into the robot's existing "scoring system."

Why This Is Special

The paper claims two main things about how this works:

1. For "Best-of-N" (Picking the Best Answer from a List)
Imagine the robot writes 20 different answers to your question.

Without LARA: You might pick the one that sounds best, but it could be dangerous.
With LARA: The system uses the calculated "Lambda" to score all 20 answers. It automatically picks the one that is the most helpful while strictly staying under the safety budget. The paper proves mathematically that this method finds the perfect balance point.

2. For "Token-Level" (Guiding the Robot Word-by-Word)
Imagine the robot is writing a sentence one word at a time.

With LARA: Instead of guessing how much to penalize a risky word, the system uses the calculated "Lambda" to adjust the probability of the next word. It's like a traffic cop who knows the exact speed limit and directs the car to stay just under it, rather than waving a hand and hoping for the best. The paper calls this a "principled heuristic"—a smart rule of thumb based on math, not a guess.

The Results

The authors tested this on two popular robot models (Llama and Qwen).

The Finding: LARA made the robots much better at balancing helpfulness and safety compared to other "guessing" methods.
The Comparison: The "Best-of-N" method using LARA performed almost as well as the expensive, retrained models (which require weeks of training), but it did it in seconds without changing the robot's brain.
The Trade-off: It successfully stopped the robots from being harmful without making them useless.

Summary Analogy

If training a safe AI is like rebuilding a house to make it fireproof, LARA is like installing a smart fire sprinkler system that automatically adjusts the water pressure based on the size of the fire. You don't need to rebuild the house; you just need to calibrate the sprinkler once, and it keeps the house safe while letting you live comfortably inside.

Technical Summary: Safe Inference-Time Alignment via Lagrangian Reward Augmentation (LARA)

Problem Statement

Reinforcement Learning from Human Feedback (RLHF) is the standard framework for aligning Large Language Models (LLMs) with human preferences. However, traditional alignment relies on policy fine-tuning (e.g., via RLHF or Direct Preference Optimization like DPO), which is computationally expensive and requires retraining when safety requirements or user populations change.

While inference-time alignment methods (such as Best-of-N reranking or token-level reward-guided decoding) have emerged to steer frozen models using auxiliary reward signals without weight updates, they currently lack a principled mechanism for safe alignment. Existing methods typically optimize a single scalar score. When helpfulness and safety (harmlessness) objectives conflict, practitioners must manually tune the trade-off between them. Furthermore, these methods do not inherently enforce explicit safety constraints (e.g., a budget on expected cost) during decoding, leading to either unsafe outputs or unnecessary reductions in helpfulness.

Methodology: Lagrangian Reward Augmentation (LARA)

The authors propose LARA, a decoder-agnostic framework that transfers the constrained optimization structure of Safe RLHF to the inference stage. Instead of updating model weights, LARA calibrates a scalar dual variable to create an augmented reward signal.

1. Theoretical Formulation

The method begins with the Safe RLHF objective, which maximizes expected helpfulness ( $r$ ) subject to an expected cost constraint ( $c \le \tau$ ), regularized by a KL-divergence term against a reference policy ( $\pi_{ref}$ ):
$\max_{\pi} \mathbb{E}[r(x,y)] - \beta D_{KL}(\pi \parallel \pi_{ref}) \quad \text{s.t.} \quad \mathbb{E}[c(x,y)] \le \tau$

Using Lagrangian duality, the authors dualize the safety constraint. This transforms the infinite-dimensional policy optimization problem into a one-dimensional convex optimization problem over a non-negative dual variable $\lambda$ (the Lagrange multiplier).

Augmented Reward: The dual variable defines an augmented reward: $r_\lambda(x, y) = r(x, y) - \lambda c(x, y)$ .
Dual Objective: The dual function $g(\lambda)$ is shown to be convex. Its derivative corresponds to the difference between the safety budget $\tau$ and the expected cost under the Gibbs-tilted policy induced by $r_\lambda$ .
Optimality: The optimal $\lambda^*$ satisfies the complementary slackness condition, ensuring the expected cost equals the budget $\tau$ (if the constraint is active) while maximizing helpfulness.

2. Empirical Calibration

LARA does not require solving the optimization analytically. Instead, it estimates $\lambda^*$ via empirical calibration on a small dataset:

Data: A small set of prompts is sampled, and multiple responses are generated from the frozen reference policy.
Estimation: The dual variable is estimated by minimizing the empirical dual objective over a compact interval $[0, \Lambda]$ .
Algorithm: Due to the convexity and monotonicity of the dual gradient, the optimal $\lambda$ is found efficiently using binary search.
Robustness: To handle distributional mismatch between calibration and deployment, the authors suggest bootstrapping the calibration set and selecting an upper quantile of the estimated $\lambda$ distribution as a conservative choice.

3. Integration with Decoding

Once calibrated, the scalar $\hat{\lambda}$ is used to compute the augmented reward $r_{\hat{\lambda}}$ . This score is passed as a "drop-in" signal to existing inference-time alignment methods:

Sequence-level (e.g., Best-of-N): The calibrated $\lambda$ yields an exact solution to the expected-cost constrained problem.
Token-level (e.g., Reward-Guided Search): The calibrated $\lambda$ provides a principled heuristic derived from the constrained objective, though it does not guarantee an exact population-level constraint satisfaction due to the greedy nature of token-level decoding.

Key Contributions

Formulation: The paper formulates safe inference-time alignment as the dual of the Safe RLHF objective, reducing the problem to a one-dimensional convex calibration task.
Geometry & Efficiency: The authors characterize the geometry of the dual objective, proving its convexity and the monotonicity of its gradient. This enables an efficient binary-search calibration procedure.
Theoretical Guarantees: A finite-sample guarantee is provided for the estimator $\hat{\lambda}_{N,K}$ , showing that the error concentrates around the population minimizer $\lambda^*$ , controlled by the number of prompts ( $N$ ) and samples per prompt ( $K$ ).
Decoder-Agnostic Framework: LARA is designed as a calibration layer compatible with any existing inference-time alignment method, replacing manual scalar tuning with a principled safety-constrained construction.

Experimental Results

The authors evaluated LARA on Llama-3.2-3B and Qwen-2.5-3B models using the BeaverTails dataset.

Constraint Satisfaction: Experiments confirmed that the calibrated $\lambda^*$ sits at the boundary of the safety constraint. Values $\lambda < \lambda^*$ violated the safety budget, while $\lambda > \lambda^*$ satisfied the constraint but unnecessarily reduced helpfulness.
Performance vs. Baselines:
- Inference-time: Among inference-time methods, Best-of-N (BoN) combined with LARA achieved the best helpfulness-harmlessness trade-off. It outperformed SafeDPO on helpfulness while remaining competitive on harmlessness, all without weight updates.
- Token-level: Token-level methods (ARGS, PARGS, CD-Q) generally performed worse on helpfulness compared to BoN, though CD-Q showed strong harmlessness improvements with increasing $\lambda$ .
- vs. Fine-tuning: While LARA-driven BoN approached the performance of fine-tuning-based baselines (Safe-RLHF, SafeDPO), the fine-tuned methods still achieved the overall best trade-off. However, LARA offers a significant advantage in deployment flexibility and cost.
Evaluation: Results were consistent across model-based evaluations (using BeaverTails reward/cost models) and GPT-4o-mini based evaluations, though GPT showed less sensitivity to $\lambda$ variations than the reward models.

Significance and Claims

The paper claims that LARA provides a principled alternative to manual penalty tuning for safe inference-time alignment. By dualizing the safety constraint, LARA translates a safety budget into a decoding objective without modifying the model weights.

For Sequence-level Samplers: LARA offers an exact expected-cost interpretation, theoretically guaranteeing the optimal trade-off under the constrained objective.
For Token-level Decoders: It offers a principled heuristic that respects the constrained objective's geometry, even if it cannot strictly enforce the population constraint.

The authors conclude that LARA successfully transfers safety constraints from training-time alignment to practical decoding-time control, improving the helpfulness-harmlessness trade-off and bringing inference-time methods closer to the performance of fine-tuning-based approaches. The work emphasizes that safety constraints can be managed effectively at inference time, avoiding the computational overhead of repeated retraining.