Source-linked AI summary
On the Representation Collapse of Sparse Mixture of Experts
Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, Xian-Ling Mao, Heyan Huang, Furu Wei
TL;DR
SMoE routing can drive token representations toward expert centroids, raising a representation-collapse concern. The paper introduces low-dimensional hyperspherical routing with normalization and temperature-controlled gating, and reports consistent multilingual gains alongside reduced collapse and more consistent routing.
Problem
SMoE routing may encourage token clustering around expert centroids, creating a representation-collapse issue that can limit representation capacity.
Method
X-MOE estimates token–expert routing scores after dimension reduction and L2 normalization, using a learnable temperature to control expert activation.
Results
X-MOE consistently improves language-modeling and fine-tuning performance across multilingual benchmarks, including an average XTREME score of 65.3 for softmax-gating X-MOE.
Takeaways & Limitations
The method alleviates representation collapse and achieves more consistent routing during both pre-training and fine-tuning.
Abstract
from arXiv · showhide
Sparse mixture of experts provides larger model capacity while requiring a constant computational overhead. It employs the routing mechanism to distribute input tokens to the best-matched experts according to their hidden representations. However, learning such a routing mechanism encourages token clustering around expert centroids, implying a trend toward representation collapse. In this work, we propose to estimate the routing scores between tokens and experts on a low-dimensional hypersphere. We conduct extensive experiments on cross-lingual language model pre-training and fine-tuning on downstream tasks. Experimental results across seven multilingual benchmarks show that our method achieves consistent gains. We also present a comprehensive analysis on the representation and routing behaviors of our models. Our method alleviates the representation collapse issue and achieves more consistent routing than the baseline mixture-of-experts methods.
1 Introduction
SMoE increases model capacity with sparse activation, but its routing mechanism can encourage representation collapse. X-MOE addresses this issue with low-dimensional hyperspherical routing and shows consistent gains across multilingual pre-training and downstream evaluation.
- Motivation: SMoE increases model capacity while keeping computational cost lower through sparse expert activation.The approach distributes tokens among several experts rather than activating all experts for every token.
- Motivation: Routing assigns each token to its best-matched experts using token–expert similarity scores.The paper frames routing as a central design problem in SMoE models.
- Method: X-MOE projects hidden vectors into a lower-dimensional space, applies L2 normalization, and adds a learnable-temperature soft expert gate.These components estimate routing scores on a low-dimensional hypersphere while controlling expert activation.
- Results: Across cross-lingual language-model pre-training and downstream fine-tuning, X-MOE consistently outperforms baseline SMoE models and improves routing consistency.The analysis also reports that X-MOE alleviates representation collapse.
- Problem: Learning routing scores encourages token clustering around expert centroids, creating a trend toward representation collapse.X-MOE is motivated by this under-explored issue in existing SMoE routing.
2 Background
SMoE routes tokens to sparsely activated feed-forward experts using dot-product similarities and gating. The paper argues that this mechanism can collapse representations toward an expert-defined low-dimensional subspace and toward selected expert embeddings.
- Sparse Mixture of Experts: An SMoE layer combines a router with several sparsely activated feed-forward expert networks.The router determines which experts process each input token.
- Routing: For token representation h and expert embedding e_i, the router computes the dot-product score s_i = h · e_i.The resulting scores feed a sparse gating function for conditional expert activation.
- Routing: Top-1 routing activates only the expert with the largest routing score, while the framework also considers softmax and sigmoid gating.Top-K routing extends the mechanism by activating up to K experts per token.
- Representation Collapse: Back-propagation updates token representations toward a linear combination of expert embeddings.The Jacobian decomposition identifies a gating-related gradient path responsible for this tendency.
- Representation Collapse: The collapse finding also holds under top-K routing, where multiple selected experts contribute to each token’s output.The paper explicitly extends the conclusion beyond top-1 routing.
- Representation Collapse: Because N experts span at most an N-dimensional subspace, hidden vectors tend toward an R^N subspace within R^d when N ≪ d.The paper presents this as a potential reduction in Transformer representation capacity.
3 Methods
X-MOE estimates token–expert routing on a low-dimensional hypersphere using projection and normalization, then adjusts gating with a learnable temperature. Its training objective combines the task loss with auxiliary load balancing, while fine-tuning freezes SMoE parameters.
- Dimension Reduction: X-MOE projects hidden states and expert embeddings into a lower-dimensional space to match the low-rank nature of SMoE routing.The projection also isolates direct hidden-state and expert-embedding interactions.
- Routing Algorithm: Figure 1 contrasts standard dot-product SMoE routing with X-MOE’s dimension reduction, L2 normalization, and gating temperature.The proposed design measures routing scores on a low-dimensional hypersphere.
- L2 Normalization: L2 normalization places token and expert representations on a common scale, stabilizing routing scores and reducing norm-driven expert dominance.Without normalization, larger expert-embedding norms can attract more tokens.
- Gating Temperature: A learnable temperature rescales normalized scores because L2 normalization restricts routing scores to [−1, 1].This prevents expert activation from becoming overly conservative.
- Training Objective: The objective jointly minimizes the task loss and an auxiliary load-balancing loss computed separately for each router.The task term can be masked language modeling for pre-training or sequence-to-sequence learning for translation.
- Fine-tuning: During fine-tuning, all expert and router parameters are frozen because small downstream datasets can lead SMoE models to overfit and route inconsistently.The load-balancing loss is still used during fine-tuning.
4 Experiments
The experiments evaluate X-MOE on cross-lingual pre-training, multilingual translation, and seven XTREME downstream benchmarks against dense and SMoE baselines. Results consistently favor X-MOE, while ablations examine routing components, dimensions, load balancing, and routing settings.
- Experimental Setup: Experiments use cross-lingual pre-training and fine-tuning, with downstream evaluation on seven XTREME understanding benchmarks.The setup includes multilingual data, Transformer-based X-MOE models, and comparisons with dense and SMoE baselines.
- Experimental Setup: Baseline models use the same pre-training data as X-MOE, while SMoE baselines include softmax- and sigmoid-gating variants.This controls the training-data comparison and distinguishes the evaluated routing implementations.
- Ablation Studies: Jointly using dimension reduction, L2 normalization, and frozen routing achieves the best ablation performance on XNLI and MLQA.Removing any one of the three routing methods lowers performance.
- Downstream Evaluation: 65.3 average score: softmax-gating X-MOE achieves the best performance across the seven XTREME downstream tasks.Results are averaged across test languages and five random seeds; X-MOE also improves consistently under top-1 and top-2 routing on XNLI.
- Upstream Evaluation: X-MOE consistently outperforms both dense and SMoE models in eight WMT-10 translation directions.The comparison uses BLEU scores for the x →en directions.
Weight XNLI MLQA
X-MOE improves multilingual language and translation performance while reducing representation collapse and producing more stable, consistent routing than SMoE baselines.
- Multilingual Translation: X-MOE consistently outperforms dense models and SMoE baselines in eight WMT-10 translation directions.The evaluation reports BLEU scores for x →en translation directions.
- Routing Ablations: Jointly using dimension reduction, L2 normalization, and routing freezing benefits X-MOE performance, while ablating any component reduces performance.The ablation study evaluates variants under matched pre-training and fine-tuning conditions on XNLI and MLQA.
- Routing Dimension: Routing dimension N/2 performs best on XNLI, whereas N/4 performs best on MLQA.The study compares routing dimensions N/4, N/2, N, 2N, and 4N.
- Representation Collapse: X-MOE achieves larger representation-collapse scores than SMoE and trends upward during pre-training, indicating less collapse.The RC metric measures within-class variability, with larger values indicating less collapse.
- Inter-run Consistency: Across 12 XNLI fine-tuning runs, X-MOE converges toward nearly identical routing behaviors, unlike the seed-sensitive SMoE baseline.X-MOE shows substantially better inter-run consistency despite differing routing early in fine-tuning.
- Routing Consistency: After 15K steps, X-MOE has a much lower routing-fluctuation ratio than SMoE, indicating more stable routing during pre-training.Lower RF ratios indicate more consistent routing.
5 Related Work
Prior work develops SMoE architectures, token-assignment algorithms, and analyses of representation collapse across neural and language-model settings.
- SMoE for Large-Scale Models: SMoE uses conditional computation to increase model capacity while maintaining computational efficiency.The related work situates SMoE as an extension of mixture-of-experts models.
- SMoE Routing Algorithms: Existing SMoE routing methods include global linear assignment, hashing, expert-selected top-k tokens, and frozen routing functions.These methods vary how tokens and experts determine assignments.
- Representation Collapse: Representation collapse describes degeneration in neural representations, including negligible within-class variation observed in classification and language-model fine-tuning.The literature also reports related collapse phenomena in visual representation learning.
6 Conclusion
The paper identifies representation collapse in SMoE models and proposes low-dimensional hypersphere routing, which improves multilingual performance and routing consistency.
- Conclusion: The proposed routing algorithm estimates token–expert scores on a low-dimensional hypersphere to alleviate representation collapse.The method combines dimensionality reduction with normalized representations for routing.
- Conclusion: Experiments across multilingual language-model benchmarks show consistent improvements over SMoE baselines in language modeling and fine-tuning.The conclusion summarizes gains across cross-lingual pre-training and downstream evaluation.
- Conclusion: The analysis finds that X-MOE reduces the trend toward representation collapse and achieves more consistent routing.The paper identifies vision and multimodal pre-training, and larger model sizes, as future evaluation directions.
- Ethical Considerations: The paper notes high computational and environmental costs as negative societal impacts of large-scale model training.It frames more efficient SMoE training as potentially reducing computation and CO2 emissions.
A Model Hyperparameters
X-MOE uses specified gating temperatures, a 250K-subword XLM-R vocabulary, and SentencePiece tokenization.
- Model Hyperparameters: The gating temperature τ0 is initialized to 0.3 for softmax gating and 0.07 for sigmoid gating.These values are listed among the X-MOE model hyperparameters.
- Model Hyperparameters: X-MOE uses the XLM-R vocabulary with 250K SentencePiece-tokenized subwords.The vocabulary follows XLM-R.
B Hyperparameters for Pre-training
This section presents the hyperparameters used for pre-training in Table 9.
- The accompanying text identifies Table 9 as the source of the pre-training hyperparameters.
- Table 9 presents the hyperparameters for pre-training.
- The pre-training hyperparameters are organized in a dedicated table.
C Hyperparameters for Fine-tuning
This section presents the hyperparameters used for fine-tuning on XTREME downstream tasks in Table 10.
- The accompanying text identifies Table 10 as the source of the fine-tuning hyperparameters.
- Table 10 presents the hyperparameters for fine-tuning on the XTREME downstream tasks.
- The fine-tuning hyperparameters are organized in a dedicated table.