AI News HubLIVE
站内改写4 分钟阅读

利用自我蒸馏推理优化Amazon Nova监督微调

本文探讨了在监督微调(SFT)中为缺乏推理轨迹的数据集生成思考令牌的方法。首先分析了推理抑制问题,然后介绍了自我蒸馏推理(SDR),并通过三个基准验证了其效果,最后提供了实用建议。SDR通过复用基础模型的思维链,在不牺牲目标性能的情况下有效缓解灾难性遗忘,甚至提升目标性能。

来源AWS Machine Learning Blog作者: Rushil Anirudh

在监督微调(SFT)过程中,为训练数据创建高质量的思维链推理轨迹往往不切实际且成本高昂。因此,您可能会选择跳过推理步骤,仅使用输入和输出进行训练。然而,推理是Amazon Nova 2系列模型的关键能力,已被证明能显著提升Nova及其他前沿模型在编码和数学等难题上的性能。借助Amazon Nova 2定制化,您可以通过SFT和强化微调(RFT)等技术,在特定领域发挥思考模型的力量。对于SFT,释放这些增益的关键要求是SFT数据集中包含黄金思维链轨迹,即经过验证和清理的强教师模型思考轨迹。

本文探讨了为缺乏推理轨迹的数据集生成思考令牌的想法。我们首先研究了推理抑制问题,然后介绍了自我蒸馏推理(SDR),并通过三个基准验证了其效果,最后提供了实用建议。SDR复用基础Amazon Nova 2 Lite模型的思维链作为非推理数据集的替代。通过实验,我们观察到SDR带来的影响,包括提升目标性能和缓解灾难性遗忘。这符合越来越多发现蒸馏信息重要性的研究,即自我蒸馏。

下图展示了SDR与普通SFT和模型合并相比,在目标性能和数学保持方面的差异。模型合并是一种保留先前习得技能的简单方法,它将SFT检查点合并回基础模型,但通常会牺牲SFT检查点的增益,因为灾难性遗忘导致目标性能与通用性能之间的权衡。而SDR在保持大部分通用性能的同时,实现了与纯SFT相当或更好的目标性能。

由于灾难性遗忘,我们发现对数据集进行SFT会导致先前知识的完全丢失。例如,普通SFT使数学性能从70%(基础)平均下降到6%。使用SDR,我们几乎可以完全恢复。自我蒸馏推理让我们能够在不牺牲目标性能的情况下保留基础模型的能力,并且在许多情况下还能提升目标性能。模型合并是缓解灾难性遗忘的常见解决方案,但SDR更有效地实现了这种保留。

SDR通过用模型自身的推理轨迹扩充数据集,提供了训练中的正则化。它不需要单独的教师模型或事后插值。这种方法无需人工标注,适用于任何领域的现有SFT数据集,并且在保留目标性能和通用性能方面比模型合并更有效。

与模型合并相比,自我蒸馏在数学性能上几乎相同(70%对68%),同时平均提升目标性能6.5%以上。这一发现与最近倡导自我蒸馏作为缓解灾难性遗忘强大机制的研究一致。

背景:无推理令牌的SFT可能导致推理能力丧失 在介绍方法之前,理解为什么对无推理轨迹的SFT会导致模型性能下降至关重要。本节研究推理抑制问题,量化推理对目标性能的影响,并讨论模型合并作为现有恢复机制。

推理抑制问题 在推理模式开启的情况下对非推理数据集进行SFT训练会导致一个关键问题:即使显式开启推理模式,模型在推理时也会失去推理能力。我们假设这种现象发生是因为训练目标同时计算推理令牌和输出令牌的损失。当训练数据只包含输入-输出对而没有中间推理步骤时,模型无法获得生成连贯推理轨迹的监督信号。损失函数本质上惩罚不直接贡献于最终输出的令牌,训练模型绕过其推理机制。这种行为体现了捷径学习,即模型依赖虚假相关性而非学习稳健的推理模式。

量化推理的影响 尽管存在抑制问题,但训练和推理时开启推理在目标性能上产生显著益处。为确保公平比较,我们保持训练和推理设置一致:比较训练和推理时都开启推理与都关闭推理。实验表明,训练和推理时都开启推理带来显著提升。下表展示了在LLaVA CoT数据集上不同LoRA合并权重下的效果。

合并权重 | 目标性能(推理开启) | 目标性能(推理关闭) | 增量 0.0(基础) | 12.30% | 12.30% | 0% 0.3 | 34.38% | 35.28% | -0.9% 0.5 | 60.51% | 48.80% | +11.7% 0.7 | 66.67% | 54.05% | +12.6% 1.0(完全合并) | 65.17% | 47.90% | +17.3% 表1:推理和合并权重对目标性能的影响。训练和推理模式保持一致。

模型合并作为恢复机制 目前针对无推理令牌的SFT用例,我们推荐的解决推理抑制问题的方法是使用模型合并作为恢复推理技能的机制。类似于模型合并如何恢复数学或编码能力,我们发现通过微调模型与基础模型之间的加权插值可以恢复推理能力。这种恢复与通用性能指标(如数学和编码基准)的恢复密切相关,表明推理是支撑多个下游技能的基础能力。

为了更好理解这一点,我们在实验中使用了以下基准:

  • ToolACE:工具调用基准(10K样本),评估模型根据用户查询和可用工具定义生成正确函数调用的能力。
  • CoCoHD:长上下文结构化提取(730样本),从美国国会听证会记录中提取元数据。
  • GovReport:长上下文抽象摘要(17K样本),生成美国政府报告摘要。
  • Invoice-OCR:多模态文档理解(约500样本),从扫描发票图像中提取结构化字段。
  • CaptionGen:视频字幕(1800样本),生成视频片段的描述性字幕。

我们研究推理是否随模型合并而出现。在一个指令跟随基准上,我们计算未经推理训练的模型生成的推理令牌数量。x轴显示定制模型与基础模型之间的合并权重,其中0.0对应基础模型,1.0是未经合并的SFT模型。

对定制化的影响 这些发现对部署定制化模型有重要影响。首先,在训练和推理时保持相同的推理模式对优化性能至关重要。其次,若不加以干预,领域特定的SFT会显著退化通用推理能力。第三,模型合并提供了一种事后解决方案,但需要仔细调整合并权重以平衡竞争目标。这些权衡促使我们开发一种无需事后干预就能保留推理能力的训练方法。

推理在SFT中的作用 SFT中的推理轨迹承担三个关键功能,将强化微调(RFT)、正则化和自我蒸馏范式联系起来。首先,推理轨迹作为隐式策略正则化,约束微调模型接近基础模型的策略,防止灾难性遗忘和能力退化,类似于RLHF中的KL正则化。其次,推理提供过程监督,类似于RFT中的逐步奖励塑造,而非仅结果监督,从而带来更稳健的泛化。第三,自我蒸馏——模型从自身预测中学习——已在从Born-Again Networks到Constitutional AI等多个领域被证明有效。推理轨迹通过捕捉模型的问题解决过程而非仅最终输出,提供了特别丰富的自我监督,如近期持续学习研究所展示的。

近期工作表明,有效的微调模型在参数空间中保持与基础模型的接近,而我们的方法利用这一洞察,将基础模型自身的推理轨迹作为训练目标,创建与模型内部表征一致的学习信号,同时在推理过程中提供密集的令牌级监督。

自我蒸馏推理 我们提出了一种有效方法,在无推理轨迹的数据集上为SFT定制化使用自我蒸馏推理(SDR)。该过程分三个阶段:首先,使用基础推理模型(如Amazon Nova 2 Lite)为SFT数据集中的每个示例生成推理轨迹。然后,通过将生成的推理轨迹添加到原始输出之前来扩充训练数据。最后,在推理模式开启的情况下对模型进行微调,对推理令牌和输出令牌都提供监督。

SDR为您的SFT工作流提供了多项优势。零标注成本,因为不需要人工创建推理轨迹。推理轨迹与基础模型自身的问题解决方法一致。该技术可扩展到任何领域的现有SFT数据集,并允许灵活使用不同的教师模型。

为SFT构建推理轨迹 我们查询Amazon Bedrock以获得给定数据集的推理轨迹。有两种方法构建思维链(以下内容因篇幅限制省略)。