Jetpack

MoE Jetpack: From Dense Checkpoints to Adaptive Mixture of Experts for Vision Tasks^[1]

作者是来自华科的Xingkui Zhu等人。论文引用[1]:Zhu, Xingkui et al. “MoE Jetpack: From Dense Checkpoints to Adaptive Mixture of Experts for Vision Tasks.” ArXiv abs/2406.04801 (2024): n. pag.

Time

• 2024.Jun ### Key Words
• pretrained dense checkpoints
• SpheroMoE
• 一句话总结：作者利用预训练的dense weights来初始化MoE models，将dense model的MLPs划分成多个experts，确保expert的多样性和适用性，同时，为了进行微调，提出了SpheroMoE layer，提出了一些改进来避免expert的over-specialized。

总结

1. 稀疏的激活的MoE相对于密集激活的models，提供了一个有效的alternative，将提高的精度和计算效率进行结合。然而，从零开始训练MoE要求大量的data和计算资源，限制了它们的广泛采用。为了解决这个问题，作者引入了MoE Jetpack，一个旨在将丰富且现成可获取的密集checkpoints微调为专家混合（MoE）模型的框架。MoE Jetpack引入了两个key techniques：(1). checkpoint recycling，用dense checkpoints对MoE models进行初始化，来加速收敛和增强精度。减小对大规模预训练的需求; (2). hyper-spherical adaptive MoE(SpheroMoE) layers，优化MoE的架构，来增强微调的性能和效率。

2. 增加model的size是提高深度学习性能的key factor，然而，model在size上expand，它们的计算需求会激增，导致在训练和推理阶段相当大的slowdown，一个有前景的、将计算成本和model size进行解耦的方式是MoE，不同于dense models，将所有网络的参数都用在输出上，MoE基于输入的tokens，动态地激活model的disinct parts，这使得不增加FLOPs的情况下实现了model scaling，然后在model upscaling阶段保持了training和推理速度，最近的进展见证了MoE在多个domain的implementations。

3. 尽管有潜力，但是MoE面临严重的challenges，由于缺乏pretrained models，不同于dense models，受益于很多可用的预训练的资源，大多数的MoE模型必须随机初始化，从零开始训练。这个过程要求大量的计算资源和大的数据集，限制了MoE的研究。因此，作者的研究旨在利用来自dense checkpoints的预训练知识，来降低MoE的训练时间和数据的要求。作者特别地研究了，**利用dense checkpoints是否能增强精度和微调期间的MoE模型的收敛速度。在本文中，作者提出了MoE Jetpack，一个将pretrained dense checkpoints微调为MoE model的方法，MoE Jetpack利用大量的dense pretraining 来提高MoE的性能和加速收敛，包含两个key techniques，首先：checkpoint recycling利用将MoE models从dense checkpoints中进行初始化，不同于sparse upcycling，简单的复制MLP来构造experts，checkpoint recycling利用多个dense checkpoints和strategic weight selection，对MoE models提供了灵活性，十个生成高质量的初始化weights，第二个technique是hyperspherical adaptive MoE layer，展示出了一个优化的MoE架构，促进了dense checkpoints的无缝集成，增强了微调的性能，现有的MoE架构，例如Switch Transformers和Soft MoE，没有采用预训练的dense checkpoints，导致微调期间的优化不高效和over-specialization，SpheroMoE layer通过normalized token mixing, expert regularization和adaptive dual-path 机制，缓解了这个问题，确保更光滑的集成和提高了性能。 通过给dense checkpoints和MoE JetPack结合，微调的MoE的model模型实现了更高的精度。作者的贡献如下：

   • 作者引入了checkpoint recycling，开拓了将dense checkpoints用于初始化MoE experts的思路，增强了初始化的灵活性，多样化了experts，消除了MoE 预训练的计算成本
   • 作者开发了一个shperoMoE layer，优化将dense checkpoints微调为MoE 架构，缓解了优化的挑战，防止了experts的over-specialization

