清华+腾讯+港大提出LLM-Neo | 通过结合KD与LoRA，提出全新构建轻量化LLM方法

本文链接：https://blog.csdn.net/m0_59235699/article/details/143779938

在本文中，作者提出了一种新颖的LLM-Neo框架，该框架高效地将大语言模型（LLM）教师的知识转移到一个紧凑的学生模型中。首先，作者重新审视了知识蒸馏（KD）和低秩适应（LoRA），并认为它们共享相同的范式。

受这一观察的启发，作者探索了一种结合LoRA和KD的策略，以增强知识转移的效率。作者首先总结了该设计的一些指南，并进一步开发了LLM-Neo。

实验结果显示，LLM-Neo在压缩Llama2和Llama3方面优于各种 Baseline 。进一步的分析证明了所提出的LLM-Neo在LoRA变体上的鲁棒性。已训练的模型可在该仓库中获取。

1 简介

大语言模型（LLM）的知识蒸馏（KD）是一种主要的模型压缩方法。核心思想是通过模仿教师模型的行为来训练一个紧凑的学生模型。主流的方式之一是对齐logits，从而将教师模型的知识转移到学生模型中。

Parameter-Efficient Fine-Tuning (PEFT) [4] 和 [5] 是另一种常用于大语言模型效率提升的技术。在各种PEFT方法中，低秩 Adapter （LoRA）自从其在推理过程中不引入额外参数以来，受到了越来越大的关注。在训练过程中，LoRA 更新一个可合并的低秩分支而不是更新原始的完整参数。因此，LoRA 可以高效地将训练示例中的知识转移到训练好的模型中。

在本文中，作者提出观点认为，知识蒸馏（KD）和低秩 Adapter （LoRA）遵循相同的范式，即旨在通过不同的源传递知识。此外，LoRA 通过低秩支路高效地传递知识，而 KD 方法则更新全部参数，并通常消耗更多资源。

因此，作者提出了一个问题：能否结合 KD 和 LoRA 来提高从教师模型向学生模型传递知识的效率？

为此，作者提出了一种名为 LLM-Neo 的新框架，将 LoRA 集成到 KD 中以实现参数高效的知识蒸馏。

具体而言，如图1 所示，作者借鉴了 LoRA 的思想，在学生模型中引入了一个低秩支路，以继承教师模型的知识。作者首先进行了全面的分析并推导出了设计 LLM-Neo 的宝贵准则。在压缩 Llama 2 和 Llama 3.1 实验中的结果表明了 LLM-Neo 的有效性和效率。进一步的分析还显示，LLM-Neo 对 LoRA 变体具有鲁棒性。

作者的贡献可总结如下：

提出了LLM-Neo，这是一种将知识蒸馏与低秩 Adapter 结合以构建轻量化大语言模型的新方法。
总结了LLM-Neo的指导原则，包括：通常来说，秩越大效果越好，128秩表现良好。
对于LoRA而言，学习率为2e-4附近较为合适。
较大的秩需要较低的学习率。

如图2所示，尝试对角平衡是一个值得考虑的方法。作者在Llama 2和Llama3.1模型上进行了广泛的实验，以证明LLM-Neo的有效性。此外，进一步的分析还证实了它与现有大语言模型优化方法的兼容性，突显了其在提高模型性能和可扩展性方面的贡献。作者发布了LLama-3.1-Neo-1B模型的权重，这是一个使用LLM-Neo方法在百万规模的BAAI对话数据集上训练的小型模型，为更广泛的科研和开发社区做出了贡献。

2 初步和方法

在本文中首先从数学角度分析了知识蒸馏（KD）和低秩 Adapter （LoRA）。随后，将LoRA梯度应用于KL散度项，扩展了KD过程，并由此提出了大语言模型-新框架（LLM-Neo）。

2.1 知识蒸馏

在知识蒸馏（KD）中，学生模型 ( S ) 从 Ground Truth 标签和教师模型的预测中学习。对于相同的输入，( S ) 模型和 ( T ) 模型都会预测下一个 Token，并生成概率分布。KD损失由两部分组成：与 Ground Truth 的交叉熵损失以及 Kullback-Leibler (KL) 散度，用于使学生模型的预测与教师模型的预测相匹配。

其中是损失权重。在该设置中，模型参数通常通过梯度进行更新：

2.2 Low-Rank 调适

LoRA 引入了一个低秩分支来建模权重更新。更新公式表示为：

其中，且和，并且。因此，更新的参数从减少到。在监督微调（SFT）过程中，模型通过最小化交叉熵损失来学习训练样本，具体定义如下：

因此，作者可以更新低秩分支通过：

2.3 将LoRA梯度扩展到KL损失

通过比较定义，可以得出结论，知识蒸馏（KD）和LoRA遵循相同的范式，即知识转移。在方程1中，第一项对应于方程(4)中的监督微调损失，两者都利用从GT标签中提取的知识来更新模型。然而，KD损失中的第二项源自教师模型的输出，代表了一种不同的知识来源。因此，可以将低秩更新扩展到项，如下所示：

这确保了交叉熵损失和KL散度项均使用LoRA的低秩参数进行优化，同时通过多种知识源保持知识传输的效率。

2.4 LLM-Neo: 结合LoRA和知识蒸馏

最后，作者将LoRA和知识蒸馏(KD)统一起来形成了LLMNeo方法。最终的损失函数是数据集的交叉熵损失和学生模型与教师模型之间KL散度的加权组合。具体来说，LLMNeo的优化过程为：

在LLM-Neo中，两个术语都贡献于更新，正如方程(4)和(6)所示。因此，LLM-Neo通过在监督学习和知识蒸馏中应用低秩更新，保留了LoRA的参数效率，并结合了数据集学习和教师模型知识迁移的优点。

