Lagrange Reward Augmentation (LARA) enables safe inference-time alignment

Researchers propose Lagrangian Reward Augmentation (LARA), a framework for steering frozen language models during decoding under explicit safety constraints without repeated weight updates. LARA dualizes the constraint to reduce optimization to a one-dimensional convex problem, producing an augmented reward that serves as a drop-in scoring signal for existing methods.

• The method derives from a KL-regularized constrained objective using reward and cost models.
• It estimates a dual variable on a small calibration set to define the augmented reward.
• For sequence-level sampling like Best-of-N reranking, the variable solves the expected-cost constrained problem.
• For token-level decoding, it provides a principled dual-calibrated heuristic.

LARA improves the helpfulness-harmlessness tradeoff, with Best-of-N performance approaching finetuning-based direct alignment baselines.