PaperlayerView exact revision on arXiv

Source-linked AI summary

Merge, Then Compress: Demystify Efficient SMoE with Hints from Its Routing Policy

Pingzhi Li, Zhenyu Zhang, Prateek Yadav, Yi-Lin Sung, Yu Cheng, Mohit Bansal, Tianlong Chen

arXiv:2310.01334v2cs.LGcs.AIcs.CL

TL;DR

SMoE's memory overhead and redundant experts raise the question of how to consolidate experts without losing vital knowledge. The paper proposes routing-guided merging followed by compression; across eight benchmarks, MC-SMoE reduces memory and FLOPs with marginal performance drops.

  • Problem

    SMoE's memory overhead and redundant experts constrain practical use, raising the question of how to consolidate experts without sacrificing vital knowledge.

    Source: §1 · p. 2

  • Method

    M-SMoE aligns expert permutations and uses routing statistics to group and frequency-weight merge experts, while MC-SMoE further compresses merged weights.

    Source: §1 · p. 2, §3.1 · p. 4, §3.2 · p. 6

  • Results

    MC-SMoE achieves up to 80% memory and 20% FLOPs reduction, with marginal performance drops across eight benchmarks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors report that merged experts are more compressible and that MC-SMoE pursues memory and parameter efficiency with competitive performance.

    Source: §5 · p. 9

    M-SMoE and MC-SMoE improve memory and TFLOPs efficiency but do not improve latency.

    Source: A1.2 EFFICIENCY DISCUSSIONS AND… · p. 16

Abstract

from arXiv · show

Sparsely activated Mixture-of-Experts (SMoE) has shown promise to scale up the learning capacity of neural networks, however, they have issues like (a) High Memory Usage, due to duplication of the network layers into multiple copies as experts; and (b) Redundancy in Experts, as common learning-based routing policies suffer from representational collapse. Therefore, vanilla SMoE models are memory inefficient and non-scalable, especially for resource-constrained downstream scenarios. In this paper, we ask: Can we craft a compact SMoE model by consolidating expert information? What is the best recipe to merge multiple experts into fewer but more knowledgeable experts? Our pilot investigation reveals that conventional model merging methods fail to be effective in such expert merging for SMoE. The potential reasons are: (1) redundant information overshadows critical experts; (2) appropriate neuron permutation for each expert is missing to bring all of them in alignment. To address this, we propose M-SMoE, which leverages routing statistics to guide expert merging. Specifically, it starts with neuron permutation alignment for experts; then, dominant experts and their "group members" are formed; lastly, every expert group is merged into a single expert by utilizing each expert's activation frequency as their weight for merging, thus diminishing the impact of insignificant experts. Moreover, we observed that our proposed merging promotes a low dimensionality in the merged expert's weight space, naturally paving the way for additional compression. Hence, our final method, MC-SMoE (i.e., Merge, then Compress SMoE), further decomposes the merged experts into low-rank and structural sparse alternatives. Extensive experiments across 8 benchmarks validate the effectiveness of MC-SMoE. For instance, our MC-SMoE achieves up to 80% memory and a 20% FLOPs reduction, with virtually no loss in performance.

1 INTRODUCTION

SMoE scales model capacity through sparse expert activation, but its memory demands and redundant experts constrain practical use. This section introduces routing-guided merging and post-merging compression as a route to compact models.

1 INTRODUCTION

  • SMoE's memory overhead and redundant experts constrain resource-restricted use, motivating consolidation without sacrificing vital knowledge.The paper attributes redundancy to representation collapse under learning-based routing and notes that conventional merging can overshadow critical experts.Source: §1 · p. 2
  • M-SMoE's merged experts occupy a lower-dimensional weight space, which MC-SMoE further compresses using low-rank decomposition.The framework targets downstream fine-tuning and zero-shot evaluation.Source: §1 · p. 2
  • MC-SMoE achieves up to 80% memory and 20% FLOPs reduction, with marginal performance drops across eight benchmarks.M-SMoE alone yields up to a 60% memory-overhead reduction with slightly improved performance.Source: §1 · p. 2
  • M-SMoE aligns expert neuron permutations, groups non-dominant experts with dominant experts, then merges each group using activation-frequency weights.Routing frequency proxies expert utilization and importance; routing similarity helps associate redundant experts with dominant ones.Source: §1 · p. 2

