---
title: EMA Routing for Mixture-of-Experts
url: https://www.emergentmind.com/topics/ema-exponential-moving-average-routing
type: topic
---

# EMA Routing for Mixture-of-Experts

Exponential Moving Average (EMA) Routing, specifically formulated as Expert Threshold (ET) routing, is a population-level sparse routing paradigm for mixture-of-experts (MoE) models. It is designed to enable dynamic computation allocation and expert load balancing without requiring explicit auxiliary losses. ET routing maintains per-expert thresholds using an exponential moving average over batch-statistics derived from router logits, facilitating a fully causal, batch-independent, and dynamically adaptive routing mechanism. This approach is targeted at scalable autoregressive language modeling and addresses limitations present in established routing methods such as Token-Choice MoE (TC-MoE) and Expert-Choice MoE [2603.11535].

## 1. EMA Threshold Computation and Maintenance

For each expert $e$, a threshold $T_e^{(t)}$ is maintained and updated at every training step $t$ using an exponential moving average (EMA). Given a batch of $N$ tokens and $E$ experts, the target per-expert token allocation is $k=N/E$. For expert $e$, the $k$-th largest router logit among all tokens (denoted $f_e^{(t)}$) is used as the statistic for threshold adjustment:
\[
T_e^{(t)} = \beta\,T_e^{(t-1)} + (1-\beta)\,f_e^{(t)}
\]
where $\beta \in [0,1)$ denotes the EMA decay factor, typically set to $0.999$. Here, $f_e^{(t)}$ is the $k$th-order statistic of $\{r_{t,i}\}_{i=1}^N$, with $r_{t,i}$ as the router logit assigned by expert $e$ to token $i$ at step $t$.

The initial threshold $T_e^{(0)}$ can be set to zero or to the appropriate order statistic from early batches; in practice, a short warm-up phase using alternative routing (e.g., Expert-Choice) is used to stabilize these threshold values before standard ET routing commences.

## 2. Token Routing Decision Rule

Each token $t$ with hidden state $x_t \in \mathbb{R}^d$ is processed by a router network producing logits $r_{t,e} = (W_r\,x_t)_e$. The routing assignment is determined by comparing each $r_{t,e}$ to threshold $T_e^{(t)}$:
\[
z_{t,e} = \mathbf{1}\left[r_{t,e} > T_e^{(t)}\right] =
\begin{cases}
1,& r_{t,e} > T_e^{(t)} \\
0,& \text{otherwise}
\end{cases}
\]
Hence, a token may be routed to any subset of experts, including none, one, or all $E$, enabling natural dynamic compute scaling per token. At inference, $T_e$ is no longer updated; routing is fully causal and depends only on each token’s representation and the frozen, EMA-learned thresholds.

## 3. Dynamic Computation Allocation and Load Balancing

The ET routing paradigm eliminates the need for explicit auxiliary loss terms for expert balancing, contrasting with the traditional TC-MoE approach which fixes the number of experts per token ($G$) and applies auxiliary objectives to maintain load balance. In ET routing, each expert $e$ raises or lowers its threshold $T_e$ to target a fixed expected token share, enforcing
\[
\mathbb{E}_t\left[z_{t,e}\right] = \Pr_t[r_{t,e} > T_e] \approx \frac{1}{E}
\]
in expectation over tokens. This self-regulating property results in emergent, population-level balance. Empirically, the per-batch load for each expert fluctuates around $N/E$; a capacity factor $C$ (e.g., $0.5$) clips per-expert loads to $[(1-C)N/E, (1+C)N/E]$ during training to mitigate out-of-memory risk, but this is rarely invoked post-warm-up.

## 4. Causality, Batch Independence, and Inference Consistency

By updating thresholds via EMA and using independent token-to-expert comparisons, ET routing achieves a fully causal workflow: each token’s routing is computed independently, and at inference the procedure requires no batch aggregation or per-step recalibration. There is no discrepancy between training and inference routing behavior, as the same per-expert thresholds are applied throughout.

Compared to batch-dependent methods like Expert-Choice (batch top-$k$), where threshold variance scales as $O(1/\sqrt{N})$, ET routing fixes thresholds and tolerates $O(\sqrt{N})$ fluctuations in per-batch expert load, trading hardware resource predictability for consistent, batch-size-independent deployment and complete train-inference alignment.

## 5. Algorithmic Implementation

The core ET routing procedure consists of the following key steps:

- For each token $t$ and each expert $e$, compute $z_{t,e} \leftarrow 1$ if $r_{t,e} > T_e$, else $0$.
- At training, for each expert $e$:
    - Calculate $f_e$ as the $k$-th largest of the current batch's $\{r_{t,e}\}$.
    - Update $T_e \leftarrow \beta \, T_e + (1-\beta) f_e$.
- At inference, $T_e$ remains unchanged; routing decisions use the fixed EMA-learned threshold.
- A capacity factor $C$ optionally clips per-batch expert assignment counts to a safe interval during training.

This algorithm operates on the assignment matrix $z \in \{0,1\}^{N \times E}$ and maintains updated per-expert thresholds throughout training.

## 6. Empirical Evaluation and Comparative Performance

In large-scale autoregressive pretraining with 2.4B parameter models on FineWeb-Edu, ET routing achieves superior perplexity relative to both dense and sparse TC-MoE baselines. Specifically, with models trained for 11.2B tokens:
- Dense baseline: cross-entropy loss 2.751
- TC-MoE (with auxiliary loss): 2.687
- ET routing: 2.620

The reduction of 0.067 nats in cross-entropy for ET versus TC-MoE translates to reaching the same loss with approximately $1.6\times$ fewer training tokens. After warm-up, expert saturation or starvation events—cases where capacity limits are exceeded or too few tokens are assigned—occur in well under 1% of training steps, indicating robust load-balancing behavior without auxiliary loss terms [2603.11535].

## 7. Design Choices and Practical Considerations

Key implementation features of ET routing include:
- EMA decay $\beta=0.999$, yielding an effective window of approximately 1000 steps.
- Thresholds initialized to zero or with early-batch statistics; a 4k-step Expert-Choice warm-up enables threshold stabilization.
- At inference, capacity-based clipping is disabled.
- No auxiliary gating or balancing losses are incorporated; all token-expert assignment balance emerges from the threshold EMA mechanism.

In sum, ET routing replaces the hard “top-$G$” per-token constraint of Token-Choice MoE and the batch-wide “top-$k$” per-expert constraint of Expert-Choice MoE with a lightweight, EMA-modulated per-expert threshold. This yields dynamic compute per token, self-balanced expert loads in expectation, and fully causal inference, providing empirically validated gains in training efficiency and modeling quality compared to prior sparse routing techniques in large-scale language modeling [2603.11535].

Source: https://www.emergentmind.com/topics/ema-exponential-moving-average-routing