模块化神经架构因其强大的泛化能力和对新领域的高效样本适应能力而越来越受到关注。然而,特别是在早期阶段,训练模块化模型面临挑战,这是由于它们固有的稀疏连接性导致的优化困难。利用单体模型的知识,如使用知识蒸馏技术,可能会促进模块化模型的训练,并使它们能够整合来自多个在不同来源预训练的模型的知识。然而,传统的知识蒸馏方法并不是为模块化模型量身定做的,当直接应用时可能会失败,这是由于独特的架构和涉及的庞大参数数量。 受到这些挑战的启发,作者提出了一种通用的模块到模块知识蒸馏(m2mKD)方法,用于在模块间传递知识。m2mKD涉及从预训练的单体模型中分离出的教师模块,以及模块化模型的学生模块。m2mKD分别将这些模块与共享元模型结合,并鼓励学生模块模仿教师模块的行为。作者评估了m2mKD在两种不同的模块化神经架构上的有效性:神经注意力电路(NACs)和视觉专家混合(V-MoE)。 通过将m2mKD应用于NACs,作者在Tiny-ImageNet上的IID准确度(高达5.6%)和Tiny-ImageNet-R上的OOD鲁棒性(高达4.2%)上取得了显著改进。平均来看,作者在ImageNet和ImageNet-R上都观察到了1%的增益。使用m2mKD训练的V-MoE-Base模型在ImageNet上的准确度也比端到端训练高了3.5%。实验结果表明,m2mKD为将模块化网络与预训练的单体模型连接起来提供了一个有前景的解决方案。 代码:https://github.com/kamanphoebe/m2mKD
尽管在计算机视觉和自然语言处理等各个领域,大型单体模型取得了显著的成功,但人们对其有限的泛化能力和不断增加的计算成本表示担忧。与此同时,模块化模型越来越受到关注,它们有望减轻单体模型的缺点。与具有固定计算图和参数的单体模型相比,模块化神经架构能够根据输入调整其参数,具有静态单体模型所缺乏的优越特性。与集体优化参数的单体模型不同,模块化模型由独立的模块组成,每个模块可以局部更新,而不会影响网络的其他部分。这些模块经过训练以专精于特定任务。在推理过程中,即使对于分布外的样本,也只有相关的模块会被激活,从而提高泛化性能。例如,DEMix Layers通过从训练数据中学习到的医学和新闻模块联合表示与COVID-19相关的数据。此外,模块化模型中的条件计算提高了计算效率。一个典型的模块化架构,即专家混合体(Mixture-of-Experts, MoE),在保持与原始模型相似的计算需求的同时,大幅提升了模型的能力。
尽管模块化架构在分布外(OOD)鲁棒性和计算效率方面超过了单体模型,但它们的算法和实现方面仍有很大的提升空间。由于优化困难,源于稀疏交互的挑战,训练模块化模型变得复杂。虽然近期的研究已经调查了模块化模型的训练不稳定性,但这一领域仍需进行彻底的研究。在工程方面,存在几种MoE的实现方式,但其中许多并没有考虑为专家进行动态工作负载分配或支持像混合精度训练这样的加速技术。尽管适应性并行已经解决了MoE的动态特性,但由于无法避免的通信开销,训练MoE仍然比其单体对应模型要慢得多。
为了缓解优化问题,利用预训练的单一模型来指导模块化模型的训练显示出潜力。知识蒸馏(KD)是一种现有技术,它将知识从预训练的教师模型转移到较小的学生模型。在单一模型的环境下,KD已被证明是有效的。然而,由于模块化模型庞大的模型尺寸,直接将传统的KD方法应用于模块化模型在计算上是昂贵的。甚至使用单一模型作为大型模块化模型的教师可能会损害性能(见表6)。此外,通过常规方法训练的单一模型可能不是作为教师的最优选择。
受到深度孵化的分割与征服训练机制的启发,作者引入了模块到模块的知识蒸馏(m2mKD),以在单体教师模型的子模块和模块化学生之间传递知识。如图1所示,首先作者采用深度孵化这一模块化训练方法,使用小型元模型孵化教师模块。接下来,作者鼓励学生模块模仿教师模块的行为。最后,用蒸馏后的学生模块初始化模块化学生模型。通过在分布式模块 Level 调整蒸馏过程,m2mKD显著降低了教师模型的容量需求,并实现了独立训练,使得蒸馏的计算成本更低。当模块化学生模型较大或在集成学习中涉及多个教师时,成本节省尤为明显。此外,教师模型以模块化方式进行训练,这可能有利于教授模块化学生。
请注意,作者的m2mKD算法对教师和学生模型的架构没有任何限制。作者在NACs和V-MoE上评估了m2mKD的性能。实验显示,使用m2mKD训练的NAC模型在Tiny-ImageNet上的IID准确率大约提高了5%,在Tiny-ImageNet-R上的分布外(OOD)鲁棒性大约增强了4%。此外,在ImageNet和ImageNet-R上平均提高了大约1%。对于V-MoE模型的实验结果表明,在小型教师的情况下,m2mKD同样有效。
作者的贡献如下:
在本节中,作者首先回顾深度孵化技术的基本方面,因为作者的工作在很大程度上依赖于它。随后,作者将介绍作者提出的方法。
深度孵化采用了一种分而治之的策略,模块化地训练大型模型,如图2所示。该过程包括三个阶段:元模型预训练、模块孵化以及组装。
值得注意的是,元模型中的层数与孵化阶段孵化的模块数相对应。
m2mKD遵循一个包含三个步骤的流程(图3):准备,模块间知识蒸馏(m2mKD),以及端到端(E2E)训练。更多细节如下所述。
准备。 在开始蒸馏过程之前,应用孵化算法来准备一个元模型和一个教师模型。这一步骤实际上是所提方法的独立部分。假设目标学生模型
总共有
个模块,元模型
应该有相同数量的元层。同样,教师模型
(由
层组成,
),被划分为
个子模块:
正如(Ni等人,2023年)提出的,初步步骤涉及以端到端的方式训练元模型,以孵化
个子模块
。随后,这些生成的子模块被重新组装形成教师模型
,并进行微调。虽然作者遵循了原始论文中描述的相同流程,但作者的关注点是子模块本身,而不是整个模型。
作者选择孵化方法,而不是简单地分离一个预训练模型,因为作者认为模块孵化阶段向子模块传授了额外的知识。这使得它们能够学会如何作为单独的模块运作,而不是作为一个完整模型的未完整片段。
模块间知识蒸馏(Module-to-module Knowledge Distillation):一旦组装模型
的微调完成,子模块或者说作者接下来称之为“教师模块”的就准备进行m2mKD。与传统的知识蒸馏方法不同,m2mKD旨在模块之间传递知识,而不是整个模型之间。
类似于(Ni et al., 2023)提出的模块孵化,作者将教师模块和学生模块分别与元模型连接起来。通过比较两个生成的混合模型的输出,鼓励学生模块模仿相应教师模块的行为。杨等人(Yang et al., 2022)的先前研究表明,不同网络中位于相似深度的块可以被认为是功能等效的。这一洞见表明神经网络在相似的网络阶段学习相似的模式。
利用这一洞见,作者在相同深度为学生分配教师模块,从而形成
个教学对
。然后每个教学对可以并行进行m2mKD。对于第
个教学对,作者用
和拼接的学生模块
替换第
个元层,从而产生两个混合网络:
修改后的学生模块,表示为
,融入了一个线性缝合层,该层直接插入在模块之前和/或之后。这个缝合层负责调整特征向量的维度,以解决元层和学生模块之间可能存在的维度不匹配问题。预缝合层的权重矩阵表示为
,而后缝合层的权重矩阵表示为
。在这里,
代表元层的维度,
代表学生模块的维度。
两个混合网络现在被认为是“完整”的模型。因此,作者可以直接应用传统的基于响应的知识蒸馏技术。具体来说,对于本文考虑的分类问题,作者通过测量给定输入
的输出logits
和
之间的Kullback-Leibler (KL)散度来进行比较。然后,总损失
被定义为分类交叉熵损失
与知识蒸馏损失
的加权和:
其中
表示平衡因子,
代表softmax的温度系数,而
表示与输入
相关联的标签。在整个蒸馏过程中,只有学生模块会被更新,而元层和教师模块都保持不变。
端到端训练。 给定一个由
个模块组成的学生模型
,作者可以并行地对所有教学对
运行 m2mKD 以获得
个提炼的学生模块
。最后一步是简单地将学习到的参数加载到
中并进行端到端(E2E)训练。注意,在这一阶段,所有学生模块的拼接层都会被丢弃。
对于 NAC 模型架构,如果在 m2mKD 和 E2E 训练中使用的数据库不同(参考第4.2节),在加载某些组件的学习参数时可能会出现维度不匹配的问题。在这种情况下,这些不兼容的元素将不会被加载。
NAC是一种新型的模块化架构,如图3(a)所示。它包括一个读入层、多个传播层和一个读出层。两个连续层内的模块通过随机核调制点积注意力(SKMDPA)稀疏连接。每个模块内的基本计算单元是ModFC层,它构成了ModFFN。大部分参数在模块间共享,每个模块使用自己的代码向量
来调节其计算。
SKMDPA机制采用稀疏注意力根据模块的签名向量
计算模块间的相似性,从而确定模块间的通信。每个模块维持一个状态向量
,作为下一层的输入。状态向量在读入层初始化,通过多个传播层更新,并最终作为读出层的输入。
数据集。 在本文中,作者专注于图像分类任务。预准备和m2mKD阶段的训练都是使用ImageNet-1k数据集进行的。对于NAC模型,作者使用ImageNet数据集及其迷你版Tiny-ImageNet进行端到端训练。对于OOD评估,作者使用了ImageNet-R(版本)数据集及其下采样子集Tiny-ImageNet-R。此外,还采用了CIFAR-100数据集进行小样本适配。
准备阶段。 由于作者的目标NAC模型总共包含10层(
),作者选择具有32层的DeiT-Huge模型作为教师模型,以确保足够的深度。DeiT-Huge模型被划分为10个子模块,第一个和最后一个子模块各包含4层,其余子模块包含3层。NAC模型的参数不超过37M,而DeiT-Huge模型包含632M参数。这使得每个学生和教师模块分别具有大约3.7M和63.2M个参数。在准备阶段开始时,在ImageNet上训练一个10层的元模型,持续300个周期。
随后,每个教师模块由元模型孵化100个周期,然后将所有模块重新组装在一起进行额外的100个周期的微调。这一阶段的训练配置与深度孵化方法相同。
m2mMKD. 当隐藏教师部分时,m2mMKD 可以被视为带有额外损失项的模块孵化。因此,作者再次采用了与模块孵化相似的实验设置用于 m2mKD。唯一的区别是每个学生模块仅被训练10个周期。作者将平衡因子
设置为 0.5,softmax 温度
设置为 1.0。鉴于元模型的维度
以及学生模块的维度
(针对Tiny-ImageNet1,或者对于ImageNet是
),每对拼接层将会有1M个参数。在用元层连接学生模块和教师模块之后,得到的混合网络
和
分别包含182.5M和242.0M个参数。
最初,NAC传播器层接收前一层的状态向量作为输入,并在SKMDPA和ModFFN的计算后更新状态。然而,作者对学生模块去除了SKMDPA部分,使得所有模块之间可以相互通信。这一修改基于两个原因:
为了确保ModFC层正常工作,作者随机初始化状态和代码向量。注意,这些向量在蒸馏过程中不会被更新,也不会被加载到目标模型中进行端到端的训练。具体来说,
包括输入分词器和读入层,
包括第
个修改后的传播器层,而
包括修改后的读出层和输出分词器。
端到端训练。在相同的数据集上进行训练时,对于复现的 Baseline 模型和精馏的NAC模型,超参数保持不变。然而,由于关于NAC的原始论文没有提供一个全面的超参数列表,且作者使用给定的超参数无法复现报告的结果,因此作者调整了一些超参数,以尽可能接近他们报告的结果。表9展示了作者选择的超参数2,包括所有更改的值。
如前所述,作者的教师模型DeiT-Huge和前两个阶段的学生模块NAC仅在ImageNet上进行训练。组装的教师模型达到了81.8%的验证准确率。使用8个A100 80GB GPU,可以在12小时内完成所有十个学生模块的m2mKD训练。对于在Tiny-ImageNet上的端到端(E2E)训练,由于维度不匹配,学生模块中的一部分输入分词器和整个输出分词器需要被丢弃。
表1比较了在Tiny-ImageNet和Tiny-ImageNet-R上复现的 Baseline 和作者蒸馏的NAC模型。复现的结果略优于原始论文中报告的值。尽管数据集与m2mKD阶段使用的数据集不同,但采用各种图先验正则化器的蒸馏模型在IID性能上平均提高了5.3%,在OOD鲁棒性上提高了3.8%。这表明添加蒸馏后的学生模块不仅增强了最终NAC在类似任务上的能力,还提高了其模块化。
为确保可复现性,作者三次重复了具有尺度无关先验的蒸馏模型的端到端训练阶段,并测量了它们的平均值和标准差。结果展示在表2中。
表3展示了与ImageNet的比较。原始论文仅提供了使用无尺度图先验在ImageNet上训练的NAC的验证准确度,为77%。作者重现了所有四种图先验的 Baseline ,最高值为76.1%。作者的精炼NAC在IID和OOD性能上分别实现了最大增益1.0%和2.4%。
在MoE模型中,一些或全部的标准单体模型中的前馈网络(FFNs)被替换为MoE层。如图3(b)所示,一个MoE层由多个专家和一个门组成,其中每个专家本质上是一个FFN。门通常实现为MLP,负责选择特定数量的专家来处理输入。然后对所选专家的输出进行汇总。一个MoE层进行的计算可以表示为:
其中
表示专家的总数,
代表第
个专家的输出,
表示门结构为第
个专家计算出的得分。
数据集。 在整个流程中(从准备到端到端训练),作者专门使用了ImageNet-1k数据集用于MoE模型的训练。在CIFAR-100和CUB-2011数据集上评估了模型的少样本适应能力。
准备. 作者选择了一个具有24层和304M参数的DeiT-Large模型作为教师模型。作者没有从零开始训练元模型和教师模型,而是使用了Deep Incubation发布的预训练权重,这些预训练权重在ImageNet-1k上的准确率达到83.9%。教师模型被均匀划分为
个子模块,每个子模块由6层组成。元模型是一个四层的DeiT模型,其嵌入维度与教师模型相同。
在的这一系列实验中,目标模型
被设定为视觉MoE基础(V-MoE-B)模型。作者没有使用精心设计的门控机制,而是采用了最早引入的门控机制,该机制仅执行一个top-
(
)操作,并且没有应用额外的约束或平衡损失:
作者的V-MoE-B模型由12个MoE层组成(即DeiT-B中的所有前馈层都变为了MoE层),每个层包含8个专家。因此,学生模型总共有483M个参数。请注意,学生模型的大小超过了教师模型的大小。V-MoE-B模型的每三层被划分为一个学生模块,总共形成了4个学生模块。
在这种情况下,作者有
和
,因此每对拼接层大约占1.6M个参数。结果是,每个学生模块大约有122.4M个参数,而每个教师模块有77.6M个参数。在这个阶段构建的两个混合网络
和
将分别包含160.5M和115.9M个参数。
与NAC模型不同,MoE的学生模块不需要修改,它们的所有学习参数都可以在端到端训练阶段加载到目标模型中。与NAC模型的设置一样,作者设置平衡因子
和softmax温度
。所有其他超参数与在Deep Incubation中用于孵化DeiT-B的超参数相同,学生模块训练100个周期。
端到端训练。 再次,作者采用了与Deep Incubation相同的超参数集来训练V-MoE-B模型,除了更新频率参数,其设置为原始值的一半。
作者将m2mKD与三种 Baseline 方法进行了比较:纯端到端训练、传统知识蒸馏和深度孵化。对于端到端训练,作者从零开始训练一个V-MoE-B模型,共300个周期,使用的超参数与DeiT-B训练中使用的相同。与在模块 Level 执行知识蒸馏的m2mKD不同,传统KD指的是在完整模型之间进行知识蒸馏。
为了进行公平比较,作者再次使用DeiT-L作为教师模型,V-MoE-B作为学生模型。它们输出逻辑之间的KL散度以
和
(见方程3)的形式加入到损失中。训练过程持续300个周期。最后,作者考虑了原始的深度孵化方法,其中作者使用他们的开源预训练权重,这个预训练权重是元模型最初为孵化DeiT-B而训练的,用于V-MoE-B实验。在ImageNet-1k上的验证结果总结在表6中。
与传统KD相比,其他方法的性能更佳,其准确率比纯端到端训练低1.9%,这是第二低的。这证实了作者的观点,即现有的知识蒸馏技术并不一定与模块化模型兼容。另一方面,深度孵化方法比端到端训练的准确率高2.97%,这表明它不仅对单体模型有效,对模块化模型也有效。作者的m2mKD方法进一步将准确率提高了0.50%,是这些方法中性能最佳的。
鉴于NACs在先前实验中在OOD(Out-of-Distribution,分布外)结果上的显著表现,作者也研究了MoE在ImageNet-R数据集上的OOD鲁棒性,这在其他与MoE相关的文献中很少讨论。令人惊讶的是,所有使用四种不同方法训练的MoE模型在ImageNet-R数据集上的准确率几乎为0%。基于这些实验,似乎NACs在OOD或零样本性能方面显著强于MoEs。
训练时间。除了准确度之外,表6还展示了相对于纯端到端训练方法所需时间计算出的训练时间比例。m2mKD的速度略慢于Deep Incubation,但在时间上大致等同于纯端到端方法,而传统KD需要1.12
的时间。进一步的实验(附录C)表明,m2mKD的比例可以减少到0.57,性能只有轻微下降。
为了公平比较,作者排除了KD和m2mKD方法中教师模型训练的时间。换句话说,在计算m2mKD的时间时,只考虑了元模型训练、蒸馏和端到端训练阶段的时间。在这一点上,自然会提出这样一个问题:是否可以简化m2mKD的准备工作阶段。当前的m2mKD流程包括在准备阶段孵化一个教师模型。作者假设这样一个模型将成为对模块化学生而言更有知识的教师。
然而,由于有许多预训练模型是公开可用的,因此使用它们作为m2mKD的教师模型而不是从零开始训练新模型将节省时间。如果采用这种方法,总的训练时间将正好是表中所示的时间。作者将在未来的工作中探讨使用不同方法训练的教师如何影响m2mKD的性能。
小样本学习。 作者评估了三种模型的小样本适应能力:一种使用深度孵化训练的DeiT-B模型,以及两种分别使用深度孵化和m2mKD训练的V-MoE-B模型。超参数可以在表10中找到。在小样本训练期间,这些模型中的所有参数都被冻结,除了新随机初始化的分类器层,其输出维度被设置为小样本类别数(所有实验中为8)。对于每个样本shot,实验都使用五个不同的种子重复进行。
在表7中,作者报告了每个实验在100个epoch内达到的最高准确度,以及基于种子的平均准确度和标准差。平均结果也描绘在图6中。
令人惊讶的是,DeiT-B模型的表现优于两个V-MoE-B模型,而后者由于模块化特性本应更有优势。这或许可以归因于相对较小的数据集大小(ImageNet-1k),导致每个专家的学习样本不足以进行全面的训练。然而,当关注两个MoE模型时,使用m2mKD训练的模型在CUB-2011的所有shots以及CIFAR-100的1-shot实验中一致地比孵化模型高出大约2%。与NACs的小样本结果(第4.2节)类似,标准差相对较大。为了更精确地评估小样本表现,有必要在更大的数据集上进行实验。
消融研究。在所提出的m2mKD流水线中,缝合层仅在蒸馏阶段使用,之后被丢弃,以确保学生模型不需要进行任何修改。尽管缝合层通常只占大型模块化模型总参数的一小部分(例如,在作者的V-MoE-B学生模型中占1%),但直接移除它们可能会导致一些从教师那里获得的知识丧失,可能会对最终模型的性能产生负面影响。为了研究这一方面,作者通过在学生模型中保留缝合层来进行额外的实验。这些实验包括减少m2mKD的蒸馏周期数量以及扩大规模。更多详细信息,请参考附录C。
在这项工作中,作者提出了一种模块间知识蒸馏(m2mKD)的通用方法,用于在模块化模型训练的背景下传递知识。m2mKD利用一个整体的教师模型来促进模块化学生模型的训练,为将预训练的整体模型转换为模块化架构以利用其优势提供了一个有前景的策略。
在NACs和MoEs上的实验结果表明,所提出的 Pipeline 有效地提高了模块化模型在IID准确性和OOD鲁棒性方面的性能,即使在学生模型大小小于教师模型大小的场景中也是如此。然而,在缩小学生模块化模型与教师模型之间性能差距方面仍有改进空间。为了进一步提高有效性,m2mKD可以与现有技术如集成学习相结合。作者希望这项工作能够推进关于模块化知识蒸馏和整体到模块化转换的研究,并期待未来在这些领域的研究。
在所提出的m2mKD Pipeline 中,缝合层仅在蒸馏阶段使用,之后被丢弃,以确保学生模型无需进行任何修改。尽管缝合层通常只占大型模块化模型总参数的一小部分(例如,在作者的V-MoE-B学生模型中占1%),但直接移除它们可能会导致从教师那里获得的一些知识丢失,可能会对最终模型的性能产生负面影响。
为了研究这个问题,作者进行了保留学生模型中缝合层的实验。然而,正如表11所示,结果表明这种方法并没有带来明显的改进。因此,作者坚持决定丢弃缝合层。
在保持缝合层设置的条件下,作者进行了一些额外的实验。之前的实验表明,缝合层的存在并不会产生太大影响,因此作者认为即使移除它们,以下结论仍然适用。首先,由于NAC模型在仅进行10个周期的学生模块蒸馏后就能实现很大改进,作者测试了在相同条件下V-MoE-B模型的性能。为此,作者调整了m2mKD阶段的一些超参数,并将学生模块蒸馏10个周期。端到端(E2E)阶段保持100个周期。
表11显示,将蒸馏周期从100减少到10,准确率下降了2.1%,但仍然比Deep Incubation Baseline 高出2.8%。作者对Deep Incubation的孵化阶段应用了相同的更改,令人惊讶的是,使用新超参数后性能并没有下降。
接下来,作者检查了m2mKD的可扩展性。选择由Deep Incubation预训练的DeiT-Huge模型(632M)作为教师模型。学生模型(
)是一个V-MoE-Large模型,每隔一层放置12个MoE层,总共包含1.0B个参数。表11将m2mKD与Deep Incubation进行了比较,前者在ImageNet-1k上的准确率高出0.27%。
接下来的消融实验是基于保留缝合层进行的。之前的实验表明,缝合层的存在并没有显著影响,作者认为即使移除它们,以下结论依然有效。首先,由于NAC模型在仅为学生模块蒸馏10个周期时就表现出显著的改进,作者在相同条件下测试了V-MoE-B模型的性能。为此,在m2mKD阶段调整了一些超参数,并为学生模块蒸馏10个周期,同时保持E2E阶段为100个周期。
如表11所示,将蒸馏周期从100减少到10,准确度下降了2.1%,但仍然比深度孵化 Baseline 高2.8%。令人惊讶的是,对深度孵化阶段的孵化过程应用相同的变化并没有导致性能下降。
接下来,作者考察了m2mKD的可扩展性。教师模型是由深度孵化预训练的DeiT-Huge模型(632M),而学生模型(
)是一个V-MoE-Large模型,每两层放置一个MoE层,总共12个MoE层,总参数量为1.0B。如表11所示,在ImageNet-1k上,m2mKD的准确度比深度孵化高出0.27%。
[1].m2mKD: Module-to-Module Knowledge Distillation for Modular Transformers.