2 RELATED WORKS

Related work covers SMoE scaling, approaches to its memory and utilization costs, and model-merging strategies. The paper highlights that common merging assumptions may not fit experts trained with distinct initializations and routing.

Sparse Mixture-of-Experts

  • SMoE scales model capacity by activating different model components for different input tokens rather than using all parameters.Examples include GShard scaling a Transformer from 2B to 600B parameters at lower training cost than a 100B dense model.Source: §2 · p. 3

Efficiency Concerns in SMoE and Existing Solutions

  • Utilization-based expert pruning saves memory but can lower performance, while tensor decomposition and distillation address other efficiency goals.These approaches are described as overlooking redundancy within SMoE layers; another cited work reports substantial expert compression without performance loss.Source: §2 · p. 3

Model Merging in Language Models

  • Many model-merging studies assume shared initialization, an assumption that may not hold for SMoE experts with distinct random initializations.Experts are also optimized on subsets of training data selected by routing networks.Source: §2 · pp. 3–4
  • Prior merging strategies include mode connectivity, permutation invariance, training regularization, and analyses of how training choices affect mergeability.SMEAR dynamically merges experts during training but does not reduce memory because the full SMoE remains at inference.Source: §2 · p. 4

3 METHODOLOGY

MC-SMoE uses routing information to align, group, and frequency-weight experts before further compressing the merged experts. Merging is followed by mostly non-positive stable-rank changes across SMoE layers, motivating post-merging compression.

3.1 ROUTING POLICY GUIDES EXPERTS MERGING

  • M-SMoE aligns expert weight permutations before merging because permutation can preserve a feed-forward expert’s function while bringing neurons into alignment.The method follows weight matching to align experts without altering their functionalities.Source: §3.1 · p. 4
  • Router-logit cosine similarity groups experts, with non-dominant experts assigned to their most similar dominant expert.The method uses routing logits because experts with similar router logits across most inputs tend to be similar.Source: §3.1 · p. 5
  • M-SMoE selects dominant experts from normalized layer-wise activation frequencies, groups other experts around them, and merges each group using usage-frequency weights.The number of dominant experts varies across layers; if routing is uniform in a layer, all its experts are treated as dominant.Source: §3.1 · p. 5

3.2 MERGING ENCOURAGES EXPERT DECOMPOSITION

  • Stable-rank change is non-positive over most SMoE layers after merging, indicating lower-dimensional merged weights and motivating additional compression.Figure 4 reports average stable-rank change ratios for dominant experts in switch-base-32.Source: §3.2 · p. 5, Figure 4 · p. 5
  • MC-SMoE decomposes each merged weight matrix into a low-rank product and an incoherent residual, then structurally prunes residual columns using cumulative importance scores.The importance score is based on each residual weight and its gradient with respect to the training objective; pruning is adaptive by layer.Source: §3.2 · p. 6

4 EXPERIMENTS

The experiments evaluate Switch and GPT-based SMoE variants across supervised fine-tuning and zero-shot NLP tasks. They report the full MC-SMoE results alongside ablations of merging ratios, grouping, distillation, alignment, and compression.

4.2 COMPETITIVE PERFORMANCE AND SUPERIOR EFFICIENCY OF MC-SMOE

  • MC-SMoE saves up to 80% in memory and 20% in FLOPs, with less than 1% performance degradation on MRPC, COPA, SQuAD, WikiQA, and HotpotQA.Table 2 evaluates supervised fine-tuning on eight datasets; all compared methods activate the same number of parameters.Source: §4.2 · p. 7, Table 2 · p. 7
  • M-SMoE reduces memory by 60% while retaining performance on five datasets and improves over full SMoE by 0.49%, 0.25%, and 0.41% on MRPC, SQuAD, and HotpotQA, respectively.It has marginal performance drops on SST2, MultiRC, and WikiQA but outperforms the other pruning and merging baselines.Source: §4.2 · p. 7