4. 现有的MoE model通过将dense ViT中的特定的MLP层替换为MoE layer得到的，每个MoE layer包含一个Router function，将输入的tokens给到一组experts，尽管MoE models中的experts是相似的，路由机制不同，对性能很重要，一些路由机制例如top-k，BASE和Sinkhorn-BASE layers，Hash layers，Expert Choice Routing，soft routing，最常用的机制是top-k 路由，通过选择性地激活和Input tokens最相关的 top-k experts来降低计算开销，路由的decision计算如下： \[ \text{Router}(x; \theta_{\text{gate}}) = \text{top-}k\left( \text{softmax}\left( \text{MLP}(x; \theta_{\text{gate}}) \right) \right) \]

   输出是结合选择的experts进行汇总： \[ y = x + \sum_{i \in E} \text{Router}(x; \theta_{\text{gate}}) \cdot \text{Expert}(x; \theta_i) \]

   \(\theta\) 表示weights，\(E\) 是激活的experts的set，\(|E| =K\)，然而，top-k 路由提出了一些挑战，例如不平衡的expert utilization，token dropping和scalability issues，

   一个更加平衡和代表性的方法，Soft MoE，通过隐式地soft assignment，有效地解决了这些挑战，不同于将tokens分配给特定的experts，Soft MoE为每个expert，计算所有输入的tokens的combinations，给定输入tokens，每个MoE layer包含 \(e\) 个expert functions，每个expert有 \(s\) 个slots，中间的输出是通过将experts用于slots得到的。最后，输出的tokens通过结合experts的输出来得到。

5. 在这个部分，作者展示了MoE Jetpack的重要的概念，划分为两个阶段：checkpoint recycling dense checkpoints来初始化MoE models，用hyperspherical adaptive MoE(SpheroMoE) layer。

6. checkpoint recycling是MoE jetpack框架的基础阶段，将预训练的dense checkpoints转换为用于MoE的高质量的初始化weights，这个方法通过weight的reuse，利用predecessors的丰富的预训练知识，提高了性能和收敛速度。这个recycling的过程涉及从predecessors的MLPs中采样一部分的weights用于构建experts，确保expert的多样性和expert size的适应性，来满足多样的需求。 为了定义checkpoint recycling的process，考虑有N 个layers的predecessors，每个layer \(L_i\) 的特征维度是 d维，hidden dimension是 4d维，目标是将这个predecessor转换成successor MoE model S，它也是有N个layers，每个表示为 \(L^`_i\)，但是特征维度是降低了，跟踪Soft MOE的架构，successor有两个segments，一个有 \(N_1\) layers的dense part和有 \(N_2\) layers的MoE part， \(N_1 = N_2 = \frac{N}{2}\)，因此，这个successor model表示为：

   \[ S = \left( \{ L'_i \}_{i=0}^{\frac{N}{2}-1}, \{ L'_t \}_{t=\frac{N}{2}}^{N-1} \right). \] 受Weight Selection的启发，作者的recycling process确保了特征维度上的一致性，作者探索了四个主要的策略来指导checkpoints的recycling。

   Importance-based weight Sampling(Default)：作者在不同的hidden units选择weights和feature dimensions，来构建初始化weights for diverse experts，对于feature dimension selection，在所有layers上保持一致性是必要的，作者通过计算所有层的layer的mean output feuatres和选择top-d` dimensions来实现： \[ \bar{O} = \frac{1}{N} \sum_{i=1}^{N} O_i, \quad \text{top-}d' = \text{argsort}(\bar{O})[:d'] \]

   \(O_i\) 是 layer \(L_i\) 的输出特征，\(\bar{O}\) 是平均输出features，对于hidden units，基于activations的magnitude，对每一layer的采样是独立的，因为units的arrangement不影响输出，来自images的activation values用于得到概率分布

7. Co-Activation Graph Partitioning：这个策略将频繁co-activated hidden units划分成一个expert，作者通过数hidden units的co-activations，构建一个co-activation graph，co-activation graph中，每个unit是一个vertex，edges表示它们的co-activation frequency，例如：\(G=(V,E)\) 是co-activation graph, \(V\) 表示hidden units， \(E\) 表示带有weights的边，表示co-activation counts，用Metis Graph partitioning，得到一些subgraphs，experts通过sub-graphs的结合形成，这个方法利用hidden units的natural grouping，确保每个expert捕捉predecessor model的特定的functional subset。

