英伟达开源新大模型：训练数据减少40倍，算力节省1.8倍

934 0 0

文章摘要

全球人工智能领域的领军企业英伟达（Nvidia）近期开源了两款基于Meta公司Llama-3.1 8B模型的大模型：Nemotron-4-Minitron-4B和Nemotron-4-Minitron-8B。这两款模型的开源地址分别为https://huggingface.co/nvidia/Nemotron-4-Minitron-4B-Base和https://huggingface.co/nvidia/Nemotron-4-Minitron-8B-Base。

英伟达采用了结构化剪枝和知识蒸馏两种高效的训练方法，显著降低了模型训练所需的数据量和算力成本。结构化剪枝通过移除整个神经元、注意力头或卷积滤波器等，保留了权重矩阵的结构，使得剪枝后的模型更适合在GPU、TPU等硬件上高效运行。这种方法不仅降低了模型的内存占用和计算需求，还可能提升模型的训练速度和推理时间。

知识蒸馏则是让剪枝后的“学生模型”模仿未剪枝的“教师模型”的行为，可以在使用极少量原始训练数据的情况下，显著提升剪枝模型的表现。英伟达在训练Nemotron-4-Minitron-4B和Nemotron-4-Minitron-8B时，使用了基于logit的知识蒸馏方法，让学生模型的输出概率分布模仿教师模型，从而学习到教师模型对数据的深层理解。

通过结构化剪枝和知识蒸馏的结合使用，Minitron-4B和Minitron-8B模型在训练数据减少了40倍，算力需求节省1.8倍的情况下，性能可媲美Mistral 7B、Gemma 7B和Llama-3 8B等知名模型。这些模型是在高达15万亿token数据训练而成，而英伟达的模型仅需大约1000亿token。

在实际应用中，剪枝策略的选择依赖于模型的具体结构和优化目标。英伟达的研究人员在剪枝的过程中发现，在不同的剪枝轴上，比如宽度和深度，采用不同的策略会产生不同的效果。初期单独剪枝神经元和注意力头比同时剪枝神经元、注意力头和嵌入通道更有效，但经过几次迭代的重新训练后，这种顺序可能会发生变化。而宽度剪枝在重新训练后通常比深度剪枝有更好的效果。

除了logits层面的蒸馏，英伟达还探索了中间层特征的蒸馏，通过将教师模型的隐藏层状态映射到学生模型，并最小化这些状态之间的差异，进一步提高了学生模型的性能。这种方法使学生模型能够学习到教师模型在不同层次上的表示能力，更好地捕捉数据的复杂特征。

在损失函数的选择上，英伟达尝试了多种损失函数，包括Kullback-Leibler散度、均方误差和余弦相似度等。他们发现，对于深度减少的学生模型，使用logits损失和中间层损失的组合可以取得更好的效果。

根据测试数据显示，经过结构化剪枝和知识蒸馏的Minitron-4B和Minitron-8B模型，在MMLU上的评分提升了16%，性能可媲美Mistral 7B、Gemma 7B和Llama-3 8B。这一成果展示了英伟达在大模型训练领域的创新能力和技术实力，为AI领域的发展和应用落地提供了新的可能性。