Mixture-of-Grouped-Experts (MoGE)

• MoGE is a framework that structures experts into explicit groups to enhance load balance, memory efficiency, and specialization in complex models.
• It employs group-balanced routing and regularization techniques to optimize compute, minimize interference, and maintain diverse expert outputs.
• Empirical evaluations show that MoGE boosts throughput, reduces catastrophic forgetting, and achieves memory and performance gains across adaptive tasks.

A Mixture-of-Grouped-Experts (MoGE) is a generalization of the mixture-of-experts paradigm in which experts are organized into explicit groups, and the routing, parameterization, or optimization processes are structured to exploit group-level characteristics. This approach has arisen in several research domains, including deep learning for LLMs, representation learning, continual learning, power systems, and statistical modeling. Key motivations include improving workload and memory efficiency, enhancing the diversity and specialization of experts, improving interpretability and controllability, mitigating cross-task interference, and enabling robust and scalable deployment across compute topologies.

1. Principles and Motivations

Mixture-of-Grouped-Experts modifies the conventional MoE architecture by introducing a grouping structure on the set of experts. In standard MoE, a gating network selects a sparse subset (Top-K) of N experts for each input, often resulting in substantial load imbalance and unstructured competition among experts. MoGE approaches address several recognized shortcomings:

• Load imbalance in distributed systems: Conventional MoE can lead to severe inter-device stragglers under Top-K routing, as certain experts are favored, making expert-parallel scaling inefficient. MoGE's explicit groupings enforce global and intra-device balance (Tang et al., 27 May 2025).
• Scalability of memory and compute: Memory overheads in Transformer decoders can be controlled by grouping attention heads and judiciously partitioning key–value storage at the expert group level, preserving high-importance features while coarsening low-importance token storage (Song et al., 16 Jun 2025).
• Expert diversity and specialization: Flat expert competition can yield collapsed or redundant representations. MoGE leverages grouping and structured regularization (e.g., group-sparsity, topographic masking) to favor spatially or semantically diverse expert use (Kang et al., 12 Apr 2025).
• Mitigation of catastrophic and forward forgetting in continual learning: By allocating separate expert groups to sequential tasks and using hierarchical routing, MoGE isolates task-specific knowledge while allowing controlled sharing and adaptation (Zhou et al., 11 Aug 2025).
• Interpretability and multi-dimensional balance: Layered or hierarchical group routing enables interpretable and controllable decomposition (e.g., by function, domain, style), facilitating analysis and calibration (Li et al., 2024).

2. Formal Architectures and Routing Mechanisms

Different instantiations of MoGE vary in architectural granularity and routing algorithms:

2.1 Group-balanced MoE for Expert Load (LLM/Hardware Context)

Let $N$ experts be partitioned into $M$ groups ($N_g = N/M$ per group).

• Group-Balanced Top-K Routing: For input $h \in \mathbb{R}^d$, affinities $S = \mathrm{Softmax}(W^\top h) \in \mathbb{R}^N$ are globally computed. Within each group $j$, exactly $K'$ ($=K/M$) experts are selected with the highest scores, enforcing $\sum_{i \in \text{Group}_j} 1_{G'(h)_i > 0} = K'$ for each group. The layer output is $y = \sum_{i=1}^N G'(h)_i \cdot E_i(h)$ (Tang et al., 27 May 2025).

2.2 Dynamic Token-wise Routing with Grouped KV (Transformer)

Each token is routed via a lightweight scoring layer $\mathsf{S}(X) = \sigma(XW_r + b_r)$ to one of $E$ head-grouping experts, where each expert $e$ defines a different attention head grouping size $2^e$.

• Training/prefill: Tokens sorted by score are greedily assigned to experts with fixed ratios $\ell_e$, yielding a one-hot assignment per token.
• Decoding: Argmax selection over scored experts per token.
• Auxiliary Consistency Loss: A cross-entropy loss aligns routing decisions during training and inference (Song et al., 16 Jun 2025).

2.3 Group-Sparse Regularized MoE (Representation Learning)

