Accepted at COLM 2026 Conference on Language Modeling · San Francisco · Oct 6–9, 2026

Safe Inference-Time Alignment via Lagrangian Reward Augmentation

Yaswanth Chittepu¹
Ativ Joshi¹
Sohini Chintala²
Scott Niekum¹

¹ University of Massachusetts Amherst · ² Independent Researcher

Find us at COLM Poster #109· Poster Session 6· Thu, Oct 8, 4:40 PM PDT· Grand Ballroom

\(\R{r} - \Lam{\lambda^\star}\,\C{c}\)

One number makes any decoder safety-aware. LARA finds the multiplier \(\Lam{\lambda^\star}\) that trades helpfulness reward against harmfulness cost exactly at your safety budget. It calibrates in minutes on a small prompt set and plugs into Best-of-\(N\), ARGS, PARGS or CD-Q, with no weight updates.

The problem

Inference-time alignment optimizes one score. Safety needs two.

Inference-time methods steer a frozen model while it decodes, which avoids expensive finetuning. But they expect one scalar score. When helpfulness and safety conflict, someone has to pick the weight between them by hand, and nothing guarantees that weight meets a safety requirement.

Read the 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.

Today: hand-tuned penalty

\(\R{r} - w\,\C{c}\)

Pick \(w\) by trial and error. Too small, and responses exceed the harm budget. Too large, and helpfulness drops for no gain in safety. Retune whenever the model or the budget changes.

LARA: derived from the constraint

\(\R{r} - \Lam{\lambda^\star}\,\C{c} \quad \text{s.t.} \quad \mathbb{E}[\C{c}] \le \tau\)

State the safety budget \(\tau\) directly. LARA solves for the smallest penalty that meets it. \(\Lam{\lambda^\star}\) is the shadow price of harm.

How it works

Move the Safe RLHF constraint from training to decoding

We start from the same KL-regularized, cost-constrained objective that Safe RLHF solves with RL. After dualizing the constraint, the search over policies collapses to a search over a single nonnegative number.

Sample and score

2,000 prompts × 20 responses

Draw responses from the frozen reference model on a small calibration set. Score each one with the reward model \(\R{r}\) and the cost model \(\C{c}\).
Calibrate λ★

binary search on a 1-D convex dual

The dual's slope, \(\tau - \mathbb{E}[\C{c}]\), is monotone in \(\lambda\), so binary search finds the root. For a safety margin we take the 97.5th percentile over 10,000 bootstrap resamples, which takes about 2 minutes.
Decode with any method

drop-in scoring signal

Give \(\R{r} - \Lam{\lambda^\star}\,\C{c}\) to Best-of-\(N\), ARGS, PARGS or CD-Q in place of their usual reward. The decoder itself does not change.

Constrained goal

\[\max_{\pi}\; \mathbb{E}_{x\sim D_x}\Big[\mathbb{E}_{y\sim\pi(\cdot\mid x)}\,\R{r(x,y)} - \beta\, D_{\mathrm{KL}}\big(\pi(\cdot\mid x)\,\|\,\pi_{\mathrm{ref}}(\cdot\mid x)\big)\Big]\]\[\text{s.t.}\quad \mathbb{E}_{x\sim D_x,\; y\sim\pi(\cdot\mid x)}\,\C{c(x,y)} \le \tau\]

Closed-form dual

\[g(\Lam{\lambda}) = \beta\, \mathbb{E}_{x\sim D_x}\!\left[\log \mathbb{E}_{y\sim\pi_{\mathrm{ref}}(\cdot\mid x)} \exp\!\left(\frac{\R{r(x,y)} - \Lam{\lambda}\,\C{c(x,y)}}{\beta}\right)\right] + \Lam{\lambda}\,\tau\]\[\pi^\star_{\Lam{\lambda}}(y\mid x) \;\propto\; \pi_{\mathrm{ref}}(y\mid x)\, \exp\!\left(\frac{\R{r(x,y)} - \Lam{\lambda}\,\C{c(x,y)}}{\beta}\right)\]

Strong duality holds, and for a fixed \(\lambda\) the optimal policy is a Gibbs tilt of the reference model by the net reward \(\R{r} - \lambda\,\C{c}\).

Why search is easy

