SoftMoE：From sparse to soft mixtures of experts
Soft MoE是一种全可微分的稀疏Transformer架构，旨在解决传统稀疏混合专家模型（Sparse Mixture of Experts, MoEs）在训练稳定性、扩展专家数量、微调效果等方面的问题。Soft MoE通过软分配机制，将输入标记的不同加权组合传递给每个专家，而不是像传统MoE那样进行硬分配。这种方法避免了传统MoE中的标记丢失和专家不平衡问题，同时保持了MoE在推理成本较低的情况下扩展模型容量的优势。
From：ICLR 2024
Paper：https://arxiv.org/pdf/2308.00951
Code：https://github.com/google-research/vmoe
文章目录
- 研究动机
- 传统稀疏混合专家模型（Sparse Mixture of Experts, MoEs）的局限性
- 全可微分模型的需求
- 推理效率的提升
- 研究方法
- 算法概述
- 算法步骤
- 1. 输入标记的加权平均（槽的生成）
- 2. 专家处理槽
- 3. 输出标记的加权平均（槽的组合）
- 算法特性
- 1. 全可微分
- 2. 无标记丢失和专家不平衡
- 3. 推理效率
- 4. 稀疏性与密集性
- 实现细节
- 实验效果
- 实验设置
- 1. 数据集
- 2. 模型架构
- 3. 训练设置
- 实验结果
- 1. 训练Pareto前沿
- 2. 长时间训练的效果
- 3. 推理时间优化
- 4. 模型扩展性测试
- 5. 多模态任务表现
- 实验分析
- 1. Soft MoE的软分配机制
- 2. 长时间训练的效果
- 3. 推理效率
研究动机
传统稀疏混合专家模型（Sparse Mixture of Experts, MoEs）的局限性
尽管稀疏MoE在扩展模型容量方面表现出色，能够在不显著增加训练或推理成本的情况下提升性能，但它们也存在一些显著的问题：
- 训练不稳定：稀疏MoE在训练过程中容易出现不稳定现象，尤其是在大规模模型中。
- 标记丢失：传统MoE在路由过程中可能会丢弃某些标记，导致模型无法处理这些标记，影响性能。
- 专家不平衡：某些专家可能会处理过多的标记，而其他专家则处理较少，导致资源分配不均。
- 扩展性问题：传统MoE在扩展专家数量时面临挑战，尤其是在需要处理大量标记的情况下。
- 微调效果不佳：传统MoE在微调过程中表现不如预期，难以充分利用预训练模型的潜力。
图1：Sparse MoE 和 Soft MoE 的区别：左：Sparse MoE，给每个 Expert 分配一定的输入 token。右：Soft MoE，给每个 Expert 分配的是所有输入 token 的加权平均值
全可微分模型的需求
传统MoE的路由机制通常是离散的，涉及复杂的优化问题（如线性规划、最优传输等），这些方法虽然有效，但往往难以实现全可微分，限制了模型的优化和扩展能力。因此，研究者希望设计一种全可微分的路由机制，以简化训练过程并提高模型的稳定性。
推理效率的提升
尽管稀疏MoE在训练过程中能够节省计算资源，但在推理阶段，传统MoE的性能提升并不总是与计算成本的增加成正比。研究者希望通过设计一种新的路由机制，能够在保持高性能的同时，进一步降低推理成本。
研究方法
算法概述
Soft MoE的核心思想是通过软分配机制，将输入标记的不同加权组合传递给每个专家，而不是像传统MoE那样进行硬分配。具体来说，Soft MoE通过计算所有输入标记的加权平均值（称为“槽”），并将这些槽分配给不同的专家进行处理。每个专家只处理一部分槽，而不是直接处理输入标记。
算法步骤
Soft MoE通过计算所有输入标记的加权平均值来实现软分配，每个专家处理这些加权平均值中的一个子集。具体来说，Soft MoE使用softmax函数计算分配权重，并将这些权重应用于输入标记和输出标记，从而实现全可微分的路由机制。
1. 输入标记的加权平均（槽的生成）
-
输入标记：假设输入标记为一个矩阵 X ∈ R m × d \mathbf{X} \in \mathbb{R}^{m \times d} X∈Rm×d，其中 m m m 是标记的数量， d d d 是每个标记的维度。
-
槽的生成：Soft MoE通过计算所有输入标记的加权平均值来生成槽。具体来说，每个槽是所有输入标记的加权平均值，权重由一个softmax函数计算得出。
计算公式如下： D i j = exp ( ( X Φ ) i j ) ∑ i ′ = 1 m exp ( ( X Φ ) i ′ j ) \mathbf{D}_{ij} = \frac{\exp((\mathbf{X} \boldsymbol{\Phi})_{ij})}{\sum_{i'=1}^{m} \exp((\mathbf{X} \boldsymbol{\Phi})_{i'j})} Dij=∑i′=1mexp((XΦ)i′j)exp((XΦ)ij)其中， D \mathbf{D} D 是分配权重矩阵， Φ ∈ R d × ( n ⋅ p ) \boldsymbol{\Phi} \in \mathbb{R}^{d \times (n \cdot p)} Φ∈Rd×(n⋅p) 是可训练的参数矩阵， n n n 是专家的数量， p p p 是每个专家处理的槽的数量。
生成的槽矩阵为 X ~ = D ⊤ X \tilde{\mathbf{X}} = \mathbf{D}^{\top} \mathbf{X} X~=D⊤X 其中， X ~ ∈ R ( n ⋅ p ) × d \tilde{\mathbf{X}} \in \mathbb{R}^{(n \cdot p) \times d} X~∈R(n⋅p)×d 是所有槽的矩阵。
2. 专家处理槽
- 专家函数：每个专家是一个多层感知机（MLP），记为 f i f_i fi，用于处理分配给它的槽。
- 槽的处理：每个专家处理其对应的槽，生成输出槽矩阵
Y
~
\tilde{\mathbf{Y}}
Y~：
Y ~ i = f ⌊ i / p ⌋ ( X ~ i ) \tilde{\mathbf{Y}}_{i} = f_{\lfloor i/p \rfloor}(\tilde{\mathbf{X}}_{i}) Y~i=f⌊i/p⌋(X~i)
其中， ⌊ i / p ⌋ \lfloor i/p \rfloor ⌊i/p⌋ 表示第 i i i 个槽对应的专家。
3. 输出标记的加权平均（槽的组合）
-
输出标记的生成：每个输出标记是所有输出槽的加权平均值，权重同样由softmax函数计算得出。
计算公式如下： C i j = exp ( ( X Φ ) i j ) ∑ j ′ = 1 n ⋅ p exp ( ( X Φ ) i j ′ ) \mathbf{C}_{ij} = \frac{\exp((\mathbf{X} \boldsymbol{\Phi})_{ij})}{\sum_{j'=1}^{n \cdot p} \exp((\mathbf{X} \boldsymbol{\Phi})_{ij'})} Cij=∑j′=1n⋅pexp((XΦ)ij′)exp((XΦ)ij) 其中， C \mathbf{C} C 是组合权重矩阵。
最终的输出标记矩阵为： Y = C Y ~ \mathbf{Y} = \mathbf{C} \tilde{\mathbf{Y}} Y=CY~ 其中， Y ∈ R m × d \mathbf{Y} \in \mathbb{R}^{m \times d} Y∈Rm×d 是所有输出标记的矩阵。
算法特性
1. 全可微分
Soft MoE的所有操作都是连续且全可微分的，这使得模型可以通过标准的反向传播算法进行训练，避免了传统MoE中离散路由机制带来的优化困难。
2. 无标记丢失和专家不平衡
由于每个槽都是所有输入标记的加权平均值，Soft MoE不会出现标记丢失或专家不平衡的问题。每个标记都会以某种权重参与到所有槽的计算中，而每个专家处理的槽数量是固定的。
3. 推理效率
Soft MoE的计算复杂度主要取决于槽的数量，而不是专家的数量。通过合理设置槽的数量，Soft MoE可以在不显著增加计算成本的情况下扩展模型容量。此外，Soft MoE避免了传统MoE中复杂的排序或top-k操作，从而提高了推理效率。
4. 稀疏性与密集性
尽管Soft MoE的每个槽是所有输入标记的加权平均值，但每个专家只处理一部分槽，因此它仍然具有稀疏性。与密集Transformer不同，Soft MoE的每个专家只处理输入标记的一个子集，而不是所有标记。
实现细节
归一化
为了避免softmax函数在模型维度较大时退化为one-hot向量，Soft MoE在计算分配权重和组合权重时，对输入标记和参数矩阵进行了L2归一化。具体来说，输入标记和参数矩阵分别被归一化为单位范数，并通过一个可训练的标量进行缩放。
算法伪代码
以下是Soft MoE层的简单JAX实现伪代码：
实验效果
Soft MoE在多个视觉任务中表现出色，尤其是在大规模数据集上的预训练和微调任务中。其全可微分的软分配机制解决了传统稀疏MoE中的训练不稳定、标记丢失和专家不平衡问题，同时保持了推理效率和模型容量扩展的优势。实验结果表明，Soft MoE为大规模模型训练提供了一种高效且稳定的解决方案，特别是在多模态学习和大规模预训练任务中具有广泛的应用潜力。
实验设置
1. 数据集
- JFT-4B：一个包含超过40亿张图像的私有数据集，涵盖29,000个类别。实验中使用该数据集进行预训练，并在预训练过程中评估模型的上游验证精度（JFT-4B Precision-at-1）和ImageNet 10-shot精度。
- ImageNet-1k：用于微调和评估模型在ImageNet验证集上的精度。
- WebLI：一个包含10亿张图像和对应文本描述的私有数据集，用于多模态任务的预训练。
2. 模型架构
- Dense Transformers (ViT)：包括ViT-Small (S), ViT-Base (B), ViT-Large (L), 和 ViT-Huge (H)。
- Sparse MoEs：包括Tokens Choice和Experts Choice路由机制的稀疏MoE模型。
- Soft MoE：本文提出的全可微分稀疏MoE模型，替换Transformer中的部分MLP块。
3. 训练设置
- 训练步数：300k步（部分模型训练更长时间，如4M步或9M步）。
- 批量大小：4096。
- 学习率调度：倒数平方根学习率调度（inverse square root schedule），带有线性预热和冷却阶段。
- 分辨率：224x224（部分实验使用更高分辨率）。
实验结果
1. 训练Pareto前沿
实验首先比较了不同模型在训练成本（FLOPs和TPUv3训练时间）与性能（JFT-4B Precision-at-1和ImageNet 10-shot精度）之间的权衡。结果表明，Soft MoE在所有训练成本预算下均优于密集Transformer和传统稀疏MoE。
- Soft MoE的优势：在相同的训练成本下，Soft MoE的性能显著优于ViT和Tokens Choice/Experts Choice稀疏MoE。例如，Soft MoE B/16在训练时间和FLOPs相近的情况下，性能超过了ViT H/14。
- 扩展性：Soft MoE能够有效利用大量专家参数，而不会显著增加计算成本。
2. 长时间训练的效果
为了进一步验证Soft MoE的性能，实验对部分模型进行了长时间训练（4M步或9M步）。结果表明，Soft MoE在长时间训练后仍然保持了显著的性能优势。
- Soft MoE S/16：在长时间训练后，性能超过了ViT B/16，且训练成本显著降低。
- Soft MoE B/16：在长时间训练后，性能接近ViT H/14，但推理速度更快，FLOPs更低。
- Soft MoE L/16：在长时间训练后，性能超过了ViT H/14，且推理速度更快。
3. 推理时间优化
实验还评估了Soft MoE在推理时间上的优势。结果表明，Soft MoE在推理速度上显著优于密集Transformer。
- Soft MoE B/16：推理速度比ViT H/14快10倍，FLOPs减少5倍。
- Soft MoE L/16：推理速度比ViT H/14快2倍，FLOPs减少2倍。
4. 模型扩展性测试
实验研究了不同专家数量和槽数量对模型性能和推理速度的影响。结果表明，Soft MoE在增加专家数量时，性能持续提升，而推理速度几乎保持不变。
- Soft MoE的扩展性：增加专家数量可以显著提升性能，而不会显著增加推理成本。
- 传统稀疏MoE的局限性：增加专家数量会导致推理速度显著下降，且性能提升有限。
5. 多模态任务表现
实验还验证了Soft MoE在多模态任务中的表现。通过在WebLI数据集上进行图像-文本对比学习，结果表明，Soft MoE在多模态任务中同样表现出色。
- Soft MoE L/16：在ImageNet和CIFAR-100的零样本评估中，性能超过了ViT L/16。
- Soft MoE的优势：在多模态任务中，Soft MoE能够更好地利用专家参数，提升图像-文本对齐的效果。
实验分析
1. Soft MoE的软分配机制
实验通过可视化和定量分析，验证了Soft MoE的软分配机制的有效性。结果表明，Soft MoE能够避免标记丢失和专家不平衡问题，且每个标记都能以某种权重参与到所有槽的计算中。
- 标记贡献分布：大多数标记对槽的贡献较为均匀，少数标记贡献较大。
- 专家贡献分布：不同专家对输出标记的贡献存在差异，表明专家参数并非冗余。
2. 长时间训练的效果
实验表明，Soft MoE在长时间训练后仍然保持了稳定的性能提升，且训练成本相对较低。这表明Soft MoE在长时间训练中具有良好的扩展性和稳定性。
3. 推理效率
实验结果表明，Soft MoE在推理速度上显著优于密集Transformer和传统稀疏MoE。通过避免复杂的排序或top-k操作，Soft MoE能够以较低的推理成本实现高性能。