MaskMoE: Token-Aware Routing in Sparse MoE

• MaskMoE is a routing method for sparse MoE models that assigns a fixed subset of experts to each token, balancing training signal and diversity.
• It adapts the number of visible experts based on token frequency, ensuring rare tokens receive concentrated supervision while frequent tokens leverage multiple experts.
• Empirical results show MaskMoE outperforms traditional dynamic and fixed routing methods in perplexity and downstream task performance.

MaskMoE is a routing method for sparse Mixture-of-Experts LLMs that enhances token-level learning by imposing a token-specific routing mask over the expert set. In the formulation introduced in "MaskMoE: Boosting Token-Level Learning via Routing Mask in Mixture-of-Experts," each vocabulary token is assigned a fixed subset of visible experts before training, and ordinary learned routing is then performed only within that subset. The method is motivated by a tension between dynamic routing, which preserves representation diversity but can disperse supervision across experts, and fixed routing, which concentrates supervision but reduces diversity. MaskMoE is designed as an intermediate regime: infrequent tokens are typically restricted to one visible expert, while frequent tokens retain access to multiple experts (Su et al., 2024).

1. Definition and conceptual scope

MaskMoE operates in the standard sparse Transformer-MoE setting in which the feed-forward network sublayer is replaced by a bank of expert FFNs, with only a small subset activated per token. Its contribution is confined to routing: it does not redesign the experts, the sparse top-$k$ activation principle, the optimizer, or the distributed MoE training setup. The method is best understood as a drop-in routing constraint layered on top of ordinary MoE routing, implemented by adding a fixed token-level mask to router logits (Su et al., 2024).

The term “MaskMoE” has acquired broader informal usage, but the precise sense associated with the named method differs from several adjacent lines of work. "Parameter-Efficient Mixture-of-Experts Architecture for Pre-trained LLMs" introduces MPOE, whose main contribution is shared expert parameterization via matrix product operators; its masking component is a gradient mask applied only to shared central tensors, so masking is auxiliary rather than architecturally constitutive (Gao et al., 2022). "Task-customized Masked AutoEncoder via Mixture of Cluster-conditional Experts" introduces MoCE rather than MaskMoE; it is a masked-autoencoder-based expert-routing method in vision, with cluster-conditional gates rather than token-frequency-conditioned routing (Liu et al., 2024). "ToMoE: Converting Dense LLMs to Mixture-of-Experts through Dynamic Structural Pruning" is closely related in spirit because it uses learned binary selection variables and sparse routing, but it discovers experts by pruning dense MLP intermediate dimensions rather than restricting access to an existing expert pool (Gao et al., 25 Jan 2025). "MergeMoE: Efficient Compression of MoE Models via Expert Output Merging" belongs to expert-merging compression rather than masking-based routing (Miao et al., 16 Oct 2025).

A common misconception is therefore to treat every MoE paper containing a mask as a MaskMoE method. In the narrow technical sense, MaskMoE refers specifically to the routing-mask mechanism of (Su et al., 2024); in broader usage, it may loosely denote mask-based MoE methods, but that broader usage conflates distinct families.

2. Motivation: routing fluctuation, underfitting, and token frequency

The central problem addressed by MaskMoE is token-level underfitting induced by routing fluctuation in dynamically routed sparse MoEs. In ordinary dynamic routing, the same token may be assigned to different experts across training iterations. Because experts do not explicitly share knowledge, the token’s supervision is fragmented across those experts. The paper argues that this fragmentation is especially harmful for infrequent tokens: frequent tokens appear often enough that dispersed updates may still suffice, whereas rare tokens may be seen too few times by any single expert, leading to underfitting (Su et al., 2024).

Fixed routing, such as random hashing-style assignment, addresses that issue by always sending a token to the same expert, thereby concentrating training signal. However, the same paper argues that fixed routing compromises representation diversity, particularly for frequent tokens that appear in many contexts and can benefit from multiple expert-specific representations. MaskMoE is explicitly framed as balancing two objectives: comprehensive or concentrated training, which favors fixed routing, and representation diversity, which favors dynamic routing (Su et al., 2024).

