NeurIPS 2024 | 收敛加快8倍，准确率提升超30%！华科大提出MoE Jetpack框架

混合专家模型（MoE, Mixture of Experts）是一种通过动态激活网络的部分结构来提升计算效率的架构，可以在保持相对稳定的计算成本的前提下大幅增加参数量，从而有效提升模型性能。

这一特性使得MoE能够兼顾模型的规模与效率，已广泛应用于各种大规模任务。

然而，MoE模型通常需要在大型数据集上预训练以获得理想性能，导致其对时间和计算资源的需求极高，这也限制了其在深度学习社区中的普及性。

为解决这一问题，华中科技大学的研究人员提出了MoE Jetpack框架，利用密集模型的预训练权重（Dense checkpoints）来微调出视觉混合专家模型（MoE，Mixture of Experts）。

目前，这项工作已被NeurIPS 2024接收。

论文标题：MoE Jetpack: From Dense Checkpoints to Adaptive Mixture of Experts for Vision Tasks

论文地址： https://arxiv.org/abs/2406.04801

代码地址： https://github.com/Adlith/MoE-Jetpack

MoE Jetpack框架的核心创新包括：

1. Checkpoint recycling：通过采样密集模型权重产生差异化的专家，组成MoE模型的初始化权重，从而加速模型收敛、提升性能，并避免大规模的MoE模型预训练。

2. SpheroMoE Layer：通过调整MoE结构，利用交叉注意力机制进行专家分配，将query和key投影到超球空间以提升微调过程的稳定性，并通过一系列专家正则化方法有效缓解MoE模型微调过程中的过拟合现象。

实验结果表明，MoE Jetpack在多个数据集和网络结构上实现了显著的性能提升。在ImageNet-1K上，模型收敛速度提升2倍，准确率提高了2.8%；在小规模数据集上，收敛速度可达8倍提升，准确率提升超过30%。

图1 （a）MoE Jetpack将密集预训练权重转化为MoE模型的初始化权重，在性能提升的同时保持等效的FLOPs。(b) 未预训练的ViT、微调的ViT、未预训练的Soft MoE 与MoE Jetpack在多个视觉数据集上的性能比较。

在多个下游数据集上的实验表明，该框架能够高效利用预训练权重，实现更快的收敛速度和更优的性能表现。

研究方法

MoE Jetpack的核心由两个阶段组成：Checkpoint Recycling（用于MoE模型的初始化）和SpheroMoE层（用于微调MoE模型），如下图所示。

图2 Checkpoint Recycling和SpheroMoE结构

Checkpoint Recycling：作为MoE Jetpack的基础阶段，Checkpoint Recycling通过将预训练的密集模型权重转换为高质量的MoE初始化权重，使新模型在性能和收敛速度上都得以提升。

具体地，Checkpoint Recycling从密集权重的多层感知器（MLP）中采样出部分权重构建专家层，以确保专家的多样性和灵活性。

本文比较了四种主要的权重回收策略：

1. 重要性采样（Importance-Based Weight Sampling）：重要性采样是 MoE Jetpack 默认的权重采样方法，通过计算输出特征的均值，选择top-d' 个最重要的特征维度，同时根据隐藏单元的激活值大小进行独立采样。这种方法确保每个专家包含关键的特征和隐藏单元，有助于提升模型的初始化质量和训练效率。

2. 图划分法（Co-Activation Graph Partitioning）：图划分法通过构建共激活图，将常一起激活的隐藏单元分组，使用图划分算法（如 Metis），将密集权重的隐藏单元划分成多个子图，组合不同的子图形成不同专家层，确保专家专注于不同的功能区域，提升模型的特征表达能力。

3. 均匀采样（Uniform Weight Selection）：均匀采样在特征维度和隐藏单元上均匀选择权重，保证每个 MoE 专家层均衡分布初始化权重。此方法实现简单，但不考虑特征重要性，因此性能提升效果较为一般。

4. 随机采样（Random Weight Sampling）：随机采样在特征维度和隐藏单元中随机抽取，生成专家层的初始化权重。该方法实现简便，但由于没有关注特征的重要性，性能通常较低。

Checkpoint Recycling引入的计算开销几乎可以忽略，同时显著提升模型性能并与收敛速度。

在MoE模型权重初始化后，SpheroMoE层进一步优化微调过程。SpheroMoE 层的引入解决了 MoE 模型在优化过程中面临的数值不稳定、专家过度专一等问题。它通过三种机制来提升模型在下游任务的性能和稳定性：

1. 超球路由机制：利用cross attention结构，将输入动态分配给 MoE 模型的不同专家。这种机制首先对随机初始化的专家查询（Q, query）和输入键（K, key）进行归一化投影（L2 Norm）至超球空间，以确保数值稳定性，并通过余弦相似度来选择输入对应的专家。最终输出由各个专家的结果组合而成，保证 MoE 模型的输出特征和密集模型之间的分布一致性。

2. 自适应双路径MoE：为提升计算效率，SpheroMoE路由将输入划分为高重要性和低重要性两类，并引导其进入不同计算路径：高重要性输入分配至包含更大参数量的核心专家；低重要性输入则进入包含较小专家的通用路径。这样的双路径结构通过划分细粒度的专家增加了专家的数量，优化了资源利用，提升了模型的性能与计算效率。

图3 自适应双路径MoE

3. 专家正则化：为避免专家层过度专注于特定输入或出现过度特化，本文引入可学习的软温度参数，用以调整softmax的平滑程度以精确控制输入的分配和输出的组合。

此外，使用专家随机失活机制能有效防止模型对特定专家的依赖。

这些设计使MoE Jetpack在下游任务微调中不仅具备了更快的收敛速度，还实现了显著的性能提升。

实验结果

本文在 ViT 和 ConvNeXt 两种典型网络结构以及八个图像分类任务上进行了广泛实验。实验结果表明，MoE Jetpack 在性能上显著优于随机初始化的 Soft MoE 模型，并且利用密集权重的微调效果明显超过直接微调经过预训练的密集模型。

表1 MoE Jetpack基于ViT和ConvNeXt在8个下游数据集上的性能表现

此外，论文还对MoE Jetpack的多种配置进行了深入研究，系统分析了不同专家数量、不同原始网络尺寸大小等因素对模型性能的影响。

表2 各种参数量的MoE Jetpack都展示出了显著的性能提升

下图展示了微调的全过程中MoE Jetpack对加速模型收敛速度和提升模型性能方面的效果，突显了其作为 MoE 模型预训练替代方法的潜力。

图4 MoE Jetpack带来了收敛速度提升

专家注意力图展示了不同的专家关注图像的不同区域，各司其职。专家贡献分布图表明，核心专家和普通专家在不同层次的贡献差异显著，展示了模型的自适应路由调度机制。

图5 专家注意力图和专家贡献分布图

总结

MoE Jetpack是一个创新框架，旨在将预训练的密集模型权重高效转换为MoE模型。通过提出Checkpoint Recycling技术，MoE Jetpack能够有效继承密集模型的知识；引入的SpheroMoE 层，显著提升微调过程的稳定性和性能。

该框架不仅降低了MoE模型的训练成本和硬件要求，还减少了对环境的影响，使得研究者在普通计算资源下也能轻松使用混合专家模型，为MoE的广泛研究与应用提供了有力支持。