DirMoE: Dirichlet-routed Mixture of Experts

Paper: 2602.09001 Authors: Amirhossein Vahidi, Hesam Asadollahzadeh, Navid Akhavan Attar, Marie Moullet, Kevin Ly, Xingyi Yang, Mohammad Lotfollahi Categories: cs.LG

Abstract

This paper proposes DirMoE, a novel routing mechanism for Mixture of Experts (MoE) models that addresses fundamental limitations of the standard Top-k + Softmax routing approach. DirMoE uses a Dirichlet variational autoencoder framework to separate two key routing decisions: which experts to activate and how to distribute contributions among them. The method employs Gumbel-Sigmoid relaxation for expert selection and implicit reparameterization for the Dirichlet distribution, achieving full differentiability. A variability-based objective provides direct sparsity control. DirMoE matches or exceeds existing routing methods while improving expert specialization.

Key Contributions

• Decoupled routing decisions: Separates expert activation (which experts) from contribution allocation (how much each contributes)
• Dirichlet-based weighting: Uses the Dirichlet distribution for principled probability allocation over active experts
• Full differentiability: Gumbel-Sigmoid and implicit reparameterization enable end-to-end gradient flow
• Direct sparsity control: Variability-based objective replaces ad-hoc load balancing losses

Limitations of Top-k + Softmax Routing

The standard MoE routing mechanism has several well-known problems:

Coupled Decisions

Top-k + Softmax entangles two distinct decisions:

1. Which experts to activate (a discrete selection problem)
2. How to weight their contributions (a continuous allocation problem)

By using the same scores for both, the routing cannot independently optimize selection and weighting.

Load Imbalance

Without auxiliary losses, Top-k routing tends to collapse to using only a few experts, wasting model capacity. The auxiliary load-balancing losses used to prevent this are heuristic and can conflict with the primary training objective.

Non-Differentiable Selection

The Top-k operation is non-differentiable, meaning gradients cannot flow through the expert selection decision. Only the selected experts receive gradient updates, creating a rich-get-richer dynamic.

The DirMoE Framework

Step 1: Expert Selection via Gumbel-Sigmoid

For each token $x$, compute activation logits for each expert $e$:

$\alpha_e = W_{\text{gate}} \cdot x + b_{\text{gate}}$

Apply Gumbel-Sigmoid relaxation to obtain differentiable binary activation decisions:

$z_e = \sigma\left(\frac{\alpha_e + g_1 - g_2}{\tau}\right)$

where $g_1, g_2 \sim \text{Gumbel}(0, 1)$ are Gumbel noise samples and $\tau$ is the temperature. As $\tau \to 0$, this approaches a hard binary decision; during training, a moderate $\tau$ allows gradient flow.

The activation mask is:

$m_e = \mathbb{1}[z_e > 0.5] \quad \text{(inference)}, \quad m_e = z_e \quad \text{(training)}$

Step 2: Contribution Weights via Dirichlet Distribution

For the active experts, DirMoE uses a Dirichlet distribution to allocate contribution weights. The Dirichlet distribution is the natural choice because:

• It produces vectors that sum to 1 (valid probability distributions)
• It can express both uniform and concentrated allocations
• Its concentration parameters control the degree of specialization

Compute Dirichlet concentration parameters for active experts:

$\beta_e = \text{softplus}(W_{\text{dir}} \cdot x + b_{\text{dir}}) \cdot m_e$

Sample contribution weights from the Dirichlet distribution:

$w \sim \text{Dir}(\beta_1, \beta_2, \ldots, \beta_E)$

The implicit reparameterization trick enables gradient computation through this sampling step.

Step 3: Expert Computation

The final output combines expert outputs weighted by the Dirichlet-sampled weights:

$y = \sum_{e=1}^{E} w_e \cdot m_e \cdot \text{Expert}_e(x)$

Variability-Based Sparsity Control

Instead of heuristic load-balancing losses, DirMoE introduces a variability-based objective that directly controls sparsity:

$\mathcal{L}_{\text{var}} = -\lambda \cdot \text{Var}(\bar{m})$

where $\bar{m}_e = \mathbb{E}_x[m_e]$ is the average activation rate of expert $e$ across the batch. Maximizing variance in activation rates encourages some experts to be frequently active and others rarely, achieving sparsity. The coefficient $\lambda$ directly controls the sparsity level.

Theoretical Advantages

Why Dirichlet?

The Dirichlet distribution $\text{Dir}(\beta_1, \ldots, \beta_K)$ has properties ideally suited for expert weighting:

• Sum-to-one constraint: $\sum_e w_e = 1$ is automatically satisfied
• Concentration control: Large $\beta_e$ values concentrate weight on expert $e$; uniform $\beta$ gives equal weights
• Sparse allocations: When some $\beta_e \ll 1$, the distribution naturally produces sparse weight vectors
• Smooth interpolation: Between fully concentrated (one expert dominates) and uniform allocation

Full Differentiability

The combination of Gumbel-Sigmoid (for selection) and implicit reparameterization (for Dirichlet sampling) ensures that gradients flow through the entire routing mechanism:

$\frac{\partial \mathcal{L}}{\partial W_{\text{gate}}}, \frac{\partial \mathcal{L}}{\partial W_{\text{dir}}} \neq 0$

This means all experts receive gradient signal, even those not selected for a particular token, through the Gumbel-Sigmoid relaxation.

Experimental Results

Performance

DirMoE matches or exceeds baseline routing methods across multiple benchmarks:

• Comparable or better perplexity on language modeling tasks
• Improved performance on downstream tasks requiring diverse expert knowledge
• More stable training dynamics without the need for auxiliary load-balancing losses

Expert Specialization

A key advantage of DirMoE is improved expert specialization:

• Experts develop more distinct specializations compared to Top-k routing
• The Dirichlet weighting allows nuanced contribution allocation (not just equal weighting of top-k)
• Reduced redundancy between experts

Sparsity Control

The variability-based objective provides precise sparsity control:

• Direct relationship between $\lambda$ and the achieved sparsity level
• No need for heuristic tuning of auxiliary loss coefficients
• Smooth trade-off between computation cost and model quality

Load Balance

Despite not using explicit load-balancing losses:

• DirMoE achieves comparable or better load balance than Top-k with auxiliary losses
• The Gumbel-Sigmoid mechanism naturally distributes activations
• No expert collapse observed during training

Significance

DirMoE provides a principled solution to the routing problem in MoE models by cleanly separating the two fundamental routing decisions and using appropriate probabilistic tools for each. The Dirichlet distribution is a natural fit for contribution allocation, and the Gumbel-Sigmoid provides differentiable expert selection. The variability-based sparsity control replaces ad-hoc load balancing with a theoretically grounded objective. Together, these innovations yield a routing mechanism that is more principled, more controllable, and empirically competitive with or superior to existing approaches.