DiEP: Adaptive Mixture-of-Experts Compression Through Differentiable Expert Pruning

作者/机构: Sikai Bai (HKUST), Haoxi Li (HKUST), Jie Zhang (HKUST), Zicong Hong (HKUST), Song Guo (HKUST)

A1 主要贡献

核心问题与研究目标

大型语言模型(LLMs)的参数规模持续增长,带来了巨大的计算和存储挑战。混合专家(MoE)架构虽然通过稀疏激活部分参数来提升计算效率,但其庞大的模型体积依然给资源受限环境下的部署带来了内存和存储方面的困难。现有的MoE剪枝方法大多采用统一的稀疏度,即在每一层剪枝掉相同比例的专家,这种策略忽略了不同MoE层中专家冗余度的差异性,常常导致次优的结果和性能下降。此外,近期研究(如图1所示)表明,MoE模型中不仅存在层内专家的相似性差异,还存在显著的层间专家相关性,这突显了现有方法(如层内归一化或暴力搜索)的局限性。例如,在有64个专家的层中仅剪枝12.5%(保留56个)就需要评估约4亿种组合,使得全局优化在计算上不可行。因此,本文的研究目标是设计一种非均匀的剪枝策略,能够自适应地调整各层的剪枝率,并联合学习层间重要性,以有效捕捉不同MoE层中变化的冗余度。

创新点(本文贡献)

为解决上述挑战,本文提出了一种新颖高效的可微分专家剪枝(Differentiable Expert Pruning, DiEP)框架,该框架将专家剪枝问题重新表述为一个连续优化问题。其主要创新点如下:
1. 连续优化框架:DiEP将离散的专家选择搜索空间松弛为连续空间,从而能够使用基于梯度的优化方法来处理指数级增长的非均匀专家组合,实现了自适应的剪枝。这与传统方法中的启发式或暴力搜索形成对比。
2. 联合优化层内与层间重要性:该方法引入了两个核心组件:层内专家重要性分数(α)和层间重要性分数(β)。通过联合优化这两个分数,DiEP能够同时确定每层内部专家的相对重要性以及不同层对全局的贡献度,从而实现对专家和层级影响的精确建模,并产生深度感知的稀疏模式。
3. 单阶段、轻量化训练:与传统的双层可微分搜索方法不同,DiEP在一个单阶段训练过程中联合优化层内和层间重要性分数。该过程由一个轻量级的重构正则化器指导,不依赖验证集,并且通过解耦α和β的梯度更新来减轻优化干扰。
4. 推理时自适应专家跳过机制:除了永久性地移除不重要的专家外,DiEP还提出了一种在线专家跳过机制。该机制在推理过程中为高度相似的专家分配衰减的权重,从而动态地绕过冗余的专家计算,进一步加速了推理速度。

实验证明,DiEP在多种自然语言处理任务和先进的MoE模型上均优于其他剪枝方法。例如,在Mixtral 8×7B模型上,DiEP仅用一半的专家就保留了约92%的原始性能,在MMLU数据集上的表现比其他剪枝方法高出7.1%。


图1:通过基于RBF核的CKA标准【20, Similarity of neural network representations revisited, 2019, PMLR】对Mixtral 8×7B模型中不同MoE层的专家对进行层内和层间相似性可视化分析。颜色越深代表专家相似度越高。

A3 背景知识

MoE模型基本构成
通常,一个混合专家(MoE)模型由L层组成,其中每一层l(l = 1, ..., L)包含N个专家。第l层所有专家的输入表示为$x^{(l)} \in R^d$,其中d是输入维度。一个路由网络为每个专家i(i = 1, ..., N)生成路由logits $\zeta_i^{(l)}$,这些logits通过一个softmax函数进行归一化,以计算路由权重$w_i^{(l)}$:

$$w_{i}^{(l)}=\frac{\exp (\zeta_{i}^{(l)})}{\sum_{j=1}^{N} \exp (\zeta_{j}^{(l)})},$$

其中$w_i^{(l)}$代表了l层中专家i的贡献度。

