通过拉格朗日奖励增强实现安全的推理时对齐

📊 论文图表(共 6 张)

展开查看 6 张图

‹ ›

TL;DR

本文提出拉格朗日奖励增强（LARA）框架，用于在推理阶段对语言模型进行安全约束下的对齐。通过将 KL 正则化约束优化对偶化，LARA 将复杂的多约束问题转化为关于对偶变量的一维凸优化，并利用小规模校准集估计最优对偶变量，从而构造可即插即用的增强奖励信号。实验表明，该方法显著改善了有用性与无害性之间的权衡，其中 Best-of-N 方法的性能接近基于微调的直接对齐基线。

研究背景与动机

大语言模型的对齐训练是使其输出符合人类偏好和 safety 约束的关键步骤。传统方法依赖于在大规模数据上对模型进行 [[RLHF]] 或 [[DPO]] 等微调，但这类方法计算代价高昂，且每次更新都需要重新训练模型。随着模型规模不断增长，重复训练的成本已成为制约对齐技术迭代的主要瓶颈。

推理时对齐作为一种轻量化替代方案应运而生。这类方法在解码阶段直接引导冻结的语言模型，通过外部奖励信号干预采样或重排序过程，从而以零训练成本实现对齐。典型的序列级方法如 [[Best-of-N]] 通过生成 N 个候选回复并选择奖励最高的那个；词元级方法则在解码每一步根据累计奖励调整采样分布。理论上，这些方法避免了权重更新，但实践中的核心挑战在于：如何在优化有用性奖励的同时，显式地满足安全约束？

现有的推理时对齐方法通常将奖励与安全惩罚编码为单一的标量分数，这种做法存在根本性缺陷。一方面，手动调节惩罚系数是一个高度不稳定的超参数搜索过程，需要反复实验才能找到合适的平衡点。另一方面，这种 ad-hoc 的组合无法从理论上保证安全约束被满足，换言之，我们缺乏对约束违反概率或期望代价的精确控制。这就好比在一艘船上同时拉两根绳子来控制方向，却没有可靠的方式知道每根绳子该用多大的力。

本文的切入点正是将优化理论中的对偶方法引入推理时对齐。与其手工调参，不如让优化问题自己“告诉”我们最优的权衡策略。

方法

LARA 的核心思想是将带安全约束的推理时对齐问题建模为一个优化问题，然后利用 [[拉格朗日对偶]] 将其转化为可解的形式。

具体而言，假设我们有一个预训练语言模型 ${\pi }_{\theta }$，一个奖励模型 $r_{\text{reward}}$ 用于评估回答的有用性，以及一个代价模型 $c_{\text{cost}}$ 用于评估回答的无害性（如越界程度）。我们的目标是：

第一个约束限制了生成分布与原始模型的偏离程度（KL 正则化），第二个约束则显式地限制了安全代价的期望值不超过阈值 $\delta$。

对第二个约束引入非负对偶变量 $\lambda$，可以构造拉格朗日函数：

其中 $\eta$ 是 KL 约束对应的拉格朗日乘子。

通过对偶变换，原问题的求解等价于先在给定 $\lambda$ 下优化 $\pi$，再关于 $\lambda$ 最大化对偶函数。这个两阶段过程揭示了一个关键洞察：对于固定的 $\lambda$，最优策略在词元级别的下一个词概率分布为：

这意味着增强奖励信号 $r_{\text{aug}}=r_{\text{reward}}-\lambda c_{\text{cost}}$ 可以直接作为即插即用的评分函数，应用于任何推理时对齐方法的外层框架中。

对偶优化问题的特殊之处在于它是一维的——只需要搜索非负标量 $\lambda$。借助 Jensen 不等式等数学工具，可以证明该问题是凸的，因此存在唯一全局最优解。作者提出在小规模校准集（几十到上百个样本）上通过二分搜索估计最优 ${\lambda }^{\ast }$，然后将增强奖励 $r_{\text{aug}}$ 替换原始奖励即可。

实验与结果

作者在两个维度上验证了 LARA 的有效性：序列级方法（Best-of-N）和词元级方法（基于奖励加权的 Nucleus Sampling）。

实验采用了 HH-RLHF 数据集，包含 36k 有用性-无害性偏好对。奖励模型和代价模型分别在 SFT 基座模型上训练而来，用于预测人类偏好和安全违规程度。

在 Best-of-N 设置下，实验对比了多种基线方法：仅使用有用性奖励（RM-only）、使用固定惩罚系数的组合奖励（RM+fixed penalty），以及 LARA（使用对偶校准的动态系数）。结果显示，LARA 在帕累托前沿上显著优于固定系数方法——在相同的有用性水平下，安全违规率降低了 15-25%。

更值得注意的是，LARA-Best-of-N 的性能曲线非常接近经过微调的 [[DPO]] 基线。具体而言，当有用性偏好率约为 85% 时，LARA 达到约 92% 的无害性保留率，而 DPO 约为 95%。这意味着推理时对齐已经可以在很大程度上“追平”需要大量训练的微调方法。

词元级实验则显示 LARA 的增强奖励在指导逐词生成时同样有效，虽然理论保证不如序列级方法严格（因为词元级解码的因果结构使得精确约束更加困难），但实践中仍能显著改善权衡曲线。

讨论与可借鉴点

LARA 的核心贡献在于为推理时对齐提供了一套原则性的安全约束处理框架。通过将问题对偶化，它将一个看似复杂的约束优化问题转化为了一个一维凸优化问题，使得最优权衡系数可以被系统性地估计而非手动调节。这种从理论出发的设计思路值得借鉴——当我们在工程实践中遇到需要平衡多个目标的场景时，首先形式化优化问题再求解，往往比直接调参更加可靠。

然而，该方法也存在局限。首先，对偶变量的校准依赖于代价模型的准确性，如果代价模型本身存在偏差，校准结果也会受影响。其次，词元级方法只能获得启发式的约束效果，无法提供与序列级方法同等的安全性保证。此外，LARA 假设安全约束可以用单一的期望代价函数概括，而实际场景中可能存在更复杂的多维约束结构。

从更宏观的角度看，这项工作反映了 [[推理时计算]] 在对齐研究中日趋重要的地位。随着模型能力增强和部署场景多样化，能够在不修改模型权重的情况下快速调整模型行为将成为必备能力。LARA 所展示的“安全约束+推理时对齐”组合，为这一方向提供了有价值的范式参考。

摘要

推理时对齐在解码过程中利用辅助奖励信号来引导一个冻结的语言模型，避免了反复更新权重的开销。然而，现有的推理时对齐方法通常优化单一标量分数，因此显式的安全约束要么被忽略，要么通过人工调节的惩罚项来编码。我们提出了拉格朗日奖励增强（LARA），一种在安全约束下进行推理时对齐的通用框架。从包含奖励模型和代价模型的 KL 正则化约束目标出发，LARA 对约束进行对偶化，并将优化问题简化为一个关于非负对偶变量的一维凸问题。该对偶变量在小型校准集上进行估计，由此定义的增强奖励可作为即插即用的评分信号，应用于现有的推理时对齐方法。对于序列级采样方法（如 Best-of-N 重排序），校准后的对偶变量对应于期望代价约束问题的解。对于词元级奖励引导的解码方法，相同的构造则给出一种有原则的对偶校准启发式方法，而非精确的约束策略保证。我们在序列级和词元级推理时对齐方法上对 LARA 进行了评估，发现 LARA 改善了有用性与无害性之间的权衡，其中 Best-of-N 在推理时方法中取得了最佳性能，逼近基于微调的直接对齐基线方法。

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.