4.3 ABLATION STUDY AND EXTRA INVESTIGATION

  • Adaptive merging ratios consistently outperform uniform ratios across four datasets, matching the variation in frequently used experts across transformer blocks.The experiments use the switch-base-32 model.Source: §4.3 · p. 8, Table 3 · p. 8
  • Router-logits similarity consistently outperforms seven alternatives for grouping experts across four datasets.The authors attribute its strength to reflecting routing decisions, which can lead similarly routed experts to train on similar data.Source: §4.3 · p. 8, Table 4 · p. 8
  • Knowledge distillation improves M-SMoE in the ablation across four datasets.The study compares fine-tuned M-SMoE with and without the distillation loss on switch-base-32.Source: §4.3 · p. 8, Table 5 · p. 8
  • Permutation alignment improves merging performance; without it, mismatched neurons can produce an inferior fusion.The ablation compares M-SMoE with and without alignment on four switch-base-32 datasets.Source: §4.3 · p. 8, Table 6 · p. 8
  • Merging maintains or boosts performance while reducing model size, whereas compression alone causes a significant performance drop; MC-SMoE also outperforms compression-only with fewer parameters.The comparison uses M-SMoE, C-SMoE, and MC-SMoE on three switch-base-32 tasks.Source: §4.3 · p. 9, Table 7 · p. 9
  • Frequency-weighted merging performs best among the tested strategies; the visualization also suggests later SMoE layers are more mergeable and compressible.The comparison includes uniform, Fisher-weighted, and frequency-weighted merging; Figure 5 shades dominant experts by remaining parameters after compression.Source: §4.3 · p. 9, Table 8 · p. 9, Figure 5 · p. 9

5 CONCLUSIONS

The conclusion presents M-SMoE as a routing-guided way to consolidate experts and reports that the merged experts are more compressible. It identifies multimodal extensions and hardware co-design as future work.

5 CONCLUSIONS

  • M-SMoE consolidates expert information into fewer, more knowledgeable experts using guidance from routing policies.The conclusion frames this as addressing SMoE memory requirements and expert redundancy.Source: §5 · p. 9
  • The authors report that MC-SMoE pursues improved memory and parameter efficiency with competitive performance, supported by comprehensive experiments.The conclusion notes that merged experts are more compressible.Source: §5 · p. 9
  • The stated future directions are extending the approach to multimodal scenarios and co-designing with hardware platforms.Source: §5 · p. 9

6 REPRODUCIBILITY STATEMENT

The paper describes the materials provided to support reproducing its MC-SMoE pipeline, including source code, hyperparameter details, pseudo-code, and an algorithm overview.

6 REPRODUCIBILITY STATEMENT

  • The source-code repository includes data preprocessing, SMoE merging, compression and pruning, and evaluation scripts.Source: §6 · p. 10
  • Hyperparameter details are provided in Appendix A2, and detailed pseudo-code for expert merging is in Appendix A3.Source: §6 · p. 10
  • Algorithm 1 provides a concise overview of the MC-SMoE pipeline.Source: §6 · p. 10

A1.1 ZERO-SHOT EVALUATION RESULTS

In zero-shot evaluation, M-SMoE consistently outperforms one-shot pruning of non-dominant experts and task-specific expert pruning. The evaluation uses fairseq-moe-15b with 512 experts per layer and retains an average of 16 experts after pruning or merging.

A1.1 ZERO-SHOT EVALUATION RESULTS

  • M-SMoE consistently outperforms one-shot pruning of non-dominant experts and task-specific expert pruning in zero-shot evaluation.The comparison uses fairseq-moe-15b, which has 512 experts per SMoE layer; pruning and merging retain an average of 16 experts.Source: A1.1 ZERO-SHOT EVALUATION RESULTS · p. 16, Table A9 · p. 16