稀疏路由与层输出
为了实现稀疏性,路由网络会选择具有最大路由权重$w_i^{(l)}$的top-k个专家。第l个MoE层的输出计算如下:

$$ \boldsymbol{y}^{(l+1)}=\sum_{i \in \operatorname{Top-} k\left(w^{(l)}\right)} w_{i}^{(l)} \cdot \operatorname{FFN}_{i}\left(\boldsymbol{x}^{(l)}\right), $$

其中$FFN_i(\cdot)$表示专家i的前馈函数,而Top-k($w^{(l)}$)指的是k个最大路由权重的索引。最终的输出$y^{(l+1)}$被传递到下一层。

A2 方法细节

4.1 稀疏专家搜索空间

稀疏专家搜索空间定义
遵循可微分架构搜索的设计原则,我们首先为混合专家(MoE)架构量身定制一个稀疏专家搜索空间,如图2所示。在此框架中,一个MoE层被建模为一个仅包含两个节点的有向无环图(DAG):一个输入节点,代表进入专家层的令牌表示;以及一个输出节点,代表所选专家变换的总和。我们不将单个专家视为独立的计算单元,而是将专家剪枝过程表述为对单个聚合专家节点的离散操作。


图2:可微分专家剪枝(DiEP)框架的示意图。(a) 初始MoE模型,存在大量专家冗余和高内存成本。(b) 可微分剪枝:通过交替更新策略联合学习层内专家分数(α)和层间重要性(β),将离散的专家搜索转化为连续优化,从而实现自适应的非均匀剪枝。(c) 最终的剪枝模型:实现了一个精简的MoE架构,在降低模型体积的同时保持了高性能。

离散选择的数学表达
根据专家剪枝的原则,一部分专家根据其重要性被保留下来,这由一个二元选择掩码$m_i^{(l)} \in \{0, 1\}$控制,其中$m_i^{(l)} = 1$表示专家i被保留,$m_i^{(l)} = 0$表示被剪枝。一个MoE层中的专家聚合过程则表示为:

$$\boldsymbol{y}^{(l+1)}=\sum_{i=1}^{N}(m_{i}^{(l)}\cdot\mathrm{FFN}_{i})(\boldsymbol{x}^{(l)}),$$

其中$FFN_i(\cdot)$表示专家i的前馈函数。

非可微问题
这种离散选择过程天生导致了一个不可微的搜索空间,使得直接优化变得棘手。为了在MoE框架内实现基于梯度的优化和结构化剪枝,我们引入了一种连续松弛机制,允许对专家选择过程进行平滑更新,同时保留模型的结构化稀疏性。

4.2 连续松弛与优化

重要性分解
具体来说,我们将专家的重要性分解为两个部分:决定每层内部专家相对重要性的层内重要性分数α,以及调节不同层在选择过程中贡献度的层间重要性分数β。这种形式使我们能够以一种数据驱动且全局优化的方式进行结构化剪枝。

重要性权重计算与层输出
我们通过对层内重要性分数$\alpha_i^{(l)}$使用softmax函数进行归一化来定义层内重要性权重$\bar{\alpha}_i^{(l)}$:

$$\bar{\alpha}_{i}^{(l)}=\frac{\exp (\alpha_{i}^{(l)})}{\sum_{j=1}^{N} \exp (\alpha_{j}^{(l)})},$$

其中$\alpha_i^{(l)}$是可学习的logits,用于确定第l层内专家的相对重要性。这种归一化确保了一个平滑且可微的选择过程。同样,层间重要性分数$\beta^{(l)}$被引入作为一个可训练的标量,用于调节第l层的整体贡献。一个MoE层l的输出随后被计算为:

$$\boldsymbol{y}^{(l+1)}=\beta^{(l)} \sum_{i=1}^{N} \bar{\alpha}_{i}^{(l)} \cdot \operatorname{FFN}_{i}\left(\boldsymbol{x}^{(l)}\right) .$$

重构正则化项
为确保剪枝后的模型能够保持对原始MoE模型$F(x)$(剪枝前)的保真度,我们引入了一个重构正则化项$\Phi(\alpha, \beta)$,定义为:

$$\Phi(\alpha, \beta)=\left\|\mathcal{F}^{\prime}(\boldsymbol{x} ; \alpha, \beta)-\mathcal{F}(\boldsymbol{x})\right\|_{F},$$

其中$\| \cdot \|_F$表示Frobenius范数。该正则化项鼓励剪枝后的模型$F'$与原始模型保持一致性。

整体目标函数
整体目标函数的公式如下:

$$\min_{\alpha, \beta} \mathcal{L}(\alpha, \beta) := \mathcal{L}_{ce}(\boldsymbol{y}, \mathcal{F}'(\boldsymbol{x}; \alpha, \beta)) + \lambda \Phi(\alpha, \beta),$$

其中λ是一个正则化系数,而$L_{ce}$是交叉熵损失。

交替更新策略
为了优化目标函数,我们采用了一种交替更新策略,即层内重要性分数α和层间重要性分数β被迭代更新:

$$\begin{aligned} \begin{aligned} \alpha^{t} & \leftarrow \alpha^{t}-\eta_{\alpha} \nabla_{\alpha} \mathcal{L}\left(\alpha^{t}, \beta^{t}\right), \\ \beta^{t} & \leftarrow \beta^{t}-\eta_{\beta} \nabla_{\beta} \mathcal{L}\left(\alpha^{t}, \beta^{t}\right). \end{aligned} \end{aligned}$$

这里,t表示迭代索引,$\eta_\alpha$和$\eta_\beta$分别是α和β的学习率,而$L(\alpha, \beta)$代表在公式7中定义的整体目标函数。从理论角度,我们在算法1中总结了优化过程,并在附录B.2中提供了详细的收敛性分析。

离散化剪枝策略
为了得到一个离散的架构,我们应用了一个结构化的剪枝机制,根据专家在所有层中的全局贡献来剔除最不重要的专家。我们不是孤立地逐层剪枝专家,而是利用学习到的层内重要性分数$\alpha_i^{(l)}$和层间重要性分数$\beta^{(l)}$来统一地确定专家的重要性。

全局重要性计算
形式上,专家i在层l的总体重要性被计算为其层内和层间重要性分数的乘积:

$$s_i^{(l)} = \alpha_i^{(l)} \cdot \beta^{(l)}.$$

全局排序与剪枝
给定专家稀疏度比率r,整个MoE模型需要剪枝的专家总数为$K = NLr$,其中N是每层的专家数,L是层数。剪枝过程通过全局排序所有专家的重要性分数$s_i^{(l)}$并移除得分最低的K个专家来完成。最终的剪枝掩码$m_i^{(l)}$定义为:

$$\begin{aligned} m_{i}^{(l)}=\begin{cases}0 & \text{if } i \in P, \\ 1 & \text{otherwise,}\end{cases} \end{aligned}$$

其中P是为剪枝选择的得分最低的K个专家的集合。

策略优势
通过联合考虑层内和层间的重要性分数,这种剪枝策略确保了对专家的全局优化选择,有效减少了计算冗余,同时保持了各层间的结构平衡。

4.3 推理时的自适应跳过

推理时自适应跳过的动机
在推理过程中,用所有选定的top-k个专家处理每个令牌会引入不必要的计算开销,但研究人员在【30, Not all experts are equal: Efficient expert pruning and skipping for mixture-of-experts large language models, 2024, ACL】中发现,并非每个被选中的专家都为令牌提供了必要的贡献。这一观察启发了自适应专家跳过的需求,即在推理期间选择性地绕过不太重要的专家以提高效率。对于MoE层中的每个令牌x,使用路由权重$w = \{w_{e_0}, w_{e_1}, \dots, w_{e_{k-1}}\}$选择top-k个专家,其输出表示为$y_{e_0}, y_{e_1}, \dots, y_{e_N}$。按照惯例,为简化起见,我们假设k=2。与以往仅依赖路由权重的方法【30】不同,我们的方法结合了专家相似性,在推理过程中动态跳过不那么重要的专家,从而提高计算效率。

跳过条件
假设索引为$e_0$和$e_1$的专家被选中,且$w_{e_1} < w_{e_0}$。为了提高推理速度,如果$w_{e_1} < \gamma w_{e_0}$,专家$e_1$将被跳过,其中$\gamma$是特定于每个MoE层和生成步骤的超参数。

$\gamma$的计算方法
在我们的实现中,$\gamma$由两个因子的乘积计算得出。首先,$\gamma_1$被确定为每个MoE层在采样校准数据上$w_{e_1}/w_{e_0}$的中值。其次,$\gamma_2$基于专家输出之间的相似性计算,使用中心核对齐(CKA)【20, Similarity of neural network representations revisited, 2019, PMLR】进行评估。具体来说,$\gamma_2$是CKA相似度$\rho(y_{e_0}, y_{e_1})$与层l中所有数据样本的平均CKA相似度$\rho(y_{e_i}, y_{e_j})$的比值。$\gamma$的最终值由下式给出:

$\gamma = \gamma_1 \times \gamma_2.$

方法效果
该方法根据专家路由权重和相似性动态调整专家跳过,显著提高了推理效率并保持了模型性能。在我们的实验中,在Mixtral 8×7B上仅使用一半专家的情况下,我们观察到推理速度提升了1.2倍至1.3倍,同时保留了约92%的平均性能。

A4 实验环境

模型设置
- 主要模型: Mixtral 8×7B。
- 泛化性验证模型: Mixtral 8×7B-Instruct, Mixtral 8×22B, Deepseek-MoE-16B, Qwen2-57B-14A。
- 架构参数:
- Mixtral 系列: 每层8个专家,每个token激活2个。Mixtral 8×7B/Instruct有32个MoE层,Mixtral 8×22B有56个MoE层。
- Deepseek-MoE-16B: 28层,每层64个专家,每个token激活2个共享专家和6个选定专家。
- Qwen2-57B-14A: 28个MoE层,每层64个专家,每个token激活8个专家。

数据集

  • 评估框架: Language Model Evaluation Harness library【11】。
  • 评估任务 (零样本): MMLU【16】, OpenBookQA【32】, BoolQ【6】, RTE【2】。
  • 校准数据集: 从C4数据集中抽取128个样本用于微调。

硬件配置
- GPU: 4块 NVIDIA GeForce A800 GPU。

软件配置与实现细节
- 训练策略: 采用参数高效的差分学习,通过交替训练周期进行,层内分数α和层间分数β的更新比例为3:1。
- 优化器: 学习率为5e-3,使用余弦学习率调度器。
- 训练设置: 训练10个epoch,批处理大小为16。
- 超参数: 正则化系数λ在Mixtral架构中设为0.01,在其他MoE模型中也设为0.01。

基线方法
- M-SMoE【22】: 基于自定义置换对齐和路由策略合并专家。
- Expert Trimming【15】: 基于激活频率移除不重要专家,或通过层和块丢弃移除结构化模块。
- NAEE【30】: 枚举专家组合,通过最小化重构损失来选择最优的剩余专家。
- S-SMOE【47】: 通过基于相似性的剪枝和合并操作来识别和解决专家冗余。

A4 实验结果

主要结果

Mixtral模型上的结果 (表1)
- Mixtral 8×7B: DiEP在所有任务上均显著优于其他剪枝策略,性能提升幅度高达7.1%。特别是在具有挑战性的MMLU数据集上,当专家稀疏度为50%时,其他方法遭遇性能瓶颈,而DiEP有效缓解了移除大量专家带来的负面影响。这表明DiEP通过在任务无关数据上进行差分优化和搜索,有效保留了关键的专家知识。
- Mixtral 8×7B-Instruct: DiEP同样大幅超越其他方法。在25%的专家稀疏度下(剪枝后移除64个专家),DiEP实现了最佳性能,与完整模型相比平均性能仅下降0.6%。
- Mixtral 8×22B: 在更大的模型上,DiEP依然表现出色。即使在移除50%专家后,仍能保留完整模型94%的性能,这揭示了MoE层中存在显著的冗余,并展示了DiEP在大规模SMoE模型上的可扩展性。

Deepseek和Qwen模型上的结果 (图3)
为了验证DiEP的泛化能力,实验扩展到了架构差异显著的Deepseek-MoE-16B和Qwen2-57B-14A模型上。
- Deepseek-MoE-16B: 基于频率的方法(M-SMoE)在MMLU上性能下降严重。相比之下,DiEP在各种剪枝率下都表现出卓越性能,比次优策略(S-MoE)平均高出约1.57%。
- Qwen2-57B-14A: DiEP在各种任务上始终能达到与完整模型相当的性能,并超越所有基线方法。从完整的Qwen-MoE模型中移除244个专家后,DiEP的平均性能仍高达68.7%,仅比完整模型低0.4%。这突显了DiEP基于通用差分优化的适应性和有效性。


表1:不同专家剪枝方法在Mixtral-8×7B、Mixtral-8×7B-Instruct和Mixtral-8×22B上的零样本性能比较。专家稀疏度r表示在所有层中剪枝掉的专家占完整模型的比例。第一和第二列分别代表专家稀疏度r=25%和r=50%的结果。


图3:在Deepseek-MoE-16B和Qwen2-57B-14A上的零样本性能比较。

消融研究

组件有效性分析 (表2)
- 实验设置: 在Mixtral 8×7B上对DiEP的不同变体进行测试。基线(NAEE)仅进行逐层搜索。变体包括仅学习层内重要性α、仅使用层间分数β重加权激活频率、以及α与固定比例的β结合。
- 实验结论: 单独使用α或β均比基线有显著性能提升,证明了两个组件的有效性。与人工固定的β相比,可学习的β能够自适应地为浅层分配更高权重,从而取得更好性能。最终,结合可学习的α和β的完整DiEP方法利用了二者的互补知识,取得了最优的性能。


表2:不同组件的性能分析。

α和β的可视化分析 (图4)
- 层内分数α (图4a): 浅层(1-15层)专家的平均分数普遍高于深层(16-32层),表明浅层在模型中扮演更重要的角色。具体到第2层,第4和第8个专家α值显著高于其他专家,说明模型持续认为它们非常重要。
- 层间分数β (图4b): β分数的变化趋势与α的观察结果一致,同样显示出浅层的重要性更高。
- 结论: 这种分布证实了交替更新策略能够有效捕捉层内和层间的依赖关系,确保模型在剪枝后保留来自浅层的关键信息。α和β之间的协同作用证明了DiEP方法的有效性。


图4:在r=50%时,Mixtral 8×7B上层内分数α和层间分数β的值分布可视化。

计算成本分析 (表3和表4)
- 剪枝时间 (表3): 基线方法NAEE采用穷举搜索,在专家数量较多的模型(如Deepseek-MoE-16B)上计算成本过高,变得不可行。相比之下,DiEP仅引入0.01%的参数开销,其剪枝时间保持稳定,且不受模型架构或专家数量影响。
- 推理成本 (表4): 在Mixtral 8×7B上,DiEP通过在线专家跳过机制提升了推理效率。在保留一半专家的情况下,DiEP实现了1.28倍的令牌生成加速和48%的GPU内存节省,同时平均性能保留了近92%。


表3:我们的DiEP和NAEE在不同模型上专家稀疏度为25%时的剪枝时间比较。


表4:专家剪枝后在Mixtral 8×7B上的推理成本分析。

A5 结论

本文提出了DiEP,一个新颖的可微分专家剪枝框架,它将专家选择问题重构为一个连续优化问题。通过启用基于梯度的优化并引入一种自适应的专家跳过机制,DiEP在保持高模型性能的同时,显著降低了内存使用并加速了推理过程。大量的实验表明,我们的DiEP在各种语言任务中均优于其他的MoE剪枝方法,为高效的稀疏MoE部署设立了新的基准。

A6 附录

A.1 专家激活频率的可视化分析

专家激活频率可视化分析
为了展示DiEP中自适应专家剪枝的有效性,我们对不同方法的专家激活频率进行了比较分析。如图5所示,虽然完整的MoE模型几乎利用了所有专家,但不同专家间的激活频率存在显著差异,导致了大量的资源浪费。之前的方法(即逐层剪枝)对所有层应用统一的专家剪枝率,忽略了不同MoE层中专家间的层内和层间变化及依赖关系。相比之下,我们的方法实现了非均匀且自适应的专家剪枝,根据专家的特定特性调整剪枝率。在Mixtral 8×7B上,我们观察到专家剪枝率从浅层到深层呈现递增趋势。我们将此现象归因于浅层主要处理多样化的低级语言特征(如词性标注和局部词序),需要更多专家来捕捉详细的语言信息。而深层主要处理全局上下文和语义信息,从细粒度细节中抽象出来,因此可以用较少的专家有效运作。


图5:Mixtral 8×7B中专家激活频率的分布。完整模型(左)在所有32层中使用所有专家,导致巨大的内存消耗。逐层剪枝(中)对每层强制实施统一的专家稀疏度。我们的DiEP(右)提供了一种更灵活的方法,根据专家的全局贡献进行跨层专家剪枝。

A.2 Deepseek-MoE-16B上的推理效率分析

Deepseek-MoE-16B上的推理效率分析
我们进一步在表5中验证了我们的自适应跳过策略在Deepseek-MoE-16B上的效率。可以观察到,我们的方法在降低模型大小和提高推理效率的同时,保持了超过95%的完整模型性能。


表5:在Deepseek-MoE-16B上的推理成本分析。

A.3 GPU内存剪枝成本的效率分析

剪枝过程中的GPU内存成本分析
为了研究DiEP在剪枝过程中内存占用的效率,我们在表6中对GPU内存成本进行了详细比较,突显了DiEP在效率和可扩展性方面的优势。
- 计算效率(时间成本):与NAEE(1.31小时)相比,DiEP的剪枝过程快了5.7倍,仅需0.23小时。DiEP的执行速度(0.23小时)也比MC-MoE(0.31小时)快了25.8%。
- 内存优化(峰值内存):DiEP使用的峰值内存(139.0GB)比MC-MoE(348.4GB)少60%,显示出优越的内存效率。虽然DiEP的峰值内存比NAEE(95.1GB)高46%,但这被其大幅缩短的剪枝时间所抵消。
- 整体资源效率(内存-小时成本):DiEP的内存-小时成本(31.97 GB·h)比MC-MoE(108.00 GB·h)低70%,比NAEE(124.58 GB·h)低74%。这些结果清楚地表明,DiEP在大幅减少运行时间的同时保持了轻量级的资源占用,是MoE剪枝中更具资源效率的选择。


表6:在Mixtral 8×7B上的GPU内存剪枝成本。

A.4 更多校准数据对适应性的验证

不同校准数据的适应性验证
为了进一步验证DiEP的适应性,我们在领域特定的GSM8K数据集上,使用两个不同的校准数据集(通用C4数据集和领域相关的Math数据集)评估了其性能,并与NAEE方法进行比较。如表7所示,实验结果系统地证明了DiEP在这些不同校准设置下的优势。具体来说,当使用通用的C4校准数据时,DiEP在50%剪枝率下比NAEE高出+3.93分,在25%剪枝率下高出+4.96分。当使用领域特定的MATH校准数据时,DiEP在50%剪枝率下保持+1.10分的优势,在25%剪枝率下将优势扩大到+2.21分。这些发现共同突显了DiEP在不同数据分布的校准数据集上更强的泛化能力和适应性。


表7:使用不同校准数据集(C4和Math)在GSM8K上的适应性验证。

A.5 合并策略

与专家合并策略的结合
受S-SMoE【47, Diversifying the expert knowledge for task-agnostic pruning in sparse mixture-of-experts, 2025, Findings of ACL】的启发,我们为DiEP引入了一种合并策略,以整合冗余专家同时保留其多样性。具体来说,被剪枝的专家会根据归一化的CKA相似度与它们最相似的保留专家分组,该相似度随后通过softmax函数归一化作为合并权重。表8显示,合并策略在25%和50%的专家稀疏度下进一步提升了性能,这突显了我们DiEP的强大可扩展性。它不仅能有效维持完整模型的性能,还能通过结合其他正交策略进一步恢复被剪枝专家的多样性。


表8:集成合并策略后的性能分析。

A.6 校准数据大小的影响

校准数据大小的影响
为了分析校准数据大小的影响,我们从C4数据集【35】中随机抽取了32, 64, 128, 256, 512和1024个序列来学习DiEP的层内分数(α)和层间分数β。如表9所示,在将Mixtral 8×7B从8个专家剪枝到6个专家时,使用128个序列达到了最佳性能。更重要的是,即使只有32个样本,DiEP也能避免性能崩溃。我们将其归因于知识蒸馏(KD)正则化强制DiEP的特征与完整模型对齐。


表9:改变校准数据集中样本数量时的专家剪枝性能。

A.7 完整的专家相似度可视化分析

完整的专家相似度可视化分析
为了验证我们关于跨层剪枝必要性的动机,我们首先使用CKA相似性度量【20】在图6中可视化了Mixtral 8×7B中每一层的层内专家相似性。分析显示,专家相似性存在显著变化,尤其是在第31层。此外,不同层之间的专家相似性也存在巨大差异,例如第28-29层的相似性高于第8-10层。我们还在图7中研究了相邻层中专家对的相似性,这表明了跨层的专家关系程度各不相同,例如第30层的专家6与第31层的专家5之间存在强相关性。这些跨层专家依赖性被以往的剪枝方法所忽略。我们的方法通过对专家权重α和层权重β的交替可微分优化,有效地捕捉了层间和层内的变化。此外,我们观察到学习到的层内和层间分数并不完全对应于可视化的专家对之间的层间相似性。这可能是因为我们只提供了相邻层之间的专家相似性进行可视化分析。然而,我们的DiEP能够学习所有MoE层中的专家冗余和依赖关系。


图6:每个MoE层内专家对的特征相似性可视化。


图7:相邻MoE层之间专家对的特征相似性可视化。

A.8 超参数分析

超参数分析
有两个关键的超参数,包括可微分搜索期间的epoch数和公式7中重构正则化项与交叉熵损失之间的权重λ值。我们首先通过改变λ的值在{5, 10, 15, 20, 30}中来分析其影响。图8a显示,在25%和50%的专家稀疏度下,当λ=0.01时性能最佳。此外,我们研究了epoch数量如何影响模型性能。如图8b所示,在25%和50%的专家稀疏度设置下,DiEP在训练10个epoch时均达到最佳结果。


图8:在Mixtral 8×7B上,在25%和50%专家稀疏度下,关于(a)权重系数λ和(b)epoch数量的超参数分析。

A.9 在更多数据集上的结果

在更多数据集上的结果
我们在更多数据集上提供了实验结果,包括ARC-c、ARC-e【7】、HellaSwag【45】和WinoGrande【36】,实验对象为Mixtral 8×7B。我们的DiEP在所有任务上均远优于NAEE。如图10所示,这些结果进一步证明了我们所提出方法的有效性。


表10:在更多数据集(包括ARC-c、ARC-e、HellaSwag、WinoGrande)上的零样本评估结果。

B.1 DiEP算法流程

DiEP算法流程
算法1展示了DiEP的详细流程。输入包括模型输入x、目标y、初始的层内分数$\alpha^0$和层间分数$\beta^0$,以及正则化系数λ。算法在未收敛前循环执行:首先固定β更新α,然后固定α更新β,最后更新迭代次数。输出为优化后的α和β分数。

Input: Model inputs $\boldsymbol{x}$, targets $\boldsymbol{y}$, initial intra-layer scores $\alpha^0$, initial inter-layer scores $\beta^0$, regularization coefficient $\lambda$.
while not converged do
| / Update intra-layer scores (fix $\beta$) /
| Update $\alpha^{t+1}$ by descending:

$$\nabla_\alpha (\mathcal{L}_{ce}(\boldsymbol{y}, \mathcal{F}'(\boldsymbol{x}; \alpha^t, \beta^t)) + \lambda \Phi(\alpha^t, \beta^t))$$


| / Update inter-layer scores (fix $\alpha$) /
| Update $\beta^{t+1}$ by descending:

$$\nabla_\beta (\mathcal{L}_{ce}(\boldsymbol{y}, \mathcal{F}'(\boldsymbol{x}; \alpha^{t+1}, \beta^t)) + \lambda \Phi(\alpha^{t+1}, \beta^t))$$
| Set $t \leftarrow t + 1$
Output: Optimized intra-layer importance scores $\alpha$ and inter-layer importance scores $\beta$.

B.2 DiEP的收敛性分析

DiEP的收敛性分析
设$\Theta := \{(\alpha, \beta)\}$为参数空间,其中$\alpha \in R^{NL}$且$\beta \in R^L$。记$\theta_1 = \alpha$和$\theta_2 = \beta$。DiEP的总体目标是:

$$\mathcal{L}(\theta_{1}, \theta_{2}) = \mathcal{L}_{\text{ce}}(\mathbf{y}, \mathcal{F}'(\mathbf{x}; \theta_{1}, \theta_{2})) + \lambda \left\| \mathcal{F}'(\mathbf{x}; \theta_{1}, \theta_{2}) - \mathcal{F}(\mathbf{x}) \right\|_{F}$$

假设条件
- A1. 有下界性: $inf_{(\alpha, \beta) \in \Theta} L(\alpha, \beta) > -\infty$。
- A2. Lipschitz平滑性: $\nabla_\theta L$是$L_i$-Lipschitz连续的。
- A3. 步长: 固定的学习率满足$0 < \eta_i < 2/L_i$。
- A4. 水平集有界性: 集合$\{(\alpha, \beta) \in \Theta : L(\alpha, \beta) \leq L(\alpha_0, \beta_0)\}$是紧致的。

算法更新
对于t = 0, 1, ...

$$\theta_{i}^{t+1}=\theta_{i}^{t}-\eta_{i} \nabla_{\theta_{i}} \mathcal{L}\left(\theta_{1}^{t}, \theta_{2}^{t}\right), \quad i \in\{1,2\}.$$

收敛性定理
- 推论1(单调性和有界迭代): 在假设A1-A4下,$\{L(\theta_1^t, \theta_2^t)\}_{t \geq 0}$是单调不增且收敛的,并且$\{(\theta_1^t, \theta_2^t)\}_{t \geq 0}$是有界的。
- 引理2(更新消失): $\lim_{t\to\infty} \|\theta_i^{t+1} - \theta_i^t\| = 0$。
- 定理1(子序列收敛到临界点): 在假设A1-A4下,序列$\{(\theta_1^t, \theta_2^t)\}_{t \geq 0}$至少有一个收敛的子序列,并且每个极限点$(\theta_1^*, \theta_2^*)$都满足$\nabla_{\theta_1} L(\theta_1^*, \theta_2^*) = 0$和$\nabla_{\theta_2} L(\theta_1^*, \theta_2^*) = 0$;即它是一个临界点。

该分析与块坐标下降法的经典结果【40】一致,保证了DiEP在标准的平滑性和步长条件下(至少在子序列上)收敛到一阶稳定点。

A7 补充细节

C 局限性

计算资源限制
尽管DiEP取得了强大的性能,但由于计算资源的限制,我们的主要实验无法在一些更大规模的MoE模型上进行,例如Deepseek V3【25】和Qwen2.5-Max【42】。

模型模态限制
此外,我们的研究主要集中在语言模型上,DiEP在多模态MoE架构中的有效性仍有待探索。研究我们的方法是否能在视觉-语言任务(如MoE-LLaVA【24】)上取得有竞争力的性能,是未来研究的一个重要方向。

D 更广泛的影响

积极影响
DiEP作为一种压缩大型MoE模型的方法,可以在计算基础设施有限的地区或场景中,部署原本不可行的高性能AI模型。这可以为服务欠缺社区的教育、医疗和公共服务等应用提供支持。

消极影响
随着AI模型变得越来越强大和高效,它们可能会自动化目前由人类执行的任务,可能导致某些行业的就业岗位被取代。虽然DiEP旨在提高效率,但AI发展的更广泛趋势加剧了这一担忧。