Technical Summary: DTOP-p MoE

Problem Statement

Sparse Mixture-of-Experts (MoE) architectures are critical for scaling foundation models efficiently by decoupling parameter count from computational cost. However, standard Top-k routing imposes a rigid sparsity pattern, selecting a fixed number of experts regardless of token difficulty or layer-specific needs. This ignores the intrinsic variance in input complexity and restricts adaptive resource allocation.

While Top-p routing (nucleus sampling) offers a more adaptive alternative by selecting experts until their cumulative probability reaches a threshold $p$, existing implementations suffer from two critical limitations:

1. Fixed Global Thresholds: Naive Top-p implementations use a static global probability threshold that fails to adapt to evolving training dynamics.
2. Uncontrolled Sparsity: These methods yield unpredictable computational costs. The number of activated experts fluctuates significantly during training, making it difficult to adhere to strict pre-training compute budgets. Furthermore, fixed thresholds are highly sensitive to hyperparameter selection, often requiring extensive tuning that is computationally expensive.

Methodology: DTOP-p MoE

The authors propose DTOP-p, a sparsity-controllable dynamic routing mechanism designed to reconcile the adaptivity of Top-p with the strict computational constraints of large-scale pre-training. The method consists of two primary components:

1. Proportional-Integral (PI) Controller

To address the non-differentiability of the threshold selection and the lack of explicit resource constraints, DTOP-p employs a feedback control loop from classical control theory.

• Mechanism: The system treats the target number of activated experts ($T$) as a setpoint and the global probability threshold ($p$) as the control variable.
• Update Rule: The controller monitors the average number of activated experts per token ($a_t$) for the current batch. It calculates a sparsity error $e_t = (T - a_t)/N$ and updates the threshold for the next batch ($p_{t+1}$) using a discrete PI control law:
  $$p_{t+1} = p_0 + K_{pro} \cdot e_t + K_{int} \cdot \sum_{i=1}^{t} e_i$$
• Effect: This creates a negative feedback loop where a positive error (insufficient activation) increases $p$ to select more experts, while a negative error reduces $p$. This ensures the model converges to the target sparsity level regardless of shifts in logit entropy during training.

2. Dynamic Routing Normalization (DRN)

Applying a single global threshold across all layers assumes uniform expert confidence distributions, which contradicts findings that different layers have varying computational needs (e.g., shallow layers for general features vs. deep layers for specialized reasoning).

• Mechanism: Instead of routing on raw logits, DTOP-p normalizes the logits within each layer $l$ using a learnable scalar $\theta_l$:
  $$P(x) = \text{softmax}\left(\theta_l \cdot \frac{z - \mu(z)}{\sigma(z)}\right)$$
  where $z$ are the raw logits, and $\mu(z)$ and $\sigma(z)$ are their mean and standard deviation.
• Effect: This decouples the global threshold from local statistical properties. A larger $\theta_l$ sharpens the distribution (activating fewer experts), while a smaller $\theta_l$ flattens it (activating more). This allows the model to learn distinct sparsity patterns across layers while maintaining a global computational budget.

Layer-wise Variant

For deployment scenarios requiring strict per-layer predictability, the authors also propose a layer-wise variant where each layer maintains its own PI controller and threshold, replacing global cross-layer budget sharing with near-deterministic layer-wise tracking.

Key Contributions

1. Analysis of Fixed-Threshold Top-p: The paper demonstrates that fixed-threshold Top-p MoE offers only marginal gains over Top-k while suffering from hypersensitivity to the threshold $p$ and unstable sparsity, rendering it unsuitable for large-scale pre-training without control mechanisms.
2. DTOP-p Framework: The proposal of a dynamic Top-p MoE that integrates a PI controller for sparsity control and Dynamic Routing Normalization for layer-wise adaptivity.
3. Comprehensive Empirical Investigation: Extensive validation across Natural Language Processing (NLP) and Computer Vision (CV) domains, demonstrating that DTOP-p outperforms both Top-k and fixed Top-p baselines while matching the average FLOPs of Top-k.

Experimental Results

The authors evaluated DTOP-p on Large Language Models (LLMs) and Diffusion Transformers (DiTs) using the DCLM-Baseline dataset (NLP) and high-quality internal image datasets (CV).

• Performance: On NLP benchmarks (13 datasets including MMLU, HellaSwag, and ARC), DTOP-p achieved an average improvement of 1.9% over the Top-k MoE baseline while strictly adhering to the target activated expert count (e.g., 8 experts).
• Stability: Unlike fixed Top-p, which fluctuates wildly in the number of activated experts, DTOP-p rapidly converges to the target sparsity (within the first 1B tokens) and maintains a low standard deviation ($\approx 1$).
• Scaling Properties: DTOP-p demonstrated robust scaling across:
  • Expert Granularity: Performance gains widened as granularity increased (e.g., from 32 experts to 128).
  • Expert Capacity: Performance improved monotonically as total expert capacity increased.
  • Model and Dataset Size: The method maintained its advantage over baselines as model size (up to 2.4B active parameters) and dataset size (up to 300B tokens) increased.
• Layer-wise Analysis: Activation distribution analysis confirmed that DTOP-p learns hierarchical sparsity patterns, activating fewer experts in shallow layers and more in deeper layers, aligning with the varying complexity of features processed at different depths.

Significance and Claims

The paper claims that DTOP-p provides a robust and efficient MoE framework for foundation model pre-training. Its primary significance lies in successfully reconciling dynamic adaptivity with strict computational constraints. By making the probability threshold effectively learnable via control theory and enabling layer-wise flexibility, DTOP-p allows models to allocate computational resources based on token difficulty and layer depth without exceeding predefined FLOP budgets.

The authors position this work as a step toward sustainable Green AI, improving the computational efficiency of sparse models. They explicitly note that while the method improves efficiency, it does not mitigate inherent biases in training data and should be coupled with rigorous safety alignment. The work is presented as a practical solution for scaling foundation models where fixed Top-k is too rigid and uncontrolled Top-p is too unpredictable.