Gate logits $z$ are smoothed spatially by mapping to a 2D grid $Z$ and applying a Gaussian $\ell_{2,1}$ penalty over overlapping $h \times h$ regions, inducing local sparsity and stabilizing expert selection under input variations (Kang et al., 12 Apr 2025).

2.4 Hierarchical Group Routing (Specialized/Task-Decomposed MoE)

Hierarchical routing is realized as:

• Group-level weights computed by softmax-normalized logits (with temperature $\tau_G$).
• Within-group normalization: Each group normalizes expert selection by group-local softmax (temperature $\tau_D$), allowing independent control over exploration-exploitation at different abstraction levels (Li et al., 2024).

2.5 Two-Level Routing for Continual Learning

A group is allocated per task; routing proceeds by (1) intra-group softmax over a fixed small expert set per group, and (2) inter-group router selecting and mixing group outputs based on task identifiers or proximity to learned prototypes. Old groups are frozen when new groups are allocated, and dynamic fusion with the base model mitigates forward forgetting (Zhou et al., 11 Aug 2025).

3. Key Algorithmic Variants and Mathematical Properties

MoGE appears in diverse domains with tailored algorithmic instantiations:

• Memory/Compute-efficient Transformer Attention: By routing tokens to experts corresponding to different granularity KV groupings and sharing projection weights, MoGE realizes proportional KV cache savings per token, with per-expert grouping sizes determining memory efficiency. Average per-token KV usage can be tightly controlled by group ratios, achieving significant compression with controlled perplexity loss (Song et al., 16 Jun 2025).
• Group-Sparse Regularization: The loss function augments the base task objective with a term:

$$\mathcal{L}_\mathrm{MoGE} = \mathcal{L}_\mathrm{task} + \lambda R(z)$$

where $R(z)$ is computed via convolutional smoothing over the softmaxed gating logits, penalizing non-sparse, spatially diffuse activation patterns and boosting invariance (Kang et al., 12 Apr 2025).

• Latent Grouped Mixture Models: In statistical modeling of grouped/clustered data, MoGE formalizes the data likelihood with cluster-wise Dirichlet-distributed mixing proportions on expert densities. The model is fit via MC-EM or Gibbs-EM, and can be extended for covariate-dependent group mixing weights (Sugasawa et al., 2017).
• Mixture-of-Gradient-Experts: In constraint screening for convex OPF, the grouping is semantic—combining an ICNN (input-convex NN) and an MGN (monotone NN) via a gating function, trained to predict dual multipliers. The convex combination is computed per constraint index, with group selection aligning with theoretical properties of the underlying optimization problem (Bose et al., 2023).

4. Empirical Evaluations and Benchmarks

MoGE methods exhibit empirically validated improvements over flat or conventional MoE in multiple axes:

• Load and Throughput: On Ascend NPUs (Pangu Pro MoE), MoGE achieves zero imbalance for batch inference, with 142–303% throughput increases versus dense baselines and tight expert utilization ($\sim 12.5\%$ usage per expert vs. 30\% in Top-K MoE) (Tang et al., 27 May 2025).
• Memory/Compute-Accuracy Trade-offs: In causal language modeling, MoGE with dynamic KV grouping outperforms static grouped query attention and CLA, yielding 5–10% lower perplexity for a fixed 50% KV cache budget, and 2–7 ROUGE-L point gains in downstream tasks (Llama3, Gemma2, OPT, TinyLlama) (Song et al., 16 Jun 2025).
• Diversity, Invariance, and Scalability: Group-sparse MoGE offers 0.3–1% higher top-1 image classification accuracy and 1–3 PPL point improvements in LM compared to MoE, with negligible memory/time overhead and better expert usage diversity (Kang et al., 12 Apr 2025).
• Multi-domain Continual Learning: In grouped continual learning, MoGE (TRGE) improves average accuracy by 1.2–1.8% and reduces catastrophic forgetting, with 61% fewer trainable parameters than flat MoE-adapter baselines (Zhou et al., 11 Aug 2025).
• Statistical Efficiency: In clustered data modeling, MoGE yields lower integrated mean squared error (MISE) in simulation and better adaptation in small-sample clusters relative to standalone or global mixture models (Sugasawa et al., 2017).
• Power System Optimization: In OPF, MoGE-based screening achieves zero false negatives in constraint selection and 20–35% reductions in solve time compared to baseline classifiers (Bose et al., 2023).