This suggests that the method is frequency-adaptive rather than purely sparsity-seeking. The role of the routing mask is not to reduce activated parameters below standard sparse-MoE levels, but to restrict the candidate expert set differently for different tokens. A plausible implication is that MaskMoE should be viewed primarily as a data-efficiency intervention for long-tail token distributions rather than as a new conditional-compute mechanism.

3. Routing mask construction and operational mechanism

MaskMoE constructs a token-specific routing mask $M_t$ before training and keeps that mask fixed during training. For a token $t$, the modified router is

$p = \mathrm{Softmax}(W_g h_t^l + M) \tag{5}$

where $h_t^l$ is the token representation at MoE layer $l$, $W_g$ is the router parameter matrix, and $M$ is the routing mask. Mask entries are $0$ for visible experts and $-\infty$ for invisible experts, so invisible experts receive zero probability after softmax (Su et al., 2024).

Mask construction begins with all experts masked out: $M_t$0 The number of visible experts $M_t$1 depends on token frequency. For infrequent tokens, $M_t$2, so each such token has exactly one visible expert and therefore behaves like fixed routing. For frequent tokens, $M_t$3, so routing remains dynamic within a restricted subset of experts (Su et al., 2024).

Token frequencies are computed on the pretraining corpus. Vocabulary items are sorted by frequency, and the top tokens covering $M_t$4 of the dataset are defined as frequent, with the remainder treated as infrequent. In the default setup, $M_t$5, frequent tokens use $M_t$6, and infrequent tokens use $M_t$7. The same masking vector for a token is reused across MoE layers in multi-layer MoE models (Su et al., 2024).

The algorithmic structure is correspondingly simple. One counts token frequencies, partitions the vocabulary into frequent and infrequent sets, samples visible experts uniformly for each token according to its assigned $M_t$8, stores the resulting vocabulary-to-mask table, and then applies the mask to router logits at each MoE forward pass. The paper does not introduce a separate inference-time routing rule, so the modified routing mechanism is implicitly used during inference as well (Su et al., 2024).

4. Mathematical formulation and training objective

The paper writes the MoE layer output as

$M_t$9

with top-$t$0 sparsity making most gate entries zero in practice. MaskMoE preserves that baseline formulation and alters only the computation of the routing distribution via the masked softmax of Eq. (5) (Su et al., 2024).

MaskMoE adopts a standard load-balancing auxiliary loss: $t$1 subject to

$t$2

where

$t$3

and

$t$4

Here $t$5 is the number of tokens in a minibatch, $t$6 is the router probability vector for token $t$7, $t$8 is the empirical fraction of tokens whose top expert is expert $t$9, and $p = \mathrm{Softmax}(W_g h_t^l + M) \tag{5}$0 is the average router probability mass assigned to expert $p = \mathrm{Softmax}(W_g h_t^l + M) \tag{5}$1 (Su et al., 2024).

The final training objective is

$p = \mathrm{Softmax}(W_g h_t^l + M) \tag{5}$2

with no additional MaskMoE-specific regularizer. The paper notes an important asymmetry: for infrequent tokens with only one visible expert, the load-balancing loss does not materially regulate routing because their assignment is fixed by the mask. In practice, the auxiliary balancing term mainly governs frequent-token routing, while infrequent-token routing relies on the random sampling of visible experts to remain relatively balanced at the population level (Su et al., 2024).

Two limit cases clarify the method’s position in the routing design space. If $p = \mathrm{Softmax}(W_g h_t^l + M) \tag{5}$3 for all tokens, MaskMoE reduces to something like Hash Layer. If $p = \mathrm{Softmax}(W_g h_t^l + M) \tag{5}$4 for all tokens, it reduces to standard dynamic routing or SMoE. MaskMoE therefore interpolates between fixed and unrestricted dynamic routing at token granularity (Su et al., 2024).

5. Architecture, implementation regime, and empirical performance

The experiments use a LLaMA-style Transformer with 24 layers and hidden size 1024. Two insertion settings are studied: a single-layer MoE setting in which only the final FFN layer is replaced by an MoE layer, and a multi-layer setting in which every other FFN layer is replaced, yielding 12 MoE layers. Pretraining is conducted on The Pile, described as over 825GB and 22 domains, with a budget of 100B tokens and a LLaMA tokenizer of vocabulary size 32k. Unless otherwise specified, top-1 expert selection is used, and the default expert count for SMoE and Hash Layer is 64 experts per MoE layer (Su et al., 2024).