8. Uniform Weight Selection：weights是在dimension和hidden units上均匀选择，对于一个特征维度是d的predeccessor和一个特征维度是 d`的successor，选择的weights是这样的： \[ W_{\text{successor}}^{(i)} = W_{\text{predecessor}}^{(k)}, \quad k = \left\lfloor \frac{i \cdot d}{d'} \right\rfloor, \quad i \in \{0, \dots, d'-1\}. \]

   这个方法确保了pretrained weights在successor MoE上的even distribution。

   通过checkpoint recycling带来的计算开销是可忽略不计的，Random Sampling、Uniform Selection有很少的额外的处理，因为它们直接从dense checkpoint中的选择experts，没有进一步的开销，然而Graph Partitioning和Importance-based Sampling设计一个preliminary inference step来决定neuron importance或者co-activation patterns。

9. SpheroMoE layer：为了增强性能和稳定性，涉及了hyperspherical adaptive MoE layer，引入了三个key improvements：SpheroMoE 路由来缓解优化挑战，Expert Regularization来防止over-specialization和Adaptive Dual-path MoE。如图所示： SpheroMoE机制利用了cross-attention将分布Inputs across experts，每个experts接受一个input slots，这是一个所有input tokens的加权，为了保持dense checkpoints和MoE layers的一致性，对输入的tokens进行normalization，Queries进行随机初始化，然后也进行normalized。为了减小数值的不稳定性，Q和K(\(X_{norm}\)) 投射一个hyperspherical space，用L2 normalization，确保resulting dot product反应余弦相似度，而不是没有边界的值，这个将values限制在一个stable range，阻止了softmax的饱和，得到一个更平衡的attention distributions，提高了model的泛化性。这个相似度 \(Q_{norm}\) 和 \(K\) 计算出来，产生similarity logits \(S = Q_{norm}K^{T}\)， 输入的slots \(X\) for experts通过在 \(S\)d n dimension进行一个softmax操作得到。 每个expert独立地处理，产生outputs，这些outputs通过 \(S\) 进行加权，来汇聚experts的contributions，产生最后的output \(Y\) of the MoE layer

总结，SpheroMoE路由利用layer normalization，hyperspherical projection和交叉注意力，来有效地在distribute inputs across experts，确保numerical stability和pretrained dense models的一致性。

10. 为了提高泛化性，SpheroMoE对路由和expert behavior进行正则化，阻止experts over-specializing，对于前者，作者引入了可学习的softmax temperatures \(T_{dispatch} 和 T_{combine}\) 来精确地控制token dispatch和slot combination，使得smooth transitions，因为token dispatch在不同的slots分配tokens，然而slot combination将slot outputs汇聚为tokens，这些dual temperatures提供了灵活的控制，temperatures初始化的时候比较高，来提升attention的broad distribution，防止early specialization。在训练的时候，这些temperatures adaptively 降低，使得experts关注相关的features，specialize where advantageous，另外，作者对similarity logits \(S\) 加了一定的normal noise，提高泛化性，对于后者，作者利用随机expert dropout，每个expert \(i\) 有概率p被随机deactivated，它确保没有任何一个专家成为整个输出的依赖，从而促进所有专家更加均衡地被利用。这些技术构成了一种专家正则化策略，保持了专家的通用性，缓解了过拟合问题，从而确保 MoE 模型在下游数据集上具有稳健的性能表现。

11. Adaptive Dual-path MoE：为了缓解less critical tokens的computational redundancy，将资源集中在essential ones，SpheroMoE路由直接将input tokens路由到core, universal slots，基于此，Adaptive Dual-path架构被设计用来将每个slot分配给distinct pathway用于优化processing，这个core pathway包含一个有限的core experts，相比之下，universal pathway包含一个larger set of small experts，每个的参数是core experts的one-fourth，由于处理less critical slots，这个dual-path configuration，通过聚焦computation on key tokens，提高了resource allocation，保持了model的精度，增强了处理的效率。

\(Fig.1^{[1]}\) a. checkpoint recycling: 从dense checkpoints MLP中选择hidden units和dimensions，将预训练的weights转换成diverse experts，来初始化MoE models，b. SpheroMoE layer： 用cross-attention将tokens给到slots，将slots结合为tokens，输入的tokens project为keys，一个随机的query初始化，来计算token-slot相似度，experts来处理它们分配的slots

\(Fig.2^{[1]}\) Adaptive Dual-path MoE 架构，通过将其弄成Dual-branch系统，增强了SpheroMoE router，这个设计用于优化计算效率和模型性能，将high-impact tokens给到core path，大量的experts将没那么重要的tokens给到universal path。