苹果发现模型蒸馏Scaling Law!教师模型并非越强越好
还给出了资源分配建议
克雷西 发自 凹非寺
量子位 | 公众号 QbitAI
大模型蒸馏也有Scaling Law了!
苹果最新研究,发现了蒸馏过程中学生模型和教师模型能力之间的幂律关系。
值得关注的是,蒸馏过程当中的教师模型,并不是越强越好。
他们发现,学生模型的损失随教师模型损失降低整体呈下降趋势,但如果教师模型能力过强,学生模型性能反而会恶化。
并且存在一个学生模型相对教师模型学习能力的转折点,在其两侧分别呈现出了不同的幂律关系。
基于一系列发现,作者还针对模型蒸馏过程中的计算资源分配给出建议,帮助根据学生模型的规模、可用资源等因素选择效果更好的分配方式。
大模型蒸馏的Scaling Law
通过对不同参数量、蒸馏数据量以及不同能力的学生模型(以监督训练时的损失衡量)和教师模型进行实验和数据拟合,作者发现蒸馏时的Scaling Law。
整体来看,学生模型的交叉熵损失LS由教师模型的交叉熵损失LT和学生模型的模仿能力决定。
而学生模型的模仿能力由三部分相乘得到,分别与教师模型损失、教师学生模型之间的能力比值(带波浪线的LS表示通过监督方式训练的学生模型损失)和数据规模NS、蒸馏数据量DS相关。
至于公式中的c0、c1、d1、f1、α′、β′和γ′,则都是需要根据实验结果拟合的参数(均为正数)。
也就是说,如果参数量和蒸馏数据量固定,对于同一个学生模型而言,影响因素就剩下教师模型了。
直观上看,学生模型损失LS总体上随着教师模型损失LT的降低而降低。
但当教师模型能力远超学生模型时,继续提升教师性能的边际效用递减,学生模型性能可能反而变差。
式中的d1就是其中的转折点,在其两侧分别服从不同形式的幂律——
- 当比值小于d1时,LS主要取决于自身规模NS和数据量DS;
- 比值大于d1时,LS主要取决于LT。
而当学生模型参数量NS和蒸馏数据量DS同时趋于无穷大时,知识蒸馏能让学生模型性能最终逼近教师模型。
蒸馏计算资源怎么分
基于以上发现,作者在论文中针对不同的蒸馏情况,给出了高效的计算资源分配建议。
其中主要是资源在教师模型训练、教师模型推理和学生模型训练之间的分配,除了可用资源量之外,主要影响因素是学生模型的大小。
- 当总预算较小(≲10^21 FLOPs)时,大部分资源应分配给教师模型的训练;
- 当预算较大时(≳10^25 FLOPs),资源应在三个部分间平均分配;
- 对于规模较小的学生模型(≲3B参数),大部分资源应分配给教师模型;
- 而对于规模较大的学生模型(≳10B),更多资源应分给学生模型本身的训练。
另外,如果教师模型需要从头训练且只用于蒸馏单个学生模型,直接用所有资源监督训练学生模型,会比蒸馏的效果更好。
以训练1.82B参数的学生模型为例,当可用数据量超过1T token时,监督学习的学生模型损失比最理想情况下的知识蒸馏更低(下图左)。
只有当总数据量/计算量低于一个随学生模型规模增大而增大的阈值,且教师模型已经存在或将被重复使用多次时,知识蒸馏才更有效(下图右)。
作者简介
这篇论文的作者主要来自苹果位于英国剑桥的实验室。
第一兼通信作者是苹果ML工程师Dan Busbridge,在校期间读的是物理,拥有理论粒子物理学博士学位。
但从第一份工作开始,Dan的工作就和计算机相关,先是在一家英国公司担任数据科学家,2016年开始研究机器学习,2020年加入苹果。
第二位作者是实习生Amitis Shidani,正在牛津大学读博,专业是统计与机器学习。
第三位Floris Weers,来自庞若鸣在AIML中领导的基础模型团队。
还有Jason Ramapuram和Etai Littwin,都是2017年就加入苹果的资深工程师。
资历最老的是Russ Webb,是苹果的高级ML研究经理,早在2010年就已经加入苹果团队。
苹果在剑桥的实验室成立时,就是由Russ出任研究主管。
整体上研究是由Dan来负责,各作者所负责的具体工作,在附录中也有介绍。
论文地址:
https://arxiv.org/abs/2502.08606