For Share-MoE and MaskMoE in the shared-expert configuration, the setup uses 1 shared expert and 128 routed experts, with each expert at 0.5x the size of a standard FFN. During both training and inference, the shared expert is always active together with the top-1 selected routed expert. The optimizer is AdamW, with cosine decay learning-rate scheduling; the learning rate is $p = \mathrm{Softmax}(W_g h_t^l + M) \tag{5}$5 for dense and single-layer MoE and $p = \mathrm{Softmax}(W_g h_t^l + M) \tag{5}$6 for 12-layer MoE. Implementation uses DeepSpeed with random token selection enabled. The paper emphasizes that activated parameters and FLOPs remain unchanged relative to matched baselines: dense models have 468M parameters and 468M activated; single-layer MoE variants have 1.3B total and 468M activated; 12-layer MoE variants have 10B total and 468M activated (Su et al., 2024).

The main perplexity results on The Pile validation set show the strongest performance for MaskMoE in both insertion regimes. In the single-layer setting, perplexities are 6.95 for Dense, 6.62 for SMoE, 6.56 for Hash Layer, 6.72 for Share-MoE, and 6.48 for MaskMoE. In the 12-layer setting, perplexities are 6.18 for SMoE, 6.16 for Hash Layer, 6.15 for Share-MoE, and 6.11 for MaskMoE. The paper reports improvements of 0.14 over SMoE, 0.08 over Hash Layer, and 0.24 over Share-MoE in the single-layer setting, and 0.07, 0.05, and 0.04 respectively in the 12-layer setting (Su et al., 2024).

Downstream evaluation covers BoolQ, HellaSwag, LAMBADA, PIQA, SIQA, StoryCloze, Arc-e, TriviaQA (5-shot), and WebQs (5-shot). In the single-layer setting, MaskMoE scores 58.38 on BoolQ, 44.47 on HellaSwag, 55.36 on LAMBADA, 68.99 on PIQA, 41.86 on SIQA, 65.15 on StoryCloze, 53.14 on Arc-e, 10.39 on TriviaQA, and 7.14 on WebQs; it is reported to outperform SMoE on all 9/9 benchmarks, and Hash Layer and Share-MoE on 7/9 each. In the 12-layer setting, MaskMoE scores 58.32, 47.46, 57.46, 70.62, 41.91, 65.69, 55.35, 15.20, and 6.74 respectively; it is better than SMoE on 8/9 tasks, Hash Layer on 8/9, and Share-MoE on 6/9 (Su et al., 2024).

Ablations over routing-mask design identify $p = \mathrm{Softmax}(W_g h_t^l + M) \tag{5}$7 as a strong configuration, with perplexity 6.506, compared with 6.558 for Hash Layer $p = \mathrm{Softmax}(W_g h_t^l + M) \tag{5}$8 and 6.618 for SMoE $p = \mathrm{Softmax}(W_g h_t^l + M) \tag{5}$9. The reported ablations support four conclusions stated in the paper: MaskMoE beats Hash Layer, concentrating routing into fewer visible experts beats SMoE, setting infrequent tokens to $h_t^l$0 is especially useful, and when $h_t^l$1, a range of $h_t^l$2 values still outperform both Hash Layer and SMoE. An additional study with and without a shared expert reports average benchmark score / Pile PPL of 43.87 / 6.62 for SMoE, 44.56 / 6.56 for Hash Layer, 44.25 / 6.72 for Share-MoE, 44.99 / 6.48 for MaskMoE, and 44.85 / 6.51 for MaskMoE without shared experts (Su et al., 2024).

6. Relations to neighboring MoE families and recurrent ambiguities

MaskMoE is best distinguished from methods whose masks operate on gradients, parameters, or cluster assignments rather than token-to-expert visibility. MPOE shares parameters across experts by decomposing each expert matrix into one central tensor and four auxiliary tensors when $h_t^l$3, then sharing the central tensor across experts at a layer. Its mask is a Bernoulli gradient mask applied only to central-tensor updates, defined by $h_t^l$4, with update

