Source-linked AI summary
DynaMoE: Dynamic Token-Level Expert Activation with Layer-Wise Adaptive Capacity for Mixture-of-Experts Neural Networks
Gökdeniz Gülmez
TL;DR
Standard MoE systems fix both the number of experts activated per token and expert capacity across layers, limiting adaptation to input complexity and network depth. DynaMoE introduces percentile-based dynamic routing and six layer-wise capacity schedules, then analyzes their expressivity, compute, and gradient behavior. Across the reported tasks, optimal schedules depend on task and model scale, with descending schedules improving image classification and scale-dependent schedules benefiting language modeling.
Problem
Standard MoE implementations use fixed Top-K routing and uniform expert allocation, despite varying computational needs across inputs and network depths.
Method
DynaMoE combines input-dependent expert activation with a layer-wise schedule mapping each layer to an expert count between Nmin and Nmax.
Results
Descending schedules improve image-classification accuracy by up to 5.47%, while language-modeling optima vary by scale: descending for Tiny, ascending for Small, and uniform for Medium.
Takeaways & Limitations
Expert distribution is task-dependent, challenging uniform allocation as a general MoE design assumption.
Takeaways & Limitations
Language-modeling experiments used only 1,000 training samples, with perplexities of 1,000–2,500 and near-identical accuracy, so larger-scale pretraining is required.
Abstract
from arXiv · showhide
Mixture-of-Experts (MoE) architectures have emerged as a powerful paradigm for scaling neural networks while maintaining computational efficiency. However, standard MoE implementations rely on two rigid design assumptions: (1) fixed Top-K routing where exactly K experts are activated per token, and (2) uniform expert allocation across all layers. This paper introduces DynaMoE, a novel MoE framework that relaxes both constraints through dynamic token-level expert activation and layer-wise adaptive capacity allocation. DynaMoE introduces a principled routing mechanism where the number of active experts per token varies based on input complexity. Concurrently, the framework implements six distinct scheduling strategies for distributing expert capacity across network depth, including descending, ascending, pyramid, and wave patterns. We theoretically analyze the expressivity gains of dynamic routing and derive bounds on computational efficiency. Through extensive experiments on MNIST, Fashion-MNIST, CIFAR-10 (image classification), and Recycling-the-Web (language modeling) across multiple model scales, we demonstrate that DynaMoE achieves superior parameter efficiency compared to static baselines. Our key finding is that optimal expert schedules are task- and scale-dependent: descending schedules (concentrating capacity in early layers) outperform uniform baselines on image classification. For language modeling, optimal schedules vary by model size, descending for Tiny, ascending for Small, and uniform for Medium. Furthermore, dynamic routing reduces gradient variance during training, leading to improved convergence stability. DynaMoE establishes a new framework for adaptive computation in neural networks, providing principled guidance for MoE architecture design.
1 Introduction
DynaMoE addresses rigid fixed-Top-K routing and uniform layer capacity by adapting expert activation per token and distributing experts across depth. It evaluates these mechanisms across image and language tasks to identify task-dependent benefits.
- MoE uses conditional computation to scale parameter counts while keeping computational costs manageable.
- Standard MoE activates exactly K experts per token, regardless of input complexity or layer characteristics.
- Uniform expert allocation gives every layer identical capacity despite potentially different needs across network depth.
- DynaMoE varies active experts per token according to input characteristics and provides theoretical computational-efficiency guarantees.
- DynaMoE formalizes six strategies for distributing expert capacity across network depth and empirically validates them.
- Experiments across image classification and language modeling identify task-dependent optimal expert distributions and improvements over static baselines.
2 Related Work
Prior work addresses MoE scalability, routing, adaptive computation, and heterogeneous capacity, while DynaMoE combines input-dependent routing with systematic layer-wise expert scheduling.
- Earlier MoE research established scalable sparse routing, including single-expert Switch Transformers and expert-choice routing.
- Dynamic-computation methods such as early exiting, adaptive width, and dynamic depth primarily modify forward execution rather than the routing mechanism itself.
- Adaptive-computation research frames compute allocation per token as proportional to input complexity, motivating DynaMoE’s variable-K routing.
- Learned thresholds or per-layer learned K values are identified as natural extensions that DynaMoE does not explore.
- Prior capacity-allocation work includes heterogeneous expert sizes, expert pruning, and evidence that converted dense layers require varying experts across depth.
- DynaMoE presents systematic predefined layer-wise expert-count scheduling in a unified framework linked to information-theoretic and optimization-theoretic principles.
3 Methodology
DynaMoE combines percentile-based token routing with layer-wise schedules that vary expert counts across depth. The implementation includes exploration noise, soft weighting, residual MLP layers, and explicit load-imbalance limitations.
- Dynamic Token-Level Routing: DynaMoE generalizes fixed Top-K routing by making the number of selected experts input-dependent.
- Dynamic Token-Level Routing: A percentile threshold selects experts per token, so active-expert count K(x) varies between Kmin = 1 and Kmax = ⌈(1 −τ) · N⌉.
- Dynamic Token-Level Routing: Temperature scaling stabilizes softmax computation, while Gaussian gate noise with σ = 0.1 encourages exploration during training.
- Capacity Control and Load Balancing: The routing design processes all selected tokens without explicit capacity constraints or auxiliary balancing losses, using minimum activation and soft weighting instead.
- Capacity Control and Load Balancing: Persistent over-selection of popular experts remains possible, especially for diverse or large-vocabulary distributions, and may worsen at scale.
- Layer-Wise Expert Distribution: DynaMoE defines an expert schedule mapping each layer to a count between Nmin and Nmax, with six depth-distribution strategies.
- Layer-Wise Expert Distribution: The schedules distribute capacity across 12 layers, including descending early-layer concentration, ascending late-layer concentration, and intermediate pyramid and uniform patterns.
- Model Architecture: The architecture applies scheduled experts in residual DynaMoE layers followed by layer normalization and a classification head.
4 Theoretical Analysis
The analysis characterizes dynamic routing as expanding routing-pattern diversity and potentially reducing gradient variance, while relating compute to active experts and schedules to capacity allocation. These claims depend on explicit assumptions and fairness constraints.
- Expressivity: Dynamic routing permits any active-expert count from 1 through Kmax, whereas fixed routing permits only K, expanding reachable activation patterns when Kmax > K.
- Expressivity: The expressivity theorem concerns routing-pattern diversity and yields a lower bound on functions implementable with fixed expert weights, not function-space volume.
- Computational Efficiency: Expected expert computation scales with E[K(x)]·O(d2), while layer parameter storage scales as O(d2 · S(ℓ)) and is not included in active-expert FLOPs.
- Computational Efficiency: Compute-efficiency comparisons use active-expert FLOPs per token rather than wall-clock throughput or total parameter counts.
- Computational Efficiency: For τ = 0.7 and S(ℓ) = 8, expected active experts are approximately 2.4 versus fixed K = 2.
- Gradient Variance: The gradient-variance bound is non-vacuous only when dynamic routing has higher entropy than fixed routing.
- Gradient Variance: The variance analysis assumes more uniform dynamic routing and conditionally independent expert outputs, and is a qualitative potential reduction rather than a tight guarantee.
- Gradient Variance: Empirically, dynamic routing is associated with faster convergence and lower-variance loss curves, consistent with more balanced expert usage.
5 Experimental Setup
The experiments evaluate DynaMoE across image-classification and language-modeling datasets, four model sizes, standardized training settings, and multiple efficiency and behavior metrics. The present architecture is MLP-based; attention probes are reserved for future Transformer extensions.
- Datasets: Experiments cover MNIST, Fashion-MNIST, CIFAR-10, and Recycling-the-Web-1k for image classification and language modeling.The image datasets contain 10 classes; Recycling-the-Web-1k contributes 1,000 text samples.
- Model Configurations: Four model sizes are evaluated to study scaling behavior.
- Training Configuration: Training uses AdamW, cosine-annealed learning rates, task-specific batch sizes and epochs, temperature T = 0.5, and noise scale σ = 0.1.AdamW uses β1 = 0.9, β2 = 0.999, and weight decay 10−4; the learning rate is 10−3.
- Evaluation Metrics: Evaluation measures accuracy, perplexity, parameters per accuracy ratio, convergence speed, and expert-utilization entropy.Language-modeling perplexity is defined as PPL = exp(L), where L is cross-entropy loss.
- Architecture Scope: The current experiments use an MLP-based architecture without self-attention, while attention-related probes are proposed only for future Transformer extensions.Proposed probes include attention entropy, effective attention distance, head specialization, and a superposition pressure proxy.
6 Results
DynaMoE’s results show that adaptive expert scheduling improves image-classification efficiency, while language-modeling benefits depend on model scale and schedule choice. Training dynamics and activation patterns further indicate schedule-specific convergence behavior and depth-wise computation allocation.
- Expert Schedule Ablation: 92.68% accuracy: descending scheduling outperforms the MLP baseline by 3.26% and uniform MoE by 1.33% on Small MNIST.It also reaches 95% of final accuracy in 6 epochs.
- Expert Count Sensitivity: E8–1 provides the best accuracy-efficiency tradeoff, while E16–1 has diminishing returns and E8–4 reduces flexibility and performance.
- Scaling Analysis: Descending scheduling consistently outperforms baselines across image-classification scales, with gaps of +1.6% for Tiny and +3.3% for Small.
- Cross-Dataset Generalization: Descending scheduling improves performance across datasets, with a larger gain on CIFAR-10 of +5.47%.
- Language Modeling: Per-token accuracy is nearly identical for Small and Medium models, making perplexity the primary differentiator in those comparisons.All configurations reach approximately 0.0505 best validation accuracy at those scales.
- Language Modeling: Language-modeling optima shift with scale: descending for Tiny, ascending for Small, and uniform for Medium.The reported best perplexities are 1011.80, 2308.29, and 2383.89, respectively.
7 Analysis and Discussion
The analysis attributes schedule performance to matching expert capacity with representational diversity across depth. Descending schedules favor image classification, while language modeling can favor ascending schedules as deeper representations become more diverse.
- Why Descending Schedules Work: Descending schedules place maximal approximation capacity in early layers, where curvature, input heterogeneity, and gradient-estimate noise are highest.Later representations are more class-conditioned and closer to linear, reducing the marginal value of additional experts.
- Why Descending Schedules Work: Ascending schedules create an early-layer gradient bottleneck by routing initial representations through only 1–2 expert pathways.The paper contrasts this with descending schedules, whose larger early expert pools provide a structurally stronger gradient pathway.
- Representational Diversity-Convergence Principle: The RDC Principle states that optimal expert count should track representational diversity, which typically decreases with depth in supervised pattern recognition.This motivates descending capacity allocation, with more experts in early layers and fewer in later layers.
- Task-Dependent Scheduling: The optimal schedule is task-dependent: descending matches image classification, whereas ascending or pyramid can suit language modeling when diversity grows with depth.The RDC framework recommends estimating each task’s representational-diversity profile before selecting a schedule.
- Uniform Allocation: The uniform schedule beats dense MLP but trails descending DynaMoE by 1.33% to 2.67% across datasets.Uniform allocation benefits from conditional computation yet wastes capacity in deep layers and under-provisions early layers.
- Uniform Allocation: Uniform DynaMoE converges faster than MLP, reaching 95% accuracy in 7 versus 8 epochs, while dynamic routing supplies adaptive complexity and lower gradient variance.These benefits explain its intermediate position between descending DynaMoE and the dense baseline.
- Parameter Efficiency: The descending schedule’s parameter efficiency exceeds uniform allocation because early-layer parameters purchase more accuracy improvement than deep-layer parameters.The reported efficiency metric is 2.44 for uniform, with descending described as higher.
8 Conclusion
DynaMoE combines dynamic token-level routing with layer-wise expert distribution and evaluates the framework across image classification and language modeling. Results show that expert schedules are task- and scale-dependent, while future work remains needed for learned schedules, transformers, and larger models.
- DynaMoE combines dynamic token-level routing, six layer-wise capacity schedules, theoretical analysis, and evaluation on image classification and language modeling.
- 5.47% accuracy improvements are achieved by descending schedules on image classification relative to uniform allocation.
- Small ascending schedules achieve the best language-modeling PPL of 2308.29, marginally outperforming the MLP baseline of 2311.02 by 0.1%.
- Medium uniform schedules achieve the strongest language-modeling improvement, reducing best validation PPL by 3.4% versus MLP.The reported PPL values are 2383.89 for uniform DynaMoE and 2468.16 for MLP.
- Expert distribution is task-dependent: descending helps image classification, while language modeling favors scale-dependent schedules.
- Future work should examine learned scheduling, transformer architectures, and scaling to billion-parameter models.
A.1 Training Curves
The cited passages identify convergence behavior and per-layer expert statistics as the focus of the training-curve materials, but do not report their numerical outcomes.
- Figure 11 shows convergence behavior across schedules.
- Table 6 provides per-layer expert statistics for the descending schedule.
B Proof of Theorem 4.1
The proof compares the number of expert activation patterns available under dynamic routing with fixed Top-K routing, using factorial approximations to derive an expressivity ratio.
- Dynamic routing analyzes activation patterns for a layer with N experts and threshold τ.
- The dynamic-routing bounds use Kmax = ⌈(1 −τ)N⌉ and Kmin = 1.
- Fixed Top-K routing has a simpler activation-pattern count for comparison.
- The expressivity ratio is obtained by comparing the dynamic and fixed-routing cardinalities.
- Stirling’s approximation for factorials is used in the derivation.
C Implementation Details
The implementation describes the DynaMoE layer forward pass, including inputs, gating, threshold-based expert selection, and hyperparameter sensitivity analysis.
- Algorithm 1 presents the DynaMoE layer computation.
- The forward pass accepts input x, expert count N, threshold τ, and temperature T.
- Gates are computed with softmax, optionally perturbed by training noise, and thresholded using a percentile rule.
- The layer guarantees at least one selected expert and aggregates selected expert outputs using temperature-scaled softmax scores.
- Table 7 reports sensitivity to key hyperparameters on MNIST.