蒸馏scaling law

摘要 我们提出了一种==蒸馏缩放定律,该定律基于计算预算及其在学生模型和教师模型之间的分配,来估计蒸馏模型的性能==。 我们的发现降低了大规模使用蒸馏的风险;现在可以进行教师模型和学生模型的计算分配,以最大限度地提高学生模型的性能。 我们提供了计算最优的蒸馏方案,适用于以下两种情况:1) 教师模型已存在;2) 教师模型需要训练。 如果需要蒸馏多个学生模型,或者已经存在教师模型,那么在计算量达到一定水平之前,蒸馏的表现优于监督预训练,而该计算水平会随着学生模型大小的增长而可预测地增长。 如果只需要蒸馏一个学生模型,并且教师模型也需要训练,那么应该改为进行监督学习。 此外,我们还提供了对大规模蒸馏研究的见解,这些见解加深了我们对蒸馏的理解,并为实验设计提供了信息。 介绍 对缩放定律的研究表明,如果先前训练的[lms]遵循计算最优的训练范式,它们可以表现得更加出色。这种范式旨在确定在给定的计算预算下,能够产生最佳性能模型的模型大小和训练 Token 数量。 许多后续研究都遵循了计算最优训练方法。 计算最优模型的大小随着计算量的增长而增长,这导致推理成本增加,从而使得它们更难被有效利用。 实际上,这意味着计算最优模型速度慢、服务成本高昂、消耗更多电池电量、提高了学术研究的门槛,并且会产生显著的碳排放。 随着推理量高达每天数十亿 Token,lm的推理成本通常远高于其预训练成本,并且在测试时计算量扩展的时代,这一成本还将进一步增加。 不可持续的推理成本催生了一种替代的训练范式,即过拟合训练,在这种范式中,所使用的训练数据量远超计算最优状态下的数据量,从而能够训练出小型且性能强大的模型。当以模型的整个生命周期而非仅仅预训练成本来衡量计算量时,过拟合训练的模型能更好地满足计算最优性。由于监督学习的扩展法则遵循模型大小和训练数据的幂律关系,因此性能的边际效益递减出现得比计算最优情况下要早得多。为了达到合理的性能水平,这些模型需要在数万亿的 Token 上进行训练,这既昂贵又耗时。我们力求找到一种模型,它能在更低的训练成本下达到小型过拟合训练模型的性能水平。一个常用的备选方案是知识蒸馏,即由一个性能强大的教师 [lm]为一个较小的学生 [lm]模型生成训练目标。当知识蒸馏被用于 [lm]预训练时,我们称之为知识蒸馏预训练。关于为什么知识蒸馏有效,存在多种解释,从暗知识迁移(即信息蕴含在不正确类别的概率比率中),到作为一种正则化手段,或是降低学习过程中的噪声等等。尽管对于知识蒸馏有效的原因尚未达成共识,但在 Gemma 和 Gemini、Minitron 和 AFM 系列的 [lms]模型中,知识蒸馏预训练在预训练损失和下游评估方面,都产生了比监督预训练更强大的模型。然而,与此同时,@DBLP:conf/icml/Liu0ILTFXCSKLC24 报告称,知识蒸馏产生的模型比监督预训练产生的模型能力更弱。 鉴于大量计算资源正被用于 [lms]的蒸馏预训练,至关重要的是要了解如何合理分配这些资源,以尽可能地训练出性能最佳的模型。同时,我们也需要了解,在相同资源条件下,蒸馏预训练相比于监督预训练是否具有优势。为了填补这一知识空白,我们对蒸馏方法进行了广泛的对照研究,学生模型和教师模型的参数规模从 1.43 亿到 126 亿不等,训练数据量从数十亿 Token 扩展到最多 5120 亿 Token。通过这些实验,我们得出了蒸馏缩放定律,该定律将学生模型的性能估计为关于资源的函数(包括教师模型、学生模型的大小以及用于蒸馏的数据量)。该定律解决了在特定资源约束下,蒸馏方法在生成具有期望能力的模型时,是否有效的问题。我们的研究发现: 大小为 $N_S$ 的学生模型,在从大小为 $N_T$ 的教师模型蒸馏得到的 $D_S$ 个 Token 上进行蒸馏时,其交叉熵可以使用我们的蒸馏缩放定律进行预测。 教师模型的大小 $N_T$ 和教师模型训练所用的 Token 数量 $D_T$,仅通过它们所决定的教师模型交叉熵 $L_T=L_T(N_T,D_T)$ 来影响学生模型的交叉熵。 教师交叉熵对学生损失的影响遵循幂律,该幂律根据学生和教师的相对学习能力在两种行为之间转换,反映了知识蒸馏中的一种称为能力差距的现象,即更强的教师模型反而会产生更差的学生模型。我们的参数化解决了关于能力差距的悬而未决的问题,表明它是教师模型和学生模型之间学习能力(包括假设空间和优化能力)的差距,而不仅仅是它们相对大小的差距,后者只是一种特殊情况。我们的结果表明,当两个学习过程都获得足够的数据或计算资源时,知识蒸馏无法产生比监督学习更低的模型的交叉熵。然而,如果以下两个条件都成立,则知识蒸馏比监督学习更有效: 用于学生的总计算量或 Token 数量不大于我们的缩放定律给出的取决于学生模型规模的阈值。 教师模型已经存在,或者要训练的教师模型除了单一知识蒸馏之外还有其他用途。 我们希望我们提供的定律和分析将指导社区生产出更强大的模型,并降低推理成本和生命周期计算成本。 背景 在模型扩展时,预测模型性能至关重要,因为它使我们能够理解:i) 增加可用计算资源 ($C$) 的价值;以及 ii) 如何分配这些计算资源,通常是在模型参数 ($N$) 和数据 ($D$) 之间,以实现具有所需属性的模型。这些属性可能包括充分预测数据分布(以交叉熵 ($L$) 衡量),或者在感兴趣的下游任务上达到一定水平的性能。幸运的是,交叉熵 是可预测的,大量的经验和理论证据表明,$L$ 遵循参数 $N$ 和数据 $D$(以 Token 衡量)的幂律: ...

February 14, 2025 · 小茄墩