$h_t^l$5

This is optimization masking rather than routing masking. The paper explicitly characterizes its own main contribution as MPO-based parameter sharing, and reports, for example, a 27.2x reduction in total parameters relative to Switch Transformers for the superior model on T5-Large, with 26,000M parameters for +Switch versus 956M for +MPOE and average GLUE scores of 88.50 versus 88.82 (Gao et al., 2022).

MoCE, by contrast, is not a language-model routing method but a masked-autoencoder pretraining paradigm with cluster-conditional experts. It routes all tokens from images in the same cluster using the cluster centroid embedding,

$h_t^l$6

and is motivated by negative transfer in generic MAE pretraining rather than rare-token underfitting. Its downstream customization selects a task-specific expert path after pretraining, and its reported results include 85.54 average classification transfer versus 83.09 for MAE*, as well as 48.3 mIoU on ADE20K versus 48.1 for MAE (Liu et al., 2024).

ToMoE occupies a different neighboring region: it converts dense decoder-only LLMs into sparse MoE-like models by dynamic structural pruning. Its MLP experts are defined by binary diagonal selection matrices over the FFN intermediate dimension,

$h_t^l$7

with top-1 token-wise routing and straight-through Gumbel relaxations. This is closer to mask-derived expert construction than to MaskMoE’s restriction of token visibility over a pre-existing expert pool (Gao et al., 25 Jan 2025).

MergeMoE is farther still from the routing-mask family. It compresses pretrained MoE models by merging experts in output space, formalizing compression through matrices $h_t^l$8, $h_t^l$9, $l$0, $l$1, and $l$2 so that $l$3, with a closed-form least-squares solution $l$4. Its use of $l$5 is ordinary MoE notation rather than a new masking contribution (Miao et al., 16 Oct 2025).

Taken together, these comparisons show that “masking” in MoE literature is polysemous. In MaskMoE proper, the mask restricts token-specific expert visibility. In MPOE, the mask stochastically skips central-tensor updates. In ToMoE, masks define subnetworks carved from dense FFNs. In MoCE, cluster-conditional gating coexists with masked autoencoding. The terminological overlap is substantial, but the mechanisms are not interchangeable.

7. Limitations, interpretation, and significance

The main limitation stated for MaskMoE is the rigidity of its vocabulary split. Tokens are divided into only two groups, frequent and infrequent, using a hard threshold $l$6. The authors suggest that a smoother scheme, in which visible expert count decreases gradually with frequency rank, might work better, but do not investigate it because of compute constraints (Su et al., 2024).

A second limitation is implicit in the design choice to use token frequency as the key proxy for routing needs. The method is built on the hypothesis that rare tokens chiefly require concentrated supervision and common tokens chiefly require representational diversity. This is an operationally effective proxy in the reported language-model experiments, but the paper does not claim that frequency exhausts semantic ambiguity or contextual variability (Su et al., 2024).

The strongest practical significance of MaskMoE lies in showing that expert visibility constraints can improve perplexity and downstream task performance without changing activated parameters or FLOPs relative to matched sparse-MoE baselines. The paper also reports a noteworthy scaling behavior: for SMoE and Share-MoE, increasing expert count does not always help and can worsen perplexity at 128 experts relative to 64, whereas for Hash Layer and MaskMoE, perplexity continues to improve as the number of experts increases, with MaskMoE best across all expert counts. This suggests that restricting token assignments to denser, frequency-aware expert subsets can mitigate the supervision fragmentation that otherwise accompanies expert proliferation (Su et al., 2024).

In the literature on expert models, MaskMoE therefore occupies a specific position. It is neither a parameter-sharing architecture, nor a masked-autoencoder expert framework, nor a pruning-derived dense-to-MoE conversion method, nor an expert-merging compression method. It is a routing method for sparse MoE LLMs in which each vocabulary token is assigned a fixed candidate expert subset, so that rare tokens receive repeated exposure to the same parameters while frequent tokens retain multi-expert representational flexibility (Su et al., 2024).