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

AIGC动态4个月前发布 AIGCOPEN
845 0 0
英伟达开源新大模型:训练数据减少40倍,算力节省1.8倍

 

文章摘要


【关 键 词】 人工智能开源模型结构化剪枝知识蒸馏性能提升

全球人工智能领域的领军企业英伟达(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领域的发展和应用落地提供了新的可能性。

“极客训练营”

原文和模型


【原文链接】 阅读原文 [ 1295字 | 6分钟 ]
【原文作者】 AIGC开放社区
【摘要模型】 moonshot-v1-32k
【摘要评分】 ★★★☆☆

© 版权声明
“绘蛙”

相关文章

暂无评论

暂无评论...