【LLM】在线合并优化器以提升奖励并减轻对齐开销

sbf_2000

发布于 2024-6-7 12:59

浏览

0收藏

【LLM】在线合并优化器以提升奖励并减轻对齐开销-AI.x社区

一、结论写在前面

在强化学习人类反馈（RLHF）中，有效对齐大型语言模型（LLMs）与以人为中心的价值，同时防止通过预训练和监督微调（SFT）获得的能力退化，是一个核心挑战。插值RLIF和SFT模型参数可以调整人类偏好与基本能力之间的权衡，从而以牺牲对齐奖励为代价减少对齐开销(Alignment Tax)。

受此启发，论文提出在RLHF的每个优化步骤中整合RL策略和SFT模型，以持续调节训练方向，引入在线合并优化器。具体而言，论文通过合并SFT与预训练模型之间的参数差异来融合梯度，有效地引导梯度向SFT优化的方向最大化奖励。

对各种骨干LLMs进行的大量实验表明，与正则化和离线合并基线相比，在线合并优化器能更好地缓解对齐成本并实现更优的对齐性能。此外，论文提出了step-K在线合并优化器，弥合了在线和离线合并之间的差距，并深入分析了超参数和消融的效果。论文证明了论文的优化器适用于不同的LLM家族，如Qwen和LLaMA，跨越从1.8B到8B的各种模型大小，以及DPO和KTO等不同的RLHF算法，以及现有的模型合并方法。它显著提升了对齐奖励，同时减轻了对齐成本开销，在14个基准测试中实现了更高的整体性能。

限制：在线合并优化器的主要限制与参数效率相关。在线合并优化器增加了内存需求，因为它们需要缓存参考模型的额外增量参数，以对应训练中的增量更新权重。同时，除非参考模型也使用LoRA适配器进行训练，否则它们不能应用于LoRA训练。然而，通过将GaLore与在线合并优化器结合使用，可以进一步消除这一限制。

二、论文的简单介绍

2.1 论文的背景

理想情况下，一个最佳的RLHF策略应在保持对齐性的奖励同时避免相关的代价，力求在最大化奖励的同时最小化遗忘。依赖于神经网络的线性模式连通性，模型能力之间的权衡可简洁地描述为模型参数的插值。研究表明，通过权重插值从同一预训练模型微调的不同模型结合，往往能在原始模型间实现更平衡的性能。

受此启发，论文初步探索了将RLHF模型与其训练自的参考SFT模型进行合并。论文的观察表明，这种离线模型合并有效地缓解了对齐成本。如表4所示，离线合并模型在语言基准和语言熟练度上恢复了与SFT模型相当的性能。然而，这种改进是以相对于RLHF模型偏好评分的减少为代价的。

鉴于单次参数插值仅允许在固定能力的模型之间进行权衡，离线合并带来的适度性能提升并不令人意外。在RLHF训练过程中，每个优化步骤都提升了模型的能力。因此，论文有机会确保这些变化的方向与参考SFT模型保持一致。本文中，论文将模型合并整合到每个RLIIF优化步骤中，并引入了在线合并优化器。这一创新的优化器在提升奖励方面比传统的优化器如AdamW更有效，同时也在减少对齐成本方面，类似于离线合并。

【LLM】在线合并优化器以提升奖励并减轻对齐开销-AI.x社区

