收敛速度最高8倍，准确率提升超30%！华科发布MoE Jetpack框架 | NeurIPS 2024|华科|算法|速度|预训练

新智元报道

编辑：LRST

【新智元导读】华中科技大学的研究人员提出了MoE Jetpack框架，通Checkpoint Recycling方法和SpheroMoE结构，将密集激活模型的预训练权重微调为混合专家（MoE）模型，从而免去了MoE模型的预训练过程，大幅提升了MoE在下游任务中的精度和收敛速度。

混合专家模型（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引入的计算开销几乎可以忽略，同时显著提升模型性能并与收敛速度。

SpheroMoE层

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

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

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