苹果提出蒸馏扩展定律:预测蒸馏模型性能,优化资源分配

苹果提出蒸馏扩展定律,可预测蒸馏模型性能并优化资源分配,为更小更强的模型提供指导。

原文标题:苹果也在蒸馏大模型,给出了蒸馏Scaling Laws

原文作者:机器之心

冷月清谈:

苹果的研究人员提出了一种蒸馏扩展定律,可以根据计算预算及其在学生和教师模型之间的分配来估算蒸馏模型的性能。

该研究发现,学生模型的交叉熵可以通过一个公式进行预测,该公式考虑了学生模型的大小、训练数据量以及教师模型的交叉熵。教师模型的大小和训练数据量仅通过确定教师模型的交叉熵来影响学生模型的性能。

教师模型的交叉熵对学生模型损失的影响遵循幂律,该幂律根据学生和教师模型的相对学习能力在两种行为之间转换,反映了蒸馏中“能力差距”现象,即更强的教师模型有时会产生较差的学生模型。

研究结果表明,当学生和教师模型都有足够的数据或计算资源时,蒸馏的性能不会优于监督学习。但是,如果满足以下两个条件,蒸馏比监督学习更有效:(1)用于学生的总计算量或数据量不超过特定阈值;(2)教师模型已经存在,或者训练的教师模型有多种用途。

这项研究为在不同计算预算下优化教师和学生模型的资源分配提供了指导,并比较了蒸馏和监督学习在不同情况下的性能。研究发现,当考虑教师模型的训练成本时,监督学习通常更有效。

怜星夜思:

1、文章中提到的“能力差距”现象,除了教师和学生的模型大小差异外,还有哪些因素会导致这种现象?
2、除了模型压缩和性能提升,知识蒸馏技术还能应用于哪些场景?
3、文章提到蒸馏比监督学习更有效需要满足两个条件。在实际应用中,如何判断这两个条件是否成立?

原文内容

机器之心报道
机器之心编辑部

蒸馏模型的性能可以量化估算了。


众所周知,知识蒸馏技术当前正被大模型领域广泛使用,它可以在大幅压缩模型体量的同时保持一定的性能、降低模型时延、提升模型精度,与此同时还能对知识域进行集成和迁移。

近日,苹果研究人员提出了一种蒸馏扩展定律(Distillation Scaling Laws),基于计算预算及其在学生和教师之间的分配,我们现在开始可以估算蒸馏模型的性能了。

图 1. 蒸馏扩展定律的外推。蒸馏扩展定律适用于一系列损失为 LT 的教师的弱学生模型(L_S > 2.3)。实线表示给定学生配置(插值)下未见过的教师的预测模型行为,虚线表示见过的教师之外和强学生区域(L_S ≤ 2.3)的预测模型行为。如图所示,学生的表现可以胜过老师(详情见图 2、3 和 41)。
 
苹果研究者认为,该发现降低了大规模使用蒸馏的风险,现在我们可以基于此优化教师和学生模型的计算分配,以最大化学生模型的性能。该工作提供的计算最优的蒸馏方案适用于两种情况:1)已有教师模型,或 2)需要训练教师模型。

如果要蒸馏多个学生模型,或者已有教师模型,蒸馏在计算水平上优于监督预训练,直到计算水平随着学生模型规模的增加而可预测地增长。如果要蒸馏一个学生模型且还需要训练教师模型,则应采用监督学习。此外,作者在大规模蒸馏研究中提供了深入的见解,这些见解增加了我们对蒸馏的理解,并为实验设计提供了信息。


  • 论文标题:Distillation Scaling Laws
  • 论文链接:https://arxiv.org/pdf/2502.08606

大模型的扩展定律(Scaling Laws)表明,如果先前训练的语言模型(LM)遵循计算最优训练范式,就可以随着计算力提升而提升。由于推理成本的持续增长,目前这种定律难以持续,人们尝试使用过度训练(其中训练数据量远大于计算最优情况)来实现小型、功能强大的模型。这些模型的构建既昂贵又耗时。

我们正在寻求与训练算力投入相匹配,但训练成本更低的模型,蒸馏是一种流行的方法。但长期以来,学界对蒸馏缺乏共识,并不了解如何分配计算资源,以产生最强大的模型。为了弥补这一知识差距,研究人员对蒸馏进行了广泛研究,学生和老师的参数范围从 1.43 亿到 126 亿,使用训练数据最多达 5120 亿 token。
 
研究发现:


1. 一个大小为 N_S 的学生模型,通过从大小为 N_T 的教师模型中蒸馏 D_S 个 token 所得到的交叉熵,可以通过蒸馏扩展定律(公式 8)进行预测。
2. 老师大小 N_T 和老师训练 token 数量 D_T 仅通过确定老师的交叉熵 L_T = L_T (N_T , D_T) 来确定学生交叉熵;
3. 老师交叉熵对学生损失的影响遵循幂律,该幂律根据学生和老师的相对学习能力在两种行为之间转换,反映了蒸馏中称为能力差距的现象,即较强的老师会产生较差的学生。该工作的参数化解决了有关能力差距悬而未决的问题,表明这是老师和学生之间学习能力(假设空间和优化能力)的差距,而不仅仅是他们的相对大小,后者其实是一种特殊情况。

该结果表明,当两个学习过程都有足够的数据或计算时,蒸馏不能产生比监督学习更低的模型交叉熵。但是,如果以下两个条件都成立,则蒸馏比监督学习更有效:

1. 用于学生的总计算或 token 不大于新扩展定律给出的学生大小相关阈值;
2. 老师已经存在,或者要训练的老师有超出单次蒸馏的用途。

