英伟达™（NVIDIA®）如何修剪和提炼 Llama 3.1 以创建 Minitron 4B 和 8B

新模型采用了最先进的剪枝和蒸馏技术。

我最近开始了一份专注于人工智能的教育通讯，目前已有超过170,000名订阅者。TheSequence是一份不做作（意味着没有炒作，没有新闻等）的机器学习导向通讯，阅读时间为5分钟。目标是让您及时了解机器学习项目、研究论文和概念。请通过下面的链接订阅试试：

我们常常被大型语言模型（LLMs）特别是那些参数数量庞大的模型的进展所震撼。然而，执行70B+参数模型进行推理的成本对于大多数组织来说是不可承受的。因此，我们看到小型语言模型（SLMs）的影响力日益增长，使得执行推理工作负载变得更具成本效益。然而，往往无法从头开始预训练SLMs，因为在数据收集、预训练管道等方面存在重大挑战。一个流行的替代方案是从更大的LLMs开始，并将其蒸馏为更小的模型。剪枝和蒸馏是该领域最流行的两种技术。最近，NVIDIA发布了两个基于Llama 3.1–450B蒸馏版本的模型，分别为Minitron-8B和Minitron-4B。

Minitron专注于通过剪枝和蒸馏来减少AI模型的大小，使其在不牺牲太多准确性的情况下更加高效。剪枝通过切割层（深度剪枝）或移除神经元、注意力头或嵌入通道（宽度剪枝）来减少模型的大小。为了恢复一些丢失的准确性，剪枝后通常需要进行再训练。

蒸馏是一种相关技术，其中一个较小的模型，称为学生，从一个较大、复杂的模型（称为教师）学习。其目标是创建一个更紧凑的模型，保留较大模型的许多预测能力，同时更加快速且对资源的要求更低。

蒸馏方法：经典与SDG微调

Minitron 确定了两种关键的蒸馏风格。一种方法是 SDG 微调，其中较小的预训练学生模型使用由较大教师模型生成的数据进行精炼。在这种方法中，学生模仿教师预测的最终标记，这在一些流行的教程和 AI 平台中可以看到。

另一种方法，经典知识蒸馏，则更为复杂。学生模型不仅仅关注预测的标记，而是尝试复制教师模型的各种内部状态。这种技术在训练过程中提供了更详细的反馈，从而提高了准确性。然而，实施这种方法需要训练框架中的特定支持，因为它涉及处理来自教师内部状态的大量数据。

这两种方法并不是互相排斥的，而是可以相辅相成。Minitron 的主要重点是经典知识蒸馏方法。

剪枝和蒸馏工作流程

为了创建更高效的模型，Minitron将剪枝与经典的知识蒸馏相结合。从一个较大的模型开始，例如一个15B参数模型，Minitron评估不同组件的重要性——层、神经元等——然后将模型缩小到更小的尺寸，比如一个8B模型。较小的模型经过轻量级的再训练过程，从原始的较大模型中学习。这个过程可以重复进行，以进一步减少模型的大小，最终生成更小的版本，例如一个4B模型。

剪枝和蒸馏过程是迭代的，每个较小的模型作为下一个压缩和再训练轮次的基础。

剪枝影响

有效地剪枝一个模型需要理解其哪些部分是至关重要的。Minitron采用了一种基于激活数据的方法，通过使用小型数据集来估计各种组件的重要性——层、神经元、注意力头和嵌入通道。该方法仅需前向传播，使其比依赖于反向传播和梯度计算的技术更简单且更具成本效益。

虽然可以在模型的不同部分之间交替进行剪枝和重要性估计，但Minitron发现，在大多数情况下，一轮重要性估计就足够了。

使用经典知识蒸馏进行再训练

在剪枝后，Minitron 使用经典知识蒸馏对较小的模型进行再训练。这涉及通过最小化模型各个阶段的损失来教导剪枝后的模型，包括嵌入输出、logits 和变换器架构中的特定损失。学生模型通过比较不同层的输出，从未剪枝的教师模型中学习。

通过大量实验，Minitron 提炼了压缩语言模型的几条最佳实践：

· 模型尺寸： 首先训练最大的模型，然后逐渐剪枝和蒸馏，创建更小的版本。

· 剪枝策略： 优先考虑宽度剪枝而非深度剪枝，尤其是对于参数量高达 15B 的模型。单次重要性估计通常是足够的。

· 再训练： 使用蒸馏损失进行再训练，而不是传统训练。当显著剪枝层时，使用来自 logits、中间状态和嵌入的损失组合。对于较小的深度减少，保持仅使用 logits 的蒸馏。

Minitron 将这些技术应用于 Llama 3.1 模型系列，该系列包括参数从 405B 到 8B 的模型。具体而言，他们专注于将 8B 模型蒸馏为更高效的 4B 版本。

微调教师模型

在剪枝之前，Minitron 对 8B 模型进行了微调，以考虑与原始训练集的数据分布变化。没有这一步，教师模型在蒸馏过程中可能无法为学生提供最佳指导。

深度剪枝

为了将8B模型减少到4B，Minitron剪除了16层，通过逐一移除它们并跟踪对性能的影响来评估它们的重要性。他们发现模型开始和结束的层对保持准确性最为关键。基于这一分析，Minitron为最终的4B模型移除了特定的一组层。

宽度修剪

除了深度修剪，Minitron 还在宽度维度上进行了修剪，目标是注意力头、嵌入通道和隐藏层。修剪后，重新训练帮助恢复了在初始修剪步骤中丢失的一些性能。有趣的是，尽管宽度修剪最初导致的损失高于深度修剪，但重新训练使模型能够随着时间的推移更有效地恢复。

结果

NVIDIA 在多个基准测试中评估了 Minitron 模型，结果与基准模型的性能相匹配。

Minitron 4B-8B 展示了蒸馏和剪枝构建更小、更高效模型的潜力。尽管这种方法也面临着重大挑战，但我认为，总体而言，它为行业设定了一个重要的基准。