SimSMoE: Solving Representational Collapse via Similarity Measure

TL;DR

SimSMoE tackles representational collapse in sparse mixture-of-experts by diversifying expert representations using a centered kernel alignment ($CKA$) based similarity loss. It introduces a Similarity Learning module that identifies collapsed experts via shared-token analysis and applies a $CKA$-based loss with a projection head to minimize inter-expert similarity. The method is compatible with existing routing policies and demonstrates improved pre-training and fine-tuning performance across multiple decoder-only LLM architectures (Brainformer, GLaM, Mistral) on Enwik8, Text8 and downstream NLP tasks, with ablations supporting robustness and scalability. These results suggest that directly addressing expert representations can enhance SMoE training efficiency and scalability for large-scale LLMs.

Abstract

Sparse mixture of experts (SMoE) have emerged as an effective approach for scaling large language models while keeping a constant computational cost. Regardless of several notable successes of SMoE, effective training such architecture remains elusive due to the representation collapse problem, which in turn harms model performance and causes parameter redundancy. In this work, we present Similarity-based Sparse Mixture of Experts (SimSMoE), a novel similarity of neural network algorithm, that guarantees a solution to address the representation collapse issue between experts given a fixed FLOPs budget. We conduct extensive empirical evaluations on three large language models for both Pre-training and Fine-tuning tasks to illustrate the efficacy, robustness, and scalability of our method. The results demonstrate that SimSMoE significantly enhances existing routing policy and outperforms other SMoE training methods in performance for the tasks.

Figures (5)

• Figure 1: Illustration of the proposed SimSMoE architecture and a SMoE architecture. (a) A SMoE architecture selectively activates experts based on dot-product token-expert routing scores, directing the selected token to the chosen experts. (b) A SimSMoE architecture mitigates the issue of representation collapse by reducing the similarity among the selected experts.
• Figure 2: A Similarity Learning Layer (ours) to minimize the similarity among experts.
• Figure 3: Bit-per-Character (BPC) on validation dataset during the training phase reported for Mistral jiang2024mixtral across the three routing mechanisms. (a) SMoE with the Balancing Loss. (b) XMoE. (c) StableMoE
• Figure 4: Analysis of the similarity index for the Sparse Mixture of Experts (SMoE). Figure (a) shows the correlation between the number of shared tokens and the similarity index. Figure (b) illustrates the similarity index by layer order.
• Figure 5: Exploration of the impact of similarity learning on diversity model representation. Figure (a) shows the heatmap of differences between the hidden representations of two experts for the SMoE layer. Figure (b) shows the heatmap of differences between the hidden representations of two experts for the SimSMoE layer.