新的定律和分析有望指导 AI 社区构建更强大的模型,实现更低的推理成本和总计算成本。

蒸馏扩展率

文章概述了他们如何得出蒸馏扩展率所采取的步骤。

首先是实验设置。本文的目标是理解教师模型在蒸馏过程中的作用,因此,该研究在纯蒸馏情况下(λ = 1,公式 7)进行蒸馏,以避免数据带来的混淆。本文验证了 λ = 1 的选择能够产生与最优 λ∗ 统计相似的结果。同样,所有实验均使用蒸馏温度(τ = 1),因为该研究发现这能产生性能最佳的学生模型。



表 1 是文中出现的符号代表内容:


对应中文版本

此外,对于实验数据的选择,本文结合了三种方案:

固定 M 的教师 / 学生 IsoFLOP 实验:

本文预测在固定教师模型的情况下,学生模型的参数(N_S)和训练 token 数量(D_S)之间会呈现出幂律行为。

为了在给定的计算预算内生成尽可能多样的教师模型,本文训练了六个 Chinchilla 最优教师模型,其参数范围从 1.98 亿到 77.5 亿。对于每一个教师模型,本文按照标准训练成本,将其蒸馏到具有四种 IsoFLOP 配置的学生模型中。最终得到的学生模型交叉熵如图 2 所示。作者还注意到,在某些情况下,学生模型能够超越教师模型的表现,即展现出弱到强的泛化能力。


注:为实现扩展系数的可靠识别,此前有研究使用了两种策略:

  • (固定模型,变化数据) 对于一个固定的模型族,变化训练 token 的数量。
  • (IsoFLOP 配置) 在总计算约束下,同时变化模型大小和训练 token 的数量。

为了确保实验能够检测到这种影响,本文设定学生(N_S,D_S)是固定的,而 N_T 和 D_T 在计算约束下变化。本文进行了蒸馏实验,将四个 Chinchilla 最优(M_S = D_S/N_S ≈ 20)的学生(其参数范围从 1.98 亿到 18.2 亿),从根据四种 IsoFLOP 配置训练的教师中蒸馏出来。最终得到的学生交叉熵如图 3 所示。


最后,本文训练了固定 M 的教师模型与固定 M 的学生模型的组合,其中包含十个教师模型(M_T ≈ 20)和五种规模的学生模型,每个学生模型至少对应四种 M_S 选择。其中两个学生模型的交叉熵结果如图 4 所示。


此外,本文还需要确定蒸馏扩展定律的函数形式。首先,本文观察到教师模型的大小 N_T 和预训练 token 数量 D_T 的贡献可以通过教师模型的交叉熵 L_T 来总结。这可以从图 1 和图 3b 中看出:


总之,本文提出,学生交叉熵在 L_T 中遵循 broken 幂律,在 N_S 和 D_S 中遵循幂律:


在此之后,论文分析了在不同计算预算下如何最优地分配教师和学生模型的资源,包括教师模型的训练成本和学生模型的蒸馏成本,并比较了蒸馏和监督学习在不同计算预算下的性能,发现当教师模型的训练成本被考虑时,监督学习通常更有效。

这项工作代表了已知最大规模的蒸馏受控实证研究,系统消融了常见的蒸馏技术。正如监督扩展减轻了监督预训练中的风险一样,新工作为生产更小、更强大的模型提供了路线图,并增强了测试时扩展的可行性。

© THE END 
转载请联系本公众号获得授权
投稿或寻求报道:[email protected]

关于“能力差距”,除了模型大小,我觉得训练数据的差异也会有影响。如果教师模型在更丰富的数据集上训练,它学到的知识更全面,但如果学生模型的数据集比较局限,可能难以吸收这些知识,反而不如用一个在相似数据集上训练的教师模型效果好。

其实,判断这两个条件是否成立,最终还是要看实验结果。可以先进行一些小规模的实验,比较蒸馏和监督学习的性能,然后根据实验结果来调整策略。实践出真知嘛。

对于第一个条件,关于计算量或数据量的阈值,文章中提到了“新扩展定律给出的学生大小相关阈值”。所以,我们需要根据这个定律来计算具体的阈值,然后根据实际情况判断是否满足。当然,这需要对定律有更深入的理解。

知识蒸馏可以用于跨模态学习,比如用一个图像分类模型作为教师,指导一个文本描述模型学习图像的语义信息,这样文本模型就能根据图像生成更准确的描述。

第二个条件,关于教师模型的可用性,我觉得需要根据项目的具体需求来判断。如果已经有现成的、合适的教师模型,或者需要训练的教师模型可以用于其他任务,那么就满足条件。否则,就需要权衡训练教师模型的成本。

我觉得还可以用于个性化推荐。可以用一个通用的推荐模型作为教师,然后根据用户的个人数据蒸馏一个个性化的学生模型,这样可以提供更精准的推荐服务。

保护模型隐私也是一个应用场景。可以用公共数据训练一个教师模型,然后用私有数据蒸馏一个学生模型,这样学生模型就能在不泄露私有数据的情况下获得类似的性能。

我觉得模型架构的差异也可能导致“能力差距”。即使教师模型很大,但如果它的架构不适合特定任务,而学生模型的架构更适合,学生模型的性能也可能超过教师模型。有点像老师是全才,学生是专才的意思。

从另一个角度看,优化算法和超参数的选择也会影响“能力差距”。教师模型的训练可能使用了更高级的优化器和精细调整的超参数,而学生模型的训练资源有限,可能无法达到同样的效果,导致性能差距。