A1.2 EFFICIENCY DISCUSSIONS AND LIMITATIONS

Efficiency outcomes depend on the inference implementation: specialized groupwise processing improves throughput and latency, while the unoptimized MC-SMoE has lower throughput than M-SMoE despite using less memory and fewer FLOPs. Merging alone does not reduce latency because the router continues to manage the same number of classes.

A1.2 EFFICIENCY DISCUSSIONS AND LIMITATIONS

  • M-SMoE and MC-SMoE do not improve latency because merged experts still leave the router managing the same number of classes.The implementation redirects group members to a representative expert, but explicitly creating a new router for merged experts is non-trivial.Source: A1.2 EFFICIENCY DISCUSSIONS AND… · p. 16
  • Latency orders as SMoE ≈ M-SMoE > M-SMoE with router pruning > dense across batch sizes 32, 256, and 512.The comparison uses the switch-base-32 backbone; the authors identify router pruning as a future direction.Source: Table A10 · p. 16, A1.2 EFFICIENCY DISCUSSIONS AND… · p. 16
  • Without specialized implementation, MC-SMoE has lower throughput than M-SMoE despite consuming less memory and fewer FLOPs.The authors attribute the throughput difference to lacking specialized sparse-matrix software or hardware support for MC-SMoE.Source: Table A11 · p. 17, A1.2 EFFICIENCY DISCUSSIONS AND… · p. 17
  • Specialized groupwise processing improves M-SMoE and MC-SMoE throughput and latency.Tokens routed to a group are gathered and processed through one shared expert; the authors report gains after this optimization.Source: Table A12 · p. 17, A1.2 EFFICIENCY DISCUSSIONS AND… · p. 17

A1.3 COMPUTATIONAL COST DISCUSSION OF M-SMOE

M-SMoE’s merging costs grow linearly with the number of SMoE layers overall, although some components remain constant-cost. Experiments identify permutation alignment as the main time bottleneck and model inference as the main memory cost.

A1.3 COMPUTATIONAL COST DISCUSSION OF M-SMOE

A1.4 COMPARISON BETWEEN DIFFERENT PRUNING RATIO SCHEDULES

The cubic pruning-ratio schedule performs best among the tested schedules on COPA and MultiRC. The accompanying technique details describe permutation alignment and distillation, implementation of group merging, and grouping results across eight tasks.

A1.4 COMPARISON BETWEEN DIFFERENT PRUNING RATIO SCHEDULES

A2 MORE TECHNIQUE DETAILS

  • For M-SMoE, 256 randomly selected training samples provide expert-use frequencies and router-logit similarities; zero-shot evaluation involves no optimization.The authors also describe supervised fine-tuning with knowledge distillation from the full SMoE teacher, combining hard-label cross-entropy and soft-label KL divergence with α fixed at 0.2.Source: A2 MORE TECHNIQUE DETAILS · pp. 18–19
  • Permutation alignment preserves an expert’s function while matching neuron permutations by minimizing weight distance as a linear assignment problem.The authors solve the assignment with the Hungarian algorithm.Source: A2 MORE TECHNIQUE DETAILS · p. 19

A3 MORE IMPLEMENTATION DETAILS

  • M-SMoE replaces each expert group’s weights with a usage-frequency-weighted average and redirects its other members to the representative expert.The implementation describes this as retaining a representative while routing the group’s remaining experts to it.Source: A3 MORE IMPLEMENTATION DETAILS · p. 20

A4.1 GROUPING RESULTS OF M-SMOE

  • Grouping results are provided for switch-base-32 across SST2, MRPC, MultiRC, COPA, WinoGrande, SQuAD, WikiQA, and HotpotQA.The results are presented in Figures A6–A13.Source: A4.1 GROUPING RESULTS OF M-SMOE · p. 21
Loading 2310.01334v2…