初次阅读这篇文章，感到非常震撼。在DeepSeek将知识蒸馏方法带入大众视野后，Apple与牛津大学的研究人员迅速提出了蒸馏缩放定律，并已于2月28日完成了所有实验及一篇长达67页的论文上传至arXiv。这种效率和深度，无疑展示了大公司的研究实力。这种高效的研究速度和成果展示，再次彰显了大厂在科研领域的强大实力。
作为算力有限的独立研究者，正好借此深入研究知识蒸馏技术，如何利用蒸馏方法最大化地利用有限的计算资源。通过学习这些前沿研究，或许能够找到一些优化的思路和方法，从而在有限的条件下实现更好的模型性能。
简单看下研究动机，总结为以下几点：
1.模型缩放法则研究现状：近年来的研究揭示了语言模型性能与模型规模、训练数据量之间的关系，但在蒸馏场景下的缩放法则尚未有系统研究。
2.模型推理成本问题：随着语言模型规模增大，推理成本显著增加，研究如何在不损失性能的前提下降低推理成本成为重要课题。
3. 蒸馏的效率与性能：蒸馏理论上可降低推理成本，然而，学界对蒸馏方法缺乏共识，尤其是在如何合理分配计算资源以产生最强大的模型方面，仍存在很大的不确定性。
1、知识蒸馏
知识蒸馏的核心思想是将大型教师模型的知识迁移到小型学生模型中。具体来说，教师模型的预测分布被用作学生模型的目标分布，学生模型通过最小化与教师模型的Kullback-Leibler散度（KL Divergence）来学习：
其中，
和
分别是教师模型和学生模型的输出logits，是蒸馏温度，用于控制教师模型输出的概率分布的熵。
学生模型的总损失函数结合了下一个token预测损失和知识蒸馏损失：
- λ：混合系数，控制知识蒸馏损失和下一个token预测损失的权重。
- ：
Z-loss的权重，用于提高训练稳定性。
2、实验设计
为了拟合蒸馏缩放法则，设计了以下三种实验方案来生成缩放定律的拟合数据：
固定M教师/学生IsoFLOP配置
固定教师模型规模和训练数据量，改变学生模型的规模和蒸馏数据量。目的是研究在固定教师模型条件下，学生模型性能如何随其规模和蒸馏数据量变化。通过这种方式，可以确定学生模型在不同规模和数据量下的最优性能。
从实验结果可以观察到：
- 在算力较大的情况下，学生模型参数规模越大，其损失函数越小，而且教师模型规模越大，这种趋势越明显。
- 在学生和教师模型规模确定的情况下，算力越多，学生模型的确表现越好，这也很符合常理。
- 在算力较小的情况下，模型性能会随着模型规模先提高，后减弱，这里很容易理解，较大的模型在较小的算力下没有得到充分训练。
- 在特殊情况下，学生模型能够优于教师模型，表现出泛化能力，个人猜测可能是教师模型训练欠拟合。
IsoFLOP教师/固定M学生配置
固定学生模型规模和蒸馏数据量，改变教师模型的规模和训练数据量。目的是研究教师模型性能如何影响学生模型的最终性能。通过这种方式，可以确定最优的教师模型规模和训练数据量，以最大化学生模型的性能。
结果可以看到，并非教师模型参数越大，学生模型的交叉熵
越低。表明教师模型的性能需要与学生模型的能力相匹配，才能达到最佳蒸馏效果。
固定M教师/固定M学生配置
在固定比例下，教师模型和学生模型的性能如何随其规模变化。这种配置虽然不是拟合缩放法则所必需的，但它提供了一个更广泛的视角，展示了学生模型性能在不同教师模型条件下的变化。
同时变化教师模型和学生模型的规模：教师模型和学生模型的参数规模
和
以及它们的训练数据量
和
同时变化，但保持固定的比例
和
。
在图4中，观察到容量差距现象，教师模型的规模对学生的性能有显著影响，但存在一个最优的教师模型规模
（但并不一定是最大的模型），使得学生模型的性能
最低。即提高教师的表现并不总是提高学生的表现，甚至最终会降低学生的表现。
2、Distillation Scaling Laws
先介绍一个预备知识，已有大量理论和实验证明了交叉熵是可预测的，在监督学习场景下，模型的性能（交叉熵损失L）在模型规模（参数数量 N）和训练数据量（token数量 D）遵循幂律关系。
可以使用教师的模型参数量和训练数据量
来拟合监督缩放定律。
基于前面实验结果，研究人员建立了蒸馏扩展定律的数学模型，用以下公式来表示：
：教师模型的验证交叉熵损失。
：学生模型的验证交叉熵损失。
：如果学生模型以监督学习方式训练，其预期的验证交叉熵损失。
：学生模型的非嵌入参数数量。
：学生模型的蒸馏token数量。
是通过实验拟合的系数
简单介绍下公式的组成：
- 教师模型的交叉熵：教师模型的性能通过其交叉熵损失
来衡量。较低的
表示教师模型具有更好的性能，教师模型的性能直接影响学生模型的最终性能。 - 容量差距（Capacity Gap）：当教师模型的性能远优于学生模型时，学生模型的性能可能会下降。这被称为容量差距现象。公式中的
和
用于捕捉这种现象。具体来说，当
远小于
时，
会变得很大，导致学生模型的性能下降。 - 学生模型的性能极限：当蒸馏数据量
足够大时，学生模型的性能会趋近于教师模型的性能
。公式中的
用于描述学生模型在不同规模和数据量下的性能变化。 - 学生模型的参数数量和蒸馏数据量：学生模型的参数数量
和蒸馏数据量
对学生模型的性能有显著影响。公式中的α′和分别描述学生模型的参数数量和蒸馏数据量对性能的影响。 - 指数项
最后两个公式的关键系数可以参考下表：
结果可以看到，监督和蒸馏缩放定律在观察到的数据上拟合得非常好，相对预测误差小于1%。反映了蒸馏缩放定律的有效性，因此基于教师和学生的模型规模和数据规模可以互相实现最优模型的选择。
3、Distillation vs Supervised
为了解蒸馏在何时有益，下图是在固定计算资源下，蒸馏和监督学习的性能对比。结果展示，监督学习在给定足够的学生计算或数据时总是优于蒸馏。对于适中的数据预算，蒸馏是有利的，然而，当有大量数据可用时，监督学习优于蒸馏。
总结即，在计算资源有限的情况下，蒸馏通常比监督学习更高效。这是因为蒸馏可以通过教师模型的指导，更快地学习到有效的特征表示，从而在较少的计算资源下达到更好的性能。
4、教师模型选择
- 学习信号强度：不同大小的教师模型可能提供不同的学习信号强度，这通常通过交叉熵损失来衡量。较大的教师模型可能会提供更强的学习信号（较低的交叉熵），从而有助于学生模型更好地学习。
- 成本增加：使用较大的教师模型会带来更高的成本，这是因为需要计算教师模型的logits。这意味着更大的教师模型不仅自身训练成本更高，在用于蒸馏时也会消耗更多的计算资源。
图7展示了在不同蒸馏数据预算下学生模型的交叉熵损失变化情况。结果显示，最优教师损失
（用红线表示）随着学生大小
的增大而按幂律递减，直到学生的损失
与最优教师损失
相匹配为止。这验证了最佳教师规模几乎始终与学生规模呈线性比例关系。
图15可以看到，随着蒸馏数据量
的增加，最优教师模型的交叉熵X逐渐降低，但这种降低的速率会随着
的增加而减缓。因此可以得到：在计算资源有限的情况下，选择较小的教师模型可以减少推理成本，同时仍然能够提供有效的学习信号给学生模型。
5、计算最优蒸馏
计算最优蒸馏的目标是，在给定计算预算C的情况下，确定如何生成具有最低交叉熵的所需大小
的学生模型。具体来说，需要找到最优的学生训练数据量
、教师模型大小
和教师训练数据量
，使得在满足计算预算约束的情况下，学生的交叉熵最小化。
从图8中可以看到：
- 监督学习在足够的计算预算下总是匹配最优蒸馏：监督学习总是在某个总计算预算下匹配最佳情况蒸馏设置。这意味着，当计算预算足够大时，监督学习可以达到与蒸馏相同的性能。
- 当教师训练包含在计算中时，学生交叉熵总是高于监督设置：这意味着，如果你的唯一目标是生成具有目标大小的最佳模型，并且你没有访问教师的权限，那么你应该选择监督学习，而不是训练教师然后进行蒸馏。相反，如果目的是蒸馏出一系列模型，或使用教师作为服务模型，那么蒸馏可能比监督学习在计算上更有益。
- 较小的模型更可能从监督预训练中受益，而较大的模型更可能从蒸馏中受益：较小的模型在较大的计算预算下更倾向于监督学习，而较大的模型在较大的计算预算下更倾向于蒸馏。
图9展示了随着计算预算的变化，最优教师大小和训练数据量的变化趋势。学生和教师模型的token按幂律缩放，学生token的增长速度更快。最佳教师模型规模最初增加，直到略大于学生，之后趋于稳定。这是因为使用大型教师模型进行推理成本很高，并且随着学生token数量的增加，过拟合教师模型更有效。
6、总结
Distillation Scaling Laws（蒸馏缩放定律）为知识蒸馏技术提供了理论基础，优化了教师与学生模型间的资源分配，使得基于计算预算更科学地预测模型性能成为可能，同时降低了大规模应用蒸馏技术的风险，并揭示了“能力差距”现象的本质，为解决蒸馏过程中的挑战提供了方向。
知识蒸馏的潜力是在有限计算成本下能够训练出性能更优的模型，因此在得到蒸馏缩放定律后，可以指导有限资源下的许多场景。
