Paper deep dive
Safe Inference-Time Alignment via Lagrangian Reward Augmentation
Yaswanth Chittepu, Ativ Joshi, Sohini Chintala, Scott Niekum
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 93%
Last extracted: 7/7/2026, 4:34:36 AM
Summary
The paper introduces Lagrangian Reward Augmentation (LARA), a framework for safe inference-time alignment of frozen language models. LARA transfers the constrained optimization structure of Safe RLHF to decoding time by dualizing the safety constraint, reducing it to a one-dimensional convex problem over a nonnegative dual variable (λ). This calibrated λ generates an augmented reward that seamlessly integrates with existing sequence-level and token-level alignment methods, improving the helpfulness-harmlessness tradeoff without requiring weight updates. Best-of-N reranking demonstrates the strongest performance among inference-time approaches, closely approaching finetuning-based baselines.
Entities (9)
Relation Signals (8)
LARA → uses → Dual variable λ
confidence 97% · LARA dualizes the constraint and reduces the optimization problem to a one-dimensional convex problem over a nonnegative dual variable.
LARA → addresses → Inference-time alignment
confidence 96% · We address this gap with Lagrangian Reward Augmentation (LARA), an inference-time alignment framework that transfers the constrained trade-off...
LARA → buildsupon → Safe RLHF
confidence 95% · LARA transfers the constrained trade-off of Safe RLHF to decoding time.
LARA → reduces → KL-regularized objective
confidence 94% · 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...
Best-of-N → isa → Sequence-level sampling method
confidence 93% · 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.
Reward Model → iscombinedwith → Cost Model
confidence 92% · Starting from a KL-regularized constrained objective with a reward model and a cost model, LARA dualizes the constraint...
Best-of-N → achieves → Best performance among inference-time methods
confidence 91% · Best-of-N achieving the best performance among inference-time methods, approaching finetuning-based direct alignment baselines.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2607.02781v1
- Canonical: https://arxiv.org/abs/2607.02781v1
PDF not stored locally. Use the link above to view on the source site.
Full Text
76,225 characters extracted from source content.
Expand or collapse full text
Safe Inference-Time Alignment via Lagrangian Reward Augmentation Yaswanth Chittepu1, Ativ Joshi1, Sohini Chintala2, Scott Niekum1 1University of Massachusetts Amherst, 2Independent Researcher Correspondence to ychittepu@umass.edu 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. 1 Introduction Reinforcement Learning from Human Feedback (RLHF) (Ouyang et al., 2022) has become a standard framework for aligning large language models with human preferences by first collecting preference data, then learning a reward model, and finally fine-tuning the model to prefer higher-reward outputs. In many deployments, however, usefulness alone is insufficient: modern assistants must also avoid harmful behavior, especially in sensitive domains such as medical consultation (Yang et al., 2022; Moor et al., 2023), legal reasoning (Katz et al., 2024), and educational support (Kasneci et al., 2023; Kung et al., 2023). Safe RLHF (Dai et al., 2023), motivated by the helpfulness and safety objectives conflict (Gehman et al., 2020; Weidinger et al., 2021; Ganguli et al., 2022), learns separate reward and cost models and optimizes under an explicit safety constraint. However, traditional alignment still relies on policy fine-tuning, whether through the classical RLHF pipeline or through direct alignment methods such as DPO (Rafailov et al., 2023) and related Direct Alignment Algorithms (Rafailov et al., 2024) that update model weights to satisfy preference data. Such retraining is expensive and may need to be repeated across applications, user populations, or changing safety requirements. This cost has motivated a growing line of work on inference-time alignment, which leaves the base model frozen and instead steers decoding using auxiliary reward signals. Representative examples include sequence-level Best-of-N (BoN) reranking (Stiennon et al., 2022; Nakano et al., 2022; Beirami et al., 2024) and token-level reward-guided search methods (Khanov et al., 2024; Rashid et al., 2024; Mudgal et al., 2023; Deng & Raffel, 2023; Rashid et al., 2025; Han et al., 2024), all of which seek to obtain better aligned generations without running a full alignment stage, that would involve expensive model fine-tuning. Despite this promise, existing inference-time alignment methods do not provide a satisfactory solution for safe alignment. They require the practitioner to specify a scalar decoding score ahead of time, so when helpfulness and safety compete, the trade-off is typically hand-tuned rather than derived from a principled constrained objective. Sequence-level reranking methods can select high-scoring responses, but they do not by themselves enforce an explicit safety budget during selection. Token-level reward-guided methods can bias decoding online, but they still optimize a pre-specified scalar score rather than the solution of a constrained safe-alignment problem. As a result, current inference-time methods offer control, but not a principled mechanism for translating a safety constraint into a decoding objective. We address this gap with Lagrangian Reward Augmentation (LARA), an inference-time alignment framework that transfers the constrained trade-off of Safe RLHF (Dai et al., 2023) to decoding. The key idea is to dualize the safety constraint, yielding a single nonnegative multiplier λ and the augmented reward rλ(x,y)=r(x,y)−λc(x,y).r_λ(x,y)=r(x,y)-λ c(x,y). Calibrating λ on a small set of prompts produces a safety-aware score that can be plugged into existing inference-time alignment methods. For sequence-level samplers such as Best-of-N (Stiennon et al., 2022; Nakano et al., 2022; Beirami et al., 2024), this construction corresponds directly to the expected-cost constrained optimum. For token-level reward-guided decoders (Khanov et al., 2024; Rashid et al., 2024; Mudgal et al., 2023; Deng & Raffel, 2023; Rashid et al., 2025; Han et al., 2024), it instead provides a principled dual-calibrated heuristic induced by the same constrained objective. Our contributions are as follows. First, we formulate safe inference-time alignment through the dual of the Safe RLHF constrained objective, showing that the original policy optimization reduces to a one-dimensional convex calibration problem whose solution defines a safety-calibrated augmented reward. Second, we characterize the geometry of the dual objective, establishing monotonicity of the dual gradient and thereby yielding an efficient binary-search calibration procedure. Third, we prove a finite-sample guarantee for estimating the optimal dual variable from a small calibration dataset. Fourth, we show how the resulting augmented reward can be used within a broad class of existing inference-time alignment methods, with an exact expected-cost interpretation for sequence-level samplers and a principled heuristic interpretation for token-level reward-guided decoders. Experiments show 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 2 Background 2.1 Reinforcement Learning from Human Feedback Reinforcement Learning from Human Feedback (RLHF) (Ouyang et al., 2022) is the dominant paradigm for aligning language models to human preferences. The standard RLHF paradigm involves three stages (Ouyang et al., 2022). The first stage is Supervised Fine Tuning (SFT), where the model is trained on high-quality instruction data provided by human or LLM annotators, through the next-token prediction objective. The second stage is Reward Modeling (RM), wherein a reward model r is trained to capture human preference. The reward model is trained on preference datasets using the Bradley-Terry preference model (Bradley & Terry, 1952). In the final stage, the policy or language model is trained to produce responses preferred by humans through Reinforcement Learning (RL) using the learned reward model. In order to preserve language quality and prevent reward overoptimization(Ouyang et al., 2022; Stiennon et al., 2022; Gao et al., 2022; Rafailov et al., 2024), RLHF optimizes a KL regularized reward objective in the final stage. maxθx∼Dx,y∼πθ(⋅∣x)[r(x,y)]−βDKL(πθ(⋅∣x)∥πref(⋅∣x)), _θ\;E_x D_x,\,y _θ(· x)\! [r(x,y) ]\;-\;β\,D_KL\! ( _θ(· x)\,\|\, _ref(· x) ), (1) 2.2 Safe RLHF Standard RLHF optimizes a single reward function, which may be inadequate in the presence of competing objectives such as helpfulness and harmlessness. Safe RLHF separates these two signals, by learning a reward model rϕ(x,y)r_φ(x,y) from helpfulness preference labels, and a cost model cψ(x,y)c_ψ(x,y) from harmfulness preference labels and then frames learning as a constrained optimization problem. maxθx∼Dx,y∼πθ(⋅∣x)[r(x,y)]−βDKL(πθ(⋅∣x)∥πref(⋅∣x))s.t.x∼Dx,y∼πθ(⋅∣x)[c(x,y)]≤τ. _θ\;E_x D_x,\,y _θ(· x)\! [r(x,y) ]-β D_KL( _θ(· x)\| _ref(· x)) .t. _x D_x,\,y _θ(· x)\! [c(x,y) ]≤τ. (2) Follow-up works, such as HC-RLHF (Chittepu et al., 2025) replace the expectation constraint with a high-confidence probabilistic safety guarantee using the Seldonian framework (Thomas et al., 2019). RAD (Chittepu et al., 2026) instead imposes a distributional dominance constraint, which they solve using optimal transport (Peyré & Cuturi, 2020). Kim et al. (2025) propose SafeDPO, a direct alignment variant of Safe-RLHF that optimizes the safety-constrained RLHF problem using supervised learning. Unlike these methods, which modify the training objective and update model weights, LARA leaves the base model frozen and transfers the reward–cost trade-off to inference time through decoder-side calibration. 2.3 Inference-time Alignment Inference-time alignment seeks to avoid weight updates altogether by steering a frozen base model during decoding. At a high level, sequence-level methods (Nakano et al., 2022; Stiennon et al., 2022; Beirami et al., 2024) sample several complete responses and select the highest scored response, by a reward model, while token-level reward-guided methods (Khanov et al., 2024; Rashid et al., 2024; Mudgal et al., 2023; Deng & Raffel, 2023; Rashid et al., 2025; Han et al., 2024) modify next-token probabilities using an auxiliary scoring function. For a reference policy πref _ref, a prompt x, and a partial response y1:t−1y_1:t-1, a token-level reward-guided rule takes the form score(yt∣x,y1:t−1)=logπref(yt∣x,y1:t−1)+wr(x,y1:t),score(y_t x,y_1:t-1)= _ref(y_t x,y_1:t-1)+w\,r(x,y_1:t), which induces the tilted distribution softmax(score(yt∣x,y1:t−1))∝πref(yt∣x,y1:t−1)exp(wr(x,y1:t)).softmax\! (score(y_t x,y_1:t-1) ) _ref(y_t x,y_1:t-1) \! (w\,r(x,y_1:t) ). Different inference-time alignment methods differ primarily in how this scoring signal is defined or approximated; we discuss these variants in Appendix A. LARA is orthogonal to this decoder design choice: it calibrates the scalar trade-off between reward and cost and then supplies the resulting augmented score to any chosen inference-time procedure, thereby replacing manual scalar trade-off tuning with a principled safety-constrained construction. 3 Method: LAgrangian Reward Augmentation (LARA) Our goal is to perform safety-aware alignment at inference time without updating model weights. We begin from the Safe RLHF objective, which separates helpfulness and safety through a reward model and a cost model, and ask whether its constrained optimization structure can be transferred from training time to decoding time. LAgrangian Reward Augmentation (LARA) does so by moving the safety trade-off into a single scalar dual variable: once this variable is calibrated, it induces an augmented reward that can be used by any standard inference-time alignment procedure. This section proceeds in two steps. Section 3.1 presents the model, the dual formulation, and the empirical calibration procedure used to estimate the dual variable from a small calibration set. The output of this calibration step is a reward of the form rλ(x,y)=r(x,y)−λc(x,y),r_λ(x,y)=r(x,y)-λ c(x,y), which can be passed as a drop-in scoring signal to any existing inference-time decoding algorithm. Thus, LARA should be viewed as a decoder-agnostic calibration layer rather than as a new decoding algorithm. Section 3.2 then gives the theoretical guarantee for this estimator under finite samples. 3.1 Model and Estimation Our starting point is the Safe RLHF formulation, which separates helpfulness and safety through a reward model r(x,y)r(x,y) and a cost model c(x,y)c(x,y), where x is a prompt and y is a response. This subsection has three steps. We first state the constrained objective and explain why strong duality allows us to work with its dual. We then derive the resulting one-dimensional dual objective and the corresponding Gibbs-tilted policy for a fixed multiplier λ. Finally, we describe how λ is estimated empirically from a small calibration set and how the resulting augmented reward is passed to an existing inference-time decoder. We assume throughout that the reward and cost models are fixed and bounded: there exist R,C>0R,C>0 such that |r(x,y)|≤R|r(x,y)|≤ R and |c(x,y)|≤C|c(x,y)|≤ C for all (x,y)(x,y). Let DxD_x denote the prompt distribution, let πref(⋅∣x) _ref(· x) denote the reference distribution over responses for prompt x, let β>0β>0 be the KL regularization coefficient, and let τ∈[−C,C]τ∈[-C,C] be the safety budget. We consider the constrained optimization problem supπx∼Dx[y∼π(⋅∣x)rϕ(x,y)−βDKL(π(⋅∣x)∥πref(⋅∣x))] s.t. x∼Dx,y∼π(⋅∣x)cψ(x,y)≤τ, _π\;E_x D_x [E_y π(· x)r_φ(x,y)-β D_KL\! (π(· x)\,\|\, _ref(· x) ) ] s.t. E_x D_x,\;y π(· x)c_ψ(x,y)≤τ, (3) where the optimization is over conditional distributions π(⋅∣x)π(· x) that are absolutely continuous with respect to πref(⋅∣x) _ref(· x) for DxD_x-almost every x, so that the conditional KL divergence is well defined. The optimization problem in (3) is a concave maximization problem in the policy variable π: the objective is concave in π, and the expected-cost constraint is affine. We assume a Slater condition, namely that there exists a feasible policy π~ π whose expected cost is strictly below τ. Under this regularity condition, strong duality holds. Consequently, solving the primal problem is equivalent to minimizing the dual function over λ≥0λ≥ 0. The key observation behind LARA is that, after dualizing the safety constraint, the infinite-dimensional optimization over policies reduces to a one-dimensional optimization over a scalar dual variable. For λ≥0λ≥ 0, define the net reward rλ(x,y):=rϕ(x,y)−λcψ(x,y),r_λ(x,y):=r_φ(x,y)-λ c_ψ(x,y), (4) and the corresponding Lagrangian ℒ(π,λ):=x∼Dx[y∼π(⋅∣x)rλ(x,y)−βDKL(π(⋅∣x)∥πref(⋅∣x))]+λτ.L(π,λ):=E_x D_x [E_y π(· x)r_λ(x,y)-β D_KL\! (π(· x)\,\|\, _ref(· x) ) ]+λτ. (5) The associated dual function is g(λ):=supπL(π,λ),λ≥0,g(λ):= _πL(π,λ), λ≥ 0, (6) where the supremum is taken over the same class of policies as in the primal problem. The dual problem is then to minimize g over λ≥0λ≥ 0: infλ≥0g(λ). _λ≥ 0g(λ). (7) The next theorem shows that this dual objective admits a closed form, and that for a fixed value of λ, the optimal policy is a Gibbs tilt of the reference model by the net reward rλr_λ. The derivation is similar in spirit to the KL-regularized optimal-policy characterization used by Rafailov et al. (2023), we provide the full proof in Appendix B for completeness. Theorem 1 (Closed-form dual objective). For every λ≥0λ≥ 0, g(λ)=βx∼Dx[logy∼πref(⋅∣x)exp(rϕ(x,y)−λcψ(x,y)β)]+λτ.g(λ)=β\,E_x D_x [ _y _ref(· x) \! ( r_φ(x,y)-λ c_ψ(x,y)β ) ]+λτ. (8) Moreover, for each prompt x, the supremum in equation (6) is attained by the Gibbs-tilted policy πλ⋆(y∣x)=πref(y∣x)exp(rλ(x,y)β)y′∼πref(⋅∣x)exp(rλ(x,y′)β). _λ (y x)= _ref(y x)\, \! ( r_λ(x,y)β )E_y _ref(· x) \! ( r_λ(x,y )β ). (9) Theorem 1 is the main structural step in our method. Equation (8) replaces the original optimization over policies by a one-dimensional minimization over λ, while Equation (9) shows that, for a fixed dual variable, the optimal policy is obtained by tilting the reference model toward responses with high net reward. This is precisely the object that LARA will later approximate. Clearly, for λ=0λ=0, the dual objective reduces to the standard KL-regularized reward maximization objective, and the Gibbs-tilted policy reduces to the standard reward-only tilt. As λ increases, the dual objective incorporates an increasing penalty on harmfulness, and the Gibbs-tilted policy incorporates an increasing penalty on responses with high cost. The optimal dual variable λ⋆λ balances these competing objectives to achieve the best helpfulness subject to the harmlessness budget. We next characterize the geometry of the dual objective, see Appendix C for the proof. Lemma 2 (Derivative and Hessian of the dual objective). The first and second derivatives of g are g′(λ)=τ−x∼Dx,y∼πλ⋆(⋅∣x)[cψ(x,y)],g (λ)=τ-E_x D_x,\;y _λ (· x) [c_ψ(x,y) ], (10) and g′(λ)=1βx∼Dx[Vary∼πλ⋆(⋅∣x)(cψ(x,y))]≥0.g (λ)= 1β\,E_x D_x [Var_y _λ (· x) (c_ψ(x,y) ) ]≥ 0. (11) In particular, g is convex on [0,∞)[0,∞). Moreover, if we restrict λ to a compact interval I:=[0,Λ]I:=[0, ] (for some Λ<∞ <∞) and the variance term in (11) is continuous and bounded away from zero on I, then there exists μ>0μ>0 such that g′(λ)≥μg (λ)≥μ for all λ∈Iλ∈ I, so g is strongly convex on I. Lemma 2 yields the main interpretation of the dual variable. Equation (10) shows that the derivative is exactly the harmlessness budget minus the expected cost under the Gibbs-tilted policy. Hence any interior minimizer λ⋆>0λ >0 satisfies x∼Dx,y∼πλ⋆(⋅∣x)[c(x,y)]=τ,E_x D_x,\;y π _λ (· x)[c(x,y)]=τ, so λ⋆λ has the usual complementary-slackness interpretation as a shadow price for harmfulness. Equation (11) further shows that g is convex, so g′g is monotone nondecreasing on [0,∞)[0,∞); equivalently, the map λ↦x∼Dx,y∼πλ⋆(⋅∣x)[c(x,y)]λ _x D_x,\;y π _λ(· x)[c(x,y)] is monotone nonincreasing. This monotonicity is the key computational advantage of the dual reduction. On a compact search interval I=[0,Λ]I=[0, ], if g′(0)≥0g (0)≥ 0, then the minimizer is the boundary point λ⋆=0λ =0; otherwise any interior minimizer is characterized by the scalar root condition g′(λ⋆)=0g (λ )=0 and can therefore be found efficiently by binary search. If, in addition, g′(λ)≥μ>0g (λ)≥μ>0 on I, then g is strongly convex on I, the minimizer is unique, and this one-dimensional search problem is especially well conditioned. Hence, we therefore estimate λ⋆λ over a compact interval I=[0,Λ]I=[0, ], which turns calibration into a stable one-dimensional search problem. Specifically, we restrict the search to a compact interval because on the unbounded domain [0,∞)[0,∞) the minimizer need not be attained at a finite λ, and a uniform strong-convexity bound need not hold. The compact restriction localizes the problem to a well-conditioned one-dimensional search region. Fix a compact interval I:=[0,Λ]⊂ℝ+I:=[0, ] _+ and let β>0β>0. For each prompt x, define Zx(λ):=Y∼πref(⋅∣x)exp(r(x,Y)−λc(x,Y)β),ϕ(x,λ):=βlogZx(λ)+λτ,Z_x(λ):=E_Y _ref(· x) \! ( r(x,Y)-λ c(x,Y)β ), φ(x,λ):=β Z_x(λ)+λτ, (12) so that g(λ):=X∼Dx[ϕ(X,λ)],λ⋆∈argminλ∈Ig(λ).g(λ):=E_X D_x[φ(X,λ)], λ ∈ _λ∈ Ig(λ). (13) In practice, we calibrate λ using a small calibration dataset of prompts and multiple samples from the reference policy for each prompt. For a candidate value of λ, we approximate the expectations in (12) and (13) by Monte Carlo, average the resulting empirical objective over calibration prompts, and then optimize over λ∈[0,Λ]λ∈[0, ]. By Lemma 2, the corresponding expected cost is monotone in λ, so this calibration reduces to a stable one-dimensional search problem. Now suppose X1,…,XN∼i.i.d.DxX_1,…,X_N .i.d. D_x, and that for each i∈1,…,Ni∈\1,…,N\, conditional on XiX_i, the responses Yi,1,…,Yi,K∼i.i.d.πref(⋅∣Xi)Y_i,1,…,Y_i,K .i.d. _ref(· X_i), with the tuples (Xi,Yi,1,…,Yi,K)(X_i,Y_i,1,…,Y_i,K) i.i.d. across i. We estimate the inner expectation by Z^i,K(λ):=1K∑k=1Kexp(r(Xi,Yi,k)−λc(Xi,Yi,k)β), Z_i,K(λ):= 1K _k=1^K \! ( r(X_i,Y_i,k)-λ c(X_i,Y_i,k)β ), (14) and define the corresponding empirical inner loss and empirical dual objective as ϕ^i,K(λ):=βlogZ^i,K(λ)+λτ,g^N,K(λ):=1N∑i=1Nϕ^i,K(λ), φ_i,K(λ):=β Z_i,K(λ)+λτ, g_N,K(λ):= 1N _i=1^N φ_i,K(λ), (15) An empirical calibration of the dual variable is then obtained by λ^N,K∈argminλ∈Ig^N,K(λ). λ_N,K∈ _λ∈ I g_N,K(λ). (16) Since g^N,K g_N,K is convex in λ, its derivative is monotone nondecreasing on [0,Λ][0, ]. Therefore, when the minimizer is interior, λ^N,K λ_N,K can be computed efficiently by binary search for the root of g^N,K′(λ)=0 g _N,K(λ)=0; otherwise, the solution lies at the boundary. Once λ^N,K λ_N,K has been computed, we define the calibrated augmented reward rλ^N,K(x,y):=r(x,y)−λ^N,Kc(x,y),r_ λ_N,K(x,y):=r(x,y)- λ_N,Kc(x,y), and pass this score to the chosen inference-time decoding procedure. In this sense, LARA calibrates a scalar trade-off parameter rather than introducing a new decoding rule: the decoder remains unchanged, and only its reward signal is replaced by the safety-calibrated augmented reward. Algorithm 1 summarizes the procedure formally. Note that the dual objective in Theorem 1 and the Gibbs policy πλ⋆π _λ are defined at the level of complete responses, so the exact expected-cost interpretation of λ⋆λ applies directly to sequence-level samplers such as Best-of-N reranking. When the calibrated augmented reward is instantiated inside a token-level reward-guided decoder, the resulting procedure need not sample exactly from πλ⋆π _λ over full responses. In that setting, λ^N,K λ_N,K should be interpreted as a principled dual-calibrated trade-off parameter rather than as a quantity that exactly enforces the population expected-cost constraint. Algorithm 1 LARA: Inference-Time Safe Alignment via Dual Calibration 1:Reference model πref _ref, reward model rϕr_φ, cost model cψc_ψ, harmlessness budget τ, KL coefficient β, calibration prompts Dcal=Xii=1ND_cal=\X_i\_i=1^N, search interval I=[0,Λ]I=[0, ], inference-time decoder algorithm DecodeAlg 2:For each i∈1,…,Ni∈\1,…,N\, sample Yi,1,…,Yi,K∼i.i.d.πref(⋅∣Xi)Y_i,1,…,Y_i,K .i.d. _ref(· X_i) and define g^N,K(λ)=1N∑i=1N[βlog(1K∑k=1Kexp(r(Xi,Yi,k)−λc(Xi,Yi,k)β))+λτ]. g_N,K(λ)= 1N _i=1^N [β ( 1K _k=1^K \! ( r(X_i,Y_i,k)-λ c(X_i,Y_i,k)β ) )+λτ ]. 3:Compute λ^N,K∈argminλ∈[0,Λ]g^N,K(λ) λ_N,K∈ _λ∈[0, ] g_N,K(λ) using binary search. 4:Define the calibrated augmented reward rλ^N,K(x,y):=r(x,y)−λ^N,Kc(x,y)r_ λ_N,K(x,y):=r(x,y)- λ_N,Kc(x,y) and pass it to DecodeAlg. Accordingly, LARA should be viewed as a decoder-agnostic layer that supplies a principled helpfulness–harmlessness trade-off to an existing inference-time alignment method. Its safety interpretation is the expected-cost, soft-constraint interpretation inherited from Safe RLHF: it calibrates a shadow price for harmfulness and then decodes with the resulting augmented reward. In practice, calibration may be sensitive to mismatch between the calibration prompts and the deployment distribution. As a robustness heuristic, one may therefore bootstrap the calibration set to obtain a distribution over λ^N,K λ_N,K and deploy an upper quantile of this distribution as a conservative choice of the trade-off parameter. We do not analyze this procedure theoretically here, so we view it as a practical robustness heuristic rather than part of the core guarantee. 3.2 Theoretical Results We now turn to the finite-sample behavior of the empirical calibration procedure defined in Section 3.1. The following theorem shows that, under boundedness and strong convexity assumptions, the estimator λ^N,K λ_N,K concentrates around the population minimizer λ⋆λ , with error controlled by the outer-sample term in N and the inner Monte Carlo term in K. Theorem 3 (Finite-sample guarantee). Assume that |r(x,y)|≤R|r_(x,y)|≤ R and |c(x,y)|≤C|c_(x,y)|≤ C for all (x,y)(x,y), that g has a unique minimizer λ⋆∈(0,Λ)λ ∈(0, ), and that g is μ-strongly convex on I=[0,Λ]I=[0, ], i.e., g′(λ)≥μg (λ)≥μ for all λ∈Iλ∈ I. Define a:=exp(−R+ΛCβ),b:=exp(R+ΛCβ),a:= \! (- R+ Cβ ), b:= \! ( R+ Cβ ), (17) BK:=CbaK+Cb22a2K,B_K:= Cba K+ Cb^22a^2 K, (18) εN,K(δ):=2C2log(2(N+1)/δ)N+C2ΛβN+BK, _N,K(δ):=2C 2 (2(N+1)/δ)N+ C^2 β N+B_K, (19) and m:=minλ⋆,Λ−λ⋆.m:= \λ , -λ \. (20) If εN,K(δ)<μm _N,K(δ)<μ m, then ℙ(|λ^N,K−λ⋆|≤εN,K(δ)μ)≥1−δ.P\! ( | λ_N,K-λ |≤ _N,K(δ)μ )≥ 1-δ. (21) Theorem 3 makes the calibration error transparent. The first term in εN,K(δ) _N,K(δ) is the concentration error from using N prompts, the second is a one-dimensional discretization term arising in the uniform derivative bound, and the third is the O(K−1/2)O(K^-1/2) Monte Carlo error from approximating the inner expectation. Thus, for fixed K, the estimator is accurate up to the residual inner-expectation error, while letting both N and K grow drives λ^N,K λ_N,K toward the population minimizer λ⋆λ . See Appendix D for the proof. 4 Experiments We investigate the following two research questions: 1. For Best-of-N sampling, does LARA identify a principled λ∗λ^* such that any smaller value violates the safety constraint and any larger value satisfies it, at unnecessary cost to helpfulness? 2. How does LARA-driven inference-time alignment compare to (i) finetuning-based baselines and (i) inference-time baselines that either forgo safety penalties (λ=0λ=0) or over-penalize harmlessness (λ>λ∗λ>λ^*)? We evaluate our research questions on two models: Llama-3.2-3B (Grattafiori et al., 2024) and Qwen-2.5-3B (Qwen et al., 2025). We follow the standard RLHF pipeline. The models are first finetuned on the Alpaca dataset (Taori et al., 2023) to obtain the SFT policy, which serves as the reference policy for all methods. The reward model and cost model are trained using the standard Bradley-Terry preference modeling objective on the BeaverTails preference dataset (Ji et al., 2023), using helpfulness and harmfulness preference labels respectively. We note that this protocol follows that of Safe-RLHF (Dai et al., 2023) and SafeDPO (Kim et al., 2025), allowing for a principled comparison. For additional implementation details and hyperparameters, refer to Appendix E. 4.1 Constraint Satisfaction and Optimality of λ∗λ^* We evaluate RQ1 using Best-of-N sampling on the BeaverTails test set (Ji et al., 2023), which was also used for evaluations in Safe-RLHF (Dai et al., 2023) and Safe-DPO (Kim et al., 2025), for which LARA is theoretically grounded. Figure 1 reports the expected helpfulness [r]E[r] and expected cost (Harmfulness) [c]E[c] of the BoN policy, which selects responses according to r−λcr-λ c, as a function of λ, for both Llama and Qwen. We observe that for λ<λ∗λ<λ^*, the cost constraint [c]≤τE[c]≤τ is violated, while for λ>λ∗λ>λ^*, the constraint is satisfied but at an unnecessary reduction in expected helpfulness. LARA’s calibrated λ∗λ^* sits at the boundary, satisfying the constraint while maximizing helpfulness, consistent with the theoretical optimality conditions derived in Section 3.1. Figure 1: Expected helpfulness [r]E[r] and expected cost [c]E[c] of the BoN policy as a function of λ, for the Llama base model on the left and Qwen base model on the right. The dashed vertical line indicates the expected cost threshold τ. 4.2 Comparison Against Finetuning and Inference-Time Baselines We evaluate RQ2 by comparing LARA-driven inference-time alignment against finetuning-based baselines (Safe-RLHF, SafeDPO, SFT) and inference-time baselines at λ=0λ=0 (no safety penalty) and λ>λ∗λ>λ^* (over-penalized). We perform both model-based evaluation and evaluation with GPT as the judge (Zheng et al., 2023). We employ the same evaluation protocol used in Kim et al. (2025). We provide results for Llama here and push those for Qwen to Appendix F. Model Evaluation We use the beaver-7b-unified-reward model111https://huggingface.co/PKU-Alignment/beaver-7b-unified-reward to assess helpfulness and the beaver-7b-unified-cost model222https://huggingface.co/PKU-Alignment/beaver-7b-unified-cost to assess harmlessness, over a set of 256 test prompts from the BeaverTails (Ji et al., 2023) test set. Looking at model-based evaluation for Llama in Figure 2(a), BoN with LARA’s calibrated λ∗λ^* achieves the best helpfulness-harmlessness tradeoff among inference-time methods, outperforming SafeDPO on helpfulness while remaining competitive on harmlessness, despite requiring no weight updates. At λ=0λ=0, BoN produces highly harmful responses, while λ>λ∗λ>λ^* over-penalizes helpfulness with marginal harmlessness gain, validating the optimality of LARA’s calibrated λ∗λ^*. All RGTG methods perform particularly poorly on helpfulness, falling below the SFT model, though they achieve higher harmlessness. Among RGTG methods, CD-Q is the most harmless, with harmlessness monotonically increasing with λ. Among all methods considered, Safe-RLHF achieves the best helpfulness-harmlessness tradeoff. (a) Model Evaluation (b) GPT Evaluation Figure 2: Helpfulness vs Harmfullness for LLama3.2-3b under both model evaluation (left) and GPT evaluation (right). λ∗λ^* for Llama is 0.60.6 and λmax _max is chosen to be 2.02.0. GPT Evaluation We used GPT-4o-mini as a judge(Zheng et al., 2023), end evaluated the helpfulness and harmlessness of the model generations, over the same set of 256 test prompts from the BeaverTails (Ji et al., 2023) test set, which were also used in model evaluations. We used the same evaluation prompts used by Kim et al. (2025), where GPT is asked to rate the helpfulness and harmlessness of model responses on a scale from 1-10. Under GPT-based evaluation for Llama in Figure 2(b), the rankings are broadly consistent with model-based evaluation, with Safe-RLHF achieving the best overall tradeoff and finetuning baselines outperforming inference-time methods. BoN achieves comparable performance to SafeDPO. However, BoN shows minimal sensitivity to λ under GPT evaluation, with helpfulness and harmlessness remaining largely stable across λ=0λ=0, λ∗λ^*, and λ>λ∗λ>λ^*, suggesting that the beaver reward and cost models are more sensitive to λ than GPT’s judgment. According to GPT evaluation, CD-Q performs much better with harmlessness improving consistently with λ, and emerges as the strongest RGTG method. ARGS and PARGS continue to perform poorly on helpfulness, though harmlessness does improve with λ for both methods. 5 Conclusion In this work, we introduced Lagrangian Reward Augmentation (LARA), a decoder-agnostic framework for safe inference-time alignment. By dualizing a KL-regularized constrained objective, LARA reduces the reward–cost trade-off to a one-dimensional calibration problem whose solution defines a safety-calibrated augmented reward. This yields a principled alternative to manual penalty tuning and can be integrated directly into existing inference-time alignment methods. For sequence-level samplers, the resulting dual variable has an exact expected-cost interpretation, while for token-level reward-guided decoders it provides a principled heuristic induced by the same constrained objective. Overall, LARA shows how safety constraints can be transferred from training-time alignment into practical decoding-time control without updating model weights. Experimentally, we 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. Acknowledgements This work has taken place in the Safe, Correct, and Aligned Learning and Robotics Lab (SCALAR) at The University of Massachusetts Amherst. SCALAR research is supported in part by the NSF (IIS-2437426) and Open Philanthropy. The computational resources for this work were provided by the University of Massachusetts Amherst’s partnership with the Unity Research Computing Platform, a multi-institutional cluster led by the University of Massachusetts and the University of Rhode Island. Scott Niekum holds concurrent appointments as an Associate Professor at the University of Massachusetts Amherst and as an Amazon Scholar. This paper describes work performed at the University of Massachusetts Amherst and is not associated with Amazon. References Beirami et al. (2024) Ahmad Beirami, Alekh Agarwal, Jonathan Berant, Alexander D’Amour, Jacob Eisenstein, Chirag Nagpal, and Ananda Theertha Suresh. Theoretical guarantees on the best-of-n alignment policy. arXiv preprint arXiv:2401.01879, 2024. Bradley & Terry (1952) Ralph Allan Bradley and Milton E Terry. Rank analysis of incomplete block designs: I. The method of paired comparisons. Biometrika, 39(3/4):324–345, 1952. Chittepu et al. (2025) Yaswanth Chittepu, Blossom Metevier, Will Schwarzer, Austin Hoag, Scott Niekum, and Philip S Thomas. Reinforcement learning from human feedback with high-confidence safety guarantees. In Reinforcement Learning Conference, 2025. Chittepu et al. (2026) Yaswanth Chittepu, Ativ Joshi, Rajarshi Bhattacharjee, and Scott Niekum. Safe rlhf beyond expectation: Stochastic dominance for universal spectral risk control. arXiv preprint arXiv:2603.10938, 2026. Christiano et al. (2017) Paul Francis Christiano, Jan Leike, Tom B. Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. ArXiv, abs/1706.03741, 2017. URL https://api.semanticscholar.org/CorpusID:4787508. Dai et al. (2023) Josef Dai, Xuehai Pan, Ruiyang Sun, Jiaming Ji, Xinbo Xu, Mickel Liu, Yizhou Wang, and Yaodong Yang. Safe rlhf: Safe reinforcement learning from human feedback. arXiv preprint arXiv:2310.12773, 2023. Deng & Raffel (2023) Haikang Deng and Colin Raffel. Reward-augmented decoding: Efficient controlled text generation with a unidirectional reward model. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, p. 11781–11791, Singapore, December 2023. Association for Computational Linguistics. doi: 10.18653/v1/2023.emnlp-main.721. URL https://aclanthology.org/2023.emnlp-main.721/. Ganguli et al. (2022) Deep Ganguli, Liane Lovitt, Jackson Kernion, Amanda Askell, Yuntao Bai, Saurav Kadavath, Ben Mann, Ethan Perez, Nicholas Schiefer, Kamal Ndousse, et al. Red teaming language models to reduce harms: Methods, scaling behaviors, and lessons learned. arXiv preprint arXiv:2209.07858, 2022. Gao et al. (2022) Leo Gao, John Schulman, and Jacob Hilton. Scaling laws for reward model overoptimization. In International Conference on Machine Learning, 2022. URL https://api.semanticscholar.org/CorpusID:252992904. Gehman et al. (2020) Samuel Gehman, Suchin Gururangan, Maarten Sap, Yejin Choi, and Noah A Smith. Realtoxicityprompts: Evaluating neural toxic degeneration in language models. In Findings of the association for computational linguistics: EMNLP 2020, p. 3356–3369, 2020. Grattafiori et al. (2024) Aaron Grattafiori et al. The llama 3 herd of models, 2024. URL https://arxiv.org/abs/2407.21783. Han et al. (2024) Seungwook Han, Idan Shenfeld, Akash Srivastava, Yoon Kim, and Pulkit Agrawal. Value augmented sampling for language model alignment and personalization, 2024. URL https://arxiv.org/abs/2405.06639. Ji et al. (2023) Jiaming Ji, Mickel Liu, Josef Dai, Xuehai Pan, Chi Zhang, Ce Bian, Boyuan Chen, Ruiyang Sun, Yizhou Wang, and Yaodong Yang. Beavertails: Towards improved safety alignment of llm via a human-preference dataset. Advances in Neural Information Processing Systems, 36:24678–24704, 2023. Kasneci et al. (2023) Enkelejda Kasneci, Kathrin Seßler, Stefan Küchemann, Maria Bannert, Daryna Dementieva, Frank Fischer, Urs Gasser, Georg Groh, Stephan Günnemann, Eyke Hüllermeier, et al. Chatgpt for good? on opportunities and challenges of large language models for education. Learning and individual differences, 103:102274, 2023. Katz et al. (2024) Daniel Martin Katz, Michael James Bommarito, Shang Gao, and Pablo Arredondo. Gpt-4 passes the bar exam. Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences, 382(2270), 2024. Khanov et al. (2024) Maxim Khanov, Jirayu Burapacheep, and Yixuan Li. Args: Alignment as reward-guided search. arXiv preprint arXiv:2402.01694, 2024. Kim et al. (2025) Geon-Hyeong Kim, Youngsoo Jang, Yu Jin Kim, Byoungjip Kim, Honglak Lee, Kyunghoon Bae, and Moontae Lee. Safedpo: A simple approach to direct preference optimization with enhanced safety. arXiv preprint arXiv:2505.20065, 2025. Kung et al. (2023) Tiffany H Kung, Morgan Cheatham, Arielle Medenilla, Czarina Sillos, Lorie De Leon, Camille Elepaño, Maria Madriaga, Rimel Aggabao, Giezel Diaz-Candido, James Maningo, et al. Performance of chatgpt on usmle: potential for ai-assisted medical education using large language models. PLoS digital health, 2(2):e0000198, 2023. Kwon et al. (2023) Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention, 2023. URL https://arxiv.org/abs/2309.06180. Moor et al. (2023) Michael Moor, Oishi Banerjee, Zahra Shakeri Hossein Abad, Harlan M Krumholz, Jure Leskovec, Eric J Topol, and Pranav Rajpurkar. Foundation models for generalist medical artificial intelligence. Nature, 616(7956):259–265, 2023. Mudgal et al. (2023) Sidharth Mudgal, Jong Lee, Harish Ganapathy, YaGuang Li, Tao Wang, Yanping Huang, Zhifeng Chen, Heng-Tze Cheng, Michael Collins, Trevor Strohman, et al. Controlled decoding from language models. arXiv preprint arXiv:2310.17022, 2023. Nakano et al. (2022) Reiichiro Nakano, Jacob Hilton, Suchir Balaji, Jeff Wu, Long Ouyang, Christina Kim, Christopher Hesse, Shantanu Jain, Vineet Kosaraju, William Saunders, Xu Jiang, Karl Cobbe, Tyna Eloundou, Gretchen Krueger, Kevin Button, Matthew Knight, Benjamin Chess, and John Schulman. Webgpt: Browser-assisted question-answering with human feedback, 2022. URL https://arxiv.org/abs/2112.09332. Ouyang et al. (2022) Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730–27744, 2022. Peyré & Cuturi (2020) Gabriel Peyré and Marco Cuturi. Computational optimal transport, 2020. URL https://arxiv.org/abs/1803.00567. Qwen et al. (2025) Qwen et al. Qwen2.5 technical report, 2025. URL https://arxiv.org/abs/2412.15115. Rafailov et al. (2023) Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in neural information processing systems, 36:53728–53741, 2023. Rafailov et al. (2024) Rafael Rafailov, Yaswanth Chittepu, Ryan Park, Harshit Sushil Sikchi, Joey Hejna, Brad Knox, Chelsea Finn, and Scott Niekum. Scaling laws for reward model overoptimization in direct alignment algorithms. Advances in Neural Information Processing Systems, 37:126207–126242, 2024. Rashid et al. (2024) Ahmad Rashid, Ruotian Wu, Julia Grosse, Agustinus Kristiadi, and Pascal Poupart. A critical look at tokenwise reward-guided text generation. arXiv preprint arXiv:2406.07780, 2024. Rashid et al. (2025) Ahmad Rashid, Ruotian Wu, Rongqi Fan, Hongliang Li, Agustinus Kristiadi, and Pascal Poupart. Towards cost-effective reward guided text generation. arXiv preprint arXiv:2502.04517, 2025. Stiennon et al. (2022) Nisan Stiennon, Long Ouyang, Jeff Wu, Daniel M. Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul Christiano. Learning to summarize from human feedback, 2022. URL https://arxiv.org/abs/2009.01325. Taori et al. (2023) Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford Alpaca: An instruction-following LLaMA model. https://github.com/tatsu-lab/stanford_alpaca, 2023. Thomas et al. (2019) Philip S Thomas, Bruno Castro da Silva, Andrew G Barto, Stephen Giguere, Yuriy Brun, and Emma Brunskill. Preventing undesirable behavior of intelligent machines. Science, 366(6468):999–1004, 2019. Van der Vaart (2000) Aad W Van der Vaart. Asymptotic statistics, volume 3. Cambridge university press, 2000. Weidinger et al. (2021) Laura Weidinger, John Mellor, Maribeth Rauh, Conor Griffin, Jonathan Uesato, Po-Sen Huang, Myra Cheng, Mia Glaese, Borja Balle, Atoosa Kasirzadeh, et al. Ethical and social risks of harm from language models. arXiv preprint arXiv:2112.04359, 2021. Yang & Klein (2021) Kevin Yang and Dan Klein. Fudge: Controlled text generation with future discriminators. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, p. 3511–3535, 2021. Yang et al. (2022) Xi Yang, Aokun Chen, Nima PourNejatian, Hoo Chang Shin, Kaleb E Smith, Christopher Parisien, Colin Compas, Cheryl Martin, Anthony B Costa, Mona G Flores, et al. A large language model for electronic health records. NPJ digital medicine, 5(1):194, 2022. Zheng et al. (2023) Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric P. Xing, Hao Zhang, Joseph E. Gonzalez, and Ion Stoica. Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena, 2023. URL https://arxiv.org/abs/2306.05685. Large Language Model Usage Large Language Models (LLMs) were used for grammatical editing and improving writing flow. Furthermore, LLMs were used by the authors to implement baselines, whose code was not available online, e.g., CD-Q (Mudgal et al., 2023), and also to create plots. The implementations were reviewed and validated by the authors. All research methodology, experimental design, data analysis, and scientific conclusions are entirely the work of the human authors. Appendix A Related Work Reinforcement Learning from Human Feedback (RLHF) is an idea that traces back to preference-based learning (Christiano et al., 2017) in reinforcement learning and was brought into the LLM setting by InstructGPT (Ouyang et al., 2022). In recent alignment work, this pipeline has been complemented by direct preference-optimization methods such as DPO (Rafailov et al., 2023), which exploit the closed-form solution of a KL-regularized reward-maximization problem to avoid performing RL. Together, these works establish the modern alignment landscape in which RLHF-style reward-based fine-tuning remains the canonical baseline, while direct alignment methods offer cheaper but still weight-updating alternatives. Safety has emerged as a central concern within this alignment pipeline because maximizing user preference alone does not reliably prevent toxic, misleading, or otherwise harmful generations. Empirical studies such as RealToxicityPrompts (Gehman et al., 2020) document toxic degeneration in neural language models, red-teaming work (Ganguli et al., 2022) systematically catalogs harmful failure modes and their scaling behavior, and broader surveys of language-model harms (Weidinger et al., 2021) emphasize that safety risks are social as well as technical. On the data side, BeaverTails (Ji et al., 2023) provides a large-scale safety-alignment dataset with explicit harmfulness annotations, enabling the separation of usefulness and safety signals. Building on this perspective, Safe RLHF (Dai et al., 2023) formalizes safety alignment by learning separate reward and cost models and solving a constrained optimization problem that maximizes helpfulness subject to a harmlessness budget, while SafeDPO (Kim et al., 2025) develops a direct-preference alternative that incorporates safety without reverting to a full RL loop. The recent Safe RLHF Beyond Expectation paper (Chittepu et al., 2026) sharpens the limitation of this expectation-based view: controlling only expected cost can miss rare but severe failures, motivating richer risk-sensitive safety criteria beyond a single expectation statistic. Inference-time alignment aims to satisfy alignment objectives without additional fine-tuning by controlling a frozen model at decode time. This setting is attractive when rapid deployment, low retraining overhead, or policy modularity is required. Best-of-N reranking (Nakano et al., 2022; Stiennon et al., 2022; Beirami et al., 2024) samples several complete responses and selects the one with the best reward score, giving a simple sequence-level test-time alignment mechanism with theoretical guarantees. ARGS (Khanov et al., 2024) instead performs reward-guided search directly during generation by modifying token selection using a reward signal at decoding time. Controlled Decoding (Mudgal et al., 2023) learns prefix scorers that bias next-token probabilities toward desired attributes. Reward-Augmented Decoding (Deng & Raffel, 2023) likewise injects reward information into token-level generation, but is designed around a unidirectional reward model to make controlled decoding more efficient. FUDGE (Yang & Klein, 2021) is an earlier controlled-generation approach that uses future discriminators to steer generation toward target attributes without retraining the base LM. Rashid et al. (2024) provide an important critique of tokenwise reward guidance, showing that full-sequence reward models are generally mismatched to prefix-level scoring and motivating partial-sequence training; this line of work underlies PARGS. FarMA (Rashid et al., 2025) then addresses the computational side of the same problem by reducing reward-model calls and making reward-guided decoding substantially cheaper at test time. FarMA also uses the optimal completion score of partial sequences to guide decoding, unlike earlier methods. Our work is orthogonal to these decoder designs: rather than proposing yet another search rule, it derives a calibrated helpfulness-versus-harmfulness objective from the dual of constrained Safe RLHF and uses that calibrated score as a drop-in signal for existing inference-time methods, thereby replacing manual scalar trade-off tuning with a principled safety-constrained construction. Appendix B Proof of Theorem 1 Proof. For notational simplicity, we present the argument in the countable-response case; the general measurable-space version follows from the same Gibbs variational identity. Fix λ≥0λ≥ 0. Starting from (5), ℒ(π,λ)=x∼Dx[y∼π(⋅∣x)rλ(x,y)−βDKL(π(⋅∣x)∥πref(⋅∣x))]+λτ.L(π,λ)=E_x D_x\! [E_y π(· x)r_λ(x,y)-β\,D_KL\! (π(· x)\,\|\, _ref(· x) ) ]+λτ. Therefore g(λ)=supπℒ(π,λ)=λτ+supπx∼Dx[y∼π(⋅∣x)rλ(x,y)−βDKL(π(⋅∣x)∥πref(⋅∣x))].g(λ)= _πL(π,λ)=λτ+ _πE_x D_x\! [E_y π(· x)r_λ(x,y)-β\,D_KL\! (π(· x)\,\|\, _ref(· x) ) ]. (22) Since the optimization over π(⋅∣x)π(· x) is pointwise in x, the supremum decouples across prompts: g(λ)=λτ+x∼Dx[supq≪πref(⋅∣x)y∼qrλ(x,y)−βDKL(q∥πref(⋅∣x))].g(λ)=λτ+E_x D_x [ _q _ref(· x) \E_y qr_λ(x,y)-β\,D_KL\! (q\,\|\, _ref(· x) ) \ ]. (23) Fix a prompt x, and abbreviate px(y):=πref(y∣x).p_x(y):= _ref(y x). For any candidate distribution q on responses, Fx(q):=y∼qrλ(x,y)−βDKL(q∥px)=∑yq(y)rλ(x,y)−β∑yq(y)logq(y)px(y).F_x(q):=E_y qr_λ(x,y)-β\,D_KL(q\|p_x)= _yq(y)\,r_λ(x,y)-β _yq(y) q(y)p_x(y). (24) We now solve the maximization of Fx(q)F_x(q). First-order characterization of the optimizer. Introduce a multiplier η for the normalization constraint ∑yq(y)=1 _yq(y)=1, and consider (q,η)=∑yq(y)rλ(x,y)−β∑yq(y)logq(y)px(y)+η(∑yq(y)−1).J(q,η)= _yq(y)\,r_λ(x,y)-β _yq(y) q(y)p_x(y)+η ( _yq(y)-1 ). Differentiate with respect to q(y)q(y): ∂q(y)=rλ(x,y)−β(logq(y)px(y)+1)+η. ∂ q(y)=r_λ(x,y)-β ( q(y)p_x(y)+1 )+η. Setting the derivative equal to zero gives rλ(x,y)−β(logq(y)px(y)+1)+η=0,r_λ(x,y)-β ( q(y)p_x(y)+1 )+η=0, hence logq(y)px(y)=rλ(x,y)+η−β. q(y)p_x(y)= r_λ(x,y)+η-β. Exponentiating both sides yields q(y)=px(y)exp(rλ(x,y)+η−β).q(y)=p_x(y) \! ( r_λ(x,y)+η-β ). Thus q(y)q(y) is proportional to px(y)exp(rλ(x,y)/β)p_x(y) (r_λ(x,y)/β), i.e. q(y)∝px(y)exp(rλ(x,y)β).q(y) p_x(y) \! ( r_λ(x,y)β ). Normalizing shows that the optimizer must be πλ∗(y∣x)=px(y)exp(rλ(x,y)/β)Zx(λ), _λ^*(y x)= p_x(y) (r_λ(x,y)/β)Z_x(λ), (25) where Zx(λ):=∑ypx(y)exp(rλ(x,y)β)=y∼πref(⋅∣x)exp(rλ(x,y)β).Z_x(λ):= _yp_x(y) \! ( r_λ(x,y)β )=E_y _ref(· x) \! ( r_λ(x,y)β ). (26) This is exactly (9). Compute the optimal value exactly. We now evaluate the supremum in closed form by completing the KL divergence. By (25), logπλ∗(y∣x)=logpx(y)+rλ(x,y)β−logZx(λ). _λ^*(y x)= p_x(y)+ r_λ(x,y)β- Z_x(λ). Therefore, for any q≪pxq p_x, DKL(q∥πλ∗(⋅∣x)) D_KL(q\| _λ^*(· x)) =y∼q[logq(y)πλ∗(y∣x)] =E_y q [ q(y) _λ^*(y x) ] =y∼q[logq(y)−logpx(y)−rλ(x,y)β+logZx(λ)] =E_y q [ q(y)- p_x(y)- r_λ(x,y)β+ Z_x(λ) ] =y∼q[logq(y)px(y)]−1βy∼qrλ(x,y)+logZx(λ) =E_y q [ q(y)p_x(y) ]- 1βE_y qr_λ(x,y)+ Z_x(λ) =DKL(q∥px)−1βy∼qrλ(x,y)+logZx(λ). =D_KL(q\|p_x)- 1βE_y qr_λ(x,y)+ Z_x(λ). Rearranging gives y∼qrλ(x,y)−βDKL(q∥px)=βlogZx(λ)−βDKL(q∥πλ∗(⋅∣x)).E_y qr_λ(x,y)-β D_KL(q\|p_x)=β Z_x(λ)-β D_KL(q\| _λ^*(· x)). (27) Since KL divergence is always nonnegative, DKL(q∥πλ∗(⋅∣x))≥0,D_KL(q\| _λ^*(· x))≥ 0, and therefore (27) implies y∼qrλ(x,y)−βDKL(q∥px)≤βlogZx(λ),E_y qr_λ(x,y)-β D_KL(q\|p_x)≤β Z_x(λ), with equality if and only if q=πλ∗(⋅∣x)q= _λ^*(· x) almost surely. Hence supq≪pxy∼qrλ(x,y)−βDKL(q∥px)=βlogZx(λ). _q p_x \E_y qr_λ(x,y)-β D_KL(q\|p_x) \=β Z_x(λ). (28) Substituting (28) into (23) yields g(λ)=λτ+x∼Dx[βlogZx(λ)].g(λ)=λτ+E_x D_x [β Z_x(λ) ]. Using (26), we obtain g(λ)=λτ+βx∼Dxlogy∼πref(⋅∣x)exp(rλ(x,y)β).g(λ)=λτ+β\,E_x D_x _y _ref(· x) \! ( r_λ(x,y)β ). Finally, substituting rλ(x,y)=r(x,y)−λc(x,y)r_λ(x,y)=r(x,y)-λ c(x,y) gives g(λ)=βx∼Dxlogy∼πref(⋅∣x)exp(r(x,y)−λc(x,y)β)+λτ,g(λ)=β\,E_x D_x _y _ref(· x) \! ( r(x,y)-λ c(x,y)β )+λτ, which is exactly (8). This completes the proof. ∎ Appendix C Proof of Lemma 2 Proof. Define Ax(λ):=βlogy∼πref(⋅∣x)exp(r(x,y)−λc(x,y)β).A_x(λ):=β _y _ref(· x) \! ( r(x,y)-λ c(x,y)β ). Then g(λ)=x∼DxAx(λ)+λτ.g(λ)=E_x D_xA_x(λ)+λτ. Fix x and write Zx(λ):=y∼πref(⋅∣x)exp(r(x,y)−λc(x,y)β).Z_x(λ):=E_y _ref(· x) \! ( r(x,y)-λ c(x,y)β ). Then Ax(λ)=βlogZx(λ).A_x(λ)=β Z_x(λ). Differentiate: Ax′(λ) A_x (λ) =βZx′(λ)Zx(λ). =β\, Z_x (λ)Z_x(λ). Now Zx′(λ) Z_x (λ) =y∼πref(⋅∣x)[exp(r(x,y)−λc(x,y)β)⋅(−c(x,y)β)]. =E_y _ref(· x) [ \! ( r(x,y)-λ c(x,y)β )· (- c(x,y)β ) ]. Therefore Ax′(λ) A_x (λ) =βy∼πref(⋅∣x)[exp(r(x,y)−λc(x,y)β)(−c(x,y)β)]y∼πref(⋅∣x)exp(r(x,y)−λc(x,y)β) =β\, E_y _ref(· x) [ \! ( r(x,y)-λ c(x,y)β ) (- c(x,y)β ) ]E_y _ref(· x) \! ( r(x,y)-λ c(x,y)β ) =−y∼πref(⋅∣x)[c(x,y)exp(r(x,y)−λc(x,y)β)]y∼πref(⋅∣x)exp(r(x,y)−λc(x,y)β). =- E_y _ref(· x) [c(x,y) \! ( r(x,y)-λ c(x,y)β ) ]E_y _ref(· x) \! ( r(x,y)-λ c(x,y)β ). By the definition of πλ∗ _λ^* in (9), the ratio above is exactly the expectation of c(x,y)c(x,y) under πλ∗(⋅∣x) _λ^*(· x), so Ax′(λ)=−y∼πλ∗(⋅∣x)c(x,y).A_x (λ)=-E_y _λ^*(· x)c(x,y). Taking expectation over x and differentiating λτλτ gives g′(λ)=x∼DxAx′(λ)+τ=τ−x∼Dxy∼πλ∗(⋅∣x)c(x,y),g (λ)=E_x D_xA_x (λ)+τ=τ-E_x D_xE_y _λ^*(· x)c(x,y), which proves (10). To compute the second derivative, write wλ(y∣x):=πref(y∣x)exp((r(x,y)−λc(x,y))/β)Zx(λ)=πλ∗(y∣x).w_λ(y x):= _ref(y x) ((r(x,y)-λ c(x,y))/β)Z_x(λ)= _λ^*(y x). Then Ax′(λ)=−wλ[c(x,y)].A_x (λ)=-E_w_λ[c(x,y)]. Differentiate once more. Using the standard derivative of an exponential-family expectation, dλwλ[c(x,y)]=−1βVarwλ(c(x,y)). ddλE_w_λ[c(x,y)]=- 1β\,Var_w_λ(c(x,y)). Hence Ax′(λ)=1βVary∼πλ∗(⋅∣x)(c(x,y)).A_x (λ)= 1β\,Var_y _λ^*(· x)(c(x,y)). Therefore g′(λ)=x∼DxAx′(λ)=1βx∼DxVary∼πλ∗(⋅∣x)(c(x,y))≥0.g (λ)=E_x D_xA_x (λ)= 1β\,E_x D_xVar_y _λ^*(· x)(c(x,y))≥ 0. This proves (11), and convexity follows immediately. ∎ Appendix D Proof of Theorem 3 The proof combines standard Hoeffding-type concentration on a finite grid, sample-average approximation arguments for stochastic programs, and classical convex M-estimation / argmin localization under strong convexity (Van der Vaart, 2000). Proof. Empirical tilted quantities. For fixed (x,y1:K,λ)(x,y_1:K,λ), define wk(λ):=exp(r(x,yk)−λc(x,yk)β),Z^x,K(λ):=1K∑k=1Kwk(λ),w_k(λ):= \! ( r(x,y_k)-λ c(x,y_k)β ), Z_x,K(λ):= 1K _k=1^Kw_k(λ), (29) and the empirical tilted weights p^λ(k∣x,y1:K):=wk(λ)∑ℓ=1Kwℓ(λ),k=1,…,K. p_λ(k x,y_1:K):= w_k(λ) _ =1^Kw_ (λ), k=1,…,K. (30) Define the empirical tilted cost c^λ,K(x,y1:K):=∑k=1Kp^λ(k∣x,y1:K)c(x,yk). c_λ,K(x,y_1:K):= _k=1^K p_λ(k x,y_1:K)\,c(x,y_k). (31) Also define ϕ^K(x,y1:K,λ):=βlogZ^x,K(λ)+λτ. φ_K(x,y_1:K,λ):=β Z_x,K(λ)+λτ. (32) Derivatives of the empirical inner loss. Differentiating (29) with respect to λ gives dλZ^x,K(λ)=−1βK∑k=1Kc(x,yk)wk(λ). ddλ Z_x,K(λ)=- 1β K _k=1^Kc(x,y_k)w_k(λ). (33) Hence, using (32), (29), and (31), ϕ^K′(x,y1:K,λ)=βZ^x,K′(λ)Z^x,K(λ)+τ=−∑k=1Kc(x,yk)wk(λ)∑k=1Kwk(λ)+τ=τ−c^λ,K(x,y1:K). φ_K (x,y_1:K,λ)=β Z_x,K (λ) Z_x,K(λ)+τ=- _k=1^Kc(x,y_k)w_k(λ) _k=1^Kw_k(λ)+τ=τ- c_λ,K(x,y_1:K). (34) Differentiating once more yields ϕ^K′(x,y1:K,λ)=1βVarJ∼p^λ(⋅∣x,y1:K)(c(x,yJ)). φ_K (x,y_1:K,λ)= 1βVar_J p_λ(· x,y_1:K)\! (c(x,y_J) ). (35) Since −C≤c(x,y)≤C-C≤ c(x,y)≤ C, (31) and (35) imply −C≤c^λ,K(x,y1:K)≤C,0≤ϕ^K′(x,y1:K,λ)≤C2β.-C≤ c_λ,K(x,y_1:K)≤ C, 0≤ φ_K (x,y_1:K,λ)≤ C^2β. (36) Combining (34) with (36) and −C≤τ≤C-C≤τ≤ C, we obtain |ϕ^K′(x,y1:K,λ)|≤2Cfor all (x,y1:K,λ)∈×K×I. | φ_K (x,y_1:K,λ) |≤ 2C all (x,y_1:K,λ) ×Y^K× I. (37) Monte Carlo population objective. Define gK(λ):=X∼DxY1,…,YK∼i.i.d.πref(⋅∣X)[ϕ^K(X,Y1:K,λ)].g_K(λ):=E_X D_xE_Y_1,…,Y_K .i.d. _ref(· X) [ φ_K(X,Y_1:K,λ) ]. (38) Since ϕ^K(⋅,⋅,λ) φ_K(·,·,λ) is differentiable in λ and the uniform bound (37) holds on I, dominated convergence applied to the difference quotients implies gK′(λ)=X∼DxY1,…,YK∼i.i.d.πref(⋅∣X)[ϕ^K′(X,Y1:K,λ)].g_K (λ)=E_X D_xE_Y_1,…,Y_K .i.d. _ref(· X) [ φ_K (X,Y_1:K,λ) ]. (39) Also, by (15), g^N,K′(λ)=1N∑i=1Nϕ^i,K′(λ). g_N,K (λ)= 1N _i=1^N φ_i,K (λ). (40) Finally, (36) implies that both gK′g_K and g^N,K′ g_N,K are Lipschitz on I with constant L:=C2β.L:= C^2β. (41) Exact tilted quantities and derivative of g. For fixed (x,λ)(x,λ), define the exact tilted distribution πλ⋆(y∣x):=πref(y∣x)exp(r(x,y)−λc(x,y)β)Zx(λ), _λ (y x):= _ref(y x) \! ( r(x,y)-λ c(x,y)β )Z_x(λ), (42) and the corresponding tilted cost c¯λ(x):=Y∼πλ⋆(⋅∣x)[c(x,Y)]. c_λ(x):=E_Y _λ (· x)[c(x,Y)]. (43) Differentiating (12) gives Zx′(λ)=Y∼πref(⋅∣x)[−c(x,Y)βexp(r(x,Y)−λc(x,Y)β)],Z_x (λ)=E_Y _ref(· x) [- c(x,Y)β \! ( r(x,Y)-λ c(x,Y)β ) ], (44) and hence ϕ′(x,λ)=βZx′(λ)Zx(λ)+τ=τ−c¯λ(x).φ (x,λ)=β Z_x (λ)Z_x(λ)+τ=τ- c_λ(x). (45) Differentiating once more yields ϕ′(x,λ)=1βVarY∼πλ⋆(⋅∣x)(c(x,Y)).φ (x,λ)= 1βVar_Y _λ (· x)\! (c(x,Y) ). (46) Since −C≤c(x,y)≤C-C≤ c(x,y)≤ C, (45) implies |ϕ′(x,λ)|≤2C.|φ (x,λ)|≤ 2C. (47) Therefore, dominated convergence applied to the difference quotients gives g′(λ)=X∼Dx[ϕ′(X,λ)]=τ−X∼DxY∼πλ⋆(⋅∣X)[c(X,Y)].g (λ)=E_X D_x[φ (X,λ)]=τ-E_X D_xE_Y _λ (· X)[c(X,Y)]. (48) Concentration on a finite grid. Fix an integer M≥1M≥ 1 and define the grid λj:=jΛM,j=0,1,…,M, _j:= j M, j=0,1,…,M, (49) with spacing Δ:=ΛM. := M. (50) For each grid point λj _j, let Ui(j):=ϕ^i,K′(λj).U_i^(j):= φ_i,K ( _j). (51) For each fixed j, the random variables U1(j),…,UN(j)U_1^(j),…,U_N^(j) are i.i.d. across i, and by (37), −2C≤Ui(j)≤2Ca.s.-2C≤ U_i^(j)≤ 2C .s. (52) Therefore, Hoeffding’s inequality gives, for every t>0t>0, ℙ(|g^N,K′(λj)−gK′(λj)|>t)≤2exp(−Nt28C2).P\! ( | g_N,K ( _j)-g_K ( _j) |>t )≤ 2 \! (- Nt^28C^2 ). (53) Applying the union bound over j=0,…,Mj=0,…,M yields ℙ(max0≤j≤M|g^N,K′(λj)−gK′(λj)|>t)≤2(M+1)exp(−Nt28C2).P\! ( _0≤ j≤ M | g_N,K ( _j)-g_K ( _j) |>t )≤ 2(M+1) \! (- Nt^28C^2 ). (54) Set tN,M(δ):=2C2log(2(M+1)/δ)N.t_N,M(δ):=2C 2 (2(M+1)/δ)N. (55) Then (54) implies that, with probability at least 1−δ1-δ, max0≤j≤M|g^N,K′(λj)−gK′(λj)|≤tN,M(δ). _0≤ j≤ M | g_N,K ( _j)-g_K ( _j) |≤ t_N,M(δ). (56) Uniform derivative control relative to gK′g_K . On the event (56), we claim that supλ∈I|g^N,K′(λ)−gK′(λ)|≤tN,M(δ)+LΔ. _λ∈ I | g_N,K (λ)-g_K (λ) |≤ t_N,M(δ)+L . (57) Fix λ∈Iλ∈ I, and choose j∈0,…,M−1j∈\0,…,M-1\ such that λ∈[λj,λj+1]λ∈[ _j, _j+1]. Since g^N,K g_N,K is convex and differentiable, its derivative g^N,K′ g_N,K is monotone nondecreasing, and therefore g^N,K′(λj)≤g^N,K′(λ)≤g^N,K′(λj+1). g_N,K ( _j)≤ g_N,K (λ)≤ g_N,K ( _j+1). (58) Using (58), (56), and the Lipschitz bound (41), we obtain g^N,K′(λ)−gK′(λ)≤tN,M(δ)+LΔ, g_N,K (λ)-g_K (λ)≤ t_N,M(δ)+L , (59) and similarly g^N,K′(λ)−gK′(λ)≥−tN,M(δ)−LΔ. g_N,K (λ)-g_K (λ)≥-\,t_N,M(δ)-L . (60) Combining (59) and (60) yields (57). Inner Monte Carlo bias. Fix x and λ∈Iλ∈ I. Let V:=exp(r(x,Y)−λc(x,Y)β),U:=c(x,Y)V,V:= \! ( r(x,Y)-λ c(x,Y)β ), U:=c(x,Y)V, (61) where Y∼πref(⋅∣x).Y _ref(· x). (62) Define u:=Y∼πref(⋅∣x)[U],v:=Y∼πref(⋅∣x)[V].u:=E_Y _ref(· x)[U], v:=E_Y _ref(· x)[V]. (63) Then (43) implies c¯λ(x)=uv. c_λ(x)= uv. (64) Let Y1,…,YK∼i.i.d.πref(⋅∣x)Y_1,…,Y_K .i.d. _ref(· x) and define u^:=1K∑k=1KUk,v^:=1K∑k=1KVk, u:= 1K _k=1^KU_k, v:= 1K _k=1^KV_k, (65) where Uk:=c(x,Yk)exp(r(x,Yk)−λc(x,Yk)β),Vk:=exp(r(x,Yk)−λc(x,Yk)β).U_k:=c(x,Y_k) \! ( r(x,Y_k)-λ c(x,Y_k)β ), V_k:= \! ( r(x,Y_k)-λ c(x,Y_k)β ). (66) Then (31) implies c^λ,K(x,Y1:K)=u^v^. c_λ,K(x,Y_1:K)= u v. (67) Because |r(x,y)|≤R|r(x,y)|≤ R, −C≤c(x,y)≤C-C≤ c(x,y)≤ C, and 0≤λ≤Λ0≤λ≤ , −(R+ΛC)≤r(x,y)−λc(x,y)≤R+ΛC,-(R+ C)≤ r(x,y)-λ c(x,y)≤ R+ C, (68) and therefore, using (17), a≤Vk≤b.a≤ V_k≤ b. (69) Hence −Cb≤Uk≤Cb,v≥a,v^≥a.s.-Cb≤ U_k≤ Cb, v≥ a, v≥ a .s. (70) Using (64), (65), (67), and (70), we have |c^λ,K(x,Y1:K)−c¯λ(x)|=|u^v^−uv|≤|u^−u|a+|u|a2|v^−v|. | c_λ,K(x,Y_1:K)- c_λ(x) |= | u v- uv |≤ | u-u|a+ |u|a^2| v-v|. (71) Since |u|≤Cb|u|≤ Cb, (71) implies |c^λ,K(x,Y1:K)−c¯λ(x)|≤|u^−u|a+Cba2|v^−v|. | c_λ,K(x,Y_1:K)- c_λ(x) |≤ | u-u|a+ Cba^2| v-v|. (72) Taking conditional expectation over Y1,…,YK∼i.i.d.πref(⋅∣x)Y_1,…,Y_K .i.d. _ref(· x), we obtain [|u^−u||X=x]≤Var(u^∣X=x)≤CbK,E\! [| u-u|\, |\,X=x ]≤ Var( u X=x)≤ Cb K, (73) because Uk∈[−Cb,Cb]U_k∈[-Cb,Cb], and [|v^−v||X=x]≤Var(v^∣X=x)≤b2K,E\! [| v-v|\, |\,X=x ]≤ Var( v X=x)≤ b2 K, (74) because Vk∈[a,b]⊆[0,b]V_k∈[a,b] [0,b]. Therefore, using (72), (73), and (74), [|c^λ,K(x,Y1:K)−c¯λ(x)||X=x]≤CbaK+Cb22a2K=BK.E\! [ | c_λ,K(x,Y_1:K)- c_λ(x) |\, |\,X=x ]≤ Cba K+ Cb^22a^2 K=B_K. (75) This bound is uniform in x and λ∈Iλ∈ I. Uniform comparison with the exact derivative. Using (39), (34), and (48), we have gK′(λ)=τ−X∼DxY1,…,YK∼i.i.d.πref(⋅∣X)[c^λ,K(X,Y1:K)].g_K (λ)=τ-E_X D_xE_Y_1,…,Y_K .i.d. _ref(· X) [ c_λ,K(X,Y_1:K) ]. (76) Hence (75) implies that, for every λ∈Iλ∈ I, |gK′(λ)−g′(λ)|≤BK.|g_K (λ)-g (λ)|≤ B_K. (77) Combining (57) with (77), we conclude that on the event (56), supλ∈I|g^N,K′(λ)−g′(λ)|≤tN,M(δ)+LΔ+BK. _λ∈ I | g_N,K (λ)-g (λ) |≤ t_N,M(δ)+L +B_K. (78) Set ε:=tN,M(δ)+LΔ+BK. :=t_N,M(δ)+L +B_K. (79) Localization of the empirical minimizer. Since g has a unique minimizer λ⋆∈(0,Λ)λ ∈(0, ), we have g′(λ⋆)=0.g (λ )=0. (80) From Lemma 2, for every λ≥λ⋆λ≥λ , g′(λ)=∫λ⋆λg′(u)u≥μ(λ−λ⋆),g (λ)= _λ ^λg (u)\,du≥μ(λ-λ ), (81) while for every λ≤λ⋆λ≤λ , g′(λ)=−∫λ⋆g′(u)u≤−μ(λ⋆−λ).g (λ)=- _λ^λ g (u)\,du≤-μ(λ -λ). (82) Therefore, if λ≥λ⋆+ε/μλ≥λ + /μ, then (81) implies g′(λ)≥εg (λ)≥ , and hence, by (78), g^N,K′(λ)≥g′(λ)−ε≥0. g_N,K (λ)≥ g (λ)- ≥ 0. (83) Likewise, if λ≤λ⋆−ε/μλ≤λ - /μ, then (82) implies g′(λ)≤−εg (λ)≤- , and hence g^N,K′(λ)≤g′(λ)+ε≤0. g_N,K (λ)≤ g (λ)+ ≤ 0. (84) If ε<μm <μ m, where m is defined in (20), then the interval [λ⋆−εμ,λ⋆+εμ] [λ - μ,λ + μ ] (85) lies strictly inside I. By convexity of g^N,K g_N,K together with (83) and (84), every minimizer of g^N,K g_N,K over I must lie in the interval (85). Hence |λ^N,K−λ⋆|≤εμ. | λ_N,K-λ |≤ μ. (86) Final choice of grid size. Choose M=NM=N. Then (50) and (55) become Δ=ΛN,tN,N(δ)=2C2log(2(N+1)/δ)N. = N, t_N,N(δ)=2C 2 (2(N+1)/δ)N. (87) Substituting (87), (41), and (18) into (79), we obtain ε=2C2log(2(N+1)/δ)N+C2ΛβN+BK=εN,K(δ). =2C 2 (2(N+1)/δ)N+ C^2 β N+B_K= _N,K(δ). (88) Combining (86) with (88), and using the probability of the event (56), yields ℙ(|λ^N,K−λ⋆|≤εN,K(δ)μ)≥1−δ.P\! ( | λ_N,K-λ |≤ _N,K(δ)μ )≥ 1-δ. (89) This is exactly (21). ∎ Appendix E Implementation Details E.1 Baselines We build on top of the Safe-RLHF codebase333https://github.com/pku-alignment/safe-rlhf. For Safe-RLHF (Dai et al., 2023), we use the hyperparameters reported in their paper. We set the cost threshold τ=−4τ=-4 for Llama and τ=0.8τ=0.8 for Qwen. This is because the Llama SFT model already achieves a low expected cost, making a stricter threshold appropriate, whereas the Qwen SFT model requires a more permissive threshold. SafeDPO (Kim et al., 2025) does not have a publicly available implementation; we implement their approach by creating the data transformation script that augments the preference data as described in their paper, and use the DPO implementation from the Safe-RLHF codebase with this transformed data, using their reported hyperparameters. For the Reward Guided Text Generation (RGTG) methods, we use the publicly available PARGS codebase444https://github.com/ahmadrash/PARGS to create the preference dataset over partial sequences, and train the partial reward model using the reward model trainer provided in the Safe-RLHF codebase. CD-Q (Mudgal et al., 2023) does not have a publicly available implementation; we implement the value function trainer from scratch. To improve efficiency, we model the value function as a causal language model, predicting the values of all possible next tokens in the vocabulary in a single forward pass, thereby speeding up both training and inference. This follows the same strategy used in FaRMA (Rashid et al., 2025) to accelerate decoding. Safe-RLHF was trained on 4 NVIDIA A100 80GB GPUs. SafeDPO, the reward and cost models, partial reward and cost models, and reward and cost value functions were all trained on 2 NVIDIA A100 80GB GPUs. E.2 LARA We use a small partition of the BeaverTails (Ji et al., 2023) training dataset as the calibration dataset for estimating the dual parameter λ∗λ^*, consisting of 2,000 prompts with 20 generated responses per prompt. Responses were generated from the base policy using vLLM (Kwon et al., 2023). We set the cost threshold τ=−4τ=-4 for Llama and τ=0.8τ=0.8 for Qwen, the same values used for the baselines. As noted in Section 3.1, to account for potential distribution mismatch between the calibration and deployment datasets, we adopt the principle of pessimism. Specifically, we compute the empirical distribution of λ∗λ^* via bootstrapping and use the 97.5th percentile (the upper end of the 95% confidence interval) as our estimate of λ∗λ^*. We generate 10,000 bootstrap samples of λ∗λ^*, which completes in approximately 2 minutes owing to the efficiency of binary search. E.3 Inference We use vLLM (Kwon et al., 2023) to generate completions for Best-of-N. For ARGS and PARGS, we use the decoding code available in the PARGS codebase. For CD-Q, we implement our own inference script, which achieves faster decoding than ARGS and PARGS owing to the causal language model structure of the value function, as described above. During inference, we use top-k sampling with k=50k=50 and temperature 1.21.2. For Best-of-N, we use N=20N=20. We note that it is important to choose N to be at least as large as the number of responses per prompt used during calibration. A smaller N would result in an overly aggressive λ∗λ^*, as BoN would not explore as far from the base policy as was observed during calibration. For RGTG methods (ARGS, PARGS, CD-Q), we consider the top k=50k=50 tokens at each decoding step with weighting w=2w=2. Appendix F Additional Results We provide the evaluation results for Qwen in Figure 3. (a) Model Evaluation (b) GPT Evaluation Figure 3: Helpfulness vs Harmfulness for Qwen2.5-3b under both model evaluation (left) and GPT evaluation (right). We observed that PARGS performed similarly to ARGS for Llama, and hence omit it from the figure. λ∗λ^* for Qwen is 2.02.0 and λmax _max is chosen to be 3.03.0.