5. Application Domains and System-level Integration

MoGE approaches have been adapted to distinct application areas:

• LLMs/Transformers: Efficient hardware utilization and sparsity for >70B parameter LLMs, via group-balanced routing and hybrid parallelism on Ascend/NPU clusters, including fused kernel and quantization designs (Tang et al., 27 May 2025).
• Memory-efficient Causal Language Modeling: MixSGA MoGE dynamically assigns per-token memory, enabling robust scaling for sequence-length-induced KV explosion (Song et al., 16 Jun 2025).
• Vision Transformers and LLMs: Group-sparse regularization is leveraged to stabilize specialization and learn invariant representations in ViT/Swin backbones (Kang et al., 12 Apr 2025).
• Continual and Multi-domain Learning: Hierarchically grouped LoRA expert routing enables robust, parameter-efficient adaptation, mitigating both catastrophic and forward forgetting on large multi-task image/text datasets (Zhou et al., 11 Aug 2025, Li et al., 2024).
• Statistical Cluster Modeling: Hierarchical MoGE generalizes finite mixture models to accommodate latent group structure and covariate-adaptive mixture weights in hierarchical/clustered data (Sugasawa et al., 2017).
• Physics-Informed and Gradient-Driven Applications: MoGE integrating convexity and monotonicity constraints provides high-accuracy constraint screening in power flow optimization (Bose et al., 2023).

6. Theoretical and Practical Considerations

• Theoretical properties: Grouping in MoGE enforces structured sparsity, diversity, and, where relevant, statistical properties such as strong duality, exact support recovery, or robust expert selection under clustered likelihoods (Sugasawa et al., 2017, Bose et al., 2023, Kang et al., 12 Apr 2025).
• Practicality: MoGE methods typically add minimal parameter/memory/compute overhead while yielding system-level efficiencies or improved accuracy/diversity. Group Regularization and grouping schedules (kernel size, filter type, temperature) must be tuned for optimal effect (Kang et al., 12 Apr 2025, Li et al., 2024).
• Extensibility: The group structure is broadly applicable and can be aligned with domain, functionality, style, device affinity, or statistical cluster structure, with empirical effectiveness proven across diverse model classes (Tang et al., 27 May 2025, Song et al., 16 Jun 2025, Li et al., 2024, Zhou et al., 11 Aug 2025, Sugasawa et al., 2017).

7. Limitations, Open Problems, and Outlook

MoGE is not without constraints:

• Its effectiveness at very large scale (billions of experts, extreme heterogeneity) remains to be investigated beyond current benchmarks (Kang et al., 12 Apr 2025).
• Generalization to new modalities (speech, video or multimodal) is under-explored.
• The optimal design of expert grouping, routing functions, and group sizes is application-specific and may require empirical exploration.
• Stability under nonstationary distributions and grouping dynamics over long-term continual learning is an ongoing research area (Zhou et al., 11 Aug 2025).
• Specialized expert architectures or group-level learning dynamics may offer further performance gains but have not been systematically benchmarked.

In summary, the Mixture-of-Grouped-Experts (MoGE) paradigm offers a robust, generalizable framework for scalable, efficient, interpretable, and diverse expert allocation across a range of modern AI, statistical, and engineering systems. Its continuing evolution reflects the increasing need for structured sparsity, workload balancing, and specialization in high-capacity models and evolving computational environments (Song et al., 16 Jun 2025, Tang et al., 27 May 2025, Kang et al., 12 Apr 2025, Li et al., 2024, Bose et al., 2023, Zhou et al., 11 Aug 2025, Sugasawa et al., 2017).