DirMoE: Dirichlet-Routed Mixture of Experts
Authors: Amirhossein Vahidi
Hesam Asadollahzadeh
Navid Akhavan Attar
Marie Moullet
Kevin Ly
Xingyi Yang
Mohammad Lotfollahi
Venue: ICLR 2026 Poster (2026)
Keywords: MoE
Mixture of experts
sparsity
intrepretability
Primary area: foundation or frontier models, including LLMs
TL;DR: An intretable probablisitc sparse mixture of expert based Dirichlet distribution.

Abstract: Mixture-of-Experts (MoE) models have demonstrated exceptional performance in large-scale language models. Existing routers typically rely on non-differentiable Top-$k$+Softmax, limiting their performance and scalability. We argue that two distinct decisions, which experts to activate and how to distribute expert contributions among them, are conflated in standard Top-$k$+Softmax. We introduce Dirichlet-Routed MoE (DirMoE), a novel end-to-end differentiable routing mechanism built on a Dirichlet variational autoencoder framework. This design fundamentally disentangles the core routing problems: expert selection, modeled by a Bernoulli component, and expert contribution among chosen experts, handled by a Dirichlet component. The entire forward pass remains fully differentiable through the use of Gumbel-Sigmoid relaxation for the expert selection and implicit reparameterization for the Dirichlet distribution. Our training objective, a variational ELBO, includes a direct sparsity penalty that precisely controls the number of active experts in expectation, alongside a schedule for key hyperparameters that guides the model from an exploratory to a definitive routing state. Moreover, our DirMoE router matches or exceeds other methods while improving expert specialization.

Published as a conference paper at ICLR 2026
DIRMOE: DIRICHLET- ROUUTED MIXTURE OF EXPERTS
Amirhossein Vahidi
Hesam Asadollahzadeh 1.2.
Navid Akhavan Attar
Marie M
Kevin Ly"
Xingyi Yang"
Mohammad Lotfollahi'
Wellcome Sanger Institute, Wellcome Genome Campus, Cambridge, UK
School of Computing and Information Systems (CIS), Faculty of Engineering and IT (FEIT),
University of Melbourne, Australia
"The Hong Kong Polytechnic University, Hong Kong [hall, av13, ml 1 9 9 sanger. ac.uk
'qqual contribution
arXiv:2602.09001v1 [cvlc...GG [cs.LG] 9 Feb 2026
ABSTRACT
Mixture-of-Experts (MoE) models have demonstrated exceptional performance in large-scale language models. Existing routers typically rely on con-differnntiabee SoffkmSof limiting their performance and scalability. We argue that two dis- tinct decisions, which experts to activate and how to distribute expert contributions among them, are conflated in standard Hop-k++Sfttmas We introduce Dirichlet- Routed MoE (DirMoE), a novel end-to-end differentiable routing mechanism built Dirichlet variational autoencoder framework. This design fundamentally dis- entangles the core routing problems: expert selection, modeled by a Bernoulli component, and expert contribution among chosen experts, handled by a Dirichlet component. The entire forward pass remains fully differentiable through the use of Gumbel-Sigmoid relaxation for the expert selection and implicit reparameter- ization for the Dirichlet distribution. Our training objective, a variational ELBO, includes direct sparsity penalty that precisely controls the number of active ex- perts in expectation, alongside schedule for key hyperparameters that guides the model from an exploratory definitive routing state. Moreover, our DirMoE router matches or exceeds other methods while improving expert specialization.
INTRODUCTION
Modern sparse Mixture-of-Experree (MoE) layers scale capacity without proportional computation by routing each token to small subset of experts (Shazeer et al. 2017 Jiang et al. 20.24 DeepSSeek 2024 Du al., 2021). The main component of MoE is the router. The router is responsible for answering two questions: (i) which experts participate and (ii) their degree of contribution, deter- mined by the probability (mass) that is distributed among them. Furthermore, the learning process should be sparse, with only a fraction of active experts, in order to decrease computational complex- ity. One of the most prominent approaches to address this is using Top-k= A persistent challenge in standard of end-to-end gradients through the discrete selection step and its requirement for secondary ob- jective (e.g., temperature tuning, auxiliary losses, straight-through estimators) to encourage sparsity during training, which complicates stability and calibration. Beyond optimization, the single Soft- max entangles two decisions: expert selection and how much their contribution is, making expert usage hard to interpret. This entanglement ties load balancing to mixture calibration via a single temperature, relying on capacity constraints to curb overload, which causes an uneven distribution across experts. These limitations motivate a routing mechanism that (i) preserves full differentiabil- ity and (ii) offers explicit, interpretable control over expert selection and probability allocation on the simplex.
routing (Lepikhin et al., 2020 Fedus et al., 2022) is the lack
ssiignment-based routers (BASE Layers) reframe dispatch an assignment problem to guar- antee balanced loads (Lewis et al. 202I): this addresses balance but not mixture calibration or interpretability of how much mass the active set receives. Most recently, ReMoE replaces Top- k+Softmax with a fully continuous gate to restore end-to-end differentiability and improve scalabil- ity (Wang al. 2025 [Puigcerver al. 2024). While this directly tackles the gradient bottleneck,

---

Published as a conference paper at ICLR 2026
Token Embedding 1
x(x1)
Output
Experti
Router
Token Embedding XT
Expert k
z(XT)
Output T
10
Ohii (x)
dP(P) prtor
Expert
X
Encoder
Sparsity Knob
Contribution
1 (x)
191
0 ~ Dir (a 2)
Gumbel
(q)
(x)
r(x)
Decoder
X
Gating Logtt
Sigmoid
x(x)
posterior
(xx)
Expert
Expert Selection
Weight
Figure Illustration of DirMoE. Given a batch X' of input embeddings, two different heads Oni()), (x) learns the active and inactive per-token expert concentration, and ((x) learns the gating logits. The routing probabilities are the normalized product of a expert selection and Dirich- let probabilities 0 (expert contribution).
the resulting "soft" routing needs auxiliary loss to enforce sparsity. In practice, these can inject interference gradients, complicate tuning, and dampen expert specialization (Wang et al. (024). Moreover, when they are combined with sepacity-basel routing (token dropping), they may intro-
duce instability (Rajbhandari et al. 2022 Gale et al. 2023).
We introduce Dirichlet variational router that factorizes routing into (i) binary selection vector over experts (the set of active experts) and (ii) a Dirichlet distribution over mixture weights on the simplex (the distribution among the active experts). As shown in Figure the router produces laxed expert selection vector via Gumbel-Sigmoid reparameterization and the expert contribution via a Dirichlet; the final routing weights are the normalized Hadamard product of the two compo- nents (see Section 44. Furthermore, the router is fully differentiable with implicit reparameterization Maddison al. 2017 2017| Jang et al. 2017, Figurnov et al., 2018). A key benefit is calibration. Under Dirichlet prior. the total mass assigned to any subset of experts follows a Beta distribution. We exploit this to derive one-parameter (sparsity knob) control that sets the expected mass on the ac- tive set, decoupled from the size of active experts (k). Thus, & cleanly determines how many experts participate, while Beta-calibrated ratio of Dirichlet shape parameters determines how much they collectively contribute. This design retains the end-to-end gradient advantage and adds transparent sparsity knob to calibrate expert selection and expert contribution.
Our contributions are:
We introduce a novel probabilistic router to improve sparsity in MoE. DirMoE the first method using Dirichlet variational autoencoder producing sparse, interpretable routes the simplex for MoE.
We use Gumbel-sigmon for expert selection and implicit reparameterization for Dirichlet samples enable end-to-end gradients. This design maintains the MoE's full differentia- bility.
We introduce a coefficient (sparsity knob) that can control the sparsity the router, pro- viding practical calibration knob for specialization and sparsity.
We demonstrate the strong performance of our method across multiple datasets and tasks: DirMoE shows competitive zero-shot results and improves expert specialization.
2

---

Published as a conference paper at ICLR 2026
2 RELATED WORK
Mixture-of-Experr and Routing GShard and Switch established practical routing via auxiliary balancing and capacity control (Lepikhin al., 2020; Fedus et al. 2022); GLaM demonstrated trillion-parameter efficiency (Du et al. 2021). To avoid discrete Top-k, recent differentiable routers restore end-to-end gradients: Soft-MoE replaces hard selection with soft assignments (Puigcerver et al. 2024); ReMoE uses ReLU routing with explicit sparsity control and balancing (Wang et al.) 2025 and Lory scales fully differentiable MoE to pretraining via segment routing and similarity- based batching (Zhong et al.) 2024). To address the limitation of independent routing each layer, Layerwise Recurrent Router (RMoE) introduces a GRU-based mechanism to propagate routing states, enabling cross-layer dependencies for improved expert selection (Qiu et al. 2025).
Sparsity in MoE A central goal in MoE is to achieve controllable sparsity between different experts. Controllable sparsity has evolved from temperature and capacity heuristies and balanc- ing losses (Fedus et al.) 2022, Zoph et al. 2022) toward learned or auto-tuned activation budgets. Variable-k routing allows each token to choose its expert count: DynMoE jointly learns per-token and the global number of experts (Guo et al. 2025). while DA-MoE derives token importance from attention to allocate a variable number of experts (Akhavan Aghdam et al., 2024). In contrast, dy- namic depth allocates compute across layers rather than experts: Mixture-of of-Detssss (MoD) learns to skip or process tokens under a compute budget (Raposo et al. (024)), A-MoD routes using atten- tion maps without extra parameters (Gadhikar et 2024), and Router-Tuning (MindSkip) enables dynamic depth by training only lightweight router (He et al. 2024).
3
PRELIMINARY
3.1 SPIKE AND SLAB ROUTING ON THE SIMPLEX
We model MoE routing with a spike and slab prior (Mitchell & Beauchamp 1988: Ishwaran & Rao 2005) that separates which experts active (expert selection) from how probability mass is distributed among the active experts (expert contribution). Let E' be the total number of experts, z E {0, _ be binary selection mask and 0 E Æ-11 simplex vector. We factorize
E
0 0 =
Bernoulli(z) mx(x)) x Dir a (P) (z))
Expert
spike: Expert selection
where n, (z) (0, 1) are gate priors (optionally with Beta hyperpriors) and is the token embedding input to the MoE layer. The slab uses a two-level concentration:
( 22) = {== Oni + (1 - 2,) o1o)
Ohi > o1o > 0, A O.
(2)
Where Chii the per-expert concentration when an expert is active (2; R 1), Oio is the per-expert concentration when inactive (z; x 0); lower number reduces leakage into inactive experts. À is global scale of the concentration of active and inactive experts, it controls dispersion in the Dirichlet parameter. A larger tightens samples around the mean so the distribution becomes more uniform. A smaller À yields higher variance and more mass near simplex vertices, so the distribution is more spiky (effectively sparser draws).
3.2 SIMPSON INDEX AS A SPARSITY METRIC
In the context of probability distributions, the Simpson Index defined as the sum of the squared probabilities: H (p) Ee p.. This metric provides an intuitive measure of a distribution's 's con- centration: perfectly sparse distribution with all probability mass one outcome provides an upper bound of = 1, while a uniform distribution across all E outcomes yields a lower bound of H = 1/E. For any other distribution, the value lies between these two bounds. We use the Simpson index as a direct and interpretable measure of sparsity for our routing distributions, where a higher value indicates greater sparsity.
3