图1：RLHF与在线合并优化器的示意图。在每个RLHF迭代中，论文首先获取更新权重A((0，然后对其进行稀疏化处理，并与参考模型的delta参数达成共识。论文使用这种合并的delta作为本次迭代中策略模型的更新。论文还对比了在线合并与离线合并，如图下方所示

对齐开销(Alignment Tax)。通常，使大型语言模型（LLMs）与人类偏好对齐涉及两个阶段：首先进行监督微调（SFT）以建立遵循指令的模型，随后通过人类反馈强化学习（RLHF）来增强人类偏好。当前主流的RLHF方法，如PPO和DPO，指导模型在优化奖励的同时，引入Kullback-Leibler（KL）散度惩罚项，该惩罚项存在于学习到的RL策略输出与参考的SFT模型之间。这一惩罚机制防止策略在追求偏好奖励时偏离其原始目标，即保持已获得的语言能力。

论文在DPO设置中尝试了不同的KL散度权重β，如图4所示。β的增加与基准性能的平均提升相关，但代价是MT-Bench和AlpacaEval上的性能下降。相反，降低β会导致模型失去其基本能力。在奖励优化与维持语言分类学之间寻求平衡已成为RLHF训练中的首要挑战。

2.2 论文的方法--在线合并优化器

受到离线合并的启发，本节论文探讨将模型合并融入到RLHF优化步骤中。论文首先审视常用的基于梯度的优化器。

2.2.1 从基于梯度的优化器到在线合并优化器

离线任务运算技术通过将LLMs的增量参数聚合来合并LLMs。相应地，论文旨在合并参考SFT模型和第t次训练步骤的政策模型。

然而，论文实证发现直接优化等式是不稳定的且难以收敛，并且等式需要额外的缓存来存储预训练模型的参数。因此，论文对等式进行了松弛处理，这种松弛仍然与论文在优化步骤中应用离线合并的动机很好地一致。这种松弛的另一个重要好处是避免了缓存额外的参数，增强了内存效率。通过这种松弛，论文表明在每个优化步骤中的在线合并可以通过基于梯度的增量权重与参考模型的增量参数之间的整合来近似。

2.2.2 实现方法

需要注意的是，论文的优化器框架具有高度的灵活性，并与现有的模型合并方法兼容。论文基于广泛使用的模型合并方法DARE和TIES开发了两种在线合并优化器：

OnDARE优化器。DARE采用随机稀疏化方法和线性组合作为共识方法。

OnTIES优化器。TIES使用top-k百分比稀疏化和基于符号的共识方法。具体来说，它从每个合并候选中保留关于绝对值的top-p百分比的参数。它根据符号和范数计算元素级多数符号，丢弃与多数符号不同的参数，并对剩余的参数进行加权求和。

2.3 论文的效果

2.3.1 实验设置

数据集。论文在广泛使用的偏好数据集ULTRAFEEDBACK上进行实验。与原始版本相比，ULTRAFEEDBACK的训练和评估部分分别包含约61K和2K个偏好对，由GPT-4进行排序，并通过人工努力进行清理，同时使用流行的基准测试(如TruthfulQA)进行了无害化处理。ULTRAFEEDBACK中的提示规模庞大、细粒度且来源多样。

训练。论文主要探索了在ULTRAFEEDBACK数据集上直接偏好优化(DPO)中的在线合并优化器，因为与近端策略优化(PPO)相比，DPO由于其较低的训练成本而在当今大规模LLM对齐中得到广泛应用。一般的DPO包括从策略模型中采样和标注响应。在这项工作中，论文使用DPO的离线策略设置，直接在ULTRAFEEDBACK数据集上训练论文的策略模型，这也被证明在提高有益性和无害性方面是有效的。

评估。评估对齐的大型语言模型是一项具有挑战性的任务。论文遵循全面评估产生可信评估的直接原则。因此，论文的评估包括7个类别中的12个公共基准，评估对齐LLM的综合能力:

数学：(1)GSM8K(2)Math401(3)Math23K;

编码:(1)HumanEval(2)MBPP(3)DS1000(4)CodeApex;

指令遵循(IF):(1)IFEval;

阅读理解(RC):(1)COQA(2)DROP;

知识：(1)MMLU;

agent：(1)NousResearch;

代码混合;

论文使用类别内的平均分数作为最终类别分数，使用所有类别中所有基准的平均分数作为整体评估。论文还使用基于GPT-4的评估器，通过长度控制评分，引入了MT-Bench 4和AlpacaEval 2.0，这是两个领先且流行的基准，用于评估LLM与人类偏好的一致性。

基线。论文方法的一个简单基线是普通的AdamW。论文进一步考虑离线合并方法，如线性合并、DARE和TIES作为论文的强基线，因为Lin et al. (2024)表明简单合并可以减轻对齐成本。由于对齐成本与遗忘有关，论文还将传统的正则化方法，如KL惩罚、EMA和ChildTuning作为论文的基线。

具体来说，DPO算法使用超参数β调整KL惩罚。论文还将LoRA作为论文的基线之一，因为参数高效方法在训练中应用了关于权重空间的正则化。

配置。论文使用三种LLM尺寸进行实验，即Qwen1.5-1.8B、Qwen1.5-7B和LLaMa3-8B系列。

具体而言，论文使用Qwen-1.8B-Base、Qwen-7B-Base和LLaMa-3-8B作为在线优化器的基础模型。论文在Qwen1.5-1.8B-SFT、Qwen1.5-7B-SFT和LLaMa-3-8B-it作为参考模型上对ULTRAFEEDBACK进行直接偏好优化。两个Qwen1.5监督微调模型在多语言指令数据上进行了训练，但与ULTRAFEEDBACK数据集没有重叠。

【LLM】在线合并优化器以提升奖励并减轻对齐开销-AI.x社区

2.3.2 主要结果

论文在表4中展示了论文的主要结果，其中展示了基线方法和论文提出的在线合并优化器在ULTRAFEEDBACK上的性能，分别在Qwen1.5-1.8BChat、Qwen1.5-7B-Chat和LLaMa-3-8B-Chat上训练。总的来说，与原始的AdamW相比，正则化和离线模型合并方法在大多数设置下并没有显著提高RLHF模型在基准测试上的平均性能，反而导致MT-Bench和AlpacaEval 2.0的偏好分数下降。这表明，仅仅依靠梯度dropout(ChildTuning)、抑制模型梯度更新的变化(EMA)或基于SFT模型对RLHF模型参数进行一次性调整(Merging)等技术，都不能有效解决对齐奖励-税收的权衡问题。

正则化基线在LLama-3-8B-It上表现特别好，所有正则化方法在平均基准测试分数以及MT-Bench和AlpacaEval分数上都取得了一致的改进。相比之下，论文提出的在线合并优化器，特别是OnDARE变体，在所有测试集上都取得了最显著的改进。OnDARE在基准测试平均分上取得了最高的改进，并在所有三个主干LLM上持续增强了MT-Bench和AlpacaEval 2.0，显著超过其他基线，尤其是在LLaMa-3-8B-Instruct实验中，分别在基准测试、MT-Bench和AlpacaEval上取得了1.3、0.19和1.57的改进。

尽管OnTIES和OnDARE在提高奖励和减轻税收方面都显示出了有效性，但在大多数情况下，OnDARE在平均基准测试分数方面略优于OnTIES，而OnTIES在AlpacaEval 2.0上有时具有更高的LC胜率。每个基准测试的详细分数在表5、表6、表7中报告。

2.3.3 超参数效应

本节分析了两个主要超参数，即参数保留率和合并权重，如何影响在线合并优化器的整体性能。

参数保留率是在线合并期间的参数保留率。论文在Qwen1.5-1.8B-Chat上探索了从1到1e-5的保留率，以在有限的计算资源内最大化搜索空间。如图2所示，即使在低至5e-4的低参数保留率下，在线合并优化器仍然保持稳健。这表明，在每个RLHF步骤中丢弃99.95%的基于梯度的参数修改仍然可以得到稳定的训练。与OnDARE相比，OnTIES对极低的参数保留率更敏感。这种敏感性是由于OnDARE采用了无偏随机稀疏化方法，而OnTIES使用的top-k稀疏化在训练过程中引入了显著的偏差。

【LLM】在线合并优化器以提升奖励并减轻对齐开销-AI.x社区

合并权重是参考(SFT)模型的梯度和增量参数的聚合权重。较大的合并权重在在线合并优化器中引入更强的正则化。论文尝试了从10−4到10−7的各种合并权重，并在表2中报告了结果。

随着合并权重的增加，由于训练过程中添加的正则化减少，MT-Bench分数上升，而平均基准测试分数先增加后减少，在α=5e−7时达到峰值。与EMA训练中的指数系数类似，较大的合并权重会导致训练不稳定。因此，论文建议从10−7这样的小数字开始合并权重的超参数搜索。值得注意的是，特殊情况合并权重=0使OnDARE等同于梯度dropout正则化方法，如论文基线中的ChildTuning。

2.3.4 RLHF算法的影响

论文进一步研究了它们在其他RLHF算法中的应用。具体而言，论文在IPO和KTO中实现了OnDARE和OnTIES。论文在这些设置下在ULTRAFEEDBACK数据集上训练Qwen1.5-7B-Chat，并在表3中展示了结果。

除了IPO算法的平均基准测试分数外，在线合并优化器OnDARE和OnTIES在MT-Bench上的表现都优于AdamW。这表明它们的有效性可以扩展到多个RLHF算法变体。

【LLM】在线合并优化器以提升奖励并减轻对齐开销-AI.x社区