Teacher-Guided Routing for Sparse Vision Mixture-of-Experts

1. Problem Statement

Sparse Mixture-of-Experts (MoE) architectures allow models to scale capacity without a proportional increase in computational cost by activating only a subset of "experts" for each input. However, training sparse MoE models, particularly in Vision (VMoE), faces significant stability challenges:

• Gradient Blocking: In standard sparse MoE, the router receives informative gradients only from the experts selected during the forward pass. Unselected experts provide no gradient signal to the router.
• Routing Fluctuation: Due to the lack of global feedback, the router struggles to learn appropriate expert-selection scores, especially in early training stages. This leads to unstable dynamics where the expert assigned to a specific input fluctuates excessively across training iterations.
• Expert Collapse & Imbalance: Frequent routing changes prevent experts from specializing effectively, leading to undertrained experts and suboptimal load balancing, even when auxiliary load-balancing losses are used.
• Transfer Instability: Routing patterns learned during pre-training often fail to persist during fine-tuning, causing performance drops on downstream tasks.

2. Methodology: TGR-MoE

The authors propose TGR-MoE (Teacher-Guided Routing), a framework that stabilizes the student router by leveraging a pretrained dense teacher model.

Core Architecture

• Teacher Model: A frozen, pretrained dense Vision Transformer (e.g., DeiT-III) serves as the teacher.
• Teacher Router: A lightweight, trainable router is attached to the teacher's intermediate feature representations. It does not perform actual MoE routing but generates a "pseudo-supervision" signal.
• Student Model: The standard sparse VMoE student model with its own trainable router.

Training Mechanism

The training process involves a joint optimization of the student and the teacher router:

1. Teacher Router Optimization:

   • The teacher router processes intermediate features from the frozen teacher backbone.
   • It is optimized using a Load-Balancing Loss (to ensure uniform expert usage) and an Entropy Loss (to prevent the router from collapsing into a uniform distribution, encouraging confident assignments).
   • Note: The teacher router is optimized to produce a stable, balanced, and semantically meaningful routing distribution, not to minimize the downstream task loss directly.

2. Student Router Distillation:

   • The student router learns to mimic the teacher's routing distribution.
   • A Kullback-Leibler (KL) Divergence term is added to the student's loss function:
     $$\mathcal{L}_{distill} = KL(\text{stopgrad}(\mathbf{p}_t) \parallel \mathbf{p}_s)$$
     where $\mathbf{p}_t$ is the teacher's routing probability and $\mathbf{p}_s$ is the student's.
   • The total student loss is: $\mathcal{L}_{student} = \mathcal{L}_{task} + \lambda_{distill} \sum \mathcal{L}_{distill}$.

3. Key Advantage: The teacher provides a global prior derived from rich, semantically structured features. This guides the student router even for experts that are not selected in the current forward pass, effectively solving the gradient blocking problem without freezing the student router.

3. Key Contributions

• Novel Routing Stabilization: Introduces a teacher-guided mechanism that supplies informative gradients to the router from unselected routes via a dense teacher's intermediate representations, mitigating gradient blocking.
• Reduced Routing Fluctuation: Demonstrates that TGR-MoE significantly reduces the frequency of expert assignment changes during training, allowing experts to specialize earlier and more consistently.
• Scalability: Shows that the method scales reliably with the number of experts ($E$), maintaining performance gains even in highly sparse regimes where standard VMoE saturates.
• Transfer Learning Preservation: Proves that applying TGR-MoE during fine-tuning preserves the routing knowledge acquired during pre-training, leading to superior downstream performance compared to standard fine-tuning.
• Zero Inference Cost: The teacher model is only used during training; the inference stage requires only the lightweight student model, incurring no additional computational overhead.

4. Experimental Results

Experiments were conducted on ImageNet-1K (pre-training) and downstream tasks (CIFAR-10/100, Oxford-IIIT Pets) using various model scales (Tiny, Small, Base).

• Accuracy Gains: TGR-MoE consistently outperforms standard VMoE and competitive variants (Expert Choice, SoftMoE, z-loss).
  • Example: On ImageNet-1K, TGR-MoE improved the Tiny model's Top-1 accuracy from 77.85% (VMoE) to 78.78%.
  • On CIFAR-100, it achieved 90.26% (Small model) vs. 88.68% for VMoE.
• Routing Consistency:
  • TGR-MoE achieves >70% agreement with the final routing assignment within the first 50 epochs, whereas VMoE remains unstable (~40% agreement) even later in training.
  • During fine-tuning, TGR-MoE maintains 73.75% routing agreement between pre-training and fine-tuning, compared to only 50.56% for standard VMoE.
• Expert Count Scaling: As the number of experts increases (up to 256), TGR-MoE continues to show steady accuracy improvements, while VMoE performance plateaus.
• Ablation Studies:
  • Distillation Timing: Applying distillation primarily in the first half of training yields the best results, suggesting teacher guidance is most critical when routing is volatile.
  • Upper Bound: Using the teacher router directly for inference (an unattainable oracle) yields the highest accuracy (80.19%), confirming the potential of teacher-informed routing.

5. Significance

TGR-MoE addresses a fundamental bottleneck in sparse MoE training: the lack of global supervision for the router. By decoupling the routing guidance from the sparse activation constraints and leveraging a dense teacher's rich feature space, the method enables:

1. Stable Training: Elimination of the "expert collapse" and fluctuation issues that plague large-scale MoE training.
2. Efficient Scaling: The ability to increase model capacity (more experts) without sacrificing training stability or inference speed.
3. Robust Transfer: A mechanism to carry over learned routing structures to new domains, which is crucial for practical deployment of large vision models.

The paper concludes that effective MoE routing requires supervisory signals beyond task-driven gradients, and teacher-guided priors offer a simple, powerful, and cost-effective solution to achieve stable, scalable, and high-performing expert architectures.