---

Published as a conference paper at ICLR 2026
4 METHOD
In the following, we introduce DirMoE, a fully differentiable MoE router that factors routing into two components: expert selection and expert contribution, and how they define the final expert weights. Then we introduce a new parameter \, the sparsity knob that can control sparsity.
4.1 PROBLEM SETUP
We consider a transformer decoder with sparse MoE layers. Given token embedding Rd let {E i= denote expert networks and let router produce logits {(x) E E We seek routing vector r(x) E A E-1 that satisfies: (i) only k < E experts are active per token; (ii) intrinsic activation sparsity for interpretability; and (iii) full differentiability so the router can learn task-dependent sparsity.
4.2 DIFFERENTIABLE ROUTER
The router has two components: (i) expert selection generated from logits ((x) E Re; combined below. The selection indicator is trained via binary Gumbel-Sigmole Jang et al. (2017) Maddison ct al. (017). Given router logits {(z), we sample an expert selection vector:
(x(x)) +
2
9i Logistic(O,)) 1),
(3)
Tz
where o is the sigmoid function, T3 is the temperature. This yields a soft mask (0, 1)
Tz 4 0, 2 becomes nearly binary.
(ii) Conditioned ž, we define Dirichlet variational posterior over latent simplex variable
4 xr 2) Dir(a ( (z, 2)),
with component-wise parameters (9) a; a (I, = ( (2, ami, (x) + (1 2i) (x)):
(
(4)
where ( > 0 scales the total concentration and (n), (1,,(x)) >> O may be learned functions
of (or shared across i) and the Dirichlet is sampled via normalized Gamma with implicit reparam- eterization (Figurnov et al. 2018).
For the prior, we use the same relaxed gates avoid mixing discrete and continuous objects and to keep a closed-form KL:
p(0) 2) (P (P) (2)),
(p) (2) xxx0 (x) 0hi (1 2,) ano)
(5)
with constants aiji > Onoo > 0 and scale (p) > 0. This is a smooth surrogate to spike and slab on faces of the simplex: ह 0, coordinate i driven toward (near-)zero probability; 2
joins the active slab.
We introduce the router normalized combination of two vectors (i) the expert selection, and (ii) the expert contribution 0 as follows:
r(c) normalize( o ((r) E A E-1 ~ Dir(a 2))
(6)
which gates the experts' outputs. Our goal is to have | lo s k and to place almost all of the total probability mass (4 on the active set S(x) = (x) = 1). Then the MoE output is
E
v(x) = Σ (x(x)) Ex(x)
This design maintains full differentiability throughout the entire forward pass: gradients flow through the Binary-Concre relaxation (Maddison et al., 2017; Jang et al. 2017) and through the Dirichlet distribution via implicit reparameterization (Figurnov al. 2018).
4

---

Published as a conference paper at ICLR 2026
4.3 TRAINING OBJECTIVE
We optimize a variational objective per token a::
LDiMMoR (1F) =
E (ilu)) log (x r(x)) + 80 ( ( D (x, 2)) Dir(a (P)
(Zeg)
reconstruction on r(z)
Dirichlet
Rpparsity (z)
(8)
sparsity
where By 2 O weights the slab KI,, B (Z). We optimize a 3-weighted variational
objective (Higgins et al. 2017] Alemi et al. 2017) (Eq. (8)). When 8g 1, the objective reduces to
the standard ELBO, The decoder is a single network Iv A R applied to the latent r. The details of the decoder and different loss terms are in Appendix A.2 following objective function:
So, we train the model with the
Lastal LIM + L
Sparsity regularization.
We encourage ks active experts in expectation per token using:
Repanity (2)
Aspursity EE= (x:) - k)
(10)
where Asparsity 2 0 controls sparsity strength. You can find the connection the KI. of the spike term in A.1 Figure shows the importance of Regulariza- tion loss on the desired sparsity (1 - k/E). Further
0.80
Aspunay 0.061
0.75
Aguusry m0
experimental results on the importance of Rsporsity Appendix
0.70
0.65
0.55
4.4 SCHEDULER
0.50
2500
5000
7500
10000
We discussed in Sec. that 0 Dir(a)) with two concentration levels: Oni for active experts and for inactive ones, scaled by ( > 0. Let S = {i 24 with |S| & denote the (idealized) active
Figure 2: Effect of sparsity regularization the sparsity, there less sparsity compared to the desired sparsity with lower Asparsity-
set Under Dirichlet, the sum over any subset is Beta-distributed; therefore
= Σ 0, Beta (Mtak (o))))
E[T] = kopi + (E k)oj0
kO
=: m.
EES
(11)
Solving for the ratio r Ohi/ /noo yields
m
E
(12)
Hence, for desired expected probability mass fraction m on the active experts, one can fix by equation 12 and then schedule the absolute scales separately
Temperature schedule. We use short schedule for the gate temperature move from ex- ploratory to decisive selection variables:
max(Fminn To p').
(13)
Large т. early encourages exploration; small 7. later yields near-binary ž.
We use this discrete view to calibrate the prior; in practice we train with relaxed gates ef. Sec. 4.2 Note that F[7] is independent of X (P)) changing (P)) alters dispersion (and thus sparsity tendencies) hut the prohability split between active and inactive when fixed.

---

Published as a conference paper at ICLR 2026
1.0
0.8
Strioon
0.6
0.7
2200
4000
5000 2kp
D00D
10000
2200
4002
6000
8000
20000
Siee
(a)
(b)
Figure 3: Effect of m and on (a) Simpson index and (b) sparsity.
À is calculated from m based on equation 16
perichlet-paramesr schedule. We choose r Obi/an from equation 12 and keep it fixed to preserve the target mean probability mass m on the active set. We then reduce the spike floor and tighten overall concentration during training:
lo
(t) max (Affoors Oio 7²)
(0)
(1) hi
r (0)
( ((x.0) n²
(14)
with decay factors 77 E (0, 1) and a small numerical floor Offorr for stability in digamma(I) evaluations.
DISCUSSION
5.1 CONTROL OF SPARSITY
We control sparsity with two orthogonal knobs: the mask temperature TN (exploration - decisive- ness) and the Dirichlet concentration scale X (dispersion on the active simplex face). Early explo- ration reduces rich-get-riche collapse without extra wadd-balancing losses; later, small TN yields near-binary expert selection vector and better specialization.
5.2 CALIBRATED SPARSITY VIA THE DIRICHLET CONCENTRATION
We introduce the Dirichlet concentration À as calibrated sparsity knob for the routing distribution in our MoE model. We use the Simpson index (p) E, p as our sparsity metric; a higher value of (p) indicates more concentrated (and thus sparser) distribution. You can find the role of À on the variance of the Dirichlet distribution (expert contribution) Appendix A.6
5.3 THEORETICAL JUSTIFICATION
We formally justify using as a sparsity knob through key theorem. The core idea is that the expected sparsity, as measured by the Simpson index, is monotonic function of the Dirichlet con- centration. We justify using the Dirichlet concentration as calibrated sparsity knob via the Simpson concentration Hp)) -E
Lemma 1 (Expected Simpson index under Dirichlet). Let p ~ Dir(A,3 with B E R S0 Σ 3ü and S2 EE= 88. Then
B
[[H(p) = AS2/B +1
XB +1
Proof in Appendix A.3
Theorem 1 (Monotone sparsity control by concentration). Fix B RE with E > 2 and let
~ Dir(A Then f()) (p)) is strictly decreasing in À > 0, with limits
E
lim f()) = x
and
lim f(A) = Σ A
m

---

Published as a conference paper at ICLR 2026
where m, B./B is the mean routing distribution. Proof is in Appendix A.4
Corollary 1 (Symmetric base). If 3 = 1, then m, = 1/E and Theorem
1 giiess f(X) strictly
A implying
decreasing from to L/E. The calibration in Lemma I reduces to E[H]
XE 1
1 h hE for target (1/E,1).
5.4 PRACTICAL CALIBRATION SPARSITY KNOB
Based on this theoretical foundation, we provide two practical ways to set the Dirichlet concentration achieve a desired level of sparsity.
(A) Symmetric Base Calibrator. For symmetric base where all experts are treated equally (3, 1), the expression simplifies, allowing us to directly map a target Simpson index h E (1/E, 1) to the required concentration A. This is useful for initial exploration where prior knowledge about expert specialization is assumed.
A 1-A h hE 1
(15)
model the total probability mass group of & experts with high shape parameters (ahi) and E - experts with low ones (Ata). The total probability mass the 'high' group follows Beta distribution. We set target specific variance, Car. for this mass, offering more nuanced control.
(B) Two-Group (Beta) Calibrator For method,
m2l (1)m)
(16)
s0hi + (E - -s)a
This approach is highly useful when we have a fixed number of active experts mind (e.g., 8 = k)) and want to control how concentrated the probability mass is that group. You can find the derivation for Equation equation 16 in Appendix A.5
Which À to tune? Routing depends the posterior draw, sparsity of is controlled by the posterior scale Ng (in a ()). The prior scale (mr+) shapes regularization via the KL. In practice, fix Ay via the calibrator (to hit target h) and mildly decay (p.()) encourage decisiveness. We
investigate the À and m that are calculated based on Equation equation 16 for variance 0.01. As m
decreases (less mean mass probabilities on the active set), the À from Equation equation [6 incressss
to keep the Beta distribution variance fixed; This preserves a similar level of sparsity (Fig. pp) while
adjusting expert contribution on the active experts (reflected by H(0) Fig. Bb).
IMPLEMENTATION
Architecture Our backbone is LLaMA (Touvron et al., 2023) with 185M parameters, configured
with 12 layers, RMSNorm (Zhang & Sennrich, 2019). SwiGLU activations (Shazeer) (020). and
rotary position embeddings (RoPE) (Su et al. 2021) and grouped-query attention (GQA) (Ainslie
et al. 2023) with 4 query groups. We replace the vanilla Top-k router with DirMoE, keeping
the same expert counts and per-expert dimensions. The training follows Algorithm implementation information is in Appendix C.I
Detailed

---

Published as a conference paper at ICLR 2026
Algorithm 1 DirMoE algorithm
Require: Batch { experts {E, router logits head ((x),
posterior base
heads
(a())),0n(())))
Require: Scales: posterior ((9), prior A temperature (8) target I
1: Update schedules: -( (
(Eqs.
13
14p
2: for each token embedding do
3: Relaxed gates via simbble-sigmod
(Equation
4: Posterior slab (99) (z, 2)
(Equation
5: Prior slab (P) (2.g)
(Equation
6: Sample 0 Dir( (9)
(implicit reparam.
7: Router r = normalize ( 0) A--1
(Equation
8: MoE output y Σ (x)
(Equation
9: end for
Objective: optimize the loss in equation 8
Training settings We train the model on The Pile dataset (Gao et al. 2021). We use Byte-Pair Encoding (BPE) tokenization (Sennrich et al. 2016) and train on ~30B tokens. The zero-shot experiments use 60k steps (covering all 30B tokens), and ablation experiments run for 10k steps. Our implementation builds on Megatron-LM (Shoeybi et al., 2019). Megatron's 's MoE stack sup- ports expert parallelism with all-to-all token dispatch and -eMP expert compute for ef- ficiency (NVIDIA 2025 Hejazi| 2024 Rajbhandari et al. 2022 Gale et al.) 2023). We use the
AdamW optimizer (oshchilov & Hutter 2019) with a learning rate schedule combining cosine an-
nealing (Loshchilov & Hutter] 2017) and linear warmup (Goyal et al. 2017). All experiments use
4-8 NVIDIA HT00 GPUs. Hyperparamtters are listed in Appendix C.3
EXPERIMENTS
7.1 SCALABILITY
In this section, we compare the scalability of the DirMoE with Vanilla MoE. In Table we show the computation cost in terms of throughput and iteration time; DirMoE has matching or exceeding the vanilla MoE baseline under compute parity. Because both systems utilize all-to-all token dispatch and grouped-GEMMM GEMM expert compute, the observed efficiency is primarily determined by the global tokens- ser-step and sequence length; our method introduces no additional bottlenecks (less than %). Figure Aeggonntdemosttttte that our method has reliable training and achieves the desired sparsity for different i and the number of experts. You can find further results on scalability Appendix B.4
Table 1: Training efficiency on LLaMA-1885 with E E-8, k=1
Method
lteration time (ms)
Throughput (TFLOP///PU))
Vanilla MoE (Switch)
431.5
138.2
DirMoE (ours)
437.3
137.2
7.2 ZERO-SHOT ACCURACY
We report the result of zero-shot accuracy in Table We evaluate our method in seven widely used benchmarks. ChChlleenee (ARC-C) and ARC-Fasy (ARC-E): grade-school science multiple- choice with adversarial hard and easier subsets (Clark et 2018); BoolQ: natural yes/no questions from real search queries (Clark et al. 2019); HellaSwag: common sense sentence completion with adversarially filtered decoys (Zellers et al. 2019); LAMBADA broad context next-word prediction requiring long range discourse understanding (Paperno et al. 2016); PIQA: physical common sense about everyday actions (Bisk et al. 2020); and RACE: reading comprehension from English ex- ams (Lai et al., 2017). DirMoE yields superior or similar results to other methods; on average, we outperform them.
8

---

Published as a conference paper at ICLR 2026
222
21.
13.0
15.0
30.0
02
72-
2000
2900
1000 Step
6092
19930
1000
4301
6000 20p
5000
1000
3000
2003
5000
4265
5003
5000
(a)
(b)
(c)
Figure 4: Effect of active experts (k) (a) LLM loss and (b) sparsity; effect of number of experts
(E) on (c) LLM loss.
Table 2: Zero-shot accuracy (%) downstream tasks (higher is better). 'deotess results reproduced by us.
Method
ARC-e
ARC-e
BoolQ
HellaSwag
LAMBADA
PIQA
RACE
Avg.
Hush (Roller 2021
19.28
45.45
54.95
69.68
31.44
63.06
27.66
38.79
Lory (Zhong et aB
0024
20.31
42.97
49.54
28.75
32.35
62.24
27.75
37.70
SparseMixer-2 010
19.80
46.72
45.96
30,24
34.12
29.00
38.39
EC (Expert Choice) 7/hn
2022
18.86
42.97
60.21
29.14
29.26
61.92
27.37
38.53
Switch MoE' (Fedus
2022
20.09
44.23
57.83
29.68
32.97
63,55
27.96
39.47
ReMoF Wang et al. 2125
20.22
46.68
54.16
30.26
35.94
63.55
29.38
40.03
DirMoE (ours)
20.57
46.20
61.52
29.93
36.44
63.75
29.52
41.13
7.3 EXPERT SPECIALIZATION
DirMoE controls routing sparsity without explicit load- load-balancing loss by using the sparsity knob (A) together with the mear-binau expert selection vector. Although this probabilistic allocation can yield modest load asymmetries, it encourages experts to concentrate distinct token subpopula- tions. As illustrated in Figure 5] our method has better expert specialization across different layers compared Vanilla MoE. The latter explicitly equalizes contributions, which stabilizes utilization but also homogenizes experts and weakens their semantic focus.
Domain specialization DirvAE
Boctod talcro hune
Books
0.-
OpenlWids
64-
Roads
Stack
Wiki
0.0
Domein specialization of MoL
Layer O
Layer 5
Leyer 11
Rozted talrro Iaci
0.8
Ancty
Books
666
Openi/ue
Recds
Stack
60
5
(b)
Figure 5: Domain specialization of (a) DirMoE and (b) Vanilla MoE in different layers and domains based average routed token. The grey dashed line indicates the uniform distribution.

---

Published as a conference paper at ICLR 2026
8 CONCLUSION AND FUTURE WORK
In this paper, we introduced DirMoE, a fully differentiable probabilistic router for Mixture-of- Experts that disentangles routing into (i) per-token expert selection and (ii) per-token probability allocation over experts (expert contribution). This white-box design gives explicit control over spar- sity, enables calibrated concentration via Dirichlet parameters, and maintains end-to-end differen- tiability. Empirically. DirMoE achieves strong few-shot performance, and we ablate its components to study calibration and sparsity control. In future works, we envision two promising directions. Robustness: The probabilistic formulation and explicit concentration control may improve model robustness. Interpretability: the disentangled routing variables naturally expose which experts are active and how probability mass is allocated. This could yield clearer, more monosemantic expert specializations.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Joshua Ainslic, James Lee- Lco-Th Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points. arXiv preprint arXiv:2305. 2023. /trl:tttp///rrxxvoor//////a
13245
Maryam Akhavan Aghdam, Hongpeng Jin, and Yanzhao Wu. DA-MoE: Towards dynamic expert allocation for mixture-of-exxprrss models. arXiv preprint arXiv:2409. 06690, 2024. URL https:
//arxiv. org/ abaaaaas2000 /2409 06669
Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy. Deep variational informa- tion bottleneck. In International Conference on Learning Representations (ICLR), 2017. URL
btttp::/aaxiv.
http://b///11... 00410
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical com- monsense in natural language. In AAAI Conference on Artificial Intelligence (AAAI), 2020. URL
https://arxiv. http://///aai....o//aaao//////// wttp//a////aooooooee/aoaa/sssaaa 1tt..////91111111...1.1.11111111
Christopher Clark, Kenton Lee, Ming-Wor Wei Chang, Tom Kwiatkowski. Michael Collins, and Kristina Exploring the surprising difficulty of natural yes/no questions. In North American Chapter of the Association for Computational Linguistics (NAACI.), 2019. URL
Toutanova. Boolq:
https: http:///laaantolo.orooo/o///////
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint
2018. hRttps:: http://axxi...rr/a///////// 114./15510330077.777 05457
DeepSeck-AL. Deepseek-Ve V2: strong, economical, and efficient mixture-of-cxperts, language model. arXiv preprint arXiv:2405.04433 04434, 2024. URL
04434
Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts. arXiv preprint arXiv:21 (tt::2112.0650 12.06905, 2021.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022.
Mikhail Figurnov, Shakir Mohamed, and Andriy Mnih. Implicit reparameterizztion gradients. In Advances in Neural Information Processing Systems (NeurIPS), 2018
Advait Gadhikar, Souptik Kumar Majumdar, Niclas Popp. Piyapat Saranrittichai, Martin Rapp. and Lukas Schott. Attention is all you need for mixture-of-deppths routing. arXiv preprint www.xxxvvvvrgabbbxoeee
arXiv:2412.20875. 2024. URL
Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. Megablocks: Efficient sparse training with mixture-oo expextssss.......oo In Proceedings of the 3rd Conference on Machine Learning and Systems (MLSys), 2023. http::tttp//aaat
org/ abs 221122211112 12211.1550 15841
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Horace Hoppe, Charles Foster, Jason Phang, Anish He, Saurav Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2 2021. URL httpt:tts//aaiiiia
org/ abs /211 00027
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Käyrälä, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training imagenet in hour. arXiv preprint arXiv:1706.00 026777 2017. URL
02677
Yongxin Guo, Zhenglin Cheng, Xiaoying Tang. Zhaopeng Tu, and Tao Lin. Dynamic mixture of ex- perts: An auto auto-tuning approach for efficient transformer models. In International Conference Learning Representations (ICLR), 2025. URL https:///////....t. sattps://openzeviie net net/forum?id
T2 9z2rFe

---

Published as a conference paper at ICLR 2026
Shwai He, Tao Ge, Guoheng Sun, Bowei Tian, Xiaoyang Wang, Ang Li, and Dong Yu. Router- tuning: A simple and effective approach for enabling dynamic-depth in transformers. arXiv preprint arXiv:2410.1318. 2024. BtRLhttps: http///axxvvv/.
bbbpg//bb/2220 NVIDIA Technical Blog. June 2024. URL htt.at/ps///xxv..rr/////////1/1
Babak Hejazi. Introducing grouped GEMM APIs in cublas and more performance updates.
http:///eeeeoooeee.viia..ooooe int rntucinn gww.pedrruuuddooooo gemm apis s- in
as aalasssaan--ae a pe
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. 3-vae: Learning basic visual concepts with a con- strained variational framework. In International Conference Learning Representations (ICLR), 2017.
Hemant Ishwaran and J. Sunil Rao. Spike and slab variable selection: Frequentist and bayesian strategies. Annals of Statistics, 2005.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In International Conference on Learning Representations (ICLR), 2017.
Albert Jiang, Yuchen Xu, Abhimanyu Dubey, and al. Mixtral of experts. arXiv preprint arXiv:2401.04088 2024. URLbtttss:/aaxxiv http://..rrraab//////a/// 04088
Bowen Jin, Tianle Luo, Yuchen Zhang, Haoyu Zeng, al. Sparse backepppppbbcppoppattiiooooooo for moe training. arXiv preprint arXiv:2310.00811,
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. Race: Large-scale reading comprehension dataset from examinations. In Empirical Methods in Natural Language Process- ing (EMNLP), 2017. BtLL:tts::/ http::tt/::/aallllee lattps::/acaaattllllgs..o wtwp:oyyy.gg//1111110011
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, Zhifeng Chen, and Jeff Dean. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 16668, 2020.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. arXiv preprint arXiv:2103. /6716, 2021. doi: 10. 48550/arXiiv2133.1677
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. In Interna- tional Conference on Learning Representations (ICLR), 2017.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations (ICLR), 2019.
Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. In International Conference on Learning Representations (ICLR), 2017.
Toby J. Mitchell and John J. Beauchamp. Bayesian variable selection in linear regression. Journal of the American Statistical Association, 1988.
NVIDIA. Mixture of experts (moe) - megatron-core developer guide. https://docs. /tepee/eeeeeeeeeeeerreeeeee/eeee /ttpee//t//pppggggd///o..... quide/noe. html 2025. Accessed 2025-09-22.
nvidia. wwl.da..commeeee com.meeattron-oom
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context. In Association for Computational Linguistics (ACL), 2016. URL httpps://alaaattlllogo/ http://awt.oooooooooooooeeeee.e wwtp/o6---0- 1144
Joan Puigcerver, Carlos Riquelme, Basil Mustafa, and Neil Houlsby. From sparse to soft mixtures of experts. In International Conference on Learning Representations (ICLR), 2024. URL https:
12

---

Published as a conference paper at ICLR 2026
Zihan Qiu, Zeyu Huang, Shuang Cheng, Yizhi Zhou, Zili Wang, Ivan Titov, and Jie Fu. Layerwise recurrent router for mixturee-f-expprr The Thirteenth International Conference on Learn- ing Representationss ICLR 2025, Singapore, April 24-28, 2025. OpenReview.net, 2025. URL
net fow.rrrummme/t htt.:/orrrmmmiieeee//FF0000
Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Deepspeed- moe: Advancing mixture-of-expert inference and training to power next-generation AI scale. In Proceedings of the 39th International Conference on Machine Learning (ICML), 2022. URL mlr. prwss v1 r2/ra jbhandariiaaa2 html
https://proceedi
David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro. ofxture-of-depths: Dynamically allocating compute in transformer-bas based lan- guage models. arXiv preprint arXiv:2404.02258, 2024. URL httth::t//aaatt httpss////xvv.orrr//// abs/
2404 02258
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason Weston, Hash layers for large sparse models. arXiv preprint arXiv:2106.04406 2021. htLL:tttp///acxx.ooogg/ 2206. 04426
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (ACL), pp. 1715-1725, 2016. hou_ltttpp///aaaaattoooooooooogp
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:20002.05202 2020. URL
https: http://aaavv..rr///////0222222 05202
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-expert layer. In International Conference on Learning Representations (ICLR), 2017.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-LM: Training multi billi-biliion parameter language models using model par- allelism. arXiv preprint (rXXiv/909.08000 2019.
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. arXiv preprint (tt::2104009644, 2021.
URL https::////aaa
org.
wtwggaabs.... wwwggbab/g/////ooo////////g /2104. 09864
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Razière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Ar- mand Joulin, Edouard Grave, and Guillaume Lample. LLaMA: Open and efficient foundation language models. arXiv preprint 232.:22032.197// 2023. http:/tttt:///aee wrw./vrr/aaaee abs/
2302. 13971
Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, and Damai Dai. Auxiliaryloss--ree load balancing strategy for mixture-of-experss arXiv preprint arXiv:2408. /5664, 2024. doi: 10.48550/ arXiv.2408.1 15664. URL /
15664
Ziteng Wang, Jianfei Chen, and Jun Zhu. Remoe: Fully differentiable mixture-of-expertt with ReLU routing. In International Conference on Learning Representations (ICLR), 2025.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? In Association for Computational Linguistics (ACL), 2019. URL htt//tttss//aaaaatoo.oo wttp://ooooyoor/111111177
Biao Zhang and Rico Sennrich. Root mean square layer normalization. 1rXiv:1910.07467, 2019.
arXiv preprint
Zexuan Zhong, Mengzhou Xia, Dangi Chen, and Mike Lewis. Lory: Fully differentiable mixture-of- experts for autoregressive language model pre-training. arXiv preprint arXiv:2405.03133 2024.
htl.https::/eet wwtp:/v..orrragee.ee abw/ 03133
13

---

Published as a conference paper at ICLR 2026
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang. Vincent Zhao, Andrew Dai, Zhifeng Chen, Quoc V. Le, and James Laudon. Mixture-of-eppere with ex- pert choice routing. In Advances in Neural Information Processing Systems (NeurIPS), 2022. URL. nget:: neurips http://seeccpppppprreeeeepepp/pp http:///ce22222233333000022226ee
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. ST-MoE: Designing stable and transferable sparse expert models. arXiv
preprint arXiv:220.08900 2022, doi: 1.485550/rXii.22222.000 URL, (tRl/ttps:: htttt://aaxxxii
org/abs/ 2202. 08906
14

---

Published as a conference paper at ICLR 2026
A
ADDITIONAL DETAILS FOR METHOD AND PROOF
A. 1 CONNECTION OF RAPARSITY TO TO THE SPIKE KL.
Let the variational spike be a mean-field Bernoulli q(z a)) = II Bern(2; px())) with pı(x)
(l((z)), and consider the spike prior Pa (z) II Bern(z; T with for (or the fixed-size
microcanonical prior uniform on {z: E, = k):)). The exact spike KL is
E
E
DKL(9(z| x) (z)
p: log Pi + (1 - ps) log 1-
Ber (pi
second-order expansion of the Bernoulli KL around pt gives
(pr )²²
(per(p)
+ (pp: o (1
2T(11 - π)
Therefore, locally,
1
DKL(allP))
Pi x)2
2T - T
i=1
By E,(, - TT2 2 Ei(p: - -x)) 2
= 1
(Zi Pi - k)) so we obtain the
E
local lower bound
Pi(x)) ( -
k)² -k)²
Dku(gllP*)
(17)
2E x( -k-
Hence the loccll sparsityinality penalty is a smooth surrogate for the spike KL:
Asparsity (
(()) -k) 2
Rsparsity (x)
with Asparsity R3
[1-(1=x)
matches the local curvature of 92 DKI. when replace the (intractable in backprop) Σ, (x)) by its low-variance stochastic estimator Σ, (x).
Remarks. Using 2 in Raparsiy gives [E,,, & k) Var Σ also discourages high-variance variance gates; as Ts J O, this gap shrinks. If desired, a deterministic variant Rsparsity (I) Asparsity Pi(x)- k:) targets only the bias term.
P: (E, so the penalty
A.2 OBJECTIVE FUNCTION.
The decoder is a single network Jv A -Rd applied to the latent
f(x) = (r()))) pe(x1 (r) = (r) o²T).
so, the reconstruction term reduces scaled MSE
- log po(1 (x(z)) = 2 I| (r()))) 2 log (ogo)
and we drop the constant and often set a²=1 in practice. With the posterior (0 Dir( (x, Z)) and the slab prior p(0 ž) Dir( (p) (Zug)), the KL term inside equation
xx 2)
is
DKL Dir (x,, 2)) (ir(a = log r
Σ (4)
Σ log (a) - log r
Σ
(p)
a
Σ log Ma (p)
Σ
(af
(
( (}(aa() (1)) <<
(18)
with ab the digamma function and where
(4) (x, 2) = ( ( (x)) + (1-zx) (1))(())
(p) (Zm)) = x1x.0 xlx.0( oni + +(1 (1- on
a
15

---

Published as a conference paper at ICLR 2026
A.3 EXPECTED SIMPSON INDEX UNDER DIRICHLET PROOF Here, we provide proof for Lemma
Proof sketch. For a Dirichlet with parameters Elp yields:
a(a+tll eg(0o+1)
with = E. 0g. Summing
B[H]-
Eola+ta)) Setting oj A3; gives B[H] (a(
S [XB](X++]
AS++B B(B++I)
A8//B41
XB+1
A.4 MONOTONE SPARSITY CONTROL BY CONCENTRATION PROOF
Here, we provide proof for Lemma
aA +
Proof. By Lemma f(A)
bā - with = S2/B and B. Differentiating gives f'(
a b
Since 3 > 0 and E' 2 2, we have B = Σ, f + 2 2 3:3, E, 3 S2, hence
(b) + 1)2"
a ST-B²
B < 0, so ('( < for all À > The limits follow directly from the rational form:
s(0+) = (0 + 1)/(0 + 1) =1 and [(00) = a/b = (BB) = S2/ B
Σ, m
A.5 TWO-GROUP BETA CALIBRATOR DERIVATION
Here, we provide details on how to derive Equation equation 16 Let be the active set with [S] = (for our setup s=k). With the two-level Dirichlet Dir X 3), where B3 Ohs for E and B an otherwise, the total mass on the active group
T = Σ 0,
ES
is Beta-distributed by equation
I1
T Beta(a, a, b):
As
s) Ogo b b= X s)
Hence
S
E[T]
m,
a+1
8 Ohi + (E - 8l 0/
ab
and using Var(Beta(a, 6))
(a + ()+
Var(T) m(1 mm
m(1 m)
C Qni + (E 8) 0%
b+b
XC+1
Solving for at a target variance thar E (0, m(1 m)) yields
m (1 - m)
A
Uuar
(19)
8 S + (E s) Olo
A.6 VARIANCE AND THE ROLE OF X.
Let ag = a() For Dirichlet,
Var(O.) a, (p) (a0 (p) 300 + 1)
Scaling all concentrations by À scales ag by X, leaves means unchanged, and shrinks variances roughly as ()))))). Hence, decreasing À pushes draws toward vertices (sparser allocations) without altering the mean split induced by the base ratios.
16

---

Published as a conference paper at ICLR 2026
B EXPERIMENTS RESULTS
B. EFFECT OF Asparsity
We further investigate the effect of Repersity LLM loss, Simpson index and in Figure
22.3
+0000 +0
Agrrs 2.000
Aavery --3 931
17.3
53
2000
2006
7500
10030
7508
1900
1931
19900
7500
1000 7900 10003
(a)
(h)
(c)
Paul LM Lom Debllkee
2.1000
2,704
2313
-78
(d)
Figure 6: Effect of Aspareity on (a) LLM loss and (b) Simpson index, and (c) (d) heatmap of Loss.
We further investigate the effect of m and Asparsity on LLM loss in Fig. Intuitively, increasing concentrates more probability mass the selected experts. If the sparsity weight Aspesity is set to zero, the router can activate more than & experts on difficult tokens. This often lowers the LM loss (more capacity per token) but reduces efficiency because the effective sparsity decreases. In practice, we use a small Aspariity > 0 so the model benefits from early accuracy gains while converging back toward the desired & for efficiency.
B.2 STABILITY ANALYSIS OF ACTIVE EXPERTS
Here, we demonstrate the average number of active experts, the maximum number of active experts, and the gradient norm throughout our training and in Figure 7
(b)
(a)
Figure 7: (a) average and maximum number of active experts (b) gradient norm.
B.3 ABLATION OF VAE LOSS
Here, we further investigate the importance of the reconstruction loss and VAE loss on LLM loss, 09 and E in Figure 8
17

---

Published as a conference paper at ICLR 2026
(a)
(b)
(c)
Figure 8: Effect of reconstruction loss on (a) LLM loss and (b) ε, and (c) KL loss effect on a².
B.4 SCALABILITY RESULTS
We explored the effect of desired active experts & and Simpson index in Figure
1.0
0.35
0.8
Simpsentndks
0.30
0.
0.15
0.2
10000
10000
Stop
Siep
(a)
(b)
Figure Effect of active experts (I)) on (a) and (b) Simpson index.
C DETAILS ON TRAINING
C.I IMPLEMENTATION DETAILS
Capacity and kernel considerations. Near-binary gates and sharp allocations increase per-step variance in expert loads; if the stack drops overflow tokens (capacity-based routing), this in- troduce training noise and dead-expert effects. Two practical mitigations are standard: dropless and block-sparse kernels (c.g., MegaBlocks), which eliminate token drops by packing variable ex- pert batches into high-utilization grouped GEMMs (Gale et al. (023), or using sufficiently high capacity factor (CF) in standard MoE stacks to reduce drops at the cost of extra padding or com- pute (Rajbhandari et al.) 2022). Both approaches stabilize optimization while preserving the benefits of sparsity control from т. and A. We used the first approach.
C.2 PRACTICAL IMPLEMENTATION NOTES
Systems. Our design keeps training fully differentiable and avoids balancing losses; scale, deploy with dropless kernels (Gale et al. 2023) or adequate capacity in expert- pprelllarrlllll engines (Rajbhandari
al. 2022).
C.3 HYPERPARAMETERS
You find Hyperparameters for training in Table and hyperparameters for router in Table
18

---

Published as a conference paper at ICLR 2026
Table 3: Optimization Hyperparameters
Component
Parameter
Default Value
Optimizer
AdamW
Learning rate
1 X 10
Weight decay
1 x 10-2
Adam betas eps
(0.9, 0.995) /1 x 10-8
LR schedule
Warmup 1% steps Cosine (min Ir = 0.1 base)
Optimizer
Global grad clip
1.0
Precision
bf16 (router logits/Dirichlet in fp32)
Micro-batch size
32
Batch
Grad accumulation
512
Seed
1234
Table Routing Hyperparameters
Component
Parameter
Default Value
Experts E' target actives k
E=8, k=1
Gate temperature тı schedule equation 13
2.0 - 0.3 (cosine)
Gates (spike)
Gate bias init (per logit)
To-logit(k/ E)
Per-token centering
On t f (
z-threshold for routing map (metrics)
0.125
Posterior scale ( (Eq. (16))
( 20
Prior mean mass on actives equation
m=0.9
Ratio r' Qt equation 12
r=39.7 (for --,,
Slab (Dirichlet)
Prior 0/1 / floor equation
14
0.005
Prior O'ni = (0)
1.99
Prior scale (P)) schedule equation 14
0.5 - 0.3 (slow exp/cosine)
KL weight 3g Recon variance o² (Sec. 4.3
1 x 10-2
1.0
Leak before normalize Epsilon clamp Expected-k penalty equation 10
S. =1 x 10-3 Off
Router output
Aspariity 0.01
D AI USAGE CLARIFICATION
We used Al-assisted tools solely for editing and syntax lookups. All ideas, methods, experiments, analyses, and conclusions were conceived and written by the authors. Every Al suggestion was manually reviewed and edited before applying. No text, figures, tables, code, or data were generated or analyzed autonomously by AI. The authors retain full responsibility for the content.
19


Reviews:
[{"review_id": "mAUsm35Scj", "reviewer": "Reviewer_qUTF", "summary": "The paper proposes DirMoE, a probabilistic, fully-differentiable MoE router that explicitly separates (i) expert selection (a Bernoulli “spike” trained via Gumbel-Sigmoid) from (ii) expert contribution (a Dirichlet “slab” sampled with implicit reparameterization). Routing weights are the normalized Hadamard product of the relaxed mask and a Dirichlet draw. Training uses a β-VAE-style ELBO with a sparsity penalty to target ~k active experts; a “sparsity knob” schedules Dirichlet concentration and gate temperature, with a Beta/Simpson-index calibration giving interpretable control over mass on the active set. Experiments (LLaMA-185M, ~30B tokens) show similar training efficiency to Switch and zero-shot gains over several baselines.", "questions": "- With the quadratic penalty on $\\sum \\tilde z - k$, how often do you deviate from exactly k actives? Does this variability harm batching efficiency or increase variance in gradient estimates?\n- For the throughput numbers in Table 1, specify the full systems setup: world size and data/tensor/pipeline/expert parallel configurations (plus capacity factor and drop vs. dropless). Under expert parallelism, does DirMoE exhibit poorer balance and lower throughput?\n- The domain specialization in Figure 5 is interesting. Please detail the protocol (corpora, measurement window...) and explain the observed patterns—for example, why some experts are entirely inactive on certain corpora? (dead experts vs. specialization)\n- Theoretically, Simpson-index control is monotone in $λ$. Empirically, how robust are downstream metrics to mis-set $m/λ$? Please provide heatmaps of $(m, λ)$ vs. zero-shot accuracy and utilization.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "-  **Interpretable sparsity control.** The Beta/Dirichlet calibration + Simpson-index theory offers a principled knob for expected mass on the active set and dispersion, beyond temperature heuristics.\n- **Light systems delta.** Reported iteration time/throughput match Switch under compute parity; the router adds negligible overhead.\n- **Empirical improvements.** On 7 zero-shot tasks, DirMoE slightly outperforms Switch/ReMoE/SparseMixer on average, and shows stronger expert specialization visualizations.", "weaknesses": "- **Load-balancing risk.** The router relies on near-binary masks and dispersion calibration without an explicit balancing mechanism; the paper itself notes potential utilization skew.\n- **Limited scale & evaluation breadth.** All results are on a ~185M-param LLaMA with ~30B tokens; there is no evidence at larger scale/more experts or on standard reasoning/coding suites (e.g., MMLU, HumanEval/HellaSwag). External validity for modern LLMs remains unclear.\n- **Questionable reuse of prior results.** Several zero-shot numbers appear copied from ReMoE rather than reproduced under your setup, making them non-comparable and weakening the empirical claim. Please re-run these baselines and report variance.\n- **Added complexity with unclear necessity.** The router introduces multiple moving parts (near-binary masks, dispersion calibration, auxiliary objectives/schedules) without a clear motivation for each. A component-wise ablation is needed to establish necessity versus incidental regularization.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jQ0z1XGOba", "reviewer": "Reviewer_xEic", "summary": "This paper introduces a novel routing mechanism for Mixture-of-Experts (MoE) models called DirMoE. It aims to solve two key problems with standard Top-k+Softmax routing: the non-differentiability of its Top-k operation and its entanglement of two distinct decisions (which experts to select vs. how to weight their contributions). The proposed method explicitly disentangles the routing process into:\n\n1. **Expert Selection:** A Bernoulli distribution determines which experts to activate for a given token.\n    \n2. **Expert Contribution:** A Dirichlet distribution determines the probability weights assigned to the *active* experts.\n    \n\nThe final routing vector is a normalized product of these two components. This entire mechanism is end-to-end differentiable, using Gumbel-Sigmoid for selection and implicit reparameterization for the Dirichlet sampling. The router is trained using a variational ELBO objective, which includes a reconstruction loss, a KL divergence, and a direct sparsity penalty loss. The authors also provide a theoretical justification for a \"calibrated sparsity,\" linking the Dirichlet concentration parameter $\\lambda$ to the expected expert contribution distribution.\n\nEmpirical results on a 185M parameter model show that DirMoE achieves better performance against baselines like Switch MoE and ReMoE. Critically, it also demonstrates significantly improved expert specialization at this scale, while introducing negligible computational overhead.", "questions": "1. **Training Objective:** Could the authors please clarify the full training objective for the entire model?\n    \n2. **Justification for VAE:** What is the intuition behind forcing the routing vector $r(x)$ to be able to reconstruct the token embedding $x$? Does this encourage the router to learn a richer representation, or is it primarily a regularization?\n    \n3. **Ablation Study:** What is the performance of a model that removes the VAE objective (reconstruction loss and slab KL) and is trained *only* with the main LM loss and the $\\mathcal{R}_{sparsity}$ penalty from Eq. 9? This would help isolate the contribution of the VAE framework.\n    \n4. **Hyperparameter Sensitivity:** The method introduces many new hyperparameters (Table 4). How sensitive is the model's final performance and specialization to these settings (e.g., the KL weight $\\beta_{\\theta}$, the posterior scale $\\lambda^{(q)}$, and the prior mean mass $m$)?\n    \n5. **Scalability (optional):** Do the authors have any insights or preliminary results on how DirMoE performs at larger scales (e.g., >1B parameters, more training tokens)?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. **Novelty:** The core idea of disentangling expert selection (Bernoulli/Gumbel) from expert contribution (Dirichlet) is well-motivated and novel. Framing this as a probabilistic spike-and-slab model within a VAE framework is a novel approach to MoE routing.\n    \n2. **Controllable Sparsity:** I really like the theoretical analysis introduced in Section 5, which connects the Dirichlet concentration parameter $\\lambda$ to the expected contribution sparsity. It provides a principled calibrated knob approach for controlling the *contribution sparsity* (i.e., the dispersion of weights among *active* experts), which is a more nuanced control than what is offered by standard temperature scaling in Softmax routing.\n    \n3. **Competitive Empirical Performance and improved specialization (at small scale):** The method achieves results that are competitive with or slightly better than baselines on the 185M model benchmark (Table 2). This demonstrates the viability of the approach in a small-scale setting. The method achieves much clearer expert specialization than the baseline (vividly shown in Figure 5).\n    \n4. **Computational Efficiency:** Despite the apparent complexity of the VAE framework, the authors demonstrate that its computational overhead is negligible (<1% slowdown in iteration time compared to a vanilla MoE, as shown in Table 1). This makes it more practical.", "weaknesses": "1. **Clarity of the Training Objective:** The paper's primary weakness is its lack of clarity regarding the overall training objective. The VAE objective (Eq. 8) is presented, but it's not explicitly stated how this loss is combined with the main LM loss. This is a critical detail for reproducibility and understanding.\n    \n2. **Justification of the VAE Objective:** The VAE's reconstruction task—reconstructing the token embedding $x$ from the routing vector $r(x)$—is non-obvious. The authors could provides more intuition and ablation about this design choice. This auxiliary reconstruction task seems not well-justified.\n    \n3. **Model Complexity:** While computationally fast, the DirMoE router is significantly more complex to implement and tune than Top-k routing. It introduces a VAE (encoder, decoder, KL loss) *for every MoE layer*, along with a host of new hyperparameters and schedules (as seen in Table 4). This complexity could be a barrier to adopt this method in large-scale MoE training.\n    \n4. **Limited Ablation and Comparison to Simpler Alternatives:** Following the previous two weaknesses, **my greatest concern** for this submission is the lack of the comprehensive ablations. For example, the VAE framework is a major part of the contribution, but the router is perfectly trainable without it (using only the main LM loss and the $\\mathcal{R}_{sparsity}$ penalty). It's plausible that the method also achieves good performance without the reconstruction loss. The paper lacks a crucial ablation study that removes the VAE objective to prove that this added complexity is necessary. Furthermore, the paper fails to compare against other, potentially simpler, methods that also aim to improve router training. For example, recent work like RMoE [1] has also shown that adding extra gradient paths and information sharing via layer-wise recurrence can be effective, suggesting that DirMoE's highly complex VAE is not the only (or necessarily best) way to solve this problem (though RMoE does not offer the controllability of DirMoE, but it is still unclear whether such controllability will improve or hurt the final performance when scaled up).\n    \n5. **Limited Scale of Experiments:** The experiments are conducted on a 185M parameter model with 8 experts. This is a significant weakness that *directly undermines* the claims of \"strong performance\" and \"improved specialization.\" The primary benefits of MoE are at much larger scales. Crucially, the paper's visual argument for specialization (Figure 5) relies on a small-scale baseline that appears to fail to specialize. Recent work[2] has demonstrated that standard Top-k+Softmax routers *can* also achieve strong specialization when properly trained at scale with appropriate load balancing implementation. This suggests the paper's central claim of superior specialization is not convincingly proven against a strong, scaled baseline. The stability and effectiveness of this complex VAE router at such scales also remain unproven.\n\n[1] https://arxiv.org/abs/2408.06793\n[2] https://arxiv.org/abs/2501.11873", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5i7hfnJt4k", "reviewer": "Reviewer_cmFj", "summary": "This paper introduces Dirichlet-Routed Mixture of Experts (DirMoE), a novel, end-to-end differentiable routing mechanism designed to address the limitations of existing methods like Top-k+Softmax, which conflate expert selection and contribution. DirMoE models routing using a Dirichlet variational autoencoder framework that fundamentally disentangles these two decisions: (i) Expert Selection ($z$) is handled by a Bernoulli component (relaxed via Gumbel-Sigmoid), and (ii) Expert Contribution ($\\theta$) among chosen experts is handled by a Dirichlet component. The training objective is a variational ELBO, which includes a direct sparsity penalty ($R_{sparsity}$) controlling the expected number of active experts $k$. A core theoretical contribution is the introduction of the sparsity knob ($\\lambda$), derived from the Dirichlet concentration, which provides explicit and calibrated control over sparsity, decoupled from the active set size $k$. The authors demonstrate that DirMoE achieves competitive zero-shot performance and improves expert specialization on a LLaMA-185M backbone.", "questions": "**Complexity Management**: The proposed method relies heavily on four distinct scheduled parameters ($\\tau_z$, $\\alpha_{lo}$, $\\alpha_{hi}$, and $\\lambda^{(p)}$). Could the authors elaborate on the practical sensitivity of DirMoE to these schedules? Is there a risk that poor schedule tuning leads to instability or poor specialization?\n\n**Efficiency vs. Other Differentiable Routers**: Please provide a comparison of training iteration time and throughput against other recent fully differentiable routers, such as ReMoE or Lory, ideally using the same LLaMA-185M setup or larger models, to fully contextualize the runtime overhead of the Dirichlet VAE structure and implicit reparameterization gradients compared to alternative continuous gating functions.\n\n**Results on O(10B) models**: It would be good to see the method is effective on larger model sizes in the order 10B. Can the authors show zero-shot results at this scale?\n\n**Role of $\\lambda_{sparsity}$**: Given that the Dirichlet concentration $\\lambda$ is theoretically presented as the primary sparsity knob for contribution, what is the minimum necessary role of the explicit cardinality penalty $R_{sparsity}$ (Eq. 9) in achieving a target expected $k$ without compromising model stability or specialization? Does the model fail to converge to sparse selections without $R_{sparsity}$? Figure 2 only shows that lower $\\lambda_{sparsity}$ yields less sparsity compared to desired $1-k/E$.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "**Originality**: The idea of factorizing routing using a spike-and-slab prior implemented via Gumbel-Sigmoid (selection) and Dirichlet (contribution) is highly original. \n\n**Quality**: The use of Gumbel-sigmoid and implicit reparameterization ensures the entire forward pass remains fully differentiable, avoiding the gradient bottlenecks of standard Top-k routing. This white-box design provides superior interpretability by explicitly controlling how many experts are active ($k$) and how concentrated their probability mass is ($\\lambda$). Empirically, DirMoE demonstrates better expert specialization across different layers compared to Vanilla MoE, attributed to its mechanism which discourages homogenization often caused by explicit load-balancing losses.\n\n**Clarity**: The paper is generally clear, well-structured, and effectively contextualizes the work relative to prior art in Top-k routing (Switch, GShard) and continuous routers (ReMoE, Soft-MoE). Figure 1 concisely illustrates the complex flow of the DirMoE router. The mathematical formulation is precise, particularly concerning the spike and slab prior and the training objective. A few inconsistencies in variable naming (listed in the Typos section) slightly detract from perfect clarity, but overall, the concepts are presented insightfully.\n\n**Significance**: Given the high importance of MoE models in industry, this is a very nice contribution to the sparse MoE literature.\nAlthough there have been fully differentiable routers proposed in literature e.g., DSelect-k [Hazimeh et al. (2021)], COMET (Ibrahim et al. (2023)], these routers do not support sparse training. DirMoE appears to be an interesting routing approach to hold end-to-end differentiability. Although the complexity of the method (in terms of tuning) makes it unlikely to be adopted by the industry to have a significant impact. \n\nSome related references to include in related work on differentiable routing:\n- Hussein Hazimeh et al. DSelect-k: differentiable selection in the mixture of experts with applications to multi-task learning. In NeurIPS'21.\n- Shibal Ibrahim et al. 2023. COMET: Learning Cardinality Constrained Mixture of Experts with Trees and Local Search. In KDD'23.\n- Michael E. Sander et al. Fast, differentiable and sparse top-k: a convex analysis perspective. In ICML'23.", "weaknesses": "**Architectural and Optimization Complexity**: DirMoE relies on a complex stack of techniques: Gumbel-Sigmoid relaxation, implicit reparameterization for Dirichlet samples, a full VAE objective, and multiple scheduled hyperparameters ($\\tau_z$, $\\alpha_{lo}$, $\\lambda^{(p,t)}$). While mathematically elegant, this complexity may lead to significant tuning overhead compared to simpler, fully continuous approaches like ReMoE, which should be explicitly discussed and benchmarked for tuning difficulty.\n\n**Persistent Need for Explicit Cardinality Penalty**: Despite the principled probabilistic design intended to control sparsity, the method still requires a direct sparsity penalty ($R_{sparsity}$) (Eq. 9) to encourage $\\sum_{i=1}^E \\tilde{z}_i(x) \\approx k$ in expectation. Although the authors justify this term as a smooth surrogate, its necessity slightly dampens the claim that the VAE framework inherently solves the controllable sparsity problem entirely through parameter calibration (like $\\lambda$).\n\n**Limited Scalability Comparison**: The scalability evaluation (Table 1) only compares DirMoE against Vanilla MoE (Switch) on a small LLaMA-185M model. Given that other recent differentiable methods (ReMoE, Soft-MoE, Lory) also address the differentiability bottleneck, a comprehensive efficiency comparison showing throughput, memory, and latency overheads against these direct architectural competitors is missing and necessary to validate the robustness of the proposed router at industrial scales. \n\n**Typos**:\n- In Line 30, the citation \"Du & et al., 2022\" (and in references) should be standardized, to \"Du et al., 2022.\"\n- Line 216, delete \"either\".\n- Reference to Eq 8 in line 10 of algorithm 1 is missing.\n- I see a number of places of the format “ Eq. equation x”. Either use equation x or Eq. x.\n- In Section 4.3 and Appendix A.1, there is inconsistent usage of the regularization term's coefficient: $\\lambda_{sparsity}$ in Equation (9) and the caption of Figure 2, but $\\lambda_{card}$ in the derivation in Appendix A.1.\n- In Figure 6 captions: $\\lambda_{sparsity}$ is consistently written as $\\lambda_{sparse}$ in the legend keys in parts (a), (b), and (c). Please make it consistent.\n- In Appendix A.2, the parameters for the prior slab (Eq. 5 revisited in Eq. 17) uses tildes on $\\alpha$'s. The Appendix uses bars on $\\alpha$'s, suggesting slight notational inconsistency between the main text and appendix. Please double check if this is not a typo.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)