3 实验

3.1 实验实现

对于训练数据，作者采用BAAI Infinity-Instruct数据集，并随机抽取1,000,000个样本作为微调数据，其中包括大约5亿个Token。由于资源有限，作者进一步从BAAI Infinity-Instruct中随机抽取了几个较小的子集，分别为10,000和100,000个样本，分别对应约500万和5000万个Token。

3.2 Guidelines

作者首先将 Llama2 7B 脱疏到 TinyLlama 1.1B，基于 100,000 个训练样本。作者对学习率（1e-4, 2e-4, 5e-4 和 1e-3）进行了网格搜索，并对（2, 4, 6, 8, 16, 32, 48, 128 和 256）进行了排名。图2 显示了所有基准上的平均评估结果。根据这些结果，作者可以得出以下指导原则：

通常情况下，秩越大效果越好。与微调不同，知识蒸馏显示了一个显著的趋势：较大的秩会带来更好的结果。
学习率接近2e-4对于LoRA而言效果良好，而过大的学习率（例如1e-3）会导致性能下降。
考虑到秩与学习率之间的关系，较大的秩需要较低的学习率。

根据指导原则，作者将秩设为128，学习率设为2e-4。对于其他设置，作者遵循最新的大语言模型剪枝方法23, 24，基于Llama 3.1构建了一个1B规模的学生模型，称为Llama 3剪枝1B。对于Llama2系列实验，作者使用Llama2-7B指令作为教师模型，Tiny Llama-1.1B作为学生模型。对于Llama3.1系列实验，作者使用Llama3.1-8B指令作为教师模型，Llama 3剪枝1B作为学生模型。在LoRA-base实验中，最小学习率设置为1e-5；在SFT和KD实验中，最小学习率设置为1e-6。作者使用5000万Token的数据集进行知识蒸馏。

如表1所示，对比分析SFT、LoRA、KD和LLM-Neo的结果表明，LLM-Neo方法在内存效率和训练时间方面优于KD，在总体性能上也超过了SFT和标准的LoRA方法。

具体而言，对于Llama 3.1，LLM-Neo获得的平均分为39.21分，比LoRA高出0.87分，比KD高出0.08分。考虑到Llama 2时，LLM-Neo也同样优于 Baseline 方法。此外，LLM-Neo可以节省约25%的GPU内存和训练时间，展示了知识迁移的高效性。在对Llama2和Llama3.1系列的实验中进一步证明了作者提出的LLM-Neo的鲁棒性。

3.3 主要结果

为了评估，作者采用了多个标准基准，包括MMLU、CMMLU、C-Eval，以及在常识推理数据集上的零样本准确性（即PlQA、HellaSwag、WinoGrande、ARCeasy、ARC-challenge 和 OpenbookQA）。作者使用lm-evaluation-harness包进行这些评估。所有实验中，作者将输入的最大长度设为512，批量大小设为4，梯度累积步数设为16。在6块Nvidia A100 40G GPU上训练2个周期大约需要一个小时。

4 实验分析

4.1 使用LoRA变体加强

作者评估了大语言模型LLM-Neo对LoRA变体的鲁棒性。具体地，作者采用了MoSLoRA的最新版本，该版本通过混合同 subspaces 来改进LoRA。如图3所示，LLM-Neo-MoSLoRA在所有情况下都优于vanilla LoRA，这表明了LLM-Neo的鲁棒性。

4.2 LLM-Neo 中的缩放定律

在Llama3报告的基础上，作者还扩大了数据集规模，并逐步将LLM-Neo应用于更大规模的数据集，包括10万、20万、50万和100万训练样本。

如图4所示，结果表明随数据规模的增加，性能呈现出一致的提升，这表明LLM-Neo能够有效利用大规模数据集并高效地进行知识转移，从而突显出其在广泛训练任务中的可扩展性和鲁棒性。

4.3 Applied on Minitron

作者进一步从Llama 3.1中对Nvidia Minitron 4B进行了知识蒸馏，使用了50M的Token。所有实验均在6块A100 40G GPU上进行。图5展示了在多个基准上的结果。可以发现LLM-Neo在精简的Minitron 4B上表现良好，这表明其鲁棒性较强。具体而言，LLM-Neo的平均得分为54.37，比 Baseline 模型高0.52。训练权重已在此仓库中提供。

为了测试将LLM-Neo与其他内存优化技术结合时的鲁棒性，作者研究了在蒸馏Minitron 4B过程中与现有大语言模型优化技术的兼容性。具体来说，作者进一步使用ZeRO-1方法对LLM-Neo进行了优化。如表2所示，可以发现LLM-Neo与ZeRO1和ZeRO2都能很好地兼容，突显了不同ZeRO Level 之间的关系及其对诸如内存消耗和时间效率等性能指标的影响。具体而言，尽管作者将批次大小减小以避免内存溢出，但KD方法仍然会遇到内存不足的问题。

5 总结

在本工作中，作者提出了一种新颖的大语言模型-Neo框架，旨在高效地将大语言模型（LLM）教师的知识转移到一个紧凑的学生模型中。作者首先回顾了知识蒸馏（KD）和低秩适应（LoRA），并发现两者具有相同的范式。

因此，作者探索了结合LoRA和KD的策略以增强知识转移的效率。首先总结了一些指南，并进一步发展了LLM-Neo，即较大的秩值（约128）更为合适。压缩Llama 2和Llama 3.1的实验结果显示，LLM-Neo优于各种 Baseline 模型。在未来的工作中，作者将探索LLM-Neo在更多训练样本和更多LoRA变体上的性能表现。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述