\[g^{\prime}(\Lam{\lambda}) = \tau - \mathbb{E}_{x\sim D_x,\; y\sim\pi^\star_{\Lam{\lambda}}(\cdot\mid x)}\big[\C{c(x,y)}\big]\]\[g^{\prime\prime}(\Lam{\lambda}) = \frac{1}{\beta}\, \mathbb{E}_{x\sim D_x}\Big[\operatorname{Var}_{y\sim\pi^\star_{\Lam{\lambda}}(\cdot\mid x)}\big(\C{c(x,y)}\big)\Big] \ge 0\]

The dual is convex, so expected cost falls monotonically as \(\lambda\) rises. At the optimum, expected cost equals the budget exactly. A finite-sample bound shows the estimated \(\hat\lambda\) concentrates around \(\lambda^\star\) as the number of prompts \(N\) and samples per prompt \(K\) grow.

Try it

Tighten the budget and watch the price of harm rise

Each position below is a real calibration run for Llama-3.2-3B with Best-of-20. Lower cost means safer. The paper's main setting is \(\tau = -4\).

Safety budget \(\tau\) (maximum allowed expected cost)

λ★0.605

𝔼[c]−4.131

Slack−0.131

Constraint met

Source: paper appendix, τ sweep table. The shaded region is over budget.

Result 1 · Constraint satisfaction

For Best-of-N, λ★ lands right on the safety boundary

We sweep \(\lambda\) and rerank with \(\R{r} - \lambda\,\C{c}\). For any \(\lambda\) below \(\Lam{\lambda^\star}\), expected cost exceeds \(\tau\). Any larger value still meets the budget, but gives up helpfulness for no reason. The calibrated value sits at the boundary, as the theory predicts. The same pattern holds for Llama-3.1-8B.

Best-of-N sweep over lambda for Llama-3.2-3B: lambda star of 0.6 sits at the threshold tau of minus 4. — Expected helpfulness 𝔼[r] against expected cost 𝔼[c] for the Best-of-20 policy at each value of λ. Points to the right of the dashed line exceed the budget.

Result 2 · Against baselines

Close to finetuning quality, without finetuning

We compare LARA-guided decoding with finetuned baselines (SFT, SafeDPO, Safe-RLHF) and with the same decoders at \(\lambda = 0\) (no safety penalty) and at \(\lambda_{\max} = 2\) (over-penalized). The evaluation uses 1,024 BeaverTails test prompts. Up and to the right is better.

Helpfulness versus harmlessness under the beaver-7b reward and cost models, Llama-3.2-3B. — Left: scores from the beaver-7b-unified reward and cost models. Right: 1–10 ratings from GPT-4o-mini, using the SafeDPO evaluation prompts.

Helpfulness versus harmlessness rated by GPT-4o-mini, Llama-3.2-3B. — Left: scores from the beaver-7b-unified reward and cost models. Right: 1–10 ratings from GPT-4o-mini, using the SafeDPO evaluation prompts.

Best-of-N
Best-of-N with λ★ is the strongest inference-time method. Under the gold models it is more helpful than SafeDPO and about as harmless, with no weight updates. At λ = 0 its responses are highly harmful, and at λ_max it loses helpfulness for little safety gain.
Token-level
Among token-level methods, CD-Q comes out ahead. It trains a prefix value function, whereas ARGS scores prefixes with a full-sequence reward model and PARGS assumes each prefix inherits its full response's label. ARGS and PARGS fall below SFT on helpfulness.
Finetuning
Safe-RLHF still has the best overall tradeoff. It pays for that with full RL training. LARA gets most of the way there using a 2-minute calibration.
Evaluators
GPT-4o-mini tends to merge helpfulness and harmlessness into one score. Its two ratings correlate at 0.73, against 0.13 for the gold models. That explains why Best-of-N looks flat across λ under GPT. For judging constraint satisfaction, use the gold models, which score the two axes independently.

Citation

BibTeX

@inproceedings{chittepu2026lara,
  title     = {Safe Inference-Time Alignment via Lagrangian Reward Augmentation},
  author    = {Chittepu, Yaswanth and Joshi, Ativ and Chintala, Sohini and Niekum, Scott},
  booktitle = {Conference on Language Modeling (COLM)},
  year      = {2026}
}