英伟达神操作！剪枝、蒸馏让 Llama 3.1 8B参数减半，同尺寸最强！

Llama 3.1 被认为是引领了开源新时代。然而，新一代的模型虽然性能强大，但部署时仍需要大量计算资源。

最近，英伟达研究表明，结构化权重剪枝与知识蒸馏相结合，可以从初始较大的模型中逐步获得较小的语言模型。

图灵奖得主、Meta 首席 AI 科学家 Yann LeCun 也点赞转帖了该研究。

经过剪枝和蒸馏，英伟达研究团队将 Llama 3.1 8B 提炼为 Llama-3.1-Minitron 4B 开源了出来。这是英伟达在 Llama 3.1 开源系列中的第一个作品。

Llama-3.1-Minitron 4B 的表现优于类似大小的最先进的开源模型，包括 Minitron 4B、Phi-2 2.7B、Gemma2 2.6B 和 Qwen2-1.5B。

这项研究的相关论文早在上个月已经放出了。

• 论文链接： https://www.arxiv.org/pdf/2407.14679

• 论文标题： Compact Language Models via Pruning and Knowledge Distillation

剪枝和蒸馏

剪枝 使模型变得更小、更精简，可以通过删除层（深度剪枝）或删除神经元和注意力头以及嵌入通道（宽度剪枝）来实现。剪枝通常伴随着一定程度的再训练，以恢复准确率。

模型蒸馏 是一种将知识从大型复杂模型（通常称为教师模型）迁移到较小、较简单的学生模型的技术。目标是创建一个更高效的模型，该模型保留了原始较大模型的大部分预测能力，同时运行速度更快且资源消耗更少。

蒸馏方式主要包括两种：SDG 微调与经典知识蒸馏，这两种蒸馏方式互补 。本文主要关注经典知识蒸馏方法。

英伟达采用将剪枝与经典知识蒸馏相结合的方式来构造大模型，下图展示了单个模型的剪枝和蒸馏过程（上）以及模型剪枝和蒸馏的链条（下）。具体过程如下：

• 英伟达从 15B 模型开始，评估每个组件（层、神经元、头和嵌入通道）的重要性，然后对模型进行排序和剪枝，使其达到目标大小：8B 模型。
• 接着使用模型蒸馏进行了轻度再训练，原始模型作为老师，剪枝后的模型作为学生。
• 训练结束后，以小模型（8B）为起点，剪枝和蒸馏为更小的 4B 模型。

从 15B 模型进行剪枝与蒸馏的过程。

需要注意的点是，在对模型剪枝之前，需要先了解模型的哪部分是重要的。英伟达提出了一种基于激活的纯重要性评估策略，该策略可以同时计算所有相关维度（深度、神经元、头和嵌入通道）的信息，使用一个包含 1024 个样本的小型校准数据集，并且只需要前向传播。这种方法相比依赖梯度信息并需要反向传播的策略更加简单且具有成本效益。

在剪枝过程中，你可以针对给定轴或轴组合在剪枝和重要性估计之间进行迭代交替。实证研究显示，使用单次重要性估计就足够了，迭代估计不会带来额外的好处。

利用经典知识蒸馏进行重新训练

下图 2 展示了蒸馏过程，其中 N 层学生模型（剪枝后的模型）是从 M 层教师模型中（原始未剪枝模型）蒸馏而来。学生模型通过最小化嵌入输出损失、logit 损失以及映射到学生块 S 和教师块 T 的 Transformer 编码器特定损失组合来学习。

图 2：蒸馏训练损失

剪枝和蒸馏最佳实践

英伟达基于紧凑语言模型中剪枝和知识蒸馏的广泛消融研究，将自己的学习成果总结为以下几种结构化压缩最佳实践。