自我蒸馏方法-减轻大模型微调过程中的灾难性遗忘

写在前面

大模型在特定任务上做微调，效果往往不错，可一个让人头疼的问题也随之而来——模型原有的通用能力会“缩水”。通俗点讲，就是学到了新技能，却忘了老本事。这次要聊的，是一篇通过自我蒸馏来缓解这种微调后灾难性遗忘的工作，名为SDFT（Self-Distillation Fine-Tuning）。

为什么微调后模型遵循通用指令的能力变弱了？说白了，就是任务数据集的信息分布，跟原始大语言模型（LLM）的信息分布之间，出现了“断层”。现在主流的应对思路，是在微调时混入一些通用指令数据来“补课”。

但SDFT换了个路子。它让模型自己给自己“上课”——用模型本身对任务数据进行生成引导，构建一个自我蒸馏数据集，相当于把任务数据的信息分布朝着原始大模型的方向“拉一拉”，缩小两者之间的差距。下图直观展示了这个思路。

方法介绍

大模型的监督微调（SFT）过程，本质上是让模型根据给定的指令和输入上下文，去学习生成对应的输出，目标是让数据的信息分布和语言模型的信息分布尽可能对齐。表达式如下：

其中，θ 是模型训练参数，x 是输入上下文，c 是指令内容，y 是模型输出。

SDFT的第一步，是让原始大模型对微调指令数据重新生成回复。这相当于做一次“内容改写”，把任务数据中的指令回复，映射到大模型自己更“舒服”的输出分布上去。这个重写过程不需要额外的复杂要求，只让模型重新输出结果就好。自我蒸馏的提示模板参考下图。

为了保证蒸馏出的回复内容靠谱，SDFT还引入了一个简单的启发式评估机制。比如在数学推理任务中，如果能从蒸馏后的回复里提取出最终答案，那就认定它有效，否则就用原始回复。这一点跟作者交流过——其实只有数学任务用了这套校验逻辑，其他不太容易判断的任务，默认蒸馏效果是准确的。

最后一步，就是用蒸馏后的回复内容，替换掉原始回复，再拿来对模型进行微调。

实验结果

所有实验都基于 Llama-2-chat-7b 模型，采用 LoRA 方法训练。学习率初始设为 1e-4，按余弦调度策略衰减到 0，训练批次大小为 8。

数据集分单任务和多任务两类：

• 单任务：OpenFunctions、GSM8K、MagiCoder；
• 多任务：Alpaca、Dolly、LIMA。

评估方面，安全性用 Advbench 榜单，实用性用 AlpacaEval 榜单，知识能力用 OpenLLM 榜单。

结果如下表所示：普通微调虽然能提升目标任务的效果，但在其他任务上的性能却明显下滑。而 SDFT 能有效缓解这种“偏科”现象，有些任务甚至还有小幅提升。

另一个值得注意的点是，在 Chat 模型上做普通任务微调，会导致模型的对齐效果变差，也就是安全性下降。而 SDFT 方法能有效保住这块能力。

有趣的是，虽然微调会影响下游任务的表现，但对模型本身的知识储备影响倒不大。

进一步分析自我蒸馏数据占比的影响，结果如下图所示：蒸馏数据的比例越高，最终效果越好。

如果将自我蒸馏数据和原始数据混合起来训练，会发现效果与按 50% 比例混合时差不多，甚至略低。这从侧面说明，数据质量的重要性往往大于数量。

最后，对比微调后模型与原始模型的指令生成结果及嵌入向量，可以发现：普通微调方法随着数据量增加，模型偏移越来越严重；而 SDFT 方法微调后的模型，嵌入空间偏移明显更小。

写在最后

自我蒸馏这个方法，不引入额外数据，就能在很大程度上缓解大模型微调后的遗忘问题。未来如果引入外部模型来完善蒸馏数据的保留机制，效果或许还会有新的突破。