Paper deep dive
MoRE: Scaling mixture of experts with hardware-aware low-rank routing
Honam Wong, Surbhi Goel, Enric Boix-Adserà
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Mixture-of-Experts (MoE) layers are central to frontier language models, and recent architectures push toward more and smaller experts. In this regime, the standard linear router becomes a bottleneck: with $M$ experts and hidden dimension $h$, its per-token cost $\Theta(Mh)$ dominates the MoE layer once $M$ is large. We introduce MoRE (Mixture of Rank-reduced-routed Experts), which factorizes the router weight matrix at rank $r$ and reduces the routing cost to $O((h + M)r)$. We prove that rank logarithmic in $M$ suffices for routing expressivity when the number of active experts is fixed, and is necessary up to precision factors. We also prove that logarithmic rank preserves load balance in a Gaussian memorization model, and training on a synthetic phonebook task shows that low rank does not hurt memorization. At matched active FLOPs, the factorization allows a factor of $\Theta(h/r)$ more experts. To realize this gain in wall-clock time, we design a fused Triton kernel at inference that avoids expensive memory operations on HBM. Empirically, MoRE improves memorization on the phonebook task and performance on knowledge-intensive Q\&A benchmarks after pretraining, while matching reasoning ability. Code available at this https URL.
Tags
Links
- Source: https://arxiv.org/abs/2609.36301v1
- Canonical: https://arxiv.org/abs/2609.36301v1
Trouble viewing inline? Open PDF directly →
Full Text
145,652 characters extracted from source content.
Expand or collapse full text
MoRE: Scaling mixture of experts with hardware-aware low-rank routing Honam Wong Affiliation: University of Pennsylvania Surbhi Goel Affiliation: University of Pennsylvania Enric Boix-Adserà Affiliation: The Wharton School, University of Pennsylvania Abstract Mixture-of-Experts (MoE) layers are central to frontier language models, and recent architectures push toward more and smaller experts. In this regime, the standard linear router becomes a bottleneck: with M experts and hidden dimension h, its per-token cost Θ(Mh) (Mh) dominates the MoE layer once M is large. We introduce MoRE (Mixture of Rank-reduced-routed Experts), which factorizes the router weight matrix at rank r and reduces the routing cost to O((h+M)r)O((h+M)r). We prove that rank logarithmic in M suffices for routing expressivity when the number of active experts is fixed, and is necessary up to precision factors. We also prove that logarithmic rank preserves load balance in a Gaussian memorization model, and training on a synthetic phonebook task shows that low rank does not hurt memorization. At matched active FLOPs, the factorization allows a factor of Θ(h/r) (h/r) more experts. To realize this gain in wall-clock time, we design a fused Triton kernel at inference that avoids expensive memory operations on HBM. Empirically, MoRE improves memorization on the phonebook task and performance on knowledge-intensive Q&A benchmarks after pretraining, while matching reasoning ability.11 1 Code available at https://github.com/Matheart/MoRE_code. 11footnotetext: Equal advising. 1 Introduction Mixture-of-experts (MoE) is now the dominant architectural paradigm for frontier language models (DeepSeek-AI, 2024a; DeepSeek-AI, 2024b; Jiang et al., 2024; xAI, 2024; Yan et al., 2026). In these architectures, the standard dense MLP is replaced by a pool of independent experts and a router layer that activates only a small subset of experts per token. Recent advances in model scaling rely heavily on increasing the number of experts in MoE layers (Abnar et al., 2025). This approach expands the total parameter count, allowing models to benefit from neural scaling laws without proportionally increasing active compute (Kaplan et al., 2020; Abnar et al., 2025). In particular, at a fixed active compute budget, MoE layers achieve significantly better memorization and factual recall than dense MLPs (Jelassi et al., 2025). Recent frontier MoE models are also increasingly operating in the fine-grained regime (Appendix A), in which the layer is divided into a larger number of smaller experts, each with intermediate width below the hidden dimension (Krajewski et al., 2024). Figure 1: MoRE unlocks more scaling with the same active FLOPs. (Top) Replacing the full-rank h×Mh× M router with a low-rank factorization R1,R2R_1,R_2 inside a fused kernel cuts routing cost from O(hM)O(hM) to O(r(h+M′))O(r(h+M )), enabling M′≫M M experts at matched active FLOPs. (Bottom Left) Inference (prefilling) wall-clock time of one layer at h=512h=512, split into the router and the rest (setup as in Figure 4). The standard router’s cost grows linearly with M and dominates at large M, while MoRE’s fused low-rank router grows far more slowly. (Bottom Right) At matched FLOPs, MoRE memorizes more on the synthetic phonebook task, and the gap widens with h, as predicted by the Θ(h/r) (h/r) scaling of Section 4.1. However, further scaling the number of experts in MoE layers hits a bottleneck: routing costs begin to dominate the expert cost as the expert pool grows. This is driven by two distinct factors: the increased active compute required to calculate the routing function, and the heavy memory-bandwidth cost of writing and reading the routing scores. These combined costs can offset the efficiency gains of sparsity in the many-expert regime (see Figure 1). This motivates our main question: Can we scale the number of experts in MoEs efficiently? 1.1 Our contributions In this work, we propose Mixture of Rank-reduced-routed Experts (MoRE), an architectural design that allows for scaling to more experts with an asymptotically lower computational cost. MoRE replaces the full-rank linear router with a low-rank router. This factorization reduces the number of parameters in the routing function and the FLOPs required for a forward pass, allowing scaling to more experts. Our router design is inspired by the low-rank parameterization of key-query (KQ) weights in recent attention architectures (DeepSeek-AI, 2024a) as well as the low-rank design of LoRA adapters (Hu et al., 2021). Expressivity guarantees. To determine whether the rank bottleneck limits MoRE’s routing capacity, we study which expert assignments it can realize and how evenly it can use the experts. In Section 3, we first prove general upper and lower bounds on the rank needed to realize any top-k assignment. Then, we prove that low-rank routing in MoRE layers preserves load-balancing ability, by constructing a router with r=Θ(logM)r= ( M) that balances expert assignment in a Gaussian memorization model extending that of Jelassi et al. (2025). By spreading inputs across experts, balanced routing helps use the memorization capacity of the whole expert pool. We also show empirically on a synthetic memorization task, Phonebook, that low rank does not hurt load balancing and memorization (Section 3.2). Scaling experts efficiently. At matched active FLOPs, a MoRE layer with rank-r routing and hidden dimension h supports a factor of Θ(h/r) (h/r) more experts than an MoE layer with the standard router (Section 4.1). However, realizing this gain in wall-clock time requires more than a FLOP reduction: materializing the score tensor in HBM (High Bandwidth Memory) becomes expensive when M is large. In this paper, we focus on inference, since a model is trained once but runs inference on every request after deployment and during RL post-training. We therefore design a fused kernel that keeps scores in on-chip SRAM (Section 4). The kernel computes scores in blocks and retains only the running top-k values and expert indices. Notably, the inference time (prefilling) of MoRE matches that of MoE at matched active FLOPs, with a much larger number of experts. Experimental validation. Finally, we validate our architecture through pretraining experiments on both the phonebook task and real-world Internet data. We show that MoRE models consistently improve performance in memorization for phonebook task and improve on knowledge-intensive tasks such as TriviaQA (Joshi et al., 2017) and Natural Questions (Kwiatkowski et al., 2019), compared with standard MoEs at equal active FLOPs. Additionally, we find that MoRE matches performance on reasoning tasks. 1.2 Related Work Routing in MoE layers. The standard router in modern MoE is a learned linear layer that scores all M experts per token (Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2022). Hash routing fixes the assignment to a deterministic hash (Roller et al., 2021), but Clark et al. (2022); Dikkala et al. (2023) show that deterministic assignment underperforms learned routing at scale, motivating the search for parameter-efficient learned routers rather than discarding learning. Boix-Adsera (2025) suggests low-rank routing but evaluates it only in a limited distillation setting. In the M≪hM h regime, X-MoE (Chi et al., 2022) and L2R (Yang et al., 2026) also introduce low-rank projections, but to address issues in token representations and routing geometry. MoRE instead targets the fine-grained regime where M is large and M≪hM h no longer holds, using low-rank routing to reduce the cost of scoring many experts. He (2024) introduces PEER, a product-key retrieval layer with cost O(hM)O(h M). At matched active FLOPs, MoRE takes much less wall-clock time than PEER (see Appendix B). Fine-grained MoE and scaling. Recent architectures are increasingly designed in the fine-grained regime, pairing a small expert intermediate dimension s with a large expert count M. Krajewski et al. (2024) introduce fine-grained MoE and show that it outperforms coarse-grained MoE at a fixed training budget, and that higher granularity is out of reach because of the routing cost. Li et al. (2026) confirm through extensive experiments that larger granularity and expert counts result in better performance. Since MoE memorization capacity scales with total rather than active parameters (Jelassi et al., 2025), designs with large M are directly motivated. Nevertheless, these works do not address the routing bottleneck that emerges as M grows, and we introduce MoRE to reduce the routing cost. Knowledge storage in language models. Two threads study how language models store factual knowledge. The first measures memorization in standard transformer and MoE architectures: Geva et al. (2021) interpret FFN layers as key-value memories, Allen-Zhu and Li (2024) establish knowledge capacity scaling laws for dense LMs, and Jelassi et al. (2025) show that for MoE, memorization capacity scales with total parameters rather than active parameters, directly motivating designs with large M. A second thread replaces the FFN with a learnable key-value lookup: Memory Layers at Scale (Berges et al., 2024) scales product-key retrieval to billions of stored embeddings, Ultra-Sparse Memory (Huang et al., 2025) pushes to ultra-sparse activation patterns, and Cheng et al. (2026) extend retrieval-style storage to conditional access. MoRE is complementary to this thread: it preserves the MLP-expert structure of MoE and inherits the memorization analysis of Jelassi et al. (2025) while reducing only the routing cost. Efficient MoE systems. Existing fused MoE kernels target expert dispatch: ScatterMoE (Tan et al., 2024), MegaBlocks (Gale et al., 2023), and follow-up IO-aware kernels (Costin et al., 2025; Guo et al., 2025) fuse scatter-gather with the per-expert grouped matmul, but leave open the question of how to implement the router efficiently in the many-expert regime. vLLM (Kwon et al., 2023) provides a FusedTopKRouter that fuses top-k selection with softmax normalization, but does not fuse the router computation with top-k. At inference, expert offloading (Hwang et al., 2024; Xue et al., 2024; He et al., 2024) reduces serving cost by moving inactive experts to CPU memory, but does not reduce the router computation. MoRE instead fuses both low-rank matmuls and top-k selection into one router kernel (Section 4.2). 2 Preliminaries Notation. We write h for the hidden dimension, s for the expert intermediate dimension, M for the number of experts per layer, k for the number of active experts per token (with k<Mk<M), r for the rank of the low-rank router, and L for the number of transformer layers. With batch size B and sequence length T, each iteration processes N=BTN=BT tokens. We write [m]=1,…,m[m]=\1,…,m\, and use P for the total parameter count and τ for the per-token active FLOPs. Mixture-of-experts layer. An MoE layer has M experts fii=1M\f_i\_i=1^M and a router g:ℝh→ℝMg:R^h ^M. Each expert is usually a SwiGLU MLP (Shazeer, 2020) with three weight matrices Wigate,Wiup∈ℝs×hW^gate_i,W^up_i ^s× h and Widown∈ℝh×sW^down_i ^h× s: fi(x)=Widown(σ(Wigatex)⊙Wiupx),fMoE(x)=∑i=1Mgi(x)fi(x),f_i(x)\;=\;W^down_i (σ(W^gate_ix) W^up_ix ), f_MoE(x)\;=\; _i=1^Mg_i(x)f_i(x), with σ the SiLU activation and ⊙ the elementwise product. Each expert thus has 3hs3hs parameters. The standard router (Shazeer et al., 2017; Fedus et al., 2022) is parametrized by a linear projection R∈ℝM×hR ^M× h and applies softmax over the k largest scores: g(x,R)=softmax(TopK(Rx))g(x;R)=softmax (TopK(Rx) ), where TopK(x)TopK(x) denotes the k largest values of x∈ℝMx ^M. We write its configuration as MoE(M,k,h,s)MoE(M,k,h,s). MoRE layer. In a MoRE layer, we replace R with R1R2R_1R_2 where R1∈ℝM×rR_1 ^M× r and R2∈ℝr×hR_2 ^r× h, scoring through an r-dimensional bottleneck: g(x,R1,R2)=softmax(TopK(R1R2x)).g(x;R_1,R_2)=softmax (TopK(R_1R_2x) ). Routing parameters drop from MhMh to (M+h)r(M+h)r. We write the low-rank configuration as MoRE(M,k,h,s,r)MoRE(M,k,h,s,r). Total parameters and active cost. The total MLP parameter count per MoE/MoRE layer is 3Mhs3Mhs, and the per-token active FLOPs per layer are τ=Cekhs+CrMh(standard router),Cekhs+Cr(M+h)r(low-rank router),τ\;=\; casesC_e\,khs+C_r\,Mh&(standard router),\\ C_e\,khs+C_r\,(M+h)\,r&(low-rank router), cases (1) where Ce,CrC_e,C_r collect per-operation FLOP factors. Following Krajewski et al. (2024, App. E), at inference we have (Ce,Cr)=(6,6)(C_e,C_r)=(6,6). Training gives a similar ratio, (Ce,Cr)=(18,14)(C_e,C_r)=(18,14). 3 Low-rank routing preserves memorization capabilities MoRE replaces the standard router with a rank-r factorization, which raises two related theoretical questions. First, can the router still express the needed expert assignments through the rank bottleneck? Second, when there are many inputs to memorize, can the low-rank router still use the experts’ parameters efficiently by balancing the inputs assigned to them? We answer both questions positively in Sections 3.1 and 3.2, and confirm the answers with experiments on the synthetic Phonebook memorization task. 3.1 Routing expressivity Here an assignment is a subset of k experts. We show that, for fixed k, a single router can select every assignment in a given family for some input, with rank growing only logarithmically in the number of experts and assignments. This dependence on the number of assignments is necessary up to a factor depending on numerical precision. Setup. The rank-r router has a simple geometric form: the rows of R1∈ℝM×rR_1 ^M× r define the M expert vectors ui∈ℝru_i ^r, and each input x∈ℝhx ^h reaches R1R_1 only through its projected representation z(x):=R2x∈ℝrz(x):=R_2x ^r where R2∈ℝr×hR_2 ^r× h. The router then selects the k experts with the largest routing scores ⟨ui,z(x)⟩ u_i,z(x) , 1≤i≤M1≤ i≤ M. For exposition, throughout this section we normalize the scale by assuming ‖ui‖2≤1\|u_i\|_2≤ 1 for every expert vector and ‖z(x)‖2≤1\|z(x)\|_2≤ 1 for every input x under consideration.22 2 We state the general result in Appendix J, which only requires them to be bounded by some constants BuB_u and BzB_z. Margin and precision. Because top-k is implemented at finite precision, selecting a subset robustly requires a positive gap between selected and unselected scores. We formalize this as follows. Definition 1. Fix a rank-r router R=R1R2R=R_1R_2. A subset S∈([M]k)S∈ [M]k is realized at input x with margin η>0η>0 if S=TopK(Rx)S=TopK(Rx) and mini∈S(Rx)i−maxj∉S(Rx)j≥η. _i∈ S(Rx)_i- _j∉ S(Rx)_j≥η. Moreover, we say a family ℱ⊆([M]k)F [M]k is realizable at margin η by R if every S∈ℱS has some input xS∈ℝhx_S ^h that realizes it at margin at least η. Note that for a p-bit implementation, a margin η is resolvable when η=Ω(2−p)η= (2^-p), which is equivalent to p=Ω(log(1/η))p= ( (1/η)). We first present an upper bound based on the Johnson–Lindenstrauss lemma (Johnson and Lindenstrauss, 1984) by providing a router construction: Theorem 1 (Upper bound on required rank). There exists a universal constant C>0C>0 such that the following holds. For 1≤k≤M−11≤ k≤ M-1, any collection of top-k assignments ℱ⊆([M]k)F [M]k, and any rank r≥Cklog(M+|ℱ|),h≥r,r≥ Ck (M+|F|), h≥ r, there exists a rank-r router R=R1R2R=R_1R_2 with ‖ui‖2≤1\|u_i\|_2≤ 1 that realizes every S∈ℱS at margin Ω(1/k) (1/ k), using inputs xSS∈ℱ\x_S\_S with ‖R2xS‖2≤1\|R_2x_S\|_2≤ 1. Proof idea. In the uncompressed case r=Mr=M, take the expert vectors to be the standard basis vectors e1,…,eM∈ℝMe_1,…,e_M ^M. The representation zS∝Sk−[M]∖SM−kz_S 1_Sk- 1_[M] SM-k is positive exactly on S, so it routes to S with margin Θ(1/k) (1/ k). The Johnson–Lindenstrauss lemma maps these M+|ℱ|M+|F| vectors to r=Θ(log(M+|ℱ|)/ε2)r= ( (M+|F|)/ ^2) dimensions while preserving inner products up to additive error ε , and ε=Θ(1/k) = (1/ k) keeps a constant fraction of the margin, which gives r=O(klog(M+|ℱ|))r=O(k (M+|F|)). The full proof is in Appendix J. Next, we use a volume-packing argument to prove a lower bound on the rank. The intuition is that if two different assignments S,T∈ℱS,T are both realized with margin η, their projected representations must be separated. Theorem 2 (Lower bound on required rank). Let ℱ⊆([M]k)F [M]k. If a rank-r router R=R1R2R=R_1R_2 with ‖ui‖2≤1\|u_i\|_2≤ 1 and ‖R2xS‖2≤1\|R_2x_S\|_2≤ 1 realizes every top-k assignment S∈ℱS at margin at least η, then |ℱ|≤(1+2/η)r,r≥log|ℱ|log(1+2/η).|F|≤(1+2/η)^r, r≥ |F| (1+2/η). Proof idea. If two assignments S≠TS≠ T are realized at xSx_S and xTx_T, pick i∈S∖Ti∈ S T and j∈T∖Sj∈ T S. Expert i beats expert j by at least η at xSx_S and loses by at least η at xTx_T, so ‖ui−uj‖2≤2\|u_i-u_j\|_2≤ 2 forces ‖R2xS−R2xT‖2≥η\|R_2x_S-R_2x_T\|_2≥η. Packing |ℱ||F| such points into the unit ball of ℝrR^r gives |ℱ|≤(1+2/η)r|F|≤(1+2/η)^r. Finally, taking ℱF to be the collection of all top-k subsets gives the following consequence: Corollary 1 (Full expressivity). At precision p=Ω(logk)p= ( k), rank r=O(k2log(M/k))r=O(k^2 (M/k)) suffices to realize every top-k assignment over M experts. Conversely, any rank-r router that realizes every top-k assignment at precision scale p must have r=Ω(klog(M/k)/p)r= (k (M/k)/p). Remark. The extra factor of k comes from preserving the routing margin in the Johnson–Lindenstrauss projection. Our bounds do not establish whether this factor is necessary. 3.2 Load balancing and memorization capacity Memorization, a core function of MoE layers (Du et al., 2022; Geva et al., 2021; Jelassi et al., 2025), concerns how many facts can be stored, not which expert configurations can be addressed. Extending the Gaussian model of Jelassi et al. (2025), we prove that logarithmic rank does not affect load balancing, and then explain how load balancing relates to memorization capacity. Proposition 1 (Gaussian load balancing with logarithmic rank). Let x1,…,xn∼(0,Ih)x_1,…,x_n (0,I_h) be independent, and suppose M=2rM=2^r (r=Θ(log2M)r= ( _2M)) with h≥rh≥ r. For any 1≤k≤M1≤ k≤ M, there exists a rank-r top-k router with M experts that selects every expert with probability exactly k/Mk/M on each example. Consequently, let NaN_a be the number of examples routed to expert a, and write L:=log2MδL:= 2Mδ. For any δ∈(0,1)δ∈(0,1), the following holds with probability at least 1−δ1-δ: maxa∈[M]|Na−nkM|≤2nkLM+23L. _a∈[M] |N_a- nkM |≤ 2nkLM\;+\; 23L. The key intuition is that, under this construction, a standard Gaussian input has no preference for one expert over another in expectation. We defer the full proof to Appendix J. Thus, the low-rank factorization preserves the load-balanced mechanism of standard MoE, while reducing routing complexity from O(Mh)O(Mh) to O((M+h)logM)O((M+h) M). Why balanced load maximizes the number of stored facts. The number of facts an MLP expert can memorize is determined by its parameter count, as shown for two-layer networks (Daniely, 2020) and for language models (Allen-Zhu and Li, 2024; Allen-Zhu and Li, 2023). An MoE layer reaches the total capacity of its experts only under balanced load: an overloaded expert receives more facts than it can memorize, and an underutilized expert leaves its parameters unused. Experimental validation with trained models. We verify in training that sufficient rank does not affect memorization, using the phonebook task of Jelassi et al. (2025). Each example is a five-letter name followed by an eight-digit number, such as “alice|12345678”. We pretrain on these pairs and report the largest number of training examples for which the model attains ≥90%≥ 90\% exact-match accuracy on the numbers given the names, evaluated on 1,0001,000 training examples sampled uniformly at random. To approach the model’s full memorization capacity, we train for up to 1,5001,500 epochs, in line with Allen-Zhu and Li (2024). Figure 2(Left) varies the router rank r while keeping the expert configuration fixed: MoRE matches the memorization capacity of standard MoE once r is sufficiently large. To compare load balancing, Figure 2(Middle, Right) tracks normalized routing entropy, Hnorm=−∑a=1Mpalogpa/logMH_norm=- _a=1^Mp_a p_a/ M, where pap_a is expert a’s share of routed tokens. Both models reach high entropy by the end of training. Figure 2: (Left) Sufficient rank does not affect memorization. (Middle, Right) MoRE with r=32r=32 and MoE have similar load balancing at the end of training. On the Phonebook task, at M=192M=192, h=128h=128, s=8s=8, k=4k=4, L=6L=6. Left: maximum phonebook size MoRE memorizes against the router rank. The dashed line is MoE (standard router) at the same configuration, which MoRE already matches at r=16r=16. Middle, Right: routing entropy HnormH_norm per layer against training step for MoRE (r=32r=32) and MoE. 1 means equal load balancing and 0 means collapse onto one expert. We define HnormH_norm and show similar results for other setups in Appendix G.1. 4 Scaling experts efficiently with fused implementation of MoRE Since logarithmic rank keeps the load balanced and does not hurt memorization (Section 3), we ask: at a fixed active-FLOP budget, how much additional capacity does the low-rank router unlock? Section 4.1 shows a Θ(h/r) (h/r) gain in the expert count.33 3 In Appendix G.2, we find that adding experts monotonically reduces MoRE’s validation perplexity. The resulting large-M regime creates a memory-bandwidth bottleneck on the routing score tensor, which we resolve with a fused Triton (Tillet et al., 2019) kernel for the router (Section 4.2). 4.1 Scaling under fixed active-FLOPs Fix the active-FLOP budget τ, the hidden dimension h, and the active expert count k, so that all non-MoE components have identical parameter counts across the two settings. We compare MoE(M,k,h,s)MoE(M,k,h,s) and MoRE(M′,k,h,s′,r)MoRE(M ,k,h,s ,r) on total MLP parameters P=3MhsP=3Mhs. The full calculation is in Proposition 2 (Appendix I). The active budget τ decomposes into an expert cost CekhsC_ekhs and a router cost. For MoE, the router cost is CrMhC_rMh, i.e. CrhC_rh FLOPs per expert. For MoRE, the router cost is CrM′r+CrhrC_rM r+C_rhr, where the per-expert term CrrC_rr is an h/rh/r reduction over MoE, and CrhrC_rhr is a fixed projection overhead independent of M′M . Since P=3Mhs∝MsP=3Mhs Ms is a product, and the budget trades off linearly between scaling M (via the router) and scaling s (via the experts), P is maximized when the budget is split equally between the router and the experts. This gives M=τ2Crh(MoE),M′=τ−Crhr2Crr(MoRE).M= τ2C_rh (MoE), M = τ-C_rhr2C_rr (MoRE). In the regime τ≫Crhrτ C_rhr (equivalently M,M′≫rM,M r), the projection overhead is negligible, and M′/M≈h/rM /M≈ h/r: each MoRE expert costs only a fraction r/hr/h of the routing cost of an MoE expert. For example, at h=2048h=2048 and r=16r=16, MoRE admits roughly 128×128× more experts than MoE at the same active cost. In Section 5, we choose configurations according to this derivation, with Ce=CrC_e=C_r.44 4 Since Ce=Cr=6C_e=C_r=6 for inference and Ce≈CrC_e≈ C_r in training, the active FLOPs roughly match in both. We find that the same Θ(h/r) (h/r) conclusion also holds under multi-GPU expert parallelism (Appendix E.3.1). 4.2 Realizing the speedup at large M When the expert count M and the number of tokens N are large, running a top-k over the full (N,M)(N,M) routing-score tensor requires writing the tensor to HBM (High Bandwidth Memory, see Appendix E.1) and reading it back, and dominates the routing time (Figure 3, Left). This requires fusing operations in addition to the low-rank factorization. Figure 3: (Left) Fusing with top-k eliminates the need to materialize the scores for MoRE. (Middle) The low-rank router is faster than the standard router, while its cost grows with r. (Right) The gap between MoRE and MoE widens as h grows. M=4,096M=4,096, N=262,144N=262,144 (B=256B=256, T=1024T=1024), single NVIDIA B200, mean of 200200 trials after 5050 warm-ups. Left and middle time the router alone at h=2048h=2048 (left: r=16r=16, middle: k=4k=4). Right times the whole layer at s=64s=64, k=4k=4, r=64r=64. Below, we sketch how MoRE can be fused and leave the full details to Appendix F.1. Instead of writing the routing scores for every token to HBM and reading them back to compute the top-k scores, we compute R1R2xR_1R_2x and maintain a running top-k buffer (Figure 3, Left)55 5 Our algorithm works well for common choices of k. For a discussion of our top-k algorithm and a comparison with other implementations, see Appendix E.2.. We find that the low-rank router in MoRE is faster than the standard router across different ranks (Figure 3, Middle). MoE can be fused in a similar way, but we find that its speedup from fusing becomes small when h is large, while the speedup for MoRE remains visible, because for MoE the cost of computing RxRx and maintaining the top-k buffer scales with h instead of r (Figure 3, Right). We also show similar speedups on multiple GPUs with expert and data parallelism, where communication costs are involved (Appendix E.3). Wall-clock time at inference in the large-M regime. As shown in Figure 4, the fused MoRE layer outperforms the standard MoE layer across the whole grid, and the speedup grows with M along every row, reaching 1.831.83 to 4.54×4.54× at M=8,192M=8,192. Equivalently, at matched wall-clock time, MoRE supports about eight times more experts: at h=512h=512, the standard layer reaches M=512M=512 in 6.906.90 ms, while MoRE reaches M=4,096M=4,096 in 5.375.37 ms, realizing the scaling advantage of Section 4.1 in practice. Fusing also reduces peak memory by up to 14×14× for the whole layer (Figure 19). We find that MoE and MoRE have similar decoding speed (see Appendix E.4). Figure 4: Inference (prefilling): At a fixed configuration the MoRE layer is 1.831.83 to 4.54×4.54× faster at M=8,192M=8,192, and at matched active FLOPs it has comparable speed to the standard MoE layer while holding h/rh/r times more experts. Forward pass of a single layer. Left: MoE layer. Middle: MoRE layer (r=64r=64). Right: speedup TMoE/TMoRET_MoE/T_MoRE. s=64s=64, k=4k=4, N=262,144N=262,144 (B=256B=256, T=1024T=1024), mean of 200200 trials after 5050 warm-ups on a single NVIDIA B200. The comparison where both layers are fused is in Figure 22. We provide additional measurements on an A100, with similar speedups, in Appendix H.2. 5 Experiments: MoRE improves performance at matched active FLOPs In this section, we show that replacing MoE layers with MoRE layers in transformers can improve performance on memorization and knowledge-intensive tasks while keeping active FLOPs fixed. 5.1 Synthetic task: increased memorization capabilities at matched active FLOPs We first test on the phonebook task (Section 3.2) whether MoRE’s parameter advantage at matched active FLOPs translates into memorization capacity. MoRE memorizes substantially more pairs than MoE at every scale we sweep. In Figure 5, we sweep the hidden dimension h∈64,128,256,512h∈\64,128,256,512\ at expert width s=8s=8, and the expert width s∈8,16,32,64s∈\8,16,32,64\ at h=64h=64. At each point, we choose the optimal MoE and MoRE configurations by Proposition 2 so that they have matched active FLOPs. Notably, the gap widens with larger h, in agreement with the Θ(h/r) (h/r) trend of Proposition 2.66 6 We stop at h=512h=512 and s=64s=64 because larger models memorize more than 2 million examples, and training on such datasets for up to 1,500 epochs exceeds our compute budget. Figure 5: At fixed active FLOPs τ, MoRE memorizes more phonebook entries than MoE, and the gap widens with h as Proposition 2 predicts. Left: sweep over hidden dimension h∈64,128,256,512h∈\64,128,256,512\ at expert width s=8s=8. Right: sweep over expert width s∈8,16,32,64s∈\8,16,32,64\ at h=64h=64. Dotted lines are log-log linear fits. Configurations and sizes are in Tables 6 and 7 (Appendix C.1). 5.2 Real-world pretraining: improved performance on knowledge-intensive tasks For the real-world pretraining, we construct a 40B-token mixture dataset following the natural-language pretraining recipe of Jelassi et al. (2025): a mixture of FineWeb-edu (HuggingFaceFW, 2024), Cosmopedia (Ben Allal et al., 2024), and Wikipedia, together with the training sets of the downstream benchmarks, so that the model is familiar with the downstream tasks at test time (Jelassi et al., 2025; Allen-Zhu and Li, 2023). We select two MoE/MoRE pairs with matched active FLOPs by Section 4.1, a third pair with wider experts (s and s′s ), and a larger pair with M=1024M=1024 and M′=hr⋅M=8192M = hr· M=8192, trained with expert parallelism. By construction, MoRE has 5 to 6× more total parameters. On the first three pairs, MoRE reduces validation perplexity by 7% to 15% (Table 1).77 7 The gap is smaller for the last pair. A possible reason is that a larger model needs >>40B tokens, which is beyond our budget. In Appendix G.3, we find that the perplexity gap increases with the token budget. Figure 6: With active FLOPs fixed, a smaller rank allows more experts, and validation perplexity falls as the rank shrinks. The number of experts M grows from 64 to 8704 as r falls from 256 to 16, at h=512h=512, k=4k=4, L=12L=12 and 10B training tokens. We evaluate 0-shot via lm-evaluation-harness (Gao et al., 2023) on knowledge-intensive tasks (TriviaQA (Joshi et al., 2017), Natural Questions (Kwiatkowski et al., 2019)) and commonsense and science reasoning (HellaSwag (Zellers et al., 2019), WinoGrande (Sakaguchi et al., 2020), ARC-Easy (Clark et al., 2018), SciQ (Welbl et al., 2017)). MoRE matches or improves knowledge-intensive performance while preserving commonsense and reasoning scores. s,hs,h Mo(R)E-M Active MoE FLOPs / layer Total Param Val PPL↓ TriviaQA F1↑ NQ-Open F1↑ Hella. acc_norm↑ ARC-e acc↑ SciQ acc↑ Wino. acc↑ 32, 1024 MoE-128 1.57M 287.0M 14.01 16.57 ± 0.24 6.88 ± 0.32 36.91 ± 0.48 56.31 ± 1.02 79.60 ± 1.27 50.20 ± 1.41 MoRE-1536 1.57M 1.50B 11.96 20.21 ± 0.25 7.63 ± 0.34 42.28 ± 0.49 63.30 ± 0.99 80.20 ± 1.26 53.51 ± 1.40 64, 512 MoE-256 1.57M 362.9M 14.01 16.29 ± 0.24 5.72 ± 0.31 36.95 ± 0.48 57.32 ± 1.01 75.80 ± 1.36 51.46 ± 1.40 MoRE-1792 1.57M 1.91B 12.66 17.12 ± 0.23 6.88 ± 0.32 40.15 ± 0.49 60.44 ± 1.00 77.70 ± 1.32 50.28 ± 1.41 128, 512 MoE-256 2.36M 664.9M 12.81 18.09 ± 0.24 5.99 ± 0.29 40.00 ± 0.49 59.68 ± 1.01 76.30 ± 1.35 51.22 ± 1.40 MoRE-1792 2.26M 3.76B 11.92 20.79 ± 0.26 7.62 ± 0.34 42.07 ± 0.49 60.27 ± 1.00 78.70 ± 1.30 54.70 ± 1.40 64, 512 MoE-1024 3.93M 1.27B 12.43 21.67 ± 0.29 8.75 ± 0.40 40.82 ± 0.49 60.86 ± 1.00 77.10 ± 1.33 53.67 ± 1.40 MoRE-8192 3.98M 7.92B 11.92 22.15 ± 0.29 8.68 ± 0.39 42.33 ± 0.49 61.74 ± 1.00 81.50 ± 1.23 51.70 ± 1.40 Table 1: Pretraining results for MoE Transformers on 40B tokens. Each MoRE configuration (r=64r=64) matches the active FLOPs of its MoE counterpart, while total parameters are scaled up. All tasks are evaluated 0-shot, and ± denotes the standard error over test samples. Active MoE FLOPs shown are at inference and we have matched FLOPs at training too. See Appendices C.2 and C.3 for details. End-to-end wall-clock cost comparison. At fixed active FLOPs (Table 1 and Table 9), MoE and MoRE have similar inference time in both prefilling and decoding (Tables 2 and 11). With the same configuration, training MoRE is 1.101.10–1.23×1.23× faster than training MoE.88 8 At matched active FLOPs, training MoRE can be slower because of work that scales with total parameters, such as the AdamW update. A more careful implementation may shorten this gap (Appendix E.5). Scale Model #Experts Total params Prefill (ms) s=32s=32, h=1024h=1024 MoE 128 0.29 B 115.4 MoRE (r=64r=64) 1,536 1.50 B 115.1 s=64s=64, h=512h=512 MoE 256 0.36 B 101.4 MoRE (r=64r=64) 1,792 1.91 B 95.3 s=128s=128, h=512h=512 MoE 256 0.66 B 102.0 MoRE (r=64r=64) 1,792 3.76 B 98.0 s=64s=64, h=512h=512 MoE 1,024 1.27 B 105.4 MoRE (r=64r=64) 8,192 7.92 B 104.0 Table 2: MoRE prefills as fast as the MoE it matches in active FLOPs, while holding 55 to 66 times as many parameters. Prefill wall-clock (ms) for the pairs of Table 1: the forward pass of the whole model on one NVIDIA B200 at the per-device batch the training runs use, T=1024T=1024, mean of 200200 iterations after 100100 warm-ups. Training time for the same pairs is in Table 12. Ablation of the router rank at matched active FLOPs. Reducing r from 256256 to 1616 at fixed active FLOPs raises MoRE’s expert count from 6464 to 8,7048,704 and lowers validation perplexity monotonically, from 18.3518.35 to 14.9014.90 (Figure 6). As the benefit plateaus at r=64r=64, we keep r=64r=64 throughout our pretraining experiments. 6 Conclusion and Discussion We introduce MoRE, a principled solution to the MoE routing bottleneck that scales experts by a factor of Θ(hr) ( hr) while matching active FLOPs. We provide theoretical guarantees for its expressivity, design a hardware-efficient implementation, and demonstrate improved performance on knowledge-intensive benchmarks such as TriviaQA (Joshi et al., 2017) and Natural Questions (Kwiatkowski et al., 2019), without degrading reasoning. Limitations and future work. Our real-world training experiments are still limited in scale (h≤2048h≤ 2048, fewer than 10B total parameters) due to our limited computational budget. We also mainly compare MoE and MoRE under the principle of Section 4.1 for improving memorization, and future work could study the scaling laws of the two architectures. While we find that matching the standard router’s initialization variance can already improve upon the default initialization of R1,R2R_1,R_2 (Appendix D), a LoRA-like (Hu et al., 2021) initialization might improve performance further. 7 Acknowledgements We thank Wentao Guo for helpful discussions and feedback. HW is supported by the gift from AWS to Penn Engineering’s ASSET Center for Trustworthy AI. References Abnar et al. (2025) S. Abnar, H. Shah, D. Busbridge, A. M. E. Ali, J. Susskind, and V. Thilak Parameters vs flops: scaling laws for optimal sparsity for mixture-of-experts language models. arXiv preprint arXiv:2501.12370. Cited by: §1. Allen-Zhu and Li (2023) Z. Allen-Zhu and Y. Li Physics of language models: part 3.1, knowledge storage and extraction. External Links: 2309.14316, Link Cited by: §C.2, §C.2, §3.2, §5.2. Allen-Zhu and Li (2024) Z. Allen-Zhu and Y. Li Physics of language models: part 3.3, knowledge capacity scaling laws. External Links: 2404.05405, Link Cited by: §1.2, §3.2, §3.2. Ben Allal et al. (2024) L. Ben Allal, A. Lozhkov, G. Penedo, T. Wolf, and L. von Werra Cosmopedia: how to create large-scale synthetic data for pre-training. Hugging Face Blog. Note: https://huggingface.co/blog/cosmopedia Cited by: §C.2, §5.2. Berant et al. (2013) J. Berant, A. Chou, R. Frostig, and P. Liang Semantic parsing on Freebase from question-answer pairs. In Proceedings of EMNLP, Cited by: §C.2. Berges et al. (2024) V. Berges, B. Oğuz, D. Haziza, W. Yih, L. Zettlemoyer, and G. Ghosh Memory layers at scale. External Links: 2412.09764, Link Cited by: §1.2. Bisk et al. (2020) Y. Bisk, R. Zellers, R. L. Bras, J. Gao, and Y. Choi PIQA: reasoning about physical commonsense in natural language. In Proceedings of AAAI, Cited by: §C.2. Boix-Adsera (2025) E. Boix-Adsera Secret mixtures of experts inside your llm. arXiv preprint arXiv:2512.18452. Cited by: §1.2. Cheng et al. (2026) X. Cheng, W. Zeng, D. Dai, Q. Chen, B. Wang, Z. Xie, K. Huang, X. Yu, Z. Hao, Y. Li, H. Zhang, H. Zhang, D. Zhao, and W. Liang Conditional memory via scalable lookup: a new axis of sparsity for large language models. External Links: 2601.07372, Link Cited by: §1.2. Chi et al. (2022) Z. Chi, L. Dong, S. Huang, D. Dai, S. Ma, B. Patra, S. Singhal, P. Bajaj, X. Song, X. Mao, H. Huang, and F. Wei On the representation collapse of sparse mixture of experts. In Advances in Neural Information Processing Systems (NeurIPS), External Links: 2204.09179, Link Cited by: Appendix B, §1.2. Clark et al. (2022) A. Clark, D. de las Casas, A. Guy, A. Mensch, M. Paganini, J. Hoffmann, B. Damoc, B. Hechtman, T. Cai, S. Borgeaud, G. van den Driessche, E. Rutherford, T. Hennigan, M. Johnson, K. Millican, A. Cassirer, C. Jones, E. Buchatskaya, D. Budden, L. Sifre, S. Osindero, O. Vinyals, J. Rae, E. Elsen, K. Kavukcuoglu, and K. Simonyan Unified scaling laws for routed language models. External Links: 2202.01169, Link Cited by: §1.2. Clark et al. (2018) P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord Think you have solved question answering? try ARC, the AI2 reasoning challenge. arXiv preprint arXiv:1803.05457. Cited by: §C.2, §5.2. Costin et al. (2025) B. Costin, T. Averbuch, D. Pai, N. Chen, and B. Keigwin MoMoE: memory-optimized mixture of experts. Note: https://blog.tilderesearch.com/blog/momoeTilde Research blog, July 25, 2025 Cited by: §E.1, §1.2. Daniely (2020) A. Daniely Memorizing gaussians with no over-parameterizaion via gradient decent on neural networks. External Links: 2003.12895, Link Cited by: §3.2. DeepSeek-AI (2024a) DeepSeek-AI DeepSeek-v2: a strong, economical, and efficient mixture-of-experts language model. External Links: 2405.04434, Link Cited by: Table 3, §1.1, §1. DeepSeek-AI (2024b) DeepSeek-AI DeepSeek-v3 technical report. External Links: 2412.19437, Link Cited by: Table 3, §1. Dikkala et al. (2023) N. Dikkala, N. Ghosh, R. Meka, R. Panigrahy, N. Vyas, and X. Wang On the benefits of learning to route in mixture-of-experts models. In The 2023 Conference on Empirical Methods in Natural Language Processing, External Links: Link Cited by: §1.2. Du et al. (2022) N. Du, Y. Huang, A. M. Dai, S. Tong, D. Lepikhin, Y. Xu, M. Krikun, Y. Zhou, A. W. Yu, O. Firat, et al. Glam: efficient scaling of language models with mixture-of-experts. In International conference on machine learning, p. 5547–5569. Cited by: §3.2. Fedus et al. (2022) W. Fedus, B. Zoph, and N. Shazeer Switch transformers: scaling to trillion parameter models with simple and efficient sparsity. External Links: 2101.03961, Link Cited by: §B.1, Appendix D, §E.3, §F.2.1, §1.2, §2. Gale et al. (2023) T. Gale, D. Narayanan, C. Young, and M. Zaharia MegaBlocks: efficient sparse training with mixture-of-experts. In Proceedings of Machine Learning and Systems, Vol. 5. Cited by: §E.1, §1.2. Gao et al. (2023) L. Gao, J. Tow, B. Abbasi, S. Biderman, S. Black, A. DiPofi, C. Foster, L. Golding, J. Hsu, A. Le Noac’h, H. Li, K. McDonell, N. Muennighoff, C. Ociepa, J. Phang, L. Reynolds, H. Schoelkopf, A. Skowron, L. Sutawika, E. Tang, A. Thite, B. Wang, K. Wang, and A. Zou A framework for few-shot language model evaluation. Zenodo. External Links: Document, Link Cited by: §5.2. Geva et al. (2021) M. Geva, R. Schuster, J. Berant, and O. Levy Transformer feed-forward layers are key-value memories. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, p. 5484–5495. Cited by: §1.2, §3.2. Guo et al. (2025) W. Guo, M. Mishra, X. Cheng, I. Stoica, and T. Dao SonicMoE: accelerating moe with io and tile-aware optimizations. External Links: 2512.14080, Link Cited by: Appendix C, §E.1, §E.2, Appendix E, §1.2. Hayou et al. (2024) S. Hayou, N. Ghosh, and B. Yu LoRA+: efficient low rank adaptation of large models. External Links: 2402.12354, Link Cited by: §D.3. He et al. (2024) X. He, S. Zhang, Y. Wang, Z. Zeng, K. Tang, S. Shi, Z. Tang, H. Yin, I. W. Tsang, Y. S. Ong, and X. Chu ExpertFlow: optimized expert activation and token allocation for efficient mixture-of-experts inference. External Links: 2410.17954, Link Cited by: §1.2. He (2024) X. O. He Mixture of a million experts. External Links: 2407.04153, Link Cited by: §B.1, Appendix B, §1.2. Hsu et al. (2024) P. Hsu, Y. Dai, V. Kothapalli, Q. Song, S. Tang, S. Zhu, S. Shimizu, S. Sahni, H. Ning, and Y. Chen Liger kernel: efficient triton kernels for LLM training. External Links: 2410.10989, Link Cited by: Figure 23. Hu et al. (2021) E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen LoRA: low-rank adaptation of large language models. External Links: 2106.09685, Link Cited by: §D.3, §1.1, §6. Huang et al. (2025) Z. Huang, Q. Min, H. Huang, D. Zhu, Y. Zeng, R. Guo, and X. Zhou Ultra-sparse memory network. External Links: 2411.12364, Link Cited by: §1.2. HuggingFaceFW (2024) HuggingFaceFW Fineweb-edu (revision 22b0aca). Hugging Face. External Links: Link, Document Cited by: §C.2, §5.2. Hwang et al. (2024) R. Hwang, J. Wei, S. Cao, C. Hwang, X. Tang, T. Cao, and M. Yang Pre-gated MoE: an algorithm-system co-design for fast and scalable mixture-of-expert inference. In Proceedings of the 51st International Symposium on Computer Architecture (ISCA), External Links: 2308.12066, Link Cited by: §1.2. Jelassi et al. (2025) S. Jelassi, C. Mohri, D. Brandfonbrener, A. Gu, N. Vyas, N. Anand, D. Alvarez-Melis, Y. Li, S. M. Kakade, and E. Malach Mixture of parrots: experts improve memorization more than reasoning. External Links: 2410.19034, Link Cited by: Appendix J, Appendix J, §C.2, §C.2, §1.1, §1.2, §1.2, §1, §3.2, §3.2, §5.2. Jiang et al. (2024) A. Q. Jiang, A. Sablayrolles, A. Roux, et al. Mixtral of experts. External Links: 2401.04088, Link Cited by: Appendix B, §1. Jiang et al. (2026) T. Jiang, B. Bordelon, C. Pehlevan, and B. Hanin Hyperparameter transfer with mixture-of-expert layers. External Links: 2601.20205, Link Cited by: §D.3. Johnson and Lindenstrauss (1984) W. B. Johnson and J. Lindenstrauss Extensions of Lipschitz mappings into a Hilbert space. In Conference in Modern Analysis and Probability (New Haven, Conn., 1982), Contemporary Mathematics, Vol. 26, p. 189–206. External Links: Document Cited by: Appendix J, §3.1. Joshi et al. (2017) M. Joshi, E. Choi, D. S. Weld, and L. Zettlemoyer TriviaQA: a large scale distantly supervised challenge dataset for reading comprehension. External Links: 1705.03551, Link Cited by: §C.2, §1.1, §5.2, §6. Kaplan et al. (2020) J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei Scaling laws for neural language models. arXiv preprint arXiv:2001.08361. Cited by: §1. Kimi Team (2025) Kimi Team Kimi K2: open agentic intelligence. arXiv preprint arXiv:2507.20534. Cited by: Table 3. Krajewski et al. (2024) J. Krajewski, J. Ludziejewski, K. Adamczewski, M. Pióro, M. Krutul, S. Antoniak, K. Ciebiera, K. Król, T. Odrzygóźdź, P. Sankowski, M. Cygan, and S. Jaszczur Scaling laws for fine-grained mixture of experts. External Links: 2402.07871, Link Cited by: §C.3, §1.2, §1, §2. Kwiatkowski et al. (2019) T. Kwiatkowski, J. Palomaki, O. Redfield, M. Collins, A. Parikh, C. Alberti, D. Epstein, I. Polosukhin, J. Devlin, K. Lee, K. Toutanova, L. Jones, M. Kelcey, M. Chang, A. M. Dai, J. Uszkoreit, Q. Le, and S. Petrov Natural questions: a benchmark for question answering research. Transactions of the Association for Computational Linguistics 7, p. 452–466. External Links: Document Cited by: §C.2, §1.1, §5.2, §6. Kwon et al. (2023) W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C. H. Yu, J. E. Gonzalez, H. Zhang, and I. Stoica Efficient memory management for large language model serving with PagedAttention. In Proceedings of the 29th Symposium on Operating Systems Principles (SOSP), External Links: 2309.06180, Link Cited by: §1.2. Lepikhin et al. (2020) D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen GShard: scaling giant models with conditional computation and automatic sharding. External Links: 2006.16668, Link Cited by: §E.3, §1.2. Li et al. (2026) M. Li, S. Kudugunta, D. Rothermel, and L. Zettlemoyer Slicing and dicing: configuring optimal mixtures of experts. External Links: 2605.11689, Link Cited by: §1.2. Loshchilov and Hutter (2019) I. Loshchilov and F. Hutter Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR), External Links: Link Cited by: §C.2. Meta AI (2025) Meta AI The Llama 4 herd: the beginning of a new era of natively multimodal AI innovation. Note: https://ai.meta.com/blog/llama-4-multimodal-intelligence/ Cited by: Table 3. Mihaylov et al. (2018) T. Mihaylov, P. Clark, T. Khot, and A. Sabharwal Can a suit of armor conduct electricity? a new dataset for open book question answering. In Proceedings of EMNLP, Cited by: §C.2. Milakov and Gimelshein (2018) M. Milakov and N. Gimelshein Online normalizer calculation for softmax. External Links: 1805.02867, Link Cited by: §F.2.1. Mistral AI (2024) Mistral AI Cheaper, better, faster, stronger: continuing to push the frontier of AI and making it accessible to all. Note: https://mistral.ai/news/mixtral-8x22b Cited by: Table 3. OpenAI (2025) OpenAI gpt-oss-120b & gpt-oss-20b model card. arXiv preprint arXiv:2508.10925. Cited by: Table 3. Pope et al. (2023) R. Pope, S. Douglas, A. Chowdhery, J. Devlin, J. Bradbury, A. Levskaya, J. Heek, K. Xiao, S. Agrawal, and J. Dean Efficiently scaling transformer inference. In Proceedings of Machine Learning and Systems (MLSys), External Links: 2211.05102, Link Cited by: §E.4. Qwen Team (2025) Qwen Team Qwen3-Next: towards ultimate training & inference efficiency. Note: https://qwenlm.github.io/blog/qwen3_next/ Cited by: Table 3. Radford et al. (2019) A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever Language models are unsupervised multitask learners. Technical report OpenAI. External Links: Link Cited by: §C.2. Roller et al. (2021) S. Roller, S. Sukhbaatar, A. Szlam, and J. Weston Hash layers for large sparse models. External Links: 2106.04426, Link Cited by: §1.2. Sakaguchi et al. (2020) K. Sakaguchi, R. Le Bras, C. Bhagavatula, and Y. Choi WinoGrande: an adversarial Winograd schema challenge at scale. In Proceedings of the AAAI Conference on Artificial Intelligence, Cited by: §C.2, §5.2. Sap et al. (2019) M. Sap, H. Rashkin, D. Chen, R. L. Bras, and Y. Choi SocialIQA: commonsense reasoning about social interactions. In Proceedings of EMNLP-IJCNLP, Cited by: §C.2. Schulman and Thinking Machines Lab (2025) J. Schulman and Thinking Machines Lab LoRA without regret. Note: https://thinkingmachines.ai/blog/lora/Thinking Machines Lab Connectionism blog, September 29, 2025 Cited by: §D.3. Shazeer et al. (2017) N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean Outrageously large neural networks: the sparsely-gated mixture-of-experts layer. External Links: 1701.06538, Link Cited by: §1.2, §2. Shazeer (2019) N. Shazeer Fast transformer decoding: one write-head is all you need. External Links: 1911.02150, Link Cited by: §E.4. Shazeer (2020) N. Shazeer GLU variants improve transformer. External Links: 2002.05202, Link Cited by: §2. Talmor and Berant (2018) A. Talmor and J. Berant The web as a knowledge-base for answering complex questions. In Proceedings of NAACL-HLT, Cited by: §C.2. Tan et al. (2024) S. Tan, Y. Shen, R. Panda, and A. Courville Scattered mixture-of-experts implementation. External Links: 2403.08245, Link Cited by: Appendix C, Figure 9, Appendix E, §E.1, §E.4, Appendix E, §1.2. Tillet et al. (2019) P. Tillet, H. T. Kung, and D. Cox Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages (MAPL), p. 10–19. External Links: Document Cited by: §4. Welbl et al. (2017) J. Welbl, N. F. Liu, and M. Gardner Crowdsourcing multiple choice science questions. In Proceedings of the 3rd Workshop on Noisy User-generated Text, p. 94–106. Cited by: §C.2, §5.2. xAI (2024) xAI Open release of Grok-1. Note: https://x.ai/blog/grok-os Cited by: Table 3, §1. Xue et al. (2024) L. Xue, Y. Fu, Z. Lu, C. Sun, L. Mai, and M. Marina MoE-Infinity: efficient MoE inference on personal machines with sparsity-aware expert cache. External Links: 2401.14361, Link Cited by: §1.2. Yan et al. (2026) Z. Yan, H. Bai, X. Yao, D. Liu, T. Liu, H. Liu, P. Li, E. Wu, S. Fan, L. Tao, R. Zhang, Y. Wang, S. Xu, J. Chang, X. Chen, K. Li, Y. Bai, G. Deng, N. Zheng, V. A. Korthikanti, A. Khattar, E. He, S. Govande, S. Lym, Z. Zhu, Q. Zhang, H. Yuan, X. Ren, D. Fu, T. Ma, S. Zhang, J. Shao, R. Wang, V. Rengasamy, R. Garg, S. Bhavani, X. Li, C. Zhou, D. Wu, Y. Wei, A. Aithal, M. Andersch, M. Shoeybi, J. Yao, and J. Yang Scalable training of mixture-of-experts models with megatron core. External Links: 2603.07685, Link Cited by: §E.3, §1. Yang et al. (2025) A. Yang, A. Li, B. Yang, B. Zhang, B. Hui, B. Zheng, B. Yu, C. Gao, C. Huang, C. Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388. Cited by: Table 3. Yang et al. (2026) M. Yang, R. Togo, G. Li, T. Ogawa, and M. Haseyama L2R: low-rank and lipschitz-controlled routing for mixture-of-experts. External Links: 2601.21349, Link Cited by: Appendix B, Appendix B, §1.2. Yang et al. (2018) Z. Yang, P. Qi, S. Zhang, Y. Bengio, W. W. Cohen, R. Salakhutdinov, and C. D. Manning HotpotQA: a dataset for diverse, explainable multi-hop question answering. In Proceedings of EMNLP, Cited by: §C.2. Zadouri et al. (2026) T. Zadouri, M. Hoehnerbach, J. Shah, T. Liu, V. Thakkar, and T. Dao FlashAttention-4: algorithm and kernel pipelining co-design for asymmetric hardware scaling. External Links: 2603.05451, Link Cited by: §H.5. Zellers et al. (2019) R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi HellaSwag: can a machine really finish your sentence?. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics (ACL), p. 4791–4800. Cited by: §C.2, §5.2. Appendix 1 Introduction 1.1 Our contributions 1.2 Related Work 2 Preliminaries 3 Low-rank routing preserves memorization capabilities 3.1 Routing expressivity 3.2 Load balancing and memorization capacity 4 Scaling experts efficiently with fused implementation of MoRE 4.1 Scaling under fixed active-FLOPs 4.2 Realizing the speedup at large M 5 Experiments: MoRE improves performance at matched active FLOPs 5.1 Synthetic task: increased memorization capabilities at matched active FLOPs 5.2 Real-world pretraining: improved performance on knowledge-intensive tasks 6 Conclusion and Discussion 7 Acknowledgements References A Frontier MoE Model Configurations B Additional Discussions in Related Work B.1 Wall-clock Comparison with PEER C Experimental Details C.1 Phonebook C.2 Real-world Pretraining C.3 FLOPs Calculation and Total parameter accounting D Initialization D.1 Setup D.2 Matching variance D.3 Discussion E Efficiency E.1 GPU Memory Hierarchy and Related Analysis E.2 Top-K Algorithm Analysis E.3 Expert Parallelism E.3.1 Scaling under Expert Parallelism E.4 Decoding E.5 Wall-clock measurements for configurations in Table F Pseudocode F.1 Inference (prefilling) F.2 Training F.2.1 Forward-Pass Pseudocode F.2.2 Backward-Pass Pseudocode G Additional Plots G.1 Routing entropy G.2 Adding More Experts help G.3 The perplexity gap grows with the token budget H Additional Plots for Efficiency H.1 Forward pass of a single layer at larger expert widths H.2 The same grid on an A100 H.3 Fusing a standard MoE layer H.4 The low-rank router once both layers are fused H.5 End-to-End Forward-Pass Breakdown at Large M H.6 Fusion benefit at different batch sizes I Scaling Total Parameters J Theoretical Results Appendix A Frontier MoE Model Configurations For reference, Table 3 summarizes the configurations of recent frontier MoE models. The trend toward higher sparsity (M≫kM k) and higher granularity (s<hs<h) motivates the fine-grained regime studied throughout this paper. Model Total Active L M k h s Mixtral 8x22B Mistral AI (2024) 141B 39B 56 8 2 6144 16384 Grok-1 xAI (2024) 314B 79B 64 8 2 6144 32768 DeepSeek-V2 DeepSeek-AI (2024a) 236B 21B 60 160 6 5120 1536 DeepSeek-V3 DeepSeek-AI (2024b) 671B 37B 61 256 8 7168 2048 Llama 4 Maverick Meta AI (2025) 400B 17B 48 128 1 5120 8192 Qwen3 235B-A22B Yang et al. (2025) 235B 22B 94 128 8 4096 1536 Kimi K2 Kimi Team (2025) 1.04T 32B 61 384 8 7168 2048 gpt-oss-120b OpenAI (2025) 117B 5.1B 36 128 4 2880 2880 Qwen3-Next-80B-A3B Qwen Team (2025) 80B 3B 48 512 10 2048 512 Table 3: Configurations of frontier MoE models. M: total experts; k: active experts per token; h: hidden dimension; s: expert FFN intermediate dimension. The trend shows greater sparsity and granularity is preferred. Appendix B Additional Discussions in Related Work Low-rank routing in the M≪hM h regime. X-MoE (Chi et al., 2022) scores experts by the dot product between a hidden vector x∈ℝhx ^h and M learned expert embeddings. The gradient pulls x toward the span of those embeddings, so after training the hidden states entering a routed layer concentrate in a subspace of ℝhR^h of dimension at most M. When M≪hM h, that subspace is much smaller than ℝhR^h, a phenomenon the authors call representation collapse. To mitigate this, Chi et al. (2022) project both expert embeddings and the hidden vector to a lower dimension. L2R (Yang et al., 2026) identifies a related but distinct issue: when M≪hM h, token representations and expert embeddings become nearly orthogonal, which narrows the margin between routing logits and produces ambiguous expert selection. L2R introduces low-rank routing to improve discrimination among experts. Our work also uses low-rank routing and derives upper and lower bounds on the rank based on the margin argument, but it serves a different purpose. The mechanisms above require M≪hM h (e.g., h=4096h=4096 with M≤64M≤ 64, as in Yang et al. (2026) and Mixtral (Jiang et al., 2024)). MoRE instead targets the fine-grained regime where M is scaled up and M≪hM h does not hold. In this regime, the low-rank router reduces routing cost rather than correcting the geometry of token representations. PEER (He, 2024) reaches a very large expert pool by shrinking each expert to a single neuron and replacing the linear router with product-key retrieval, which has complexity O(hM)O(h M) and is asymptotically better than our O((M+h)r)O((M+h)r) where r=Ok(logM)r=O_k( M). However, in the following comparison, we find that MoRE improves upon PEER based on the wall-clock speed in practice, with matched active FLOPs. B.1 Wall-clock Comparison with PEER A PEER layer holds M experts, each a single neuron ei(x)=σ(ui⊤x)vie_i(x)=σ(u_i x)\,v_i with ui,vi∈ℝhu_i,v_i ^h, so an expert has width 11 and 2h2h parameters. Routing uses H query networks, each a matrix in ℝdq×hR^d_q× h, together with two sub-key tables of shape M×dq/2 M× d_q/2. Each head splits its query in half, matches the halves against the two tables, keeps the top κ on each axis, and takes the top κ of the κ2κ^2 candidate sums, so the layer activates k=Hκk=Hκ experts per token. The per-layer parameter counts and per-token active FLOPs (we count the query projections and the sub-key tables as active FLOPs too) are Pexperts=2Mh,Prouter=Hhdq+Mdq,τPEER=12kh+6Hdq(h+M),P_experts=2Mh, P_router=Hhd_q+ M\,d_q, _PEER=12\,kh+6\,Hd_q (h+ M ), (2) to be read against 3Mhs3Mhs, (M+h)r(M+h)r and Cekhs+Cr(M+h)rC_ekhs+C_r(M+h)r for MoRE. For comparison, we implement PEER with reference to the public PEER-pytorch repository,99 9 https://github.com/lucidrains/PEER-pytorch and change the code to make the heads share one set of product keys, following He (2024). We compare MoRE against it at matched total parameters and matched active FLOPs. We match against the (s=64,h=512)(s=64,\,h=512) row of Table 1, where MoRE-1792 uses M′=1792M =1792, s′=56s =56, k=4k=4 and r=64r=64 over L=12L=12 layers. Matching expert parameters gives M=32M′s′=150,528M= 32M s =150,528, and since product-key retrieval needs M to be a perfect square we take M=3882=150,544M=388^2=150,544. Matching active FLOPs at H=8H=8 then gives dq=32d_q=32 and κ=56κ=56, so k=448k=448 experts are active per token. We choose dq=32d_q=32 rather than the dq=hd_q=h of the default implementation in order to match the router FLOPs. Table 4 measures one layer at the matched configuration. At matched active FLOPs, a MoRE layer takes substantially less wall-clock time than a PEER layer. Because M is large for PEER, we additionally implement a fused kernel for the expert application that never materializes the gathered weights. This fused kernel runs 1.90×1.90× faster on that stage, but the wall-clock ratio remains large. Forward Forward ++ backward Tokens per step MoRE PEER MoRE PEER 4,0964,096 1.061.06 5.565.56 (5.2×5.2×) 2.942.94 16.3016.30 (5.5×5.5×) 16,38416,384 1.231.23 22.0622.06 (18.0×18.0×) 3.013.01 60.9660.96 (20.3×20.3×) 65,53665,536 1.931.93 87.7987.79 (45.6×45.6×) 4.754.75 244.01244.01 (51.4×51.4×) 262,144262,144 4.924.92 350.86350.86 (71.4×71.4×) 14.1914.19 974.41974.41 (68.7×68.7×) Table 4: Single-layer wall-clock time at matched active FLOPs. The ratio grows with batch size. Single-layer wall-clock in milliseconds at the matched configuration, on one NVIDIA B200 in bf16, median of 2020 iterations after 55 warm-ups. Table 5 reports validation perplexity for the three architectures trained for 11B tokens on mixture dataset specified in Appendix C. MoE and MoRE use the Switch load-balancing loss (Fedus et al., 2022). Architecture Validation PPL Seconds per step Training time MoE-256 24.2624.26 0.3210.321 20.420.4 min MoRE-1792 24.1324.13 0.4050.405 25.825.8 min PEER 24.1724.17 6.1356.135 390.1390.1 min Table 5: Validation perplexity at 11B tokens, h=512h=512, L=12L=12, four B200s, 262,144262,144 tokens per optimizer step. PEER and MoRE reach similar perplexity under the 11B-token budget. The wall-clock cost differs by 15×15×: MoRE trains in 25.825.8 minutes compared to 390.1390.1 minutes for PEER (Table 5). Appendix C Experimental Details Router initialization. We initialize each entry of the low-rank router factors as (0,σ′2)N(0,σ 2) with σ′=σ/r1/4σ = σ/r^1/4 and σ=0.02σ=0.02, so that the variance of the router logits matches the standard router. See Appendix D for the derivation and an empirical comparison to the default initialization. To make the speedup by the low-rank router more apparent, we seek an efficient implementation of the Mixture-of-Experts architecture and therefore build upon ScatterMoE (Tan et al., 2024) (refer to Appendix E for more details). We believe that further speedups are achievable on specific hardware by applying more careful optimization as in Guo et al. (2025). We use load balancing coefficient as 0.001 for both MoE and MoRE setups and do not have evidence that MoRE is sensitive to its choice. C.1 Phonebook Tokenization. Each example is tokenized as “$abcde|12345678.”, with a five-letter lowercase name, the separator “|”, and an eight-digit number, padded with “*” to a fixed sequence length of 1616. The vocabulary has V=40V=40 symbols: 2626 lowercase letters, 1010 digits, and 44 specials (“$” BOS, “.” EOS, “|” separator, “*” pad). Architecture and sweep. All models are L=6L=6-layer Transformers with top-k routing at k=4k=4 and low-rank rank r=16r=16. The active budget τ has two natural scaling axes at fixed h: the hidden dimension itself and the expert width s. We sweep each in turn: (i) h∈64,128,256,512h∈\64,128,256,512\ at fixed s=8s=8, and (i) s∈8,16,32,64s∈\8,16,32,64\ at fixed h=64h=64. For every (h,s)(h,s), the standard baseline takes M=4sM=4s experts, and the low-rank counterpart takes the largest M′M that matches the standard active FLOPs at r=16r=16, following Proposition 2. Choice of s′s . Under our convention Ce=CrC_e=C_r, Proposition 2 prescribes the optimal MoRE expert width as s′∗=s−r/(2k)=s−2s *=s-r/(2k)=s-2 at r=16,k=4r=16,k=4. For s∈8,16,32,64s∈\8,16,32,64\ this gives s′∗∈6,14,30,62s *∈\6,14,30,62\, none of which are multiples of 88 – the kernel-friendly width for ScatterMoE. We round each s′s to the nearest multiple of 88, which collapses to s itself across the entire sweep, and we use the prescribed M′∗=h(2ks−r)/(2r)=h(s−2)/4M *=h(2ks-r)/(2r)=h(s-2)/4 unchanged. The rounding inflates the total MLP parameter count by a factor s/(s−2)s/(s-2), which is largest at the smallest scale (8/6≈1.33×8/6≈ 1.33× at s=8s=8) and shrinks to ≈1.03×≈ 1.03× at s=64s=64, it does not change the qualitative trends of Figure 5. Training and learning-rate selection. Each configuration is trained with batch size 1,0241,024 under a constant-with-warmup schedule (warmup ratio 0.050.05, weight decay 0.010.01, β2=0.95 _2=0.95), with early stopping once exact-string accuracy on the digit numbers crosses 0.90.9. The learning rate is selected per configuration by a short calibration phase over ×10−4,×10−4,×10−3\1\!×\!10^-4,\,3\!×\!10^-4,\,1\!×\!10^-3\, picking the value with the highest validation character accuracy on a brief preliminary run; the calibrated rate is then used for the full 1,5001,500-epoch run. Per-dot configurations. Tables 6 and 7 list, for every dot in Figure 5, the full (M,k,h,s)(M,k,h,s) and (M,k,h,s,r)(M,k,h,s,r) configuration and the maximum memorizable phonebook size at the calibrated learning rate. Table 6 covers the vary-h sweep at s=8s=8, and Table 7 covers the vary-s sweep at h=64h=64; the (h=64,s=8)(h=64,\,s=8) corner appears in both as the shared anchor between the two sweeps. The number of experts M are small for Table 6 as the optimal M=ksM=ks in our setup, when s=8s=8, M=ks=32M=ks=32. Configuration M/M′M/M k h s r Max memorizable MoE 32 4 64 8 — 10K MoE 32 4 128 8 — 30K MoE 32 4 256 8 — 100K MoE 32 4 512 8 — 200K MoRE 96 4 64 8 16 30K MoRE 192 4 128 8 16 100K MoRE 384 4 256 8 16 300K MoRE 768 4 512 8 16 1.2M Table 6: Phonebook per-dot configurations for the vary-h sweep at s=8s=8. M (M′M for MoRE) is the number of experts; k active per token; h hidden dim; s expert intermediate dim; r router rank (MoRE only). “Max memorizable” is the largest phonebook size at which the calibrated training run reached the early-stopping threshold. Configuration M/M′M/M k h s r Max memorizable MoE 32 4 64 8 — 10K MoE 64 4 64 16 — 30K MoE 128 4 64 32 — 150K MoE 256 4 64 64 — 500K MoRE 96 4 64 8 16 30K MoRE 224 4 64 16 16 100K MoRE 480 4 64 32 16 400K MoRE 992 4 64 64 16 1.2M Table 7: Phonebook per-dot configurations for the vary-s sweep at h=64h=64. Columns as in Table 6. The (h=64,s=8)(h=64,\,s=8) row is shared with Table 6. C.2 Real-world Pretraining Data mixture. We roughly follow the natural-language pre-training recipe of Jelassi et al. (2025, Appendix B.2), which consists of a ≈60.4≈ 60.4 B token mixture: (i) FineWeb-edu deduplicated (HuggingFaceFW, 2024) (37 B); (i) Cosmopedia synthetic textbooks (Ben Allal et al., 2024) (14 B); (i) Wikipedia (4 B unique, looped 3×3× for 12 B); and (iv) the training splits of 12 downstream knowledge-intensive and commonsense benchmarks (TriviaQA (Joshi et al., 2017), Natural Questions (Kwiatkowski et al., 2019), HotpotQA (Yang et al., 2018), WebQuestions (Berant et al., 2013), ComplexWebQuestions (Talmor and Berant, 2018), ARC-Easy (Clark et al., 2018), HellaSwag (Zellers et al., 2019), OpenBookQA (Mihaylov et al., 2018), PIQA (Bisk et al., 2020), SciQ (Welbl et al., 2017), SocialIQA (Sap et al., 2019), and Winogrande (Sakaguchi et al., 2020)), each looped 3×3× to familiarize the model with the evaluation format (Allen-Zhu and Li, 2023). Every source uses only its train split. We include the benchmark training splits in line with Jelassi et al. (2025), since Allen-Zhu and Li (2023) show that mixed training is required for successful knowledge extraction. We sample disjoint training (40 B tokens) and 5,000-sequence validation sets uniformly from this pool, following the convention of Jelassi et al. (2025). We conduct zero-shot evaluations only, since we did not include few-shot demonstrations in training. Tokenization. We use the GPT-2 byte-pair-encoding tokenizer (Radford et al., 2019) with vocabulary size V=50,257V=50,257. Documents from each mixture source are concatenated and chunked to a fixed sequence length of T=1,024T=1,024 tokens, then materialized as a single uint16 memory-mapped NumPy file for efficient reuse across runs. Architecture. All models are L=12L=12-layer Transformers with top-k routing at k=4k=4, head dimension 6464, and untied input/output embeddings. We compare three reference points (h,s)∈(1024,32),(512,64),(512,128)(h,s)∈\(1024,32),\,(512,64),\,(512,128)\. For the first two, the standard baseline takes M=4sM=4s experts, and the low-rank counterpart fixes router rank r=64r=64 and takes the largest M′M that matches the standard active FLOPs, following Proposition 2. The third point (512,128)(512,128) is constructed by scaling both s and s′s of (512,64)(512,64) by 22 while keeping M=256M=256 and M′=1792M =1792 fixed. Full configurations are listed in Table 8. Training and learning-rate selection. Each configuration is trained for ≈40≈ 40 B tokens, with sequence length 1,0241,024 on 44 GPUs in DDP, using AdamW (Loshchilov and Hutter, 2019) (β1=0.9 _1=0.9, β2=0.95 _2=0.95, ϵ=10−8ε=10^-8), weight decay 0.010.01, max grad norm 1.01.0, and a cosine learning-rate schedule with warmup ratio 0.050.05, in bf16 precision. The peak (global) learning rate is selected per configuration from ×10−4,×10−3,×10−3,×10−3\5\!×\!10^-4,\,1\!×\!10^-3,\,3\!×\!10^-3,\,5\!×\!10^-3\ by best validation perplexity on a 2B-token run, which selects ×10−31\!×\!10^-3 for every configuration. We set the router learning rate, for both R1R_1 and R2R_2 of the low-rank router, equal to the global learning rate, as we find that validation loss is not very sensitive to the router learning rate (Figure 7). Figure 7: Expert coverage is insensitive to the router learning rate, and the optimal router learning rate is the global learning rate at every rank. We therefore do not tune the router learning rate separately, and set it equal to the global learning rate. Switch balancing loss, M=1792M=1792, s=56s=56, h=512h=512, k=4k=4, L=12L=12, trained on 500M tokens of the mixture corpus at global learning rate η=3×10−3η=3× 10^-3. The dashed vertical line in every panel marks that global rate, and stars mark each rank’s best cell. Top: the optimum coincides with the global learning rate at all three ranks. The 5×10−45× 10^-4 cell is omitted at every rank because its apparent advantage at this budget does not survive to 2B tokens. Bottom: coverage holds at 1.0001.000 at every rank up to 5×10−35× 10^-3, and at r=32r=32 across the entire grid. It falls only at 10−210^-2, and only for r=64r=64 and r=128r=128. Scale Model M k h s r s=32, h=1024 Standard 128 4 1024 32 – Low-rank (r=64) 1536 4 1024 24 64 s=64, h=512 Standard 256 4 512 64 – Low-rank (r=64) 1792 4 512 56 64 s=128, h=512 Standard 256 4 512 128 – Low-rank (r=64) 1792 4 512 112 64 Table 8: Architectural configurations for each variant in Table 1. M: number of experts. k: experts active per token. h: model hidden dimension. s: expert FFN hidden dimension. r: router rank (Low-rank only). All variants use the GPT-2 tokenizer. Within each Scale group, the Low-rank variant scales M up while reducing s slightly so that active FLOPs approximately match the Standard reference; the (s=128,h=512)(s=128,h=512) group is the exception, obtained by scaling s and s′s of (s=64,h=512)(s=64,h=512) by 22 with M and M′M fixed. C.3 FLOPs Calculation and Total parameter accounting For consistency with Proposition 2, we report active FLOPs per token (rather than active parameters) in Table 1, since the matched-active configuration grid is selected to fix this quantity. We use the standard convention that one fused multiply–add counts as 22 FLOPs. All h=512h=512 models use nh=8n_h=8 attention heads; h=1024h=1024 uses nh=16n_h=16, nkv=4n_kv=4. Head dimension is d=h/nh=64d=h/n_h=64 throughout. Vocabulary size is V=50,257V=50,257 (GPT-2 BPE), and input/output embeddings are untied. Active MoE FLOPs in Table 1. Per token and per MoE layer, the active FLOPs in training (forward and backward) are 18khs+14Mh18khs+14Mh for MoE and 18khs′+14(M′+h)r18khs +14(M +h)r for MoRE. At inference (forward only), they are 6khs+6Mh6khs+6Mh for MoE and 6khs′+6(M′+h)r6khs +6(M +h)r for MoRE. The constants are derived below. In the computation of the configurations, we treat Ce≈CrC_e≈ C_r. Table 1 reports the inference FLOPs, and Table 9 lists both the training and the inference FLOPs. (s,h)(s,h) Model Training FLOPs Inference FLOPs (32, 1024) MoE-128 4.19M 1.57M MoRE-1536 4.06M 1.57M (64, 512) MoE-256 4.19M 1.57M MoRE-1792 4.13M 1.57M (128, 512) MoE-256 6.55M 2.36M MoRE-1792 6.19M 2.26M (64, 512) MoE-1024 9.70M 3.93M MoRE-8192 9.72M 3.98M Table 9: Active MoE FLOPs per token per layer for the configurations in Table 1. Training FLOPs are 18khs+14Mh18khs+14Mh for MoE and 18khs′+14(M′+h)r18khs +14(M +h)r for MoRE. Inference FLOPs are 6khs+6Mh6khs+6Mh and 6khs′+6(M′+h)r6khs +6(M +h)r. FLOPs constants in Proposition 2. The forward-pass formula above tabulates the empirical FLOPs we report. For the analytical scaling argument, Proposition 2 writes the active-FLOP budget as τ=Ce⋅khs+Cr⋅Mhτ=C_e· khs+C_r· Mh, with Crr(h+M)C_rr(h+M) in place of CrMhC_rMh for the rank-r low-rank router, abstracting per-operation FLOPs into two scalar constants. We follow the full-training (forward + backward) accounting of Krajewski et al. (2024, Appendix E), which charges 66 FLOPs per pair of (active parameter, processed token) for each linear projection (forward 22 + backward input-grad 22 + backward weight-grad 22). A SwiGLU expert performs three projections, giving Ce=3×6=18C_e=3× 6=18. The standard router charges 66 for its linear projection plus 22 FLOPs per pair for each of four token-routing operations (forward and backward dispatches of inputs and outputs), giving Cr=6+4×2=14C_r=6+4× 2=14 as in Krajewski et al. (2024). The forward-only ratio is similar: Cefwd/Crfwd=6/6=1C_e^fwd/C_r^fwd=6/6=1, compared to Ce/Cr=18/14≈1.29C_e/C_r=18/14≈ 1.29 for full training. We set Ce=CrC_e=C_r, which simplifies the configurations (for example, the optimal MoE then has M=ksM=ks), and use the same CrC_r for the low-rank router. Only the ratio Ce/CrC_e/C_r enters the matching of an MoE and a MoRE configuration. The token-routing operations cost O(kh)O(kh) per token in ScatterMoE, so a stricter accounting charges each router only for its matrix multiplications (Cr=2C_r=2 at inference). Under this accounting, MoRE uses 4.5%4.5\% to 12.5%12.5\% fewer active FLOPs than its MoE counterpart in every pair of Table 1. The Θ(h/r) (h/r) ratio in Proposition 2 depends only on the leading CrMhC_rMh term, so it is insensitive to the precise value of CrC_r within the defensible range [6,20][6,20] noted by Krajewski et al. (2024). Total parameter count. Total parameters include the embedding matrix, all M experts (active and inactive), the router, attention, layer norms, and the unembedding head: Ptotal=2Vh⏟embed+LM head+h⏟final LN+L⋅[router⏟router+⏟all experts+2h2+2hnkvd⏟attention+2h⏟2 RMSNorm],P_total\;=\; 2\,V\,h_embed+LM head\;+\; h_final LN\;+\;L· [ P_router_router\;+\; 3\,M\,h\,s_all experts\;+\; 2h^2+2\,h\,n_kv\,d_attention\;+\; 2\,h_2 RMSNorm ], (3) where Prouter=hMP_router=h\,M for the standard router and r(h+M)r\,(h+M) for the low-rank variant. (RMSNorm has no bias, so each contributes h scale parameters.) The two bold terms are the only ones that depend on the expert configuration: once h is fixed, the embedding, attention, and norm terms are also fixed, and the parameter count varies only through the router and expert terms, which are set by M and s. The depth-independent embedding and LM head term is amortized over layers, so as the depth L grows the router and the experts account for a larger fraction of the total parameters. Appendix D Initialization The Switch Transformer (Fedus et al., 2022) sets the router weights elementwise to (0,σ2)N(0,σ^2) with σ=0.02σ=0.02, and prior work has reported that the precise router initialization variance has limited impact on standard-router performance. The natural default for the low-rank router carries the same per-entry standard deviation over to both factors, i.e. each entry of R1∈ℝM×rR_1 ^M× r and R2∈ℝr×hR_2 ^r× h is drawn i.i.d. from (0,σ2)N(0,σ^2) with the same σ=0.02σ=0.02. We refer to this as the default (or standard) initialization. Empirically, however, we find that the default initialization noticeably degrades validation perplexity and shifts the optimal router learning rate significantly. We instead choose σ1,σ2 _1, _2 so that the variance of the router logits matches the standard baseline. Lemma 1 makes this matching condition precise. D.1 Setup We treat x∈ℝhx ^h as a fixed nonzero input and condition on it throughout. The standard router R∈ℝM×hR ^M× h has entries i.i.d. (0,σ2)N(0,σ^2), and the low-rank factors R1∈ℝM×rR_1 ^M× r, R2∈ℝr×hR_2 ^r× h have independent entries i.i.d. (0,σ12)N(0, _1^2) and (0,σ22)N(0, _2^2) respectively, with R1R_1 independent of R2R_2. We denote the corresponding learning rates by η for the standard router, and η1,η2 _1, _2 for R1,R2R_1,R_2. Our goal is to derive σ1,σ2 _1, _2 (and, in future work, η1,η2 _1, _2) by matching the variance and the logit update of the standard router. D.2 Matching variance Lemma 1 (Variance matching). Under the setup above, for any fixed x≠0x≠ 0, Var([R1R2x]k)=Var([Rx]k)Var\! ([R_1R_2x]_k )=Var\! ([Rx]_k ) for every k∈[M]k∈[M] if and only if σ1σ2=σr. _1 _2\;=\; σ r. The matching condition does not depend on x. Proof. All variances below are conditional on the fixed input x. For the standard router, [Rx]k=∑j=1hRkjxj[Rx]_k= _j=1^hR_kj\,x_j is a linear combination of i.i.d. (0,σ2)N(0,σ^2) entries with deterministic coefficients xjx_j, so Var([Rx]k)=σ2‖x‖22.Var ([Rx]_k )\;=\;σ^2\|x\|_2^2. For the low-rank router, write z=R2x∈ℝrz=R_2x ^r, so [R1R2x]k=∑i=1r(R1)kizi[R_1R_2x]_k= _i=1^r(R_1)_ki\,z_i. The same linear-combination argument applied to R2R_2 gives Var(zi)=σ22‖x‖22Var(z_i)= _2^2\|x\|_2^2. Since R1R_1 is independent of R2R_2 and zero-mean, each summand satisfies Var((R1)kizi)=[(R1)ki2][zi2]=σ12σ22‖x‖22,Var ((R_1)_ki\,z_i )\;=\;E [(R_1)_ki^2 ]\,E [z_i^2 ]\;=\; _1^2\, _2^2\,\|x\|_2^2, and the r summands are mutually independent because they share no entry of R1R_1. Summing, Var([R1R2x]k)=rσ12σ22‖x‖22.Var ([R_1R_2x]_k )\;=\;r\, _1^2\, _2^2\,\|x\|_2^2. Equating the two variances gives rσ12σ22‖x‖22=σ2‖x‖22r\, _1^2\, _2^2\,\|x\|_2^2=σ^2\|x\|_2^2, and since ‖x‖22>0\|x\|_2^2>0 this is equivalent to σ1σ2=σ/r _1 _2=σ/ r, independent of x. ∎ In all the experiments in this paper, we set σ1=σ2 _1= _2 and η1=η2 _1= _2. Combined with Lemma 1, this gives σ1=σ2=σr1/4,σ=0.02. _1\;=\; _2\;=\; σr^1/4, σ=0.02. Figure 8: Router learning-rate sweep on FineWeb-edu at the optimal global learning rate η=10−3η=10^-3, comparing the default per-entry initialization (left) against the variance-matching initialization of Lemma 1 (right) for low-rank ranks r∈16,32,64r∈\16,32,64\. Each curve trains for 400400M FineWeb-edu tokens. The variance-matching initialization achieves lower validation perplexity across ranks, and also has stabler optimum across this range, although it might also shift leftwards when r is larger. D.3 Discussion Figure 8 presents our empirical comparison. We hypothesize that a theoretical analysis analogous to those developed for LoRA (Hu et al., 2021) and LoRA+ (Hayou et al., 2024; Schulman and Thinking Machines Lab, 2025) could provide a principled approach to selecting both the initialization and learning rates for routers in the low-rank setting, potentially eliminating the need to tune η1 _1 and η2 _2 individually for each rank. Although a key distinction from LoRA is that LoRA initializes its output to zero to preserve finetuning performance, whereas initializing the router’s output to zero is undesirable in our context: doing so would impede the ability to break ties when selecting the top k experts. In a complementary direction, Jiang et al. (2026) derive a hyperparameter-transfer parameterization for MoE routing weights via a dynamical mean-field analysis; we believe their framework can be adapted to the low-rank router setting as well. Appendix E Efficiency For the MoE architecture, we build upon ScatterMoE (Tan et al., 2024) and implement a new fused Triton kernel for the router (refer to Figure 9 for the detailed operations in the forward pass). We believe that further speedups are achievable on specific hardware by applying more careful optimization as in Guo et al. (2025). Why we build on ScatterMoE. Figure 9 illustrates the operations in one MoE layer forward pass. Beyond the router, the expert MLP involves token dispatch (sorting by expert assignment), two grouped linear transformations with a SwiGLU activation, and a scatter-back step that combines expert outputs. Each of these steps is a potential source of memory traffic. For example, a naive grouped MoE implementation pays O(Nkh)O(Nkh) memory traffic to materialize the gathered tokens (sorted by expert assignment) before the per-expert grouped GEMM, and pays the same again on the scatter-back. ScatterMoE (Tan et al., 2024) fuses scatter-gather into the per-expert grouped matmul, removing both round-trips and producing a kernel whose only HBM traffic is the input X, the weight stack, and the output Y. With this dispatch overhead eliminated, the standard router’s O(NhM)O(NhM) matmul is the only remaining cost that grows linearly with M. Our low-rank router is designed precisely to address this last bottleneck, while reusing the rest of ScatterMoE’s machinery unchanged. E.1 GPU Memory Hierarchy and Related Analysis Bandwidth tiers. Modern GPU architectures expose a multi-level memory hierarchy with sharply varying bandwidth. Off-chip high-bandwidth memory (HBM) provides the largest capacity but the lowest bandwidth. On-chip, the L2 cache offers an intermediate level, while each streaming multiprocessor (SM) contains a small pool of SRAM (registers and shared memory) with the highest bandwidth. Any data that does not fit in on-chip memory must be materialized to HBM, incurring an order-of-magnitude bandwidth penalty. Our fusing avoids writing the full (N,M)(N,M) logits and the intermediate tensors of low-rank router to HBM, and computes the Top-K operation in SRAM, which is much faster than the HBM round-trip. Input X∈ℝN×hX ^N× hStandard Router(N,h)×(h,M)(N,h)×(h,M)++ TopK ++ Softmax ++ Aux Low-Rank Router(N,h)×(h,r)(N,h)×(h,r)(N,r)×(r,M)(N,r)×(r,M)++ TopK ++ Softmax ++ Aux Sort ++ Count (N,k)→(Nk)(N,k)→(Nk) sorted indices ++ offsets Expert Linear 1st (N,h)→(Nk,2s)(N,h)→(Nk,2s) SwiGLU (Nk,2s)→(Nk,s)(Nk,2s)→(Nk,s) Expert Linear 2nd (Nk,s)→(N,h)(Nk,s)→(N,h) Output Y∈ℝN×hY ^N× hweights ++ indices (N,k)(N,k)RoutingExpert MLP Figure 9: Operations in the forward pass of one MoE layer. The first expert linear projects to 2s2s because SwiGLU uses separate gate and up projections: SiLU(Xwgate)⊙XwupSiLU(Xw_gate) Xw_up. Expert linear layers use fused scatter-gather kernels (Tan et al., 2024) that eliminate O(Nkh)O(Nkh) memory traffic from token dispatch, making the standard router the dominant remaining overhead and motivating the low-rank factorization. Comparison to other fused-MoE kernels. Existing efficient MoE kernels, ScatterMoE (Tan et al., 2024), MegaBlocks (Gale et al., 2023), MoMoE (Costin et al., 2025) and SonicMoE (Guo et al., 2025), do top-k computation on the materialized (N,M)(N,M) logits on the HBM and have to pay expensive read and write costs. Our kernel fuses the router with top-k and do sequential top-k operation on each tile, this avoid expensive memory operations between HBM and SRAM, but the cost would also scale with k, making large k infeasible in our setup. E.2 Top-K Algorithm Analysis Our implementation is faster than torch.topk when K<16K<16 and is common across frontier model choices (Table 3). But it is slower than torch.topk for K≥16K≥ 16. For our fused kernel, the per-token cost is Θ(K⋅M)+Θ(K2⋅M/BLOCK_M) (K· M)+ (K^2· M/ BLOCK\_M). For each of the M/BLOCK_M/ BLOCK\_M tiles, we run K iterations to find the K largest values and insert them to the K-element running buffer. The Θ(K⋅M) (K· M) term aggregates the K iterations of max/argmax reductions over the BLOCK_M new candidates, across M/BLOCK_M/ BLOCK\_M tiles. The Θ(K2⋅M/BLOCK_M) (K^2· M/ BLOCK\_M) term is for computing the minimum over the K-element buffer that happen inside the same K iterations. We suspect that a more carefully engineered fused topk along the lines of SonicMoE (Guo et al., 2025) (Section D and Appendix F.4), could outperform torch.topk at larger K and close this gap, which we leave to future work. E.3 Expert Parallelism We explained how MoRE can be implemented efficiently on a single GPU, here we describe how we implement it under expert parallelism, a common setting in large-scale MoE training. We consider there are E GPUs where each GPU contains M/EM/E experts, for simplicity we assume all GPUs are inside one node, and we defer inter-node communication analysis to future work. Following popular design choices (Lepikhin et al., 2020; Fedus et al., 2022; Yan et al., 2026), we replicate non-expert components including router in all GPUs, which our fusing technique still carries over as every GPU has a local router. We could also shard the router across GPUs, i.e. each GPU holds M/EM/E experts and a 1/E1/E fraction of the router, by computing a top-k locally inside each GPU and then a global top-k across GPUs, which is a hierarchical top-k which we leave to future work. It has the same scaling argument as in the main paper. (a) expert parallelism only batch replicated on every GPU (b) ++ data parallelism each GPU routes only its N/EN/E tokens X (N×h)(N× h)router over all N tokensexperts on my M/EM/E shardall-reduce (SUM)Y (N×h)(N× h)x_loc (N/E×h)(N/E× h)router over my N/EN/E tokensall-to-all (dispatch)experts on my M/EM/E shardall-to-all (combine)y_loc (N/E×h)(N/E× h) Figure 10: The two implementations measured in Table 10. Both shard the experts M/EM/E per GPU and hold a full copy of the router weights. In (a) the batch is replicated, so every GPU scores all N tokens and one all-reduce sums the partial outputs; the routing work is therefore duplicated E times. In (b) each GPU keeps only its own N/EN/E tokens and scores only those, so the routing work is divided by E, at the cost of shipping each token to the GPU owning its selected expert and back with two all-to-alls. E.3.1 Scaling under Expert Parallelism Here we show that the Θ(h/r) (h/r) advantage of Section 4.1 still holds under expert parallelism, and that replicating the router makes the routing bottleneck more severe. Active cost with a replicated router. Because the router is replicated, every device scores all tokens against all experts and computes the global top-k locally. Each token-expert pair is therefore still computed exactly once across the system, while the router runs E times. Per-token active FLOPs across GPUs become τ=Cekhs+ECrMh(standard router),Cekhs+ECr(M+h)r(low-rank router).τ\;=\; casesC_e\,khs+E\,C_r\,Mh&(standard router),\\ C_e\,khs+E\,C_r\,(M+h)\,r&(low-rank router). cases (4) Equivalently, since each device holds M/EM/E experts and therefore computes k/Ek/E of the token-expert pairs in expectation, the active FLOPs per GPU are Ce(k/E)hs+CrMhC_e(k/E)hs+C_rMh for the standard router. Either way the factor E falls on the router term and not the expert term, so routing occupies a larger share of the budget than on a single device: replicating the router sharpens exactly the bottleneck that MoRE is designed to remove. The ratio is unchanged. Applying the budget allocation of Proposition 2 to (4) gives M∗=τ2CrEh,M′=τ−ECrhr2CrEr,M′M∗=hr(1−ECrhrτ).M^*\;=\; τ2C_rEh, M \;=\; τ-E\,C_rhr2C_rEr, M M^*\;=\; hr (1- E\,C_rhrτ ). The factor E enters both configurations identically and cancels, so the Θ(h/r) (h/r) advantage carries over. The correction term is close to 11 whenever τ>ECrMh≫ECrrhτ>E\,C_rMh E\,C_rrh, that is whenever M≫rM r, which is the same condition as on a single device. Table 10 measures both implementations of Figure 10 while sweeping M from 2,0482,048 to 16,38416,384. The standard MoE router dominates the layer time under expert parallelism, and communication does not, so replacing it with the low-rank router of MoRE speeds the layer up by up to 3.1×3.1×. Router Shared phases Layer total E M MoE MoRE experts comm. MoE MoRE speedup (a) expert parallelism only 2 2,048 8.20 1.68 3.86 0.61 12.68 6.15 2.1×2.1× 2 8,192 27.90 6.09 4.15 0.61 32.66 10.85 3.0×3.0× 2 16,384 45.30 11.95 4.41 0.62 50.33 16.98 3.0×3.0× 4 2,048 9.01 1.68 3.81 0.96 13.77 6.45 2.1×2.1× 4 8,192 29.01 6.09 4.05 0.86 33.93 11.00 3.1×3.1× 4 16,384 45.34 11.95 4.25 0.79 50.38 17.00 3.0×3.0× (b) ++ data parallelism, each GPU routes its own N/EN/E tokens 2 2,048 4.25 0.97 8.59 4.12 16.96 13.68 1.2×1.2× 2 8,192 13.04 3.30 8.70 4.42 26.17 16.43 1.6×1.6× 2 16,384 22.81 6.60 9.08 4.43 36.32 20.11 1.8×1.8× 4 2,048 2.31 0.65 4.73 2.23 9.27 7.61 1.2×1.2× 4 8,192 6.73 1.93 4.75 2.22 13.70 8.89 1.5×1.5× 4 16,384 10.55 3.62 4.85 2.18 17.59 10.66 1.7×1.7× Table 10: The standard MoE router is the bottleneck of one layer under expert parallelism and dominates communication, and MoRE achieves a notable speedup by replacing it with a low-rank router. Forward pass of one MoE layer. The experts column includes the local dispatch bookkeeping, which moves nothing between devices: it turns expert ids into a sorted pair list, and in (b) it also gathers the rows that each of the other GPUs will receive. comm. is the collectives themselves, a single all-reduce of the (N,h)(N,h) output in (a) against two all-to-alls in (b), the second of which also carries the (M)(M) balancing statistic. All runs use the Switch load-balancing loss. h=512h=512, s=64s=64, k=4k=4, r=64r=64, bf16, N=262,144N=262,144, mean of 3030 trials on B200 GPUs in one node. E.4 Decoding In terms of decoding, it is memory-bound rather than compute-bound, since each step processes a single token per sequence and the cost is dominated by reading the weights and the KV cache (Shazeer, 2019; Pope et al., 2023). We initially found the decoding speedup to degrade sharply as the expert count M grows, which we found the reason is the scatter2scatter kernel of ScatterMoE (Tan et al., 2024): it loops over the entire contiguous range of expert ids the block spans from the smallest to the largest. A decoding step carries so few tokens whose expert ids are scattered across the pool, and the loop then walks almost the entire expert range, which makes decoding significantly slower at large M. We therefore modify the kernel to step directly from one expert present in the block to the next and make it significantly faster. Table 11 and Figure 11 show the resulting decode times. Figure 11: MoRE has comparable decoding time with MoE for the same configuration. Decode wall-clock of one MoE layer on the (h,M)(h,M) grid of Figure 4, one row per decode batch size B. Left: MoE layer, standard router. Middle: MoRE layer, fused low-rank router. Right: the ratio of the two. Both arms use the dispatch of this section. One NVIDIA B200 in bf16, s=64s=64, k=4k=4, r=64r=64, one token per sequence, median of ten timed calls after five warm-ups. The h=4096h=4096, M=8192M=8192 cell is blank because the expert stride overflows the int32 offsets ScatterMoE uses. Scale Model #Experts B=4B=4 B=16B=16 B=64B=64 B=256B=256 s=32s=32, h=1024h=1024 MoE 128 17.9 17.9 17.9 18.2 MoRE (r=64r=64) 1,536 18.5 18.3 18.4 18.6 s=64s=64, h=512h=512 MoE 256 18.0 17.8 17.9 18.1 MoRE (r=64r=64) 1,792 18.5 18.5 18.4 18.6 s=128s=128, h=512h=512 MoE 256 17.9 17.8 17.9 18.0 MoRE (r=64r=64) 1,792 18.5 18.3 18.5 18.5 Table 11: Decoding: MoRE matches MoE in wall-clock time even with a much larger number of experts. Per-step decode wall-clock (ms) for MoE/MoRE pairs used in our pretraining experiments with matched active FLOPs (Table 1), across decode batch size B. Each MoRE row holds 66–12×12× the experts of the MoE it is matched against, and stays within 3%3\% of it at every scale and batch size. One NVIDIA B200 in bf16, 1212 layers, k=4k=4, prefill 128128, median of the last six of twelve decode steps. Both use the expert dispatch of this section. E.5 Wall-clock measurements for configurations in Table 1 At inference, which runs only the forward pass, MoRE prefills as fast as its MoE counterpart or faster (Table 2). Training is the one place where the extra parameters cost time. Table 12 reports the training step time for the same pairs, measured over the 4040 B-token runs: MoRE costs 1515 to 42%42\% more per step, although it has 55 to 66 times as many parameters and 77 to 1212 times as many experts. Since the two models have the same active FLOPs, we attribute the extra time to work whose cost grows with the number of parameters and experts rather than with active FLOPs. First, the AdamW optimizer updates every parameter and its first and second moments at every step, and gradient clipping and the data-parallel gradient all-reduce also process every parameter. Second, the backward pass computes a weight gradient for every expert by summing over the tokens routed to that expert. We expect that storing the AdamW states and the backward gradients in lower precision could help close this gap. Scale Model #Experts Total params Training (ms) Training time ratio s=32s=32, h=1024h=1024 MoE 128 0.29 B 241.0 MoRE (r=64r=64) 1,536 1.50 B 312.5 1.30× s=64s=64, h=512h=512 MoE 256 0.36 B 165.0 MoRE (r=64r=64) 1,792 1.91 B 221.1 1.34× s=128s=128, h=512h=512 MoE 256 0.66 B 182.7 MoRE (r=64r=64) 1,792 3.76 B 259.3 1.42× s=64s=64, h=512h=512 MoE 1,024 1.27 B 313.8 MoRE (r=64r=64) 8,192 7.92 B 360.6 1.15× Table 12: MoRE costs 1515 to 42%42\% more per training step than the MoE it matches in active FLOPs. Training wall-clock (ms) for the pairs of Table 1, measured over the 4040 B-token runs themselves, on 4×4×NVIDIA B200. Prefill time for the same pairs is in Table 2. In Table 13, we also compare MoRE with a standard MoE of the same configuration, that is, with the same number and width of experts and therefore about the same total parameters. We find MoE trains 1.101.10 to 1.23×1.23× slower than MoRE. This implies that the extra training time of MoRE in Table 12 comes from its larger total number of parameters, not from the low rank. Configuration Model Total params Training (ms) Ratio M=1,536M=1,536, s=24s=24, h=1024h=1024 MoRE (r=64r=64) 1.50 B 312.5 MoE 1.51 B 344.2 1.10× M=1,792M=1,792, s=56s=56, h=512h=512 MoRE (r=64r=64) 1.91 B 221.1 MoE 1.92 B 249.6 1.13× M=1,792M=1,792, s=112s=112, h=512h=512 MoRE (r=64r=64) 3.76 B 259.3 MoE 3.77 B 289.6 1.12× M=8,192M=8,192, s=52s=52, h=512h=512, EP MoRE (r=64r=64) 7.92 B 383.8 MoE 7.96 B 472.2 1.23× Table 13: MoRE and MoE at the same configuration: MoE is 1.101.10 to 1.23×1.23× slower than MoRE. Training step time (ms) of each MoRE model of Table 12 and of a standard MoE with the same number M and width s of experts, on 4×4×NVIDIA B200. The last pair is trained with expert parallelism (EP). Appendix F Pseudocode We give Triton-style pseudocode for the two fused router kernels. Neither materializes the (N,M)(N,M) logits tensor in HBM. The inference kernel returns the top-k selection alone. The training kernel also returns the statistic that the load-balancing loss needs, and pays a second sweep over the expert dimension for it. Notation. Write the router logits as ℓ=XR2⊤R1⊤∈ℝN×M,ℓn,j∈ℝ \;=\;XR_2 R_1 \;∈\;R^N× M, _n,j (5) for token n∈[N]n∈[N] and expert j∈[M]j∈[M], where [M]=1,…,M[M]=\1,…,M\, with R2∈ℝr×hR_2 ^r× h and R1∈ℝM×rR_1 ^M× r as in Section 2. Both kernels evaluate (5) one BN×BMB_N× B_M tile at a time, so the (N,M)(N,M) logits exist only in on-chip SRAM. An overbar denotes the gradient of the total loss with respect to that tensor and always has the same shape as the tensor: A¯=∂ℒ/∂A A= /∂ A. F.1 Inference (prefilling) At inference the router must name the k experts each token will use and supply the k logits that produce the gate weights. There is no loss, so nothing reads a softmax over all M experts. Algorithm 1 therefore walks the expert dimension once, keeping a running top-k buffer in SRAM, and writes 2Nk2Nk numbers instead of NMNM. In Figure 20, we fuse the standard MoE using a similar way, but does not have intermediate z which has size in terms of r. This makes the complexity of fusing MoE scale with h instead with r, and makes the fusing benefit for MoRE diminishes at large h. X(N,h)(N,h) z=XR2⊤z=XR_2 in SRAM tile Y=zR1[blk]⊤Y=z\,R_1[blk] one sweep: running top-k TopVals, TopIdxs(N,k)(N,k) softmax over the k scores gate weights (N,k)(N,k) combine the k expert outputs one Triton kernel Figure 12: Flowchart for inference (prefilling) for low-rank router. The blue block is one Triton kernel, holding one BN×BMB_N× B_M tile of ℓ at a time, so neither the (N,r)(N,r) intermediate z nor the (N,M)(N,M) logits reaches HBM. TopIdxs names the k experts to run for each token; a softmax over the k entries of TopVals gives the weights that combine their outputs. Nothing here depends on all M experts, which is why one sweep is enough. Algorithm 1 Fused low-rank router + Top-k (inference, prefilling) 1: Input X∈ℝN×hX ^N× h, router weights R1∈ℝM×rR_1 ^M× r, R2∈ℝr×hR_2 ^r× h, top-k size k, block sizes BN,BM,BHB_N,B_M,B_H 2: TopVals∈ℝN×kTopVals ^N× k, TopIdxs∈ℤN×kTopIdxs ^N× k in HBM 3: for each token block b∈0,…,⌈N/BN⌉−1b∈\0,…, N/B_N -1\ in parallel do 4: ⊳ Phase 1: z=XblockR2⊤z=X_blockR_2 accumulated in SRAM 5: z←0z← 0 ⊳ (BN,r)(B_N,r) in fp32 registers 6: for each hidden tile i∈0,BH,2BH,…i∈\0,B_H,2B_H,…\ do 7: Load Xtile←X[bBN:(b+1)BN,i:i+BH]X_tile← X[bB_N:(b+1)B_N,\,i:i+B_H] from HBM ⊳ (BN,BH)(B_N,B_H) 8: Load R2R_2 tile from HBM ⊳ (r,BH)(r,B_H) 9: z+=Xtile⋅R2,tile⊤z +=X_tile· R_2,tile ⊳ tile matmul in SRAM 10: Narrow z to bf16 ⊳ tensor-core operands are narrow; the accumulator is fp32 11: ⊳ Phase 2: one sweep of zR1⊤zR_1 over M with a running top-k 12: Allocate top_vals←−∞top\_vals←-∞, top_idxs←0top\_idxs← 0 ⊳ (BN,k)(B_N,k) in SRAM 13: for each expert tile m0∈0,BM,2BM,…m_0∈\0,B_M,2B_M,…\ do 14: Load R1R_1 tile from HBM ⊳ (BM,r)(B_M,r), cached in L2 across blocks 15: ytile←z⋅R1[m0:m0+BM,:]⊤y_tile← z· R_1[m_0:m_0+B_M,:] , out-of-range experts set to −∞-∞ ⊳ (BN,BM)(B_N,B_M) in SRAM 16: for c=1,…,kc=1,…,k do ⊳ merge tile logits into the running top-k 17: (vmax,jmax)←(max,argmax)(v_ ,j_ )←( , ) of ytiley_tile along axis 11 18: (vmin,jmin)←(min,argmin)(v_ ,j_ )←( , ) of top_valstop\_vals along axis 11 19: If vmax>vminv_ >v_ : write (vmax,m0+jmax)(v_ ,\,m_0+j_ ) into slot jminj_ 20: Mask ytile[:,jmax]←−∞y_tile[:,j_ ]←-∞ ⊳ prevent re-selection 21: Write top_vals,top_idxstop\_vals,top\_idxs to HBM The kernel reads X once, R1R_1 once and R2R_2 once, and writes 2Nk2Nk numbers, so its HBM traffic is O(Nh+Mr+Nk)O(Nh+Mr+Nk) rather than O(NM)O(NM). The inner merge loop runs k times per expert tile and is sequential, which is why the cost grows linearly in k in Figure 3(Left) and the cost is mild within the range of commonly used k. F.2 Training Training adds the load-balancing loss, and with it a quantity that depends on all M experts rather than on the k selected ones. Figure 13 shows its effect on the forward and on the backward. X(N,h)(N,h)z=XR2⊤z=XR_2 (N,r)(N,r) tile Y=zR1[blk]⊤Y=z\,R_1[blk] sweep 1: m,Zm,Z, top-k sweep 2: P=eY−m/ZP=e^Y-m/Z TopVals, TopIdxs(N,k)(N,k)psumpsum(M)(M)ℒpredL_predℒbalL_bal sweep 1: g (N)(N) sweep 2: Y¯ Y tile =TopVals¯+P⊙(p¯−g)= TopVals+ [rgb]0.75,0,0.25P ( p-g) TopVals¯ TopVals(N,k)(N,k)p¯ p(M)(M)z¯ z(N,r)(N,r)X¯ X (N,h)(N,h)to previous layerR¯2 R_2 (r,h)(r,h)R¯1 R_1 (M,r)(M,r)forwardbackward Figure 13: Fused low-rank router for training. Blue blocks run inside the kernel, which holds one BN×BMB_N× B_M tile of ℓ at a time and never forms the full N×MN× M matrix. Purple marks what the Switch loss adds. Sweep 2 computes the quantity required for the load-balancing loss: psumi=∑nPn,ipsum_i= _nP_n,i, the softmax probability of expert i summed over tokens. TopVals¯ TopVals is the gradient of the prediction loss with respect to the k selected scores, and is sparse in M. p¯ p is the gradient of the balancing loss with respect to psumpsum, one number per expert and dense in M. F.2.1 Forward-Pass Pseudocode At forward pass in training, we leverage the techniques in online softmax of Milakov and Gimelshein (2018) to compute the load-balancing loss. Let P∈ℝN×MP ^N× M be the row-wise softmax of the logits, Pn,j=eℓn,jZn,Zn=∑j′=1Meℓn,j′∈ℝ.P_n,j\;=\; e^\, _n,jZ_n, Z_n\;=\; _j =1^Me^\, _n,j . (6) The Switch load-balancing loss (Fedus et al., 2022) is ℒbal=M∑i=1MfipiL_bal=M _i=1^Mf_ip_i, where fif_i is the fraction of dispatch slots that went to expert i and carries no gradient, and pi=1N∑nPn,ip_i= 1N _nP_n,i is that expert’s mean routing probability over all M experts. The only quantity in it that a tiled kernel cannot produce for free is the column sum psumi=∑n=1NPn,i∈ℝM,pi=psumi/N,psum_i\;=\; _n=1^NP_n,i\;∈\;R^M, p_i=psum_i/N, (7) because Pn,iP_n,i needs the normalizer ZnZ_n, which is not known until every expert has been visited. Returning psumpsum is therefore the whole of what the kernel must add, and it is M numbers rather than NMNM. Building the softmax denominator without the (N,M)(N,M) tensor. The kernel maintains the row maximum mn=maxj∈[M]ℓn,jm_n= _j∈[M] _n,j and the denominator Zn=∑j′eℓn,j′−mnZ_n= _j e _n,j -m_n across the expert tiles, and uses Pn,j=eℓn,j−mn/ZnP_n,j=e _n,j-m_n/Z_n. Neither mnm_n nor ZnZ_n is final until every expert has been seen, so both are built incrementally. After t tiles the kernel holds m(t)m^(t), the largest logit so far, and Z(t)Z^(t), the sum over those experts. When tile t+1t+1 raises the maximum to m′m , every term already in Z(t)Z^(t) carries the factor e−m(t)e^-m^(t) and now needs e−m′e^-m . That correction is one constant shared by all of them, so a single multiply rescales the whole accumulated sum, and the calculation is shown as (8): Z(t+1)=Z(t)em(t)−m′+rowsum(eytile−m′),m(t+1)=m′.Z^(t+1)=Z^(t)\,e^\,m^(t)-m \;+\; rowsum (e^\,y_tile-m ), m^(t+1)=m . (8) Here eytile−m′e^\,y_tile-m is elementwise on the (BN,BM)(B_N,B_M) tile and rowsum adds each row across the BMB_M experts of this tile, so both sides of (8) are vectors of length BNB_N: one running normalizer per token. Once m and Z are final, a second sweep regenerates each tile and accumulates colsum(eytile−m/Z) colsum(e^\,y_tile-m/Z) into psumpsum, where colsum reduces along tokens and so returns one number per expert. The second sweep is needed because the normalizer is only complete at the end of the first, and because the top-k merge overwrites ytiley_tile with −∞-∞ as it consumes columns. It costs one extra zR1⊤zR_1 matmul, 2NMr2NMr FLOPs, and no extra HBM traffic. In the implementation the first matmul z=XR2⊤z=XR_2 is left to cuBLAS and z is passed into the kernel, so the training kernel fuses the second matmul, the top-k and the softmax statistics, while the inference kernel of Algorithm 1 fuses all three. Algorithm 2 Fused low-rank router + Top-k (training, forward). Purple lines are what the Switch load-balancing loss adds. 1: z=XR2⊤∈ℝN×rz=XR_2 ^N× r, router weight R1∈ℝM×rR_1 ^M× r, top-k size k, block sizes BN,BMB_N,B_M 2: TopVals∈ℝN×kTopVals ^N× k, TopIdxs∈ℤN×kTopIdxs ^N× k, psum∈ℝMpsum ^M, m∈ℝNm ^N, Z∈ℝNZ ^N in HBM 3: for each token block b∈0,…,⌈N/BN⌉−1b∈\0,…, N/B_N -1\ in parallel do 4: Load zblockz_block from HBM ⊳ (BN,r)(B_N,r) in SRAM registers 5: ⊳ Sweep 1: tile zR1⊤zR_1 over M with running top-k and running m,Zm,Z 6: Allocate top_vals←−∞top\_vals←-∞, top_idxs←0top\_idxs← 0 ⊳ (BN,k)(B_N,k) in SRAM 7: Allocate m←−∞m←-∞, Z←0Z← 0 ⊳ (BN)(B_N) each, in SRAM 8: for each expert tile m0∈0,BM,2BM,…m_0∈\0,B_M,2B_M,…\ do 9: Load R1R_1 tile from HBM ⊳ (BM,r)(B_M,r), cached in L2 across blocks 10: ytile←z⋅R1⊤y_tile← z· R_1 , out-of-range experts set to −∞-∞ ⊳ (BN,BM)(B_N,B_M) in SRAM 11: m′←max(m,rowmax(ytile))m ← (m,\ rowmax(y_tile) ) ⊳ new running maximum, (BN)(B_N) 12: Z←Zem−m′+rowsum(eytile−m′)Z← Z\,e^\,m-m + rowsum (e^\,y_tile-m ); m←m′m← m ⊳ eq. (8) 13: for c=1,…,kc=1,…,k do ⊳ merge tile logits into running top-k 14: (vmax,jmax)←(max,argmax)(v_ ,j_ )←( , ) of ytiley_tile along axis 11 15: jmin←argminj_ ← of top_valstop\_vals along axis 11 16: If vmax>top_vals[:,jmin]v_ >top\_vals[:,j_ ]: write (vmax,m0+jmax)(v_ ,m_0+j_ ) into slot jminj_ 17: Mask ytile[:,jmax]←−∞y_tile[:,j_ ]←-∞ ⊳ destroys ytiley_tile; prevent re-selection 18: ⊳ Sweep 2: normalized column sums; m and Z are still in registers 19: for each expert tile m0∈0,BM,2BM,…m_0∈\0,B_M,2B_M,…\ do 20: ytile←z⋅R1⊤y_tile← z· R_1 ; P←eytile−m/ZP← e^\,y_tile-m/Z ⊳ (BN,BM)(B_N,B_M), recomputed 21: atomicAdd(psum[m0:m0+BM],colsum(P)) atomicAdd (psum[m_0:m_0+B_M],\ colsum(P) ) 22: Write top_vals,top_idxstop\_vals,top\_idxs, m, Z to HBM Forming the loss. The balancing loss itself stays outside the kernel and costs O(M)O(M): ℒbal=M∑i=1Mfipi=MNf⋅psum,fi=1Nk∑n=1N∑j=1k[TopIdxs[n,j]=i],L_bal\;=\;M _i=1^Mf_i\,p_i\;=\; MN\,f·psum, f_i= 1Nk _n=1^N _j=1^k1 [TopIdxs[n,j]=i ], (9) with f obtained by a bincount on TopIdxsTopIdxs and detached, as is standard. Because the kernel returns psumpsum rather than a loss value, the same outputs serve the Switch form of (9) or a different balancing objective, without touching the kernel. F.2.2 Backward-Pass Pseudocode Two gradients reach the kernel, and they have different shapes. TopVals¯∈ℝN×k TopVals ^N× k arrives from the prediction loss on the k selected scores; it is sparse in M, touching only the selected experts. p¯∈ℝM p ^M arrives from ℒbalL_bal through psumpsum, and by (9) it is p¯i=(M/N)fi p_i=(M/N)f_i, one number per expert and dense in M. The kernel must return X¯∈ℝN×h X ^N× h, R¯1∈ℝM×r R_1 ^M× r and R¯2∈ℝr×h R_2 ^r× h, as in Figure 13. Everything else follows from the chain rule on ℓ=zR1⊤ =zR_1 and z=XR2⊤z=XR_2 , so the only new quantity to derive is the logit gradient y¯ y. The logit gradient of the balancing loss. Fix a token n and an expert j∈[M]j∈[M]. Since psumi=∑n′Pn′,ipsum_i= _n P_n ,i and only row n of P depends on ℓn _n, the chain rule contracts over the expert index alone: ∂ℒbal∂ℓn,j=∑i=1Mp¯i∂Pn,i∂ℓn,j. _bal∂ _n,j= _i=1^M p_i\, ∂ P_n,i∂ _n,j. (10) The softmax Jacobian follows from the quotient rule on Pn,i=eℓn,i/ZnP_n,i=e _n,i/Z_n with Zn=∑j′eℓn,j′Z_n= _j e _n,j . The numerator depends on ℓn,j _n,j only when i=ji=j, contributing δijeℓn,i _ije _n,i with δij=1 _ij=1 if i=ji=j and 00 otherwise; the denominator always does, contributing ∂Zn/∂ℓn,j=eℓn,j∂ Z_n/∂ _n,j=e _n,j. Hence ∂Pn,i∂ℓn,j=δijeℓn,iZn−eℓn,ieℓn,jZn2=δijeℓn,iZn−eℓn,iZneℓn,jZn=Pn,i(δij−Pn,j). ∂ P_n,i∂ _n,j= _ij\,e _n,iZ_n-e _n,ie _n,jZ_n^2= _ij e _n,iZ_n- e _n,iZ_n e _n,jZ_n=P_n,i ( _ij-P_n,j ). (11) Substituting (11) into (10), the δij _ij collapses the first sum to its i=ji=j entry and the second factors out Pn,jP_n,j: ∂ℒbal∂ℓn,j=p¯jPn,j⏟from δij−Pn,j∑i=1Mp¯iPn,i=Pn,j(p¯j−gn),gn≡∑i=1Mp¯iPn,i∈ℝ. _bal∂ _n,j= p_jP_n,j_from _ij\;-\;P_n,j _i=1^M p_iP_n,i=P_n,j ( p_j-g_n ), g_n≡ _i=1^M p_iP_n,i . (12) So g∈ℝNg ^N is one scalar per token: the mean of p¯ p under that token’s own routing distribution. Because gng_n sums over all M experts, it must be complete before any single expert’s gradient can be written, so the backward takes two sweeps for the same reason the forward does. Algorithm 3 Low-rank router (training, backward) with activation recompute. Purple lines are what the Switch load-balancing loss adds. 1: X, R1∈ℝM×rR_1 ^M× r, R2∈ℝr×hR_2 ^r× h, TopIdxs∈ℤN×kTopIdxs ^N× k, TopVals¯∈ℝN×k TopVals ^N× k (top-k logit gradients), p¯∈ℝM p ^M, saved m,Z∈ℝNm,Z ^N 2: X¯∈ℝN×h X ^N× h, R¯1∈ℝM×r R_1 ^M× r, R¯2∈ℝr×h R_2 ^r× h in HBM 3: for each token block b in parallel do 4: ⊳ Phase 1: recompute z=XblockR2⊤z=X_blockR_2 in SRAM (saves (N,r)(N,r) activation memory) 5: z←0z← 0 ⊳ (BN,r)(B_N,r) in SRAM 6: for each hidden tile i do 7: z+=Xtile⋅R2⊤z +=X_tile· R_2 8: ⊳ Phase 2, sweep 1: the per-token scalar gng_n of eq. (12) 9: g←0g← 0 ⊳ (BN)(B_N) in SRAM 10: for each expert tile m0∈0,BM,2BM,…m_0∈\0,B_M,2B_M,…\ do 11: ytile←z⋅R1⊤y_tile← z· R_1 ; P←eytile−m/ZP← e^\,y_tile-m/Z 12: g+=rowsum(P⊙p¯[m0:m0+BM])g += rowsum (P p[m_0:m_0+B_M] ) 13: ⊳ Phase 2, sweep 2: the logit gradient, contracted into z¯ z and R¯1 R_1 14: z¯←0 z← 0 ⊳ (BN,r)(B_N,r) in SRAM 15: for each expert tile m0∈0,BM,2BM,…m_0∈\0,B_M,2B_M,…\ do 16: y¯←0 y← 0; scatter TopVals¯[n,j] TopVals[n,j] into column TopIdxs[n,j]TopIdxs[n,j] of this tile ⊳ sparse: k of BMB_M 17: ytile←z⋅R1⊤y_tile← z· R_1 ; P←eytile−m/ZP← e^\,y_tile-m/Z 18: y¯+=P⊙(p¯[m0:m0+BM]−g) y +=P ( p[m_0:m_0+B_M]-g ) ⊳ eq. (12), dense in M 19: z¯+=y¯⋅R1[m0:m0+BM,:] z += y· R_1[m_0:m_0+B_M,:]; R¯1[m0:m0+BM,:]+=y¯⊤z R_1[m_0:m_0+B_M,:] += y z 20: ⊳ Phase 3: X¯block=z¯⋅R2 X_block= z· R_2 tiled over h 21: for each hidden tile i do 22: Load R2R_2 tile, write X¯tile←z¯⋅R2 X_tile← z· R_2 23: R¯2+=z¯⊤Xblock R_2 += z X_block ⊳ (r,BN)×(BN,h)→(r,h)(r,B_N)×(B_N,h)→(r,h) Appendix G Additional Plots G.1 Routing entropy We report the routing entropy of the per-expert token counts c1,…,cMc_1,…,c_M of one layer. Writing pa=ca/∑bcbp_a=c_a/ _bc_b for the share of the layer’s routed tokens that expert a receives, the normalized entropy is Hnorm=(−∑apalogpa)/logMH_norm= (- _ap_a p_a )/ M, accumulated over the full training set. It is 11 when every expert carries the same load, and falls towards 00 as routing concentrates on fewer experts. The counts are read from saved checkpoints: for each checkpoint we pass the full training set through that fixed model once and count the tokens each expert receives. Figure 14 repeats Figure 2 at h=64h=64, s=8s=8, M=1024M=1024, where the expert pool is more than five times larger, and we find similar results. Figure 14: Similar results at a larger expert pool. Phonebook, h=64h=64, s=8s=8, M=1024M=1024, k=4k=4, L=6L=6, panels as in Figure 2. Left: largest phonebook memorized, searched on a 2525K grid. Middle, Right: HnormH_norm per layer against training step, for r=32r=32 and for MoE, both trained on the 200200K phonebook. This configuration is off the matched-FLOP curve of Proposition 2: it places several ranks below log2M _2M, and both arms hold the same experts, so only the router differs. G.2 Adding More Experts help We sweep the size of the expert pool over M∈112,224,448,896,1792M∈\112,224,448,896,1792\ while holding the rest of the architecture fixed at h=512h=512, s=56s=56, r=64r=64, k=4k=4 and L=12L=12. All five configurations train on the same 55B tokens. Figure 15 shows that validation perplexity decreases monotonically over the whole range, from 19.1019.10 at M=112M=112 to 17.0517.05 at M=1792M=1792, demonstrating the benefits of using more experts. Figure 15: Adding more experts on MoRE keeps improving quality. Validation perplexity of MoRE (r=64r=64) after 55B training tokens, against the number of experts M. The curve is monotone across a 16×16× increase in the pool size. G.3 The perplexity gap grows with the token budget In Figure 16, we find that the perplexity gap between MoRE and MoE grow with token budget, which we think the gap would grow larger if more tokens are trained. This also matches our intuition, since MoRE’s benefits of memorizing more information would be more apparent when more tokens are trained. Figure 16: The perplexity gap between the two configurations grows with the token budget. Validation perplexity of MoE-1024 minus that of MoRE-8192 on the 5,0005,000 held-out mixture sequences, against training tokens. Positive values mean MoRE has the lower perplexity. Both models have matched active FLOPs and 1.271.27B against 7.927.92B total parameters. We plot the first 2525B tokens before the cosine schedule enters its annealing tail. Appendix H Additional Plots for Efficiency H.1 Forward pass of a single layer at larger expert widths Figure 17 repeats Figure 4 at s∈128,256,512s∈\128,256,512\ with the same protocol: inference (prefilling) of the whole layer, k=4k=4, r=64r=64, N=262,144N=262,144 (B=256B=256, T=1024T=1024), bf16, mean of 30 trials after 10 warm-ups on a single NVIDIA B200. The fused MoRE layer is faster in every cell that fits in memory. The speedup still grows with M along every row, and it shrinks as s grows, because wider experts make the expert computation, which both layers share, a larger part of the layer. At M=8,192M=8,192 it is 1.71.7 to 4.2×4.2× at s=128s=128, 1.61.6 to 3.6×3.6× at s=256s=256 and 1.61.6 to 3.0×3.0× at s=512s=512. Figure 17: The fused MoRE layer stays faster at larger expert widths. Inference (prefilling) of a single layer at s=128s=128, 256256 and 512512 (rows). (a) MoE Layer. (b) MoRE Layer (r=64r=64). (c) speedup TMoE/TMoRET_MoE/T_MoRE. The blank cell at s=512s=512 ran out of memory on one B200. k=4k=4, r=64r=64, N=262,144N=262,144 (B=256B=256, T=1024T=1024), bf16, mean of 30 trials after 10 warm-ups on a single NVIDIA B200. H.2 The same grid on an A100 Figure 18 repeats the inference (prefilling) measurement of Figure 4 on a single NVIDIA A100-80GB, over the same (h,M)(h,M) grid: hidden dimension h∈64,128,256,512,1,024,2,048,4,096h∈\64,128,256,512,1,024,2,048,4,096\ against expert count M∈64,128,256,512,1,024,2,048,4,096,8,192M∈\64,128,256,512,1,024,2,048,4,096,8,192\, at s=64s=64, k=4k=4 and r=64r=64. The fused MoRE layer is faster in every cell, by 1.041.04 to 3.07×3.07×, and the speedup grows with M along every row, reaching 2.602.60 to 3.07×3.07× at M=8,192M=8,192. Figure 18: The A100 reproduces the B200 result, with the fused MoRE layer 2.602.60 to 3.07×3.07× faster at M=8,192M=8,192. Inference (prefilling) of a single layer, the two arms at equal M. (a) MoE Layer. (b) MoRE Layer (r=64r=64). (c) speedup TMoE/TMoRET_MoE/T_MoRE. Single NVIDIA A100-80GB, s=64s=64, k=4k=4, N=262,144N=262,144 (B=256B=256, T=1024T=1024), bf16, mean of 3030 trials after 1010 warm-ups. Figure 19: Fusing the low-rank router reduces peak memory by up to 14×14× for the whole layer. Ratio of peak memory over a forward and a backward, unfused MoRE against fused MoRE, on the same grid. H.3 Fusing a standard MoE layer The kernel of Section 4.2 can also be written for a standard router. In Figure 20, we find the benefit of fusing diminishes as h grows larger. Figure 20: On both A100 and B200, fusing Standard MoE does not get much acceleration at large h. Inference (prefilling) of a single layer, both arms a standard MoE layer at the same M. (a) unfused router. (b) fused router. (c) speedup. (Top) A100, 0.860.86 to 2.83×2.83×. (Bottom) B200, 1.011.01 to 4.21×4.21×. At h=4,096h=4,096, M=8,192M=8,192 fusing takes the layer from 174.8174.8 to 202.5202.5 ms on the A100 and from 61.161.1 to 60.460.4 ms on the B200. s=64s=64, k=4k=4, N=262,144N=262,144 (B=256B=256, T=1024T=1024), bf16, mean of 3030 trials after 1010 warm-ups. Figure 20 holds M fixed along each row. Reading one row of it, M=4,096M=4,096, against the low-rank router at the same shape gives Figure 21: fusing takes the same work off both layers, and only the standard one gives it back as h grows. Figure 21: Fusing keeps saving time for MoRE as h grows, but stops saving time for MoE. Forward pass of a single layer, the unfused-router time minus the fused-router time, against the hidden dimension h at a fixed pool of M=4,096M=4,096 experts. What fusing removes is the round trip of the (N,M)(N,M) scores, which depends on N and M but not on h. The standard router gives it back because its fused kernel reloads a (BM,h)(B_M,h) tile of the router weight for every token block. s=64s=64, k=4k=4, r=64r=64, N=262,144N=262,144 (B=256B=256, T=1024T=1024), bf16, mean of 3030 trials after 1010 warm-ups on a single NVIDIA B200. H.4 The low-rank router once both layers are fused Figure 22 compares MoE and MoRE (r=64r=64) where both are fused. Figure 22: Comparing fused MoE and fused MoRE, the low-rank factorization gives speedup up to 2.82×2.82× on an A100 and 1.72×1.72× on a B200. Inference (prefilling) of a single layer, both arms fused, at the same M. (a) MoE Layer with the fused router. (b) MoRE Layer (r=64r=64) with the fused router. (c) speedup. (Top) A100, 1.001.00 to 2.82×2.82×. (Bottom) B200, 0.990.99 to 1.72×1.72×. s=64s=64, k=4k=4, N=262,144N=262,144 (B=256B=256, T=1024T=1024), bf16, mean of 3030 trials after 1010 warm-ups. H.5 End-to-End Forward-Pass Breakdown at Large M Figure 23 measures the end-to-end wall-clock time of a 12-layer transformer with 8 attention heads (FlashAttention-4 (Zadouri et al., 2026)), vocabulary size V=50,257V=50,257, h=512h=512, s=64s=64, k=4k=4, B=64B=64, T=1024T=1024, sweeping M∈1,024,2,048,4,096,8,192M∈\1,024,2,048,4,096,8,192\. Figure 23: End-to-end forward-pass breakdown at h=512h=512, sweeping M. Each bar decomposes the wall-time of a 12-layer transformer forward pass into Attention, Router, Experts (dispatch + grouped MLP), RMSNorm, and the Liger-fused (Hsu et al., 2024) LM head + cross-entropy. (Left) Standard MoE router (unfused h×Mh×M matmul + top-k + softmax + load-balancing-loss as separate kernels): the router slice grows from 16% of total at M=1,024M=1,024 to 38% at M=8,192M=8,192, becoming the largest single component and the dominant driver of the rising end-to-end cost. (Right) Fused low-rank MoRE router (r=64r=64): the router slice stays at 5–12% across the same sweep, and total wall-time grows much more slowly with M. At M=8,192M=8,192, MoRE is 30% faster end-to-end (108.2108.2 ms vs. 152.8152.8 ms), with router cost reduced from 57.957.9 ms to 13.413.4 ms. Mean of 200 trials with 50 warm-ups, 95% C.I. on a single NVIDIA B200. H.6 Fusion benefit at different batch sizes In Figure 24, we find that the benefit of fusing grows with the batch size. The fully fused MoRE router is 1.4×1.4× faster than the unfused one at B=2B=2 and 6.3×6.3× faster at B=256B=256. At B=2B=2 both routers take about 0.10.1 ms. As B grows, most of the unfused router’s time goes to the top-k over the (N,M)(N,M) score tensor that it stores in HBM, whose size grows with the number of tokens N=B⋅TN=B· T, while the fused kernel never stores this tensor. Figure 24: The benefit of fusing the MoRE router grows with the batch size. Forward-pass time of the low-rank router, unfused and fully fused, with the speedup of the fused kernel above each pair. h=2048h=2048, M=4,096M=4,096, r=16r=16, k=4k=4, T=1024T=1024, bf16, mean of 200200 trials after 5050 warm-ups on a single NVIDIA B200. Appendix I Scaling Total Parameters Proposition 2 (Optimal scaling under fixed active-FLOPs). Fix the active-FLOP budget τ, the hidden dimension h, and the active expert count k. The MoE configuration (M,k,h,s)(M,k,h,s) and the MoRE configuration (M′,k,h,s′,r)(M ,k,h,s ,r) that maximize total MLP parameters under matched τ are M∗=τ2Crh,s∗=τ2Cekh,TMoE∗=3τ24CeCrhk,M^*\;=\; τ2C_rh, s^*\;=\; τ2C_ekh, T_MoE^*\;=\; 3τ^24C_eC_r\,hk, M′=τ′2Crr,s′=τ′2Cekh,TMoRE∗=3(τ′)24CeCrrk,M \;=\; τ 2C_rr, s \;=\; τ 2C_ekh, T_MoRE^*\;=\; 3(τ )^24C_eC_r\,rk, where τ′:=τ−Crhrτ :=τ-C_rhr. Furthermore, in the regime M,M′≫rM,M r (equivalently τ≫Crhrτ C_rhr), M′M∗→hr,TMoRE∗TMoE∗→hr. M M^*\;→\; hr, T_MoRE^*T_MoE^*\;→\; hr. Proof. We use the same budget-allocation argument in both settings. MoE. Allocate fraction α∈[0,1]α∈[0,1] of the budget to the router and 1−α1-α to the experts: CrMh=ατ,Cekhs=(1−α)τ,C_rMh\;=\;ατ, C_ekhs\;=\;(1-α)τ, so M=ατ/(Crh)M=ατ/(C_rh) and s=(1−α)τ/(Cekh)s=(1-α)τ/(C_ekh). Substituting into TMoE=3MhsT_MoE=3Mhs, TMoE(α)=3τ2CeCrhk⋅α(1−α),T_MoE(α)\;=\; 3τ^2C_eC_r\,hk·α(1-α), which is maximized at α=1/2α=1/2, yielding the stated M∗,s∗,TMoE∗M^*,s^*,T_MoE^*. MoRE. The router cost expands as Cr(M′+h)r=CrM′r+CrhrC_r(M +h)r=C_rM r+C_rhr, where the second term is a fixed overhead independent of M′M . Subtracting it leaves a variable budget τ′:=τ−Crhrτ :=τ-C_rhr, which we split the same way: CrM′r=ατ′,Cekhs′=(1−α)τ′,C_rM r\;=\;ατ , C_ekhs \;=\;(1-α)τ , so M′=ατ′/(Crr)M =ατ /(C_rr) and s′=(1−α)τ′/(Cekh)s =(1-α)τ /(C_ekh). Substituting into TMoRE=3M′hs′T_MoRE=3M hs , TMoRE(α)=3(τ′)2CeCrrk⋅α(1−α),T_MoRE(α)\;=\; 3(τ )^2C_eC_r\,rk·α(1-α), again maximized at α=1/2α=1/2, yielding the stated M′,s′,TMoRE∗M ,s ,T_MoRE^*. Limit. The ratios are M′M∗=hr(1−Crhrτ),TMoRE∗TMoE∗=hr(1−Crhrτ)2. M M^*\;=\; hr (1- C_rhrτ ), T_MoRE^*T_MoE^*\;=\; hr (1- C_rhrτ )^2. Since τ=Cekhs+CrMh>CrMh≫Crhrτ=C_ekhs+C_rMh>C_rMh C_rhr when M≫rM r, the correction Crhr/τ≪1C_rhr/τ 1, and both ratios converge to h/rh/r. ∎ Appendix J Theoretical Results Section 3 states the rank bounds under the normalization ‖ui‖2≤1\|u_i\|_2≤ 1 and ‖R2x‖2≤1\|R_2x\|_2≤ 1. The proofs below are given in a slightly more general form with arbitrary norm bounds BuB_u and BzB_z. In this form, it is convenient to measure margins in normalized units: a normalized margin η corresponds to a raw logit margin ηBuBzη B_uB_z. Setting Bu=Bz=1B_u=B_z=1 recovers the statements used in Section 3. We first prove the lower bound. The argument shows that realizing many distinct top-k assignments forces many separated projected representations to fit inside an r-dimensional ball. Theorem 3 (Lower bound on required rank, general norm bounds). Let ℱ⊆([M]k)F [M]k and let Bu,Bz>0B_u,B_z>0. If a rank-r router R=R1R2R=R_1R_2 with ‖ui‖2≤Bu\|u_i\|_2≤ B_u and ‖R2xS‖2≤Bz\|R_2x_S\|_2≤ B_z realizes every S∈ℱS at margin at least ηBuBzη B_uB_z, then |ℱ|≤(1+2/η)r,r≥log|ℱ|log(1+2/η).|F|≤(1+2/η)^r, r≥ |F| (1+2/η). Proof. The proof is a volume packing argument: distinct assignments force their projections zSz_S to be far apart in ℝrR^r, and we count how many such well-separated points fit in a ball. For each S∈ℱS , choose an input xSx_S realizing S with margin ηBuBzη B_uB_z, and write zS=R2xSz_S=R_2x_S. Take any two distinct S,T∈ℱS,T : since |S|=|T|=k|S|=|T|=k, there exist i∈S∖Ti∈ S T and j∈T∖Sj∈ T S. At xSx_S expert i wins over j, while at xTx_T expert j wins over i, so ⟨ui−uj,zS⟩≥ηBuBz,⟨uj−ui,zT⟩≥ηBuBz. u_i-u_j,\,z_S \;≥\;η B_uB_z, u_j-u_i,\,z_T \;≥\;η B_uB_z. Adding the inequalities and applying Cauchy–Schwarz with ‖ui−uj‖2≤2Bu\|u_i-u_j\|_2≤ 2B_u gives ‖zS−zT‖2≥ηBz\|z_S-z_T\|_2≥η B_z. Hence zS:S∈ℱ\z_S:S \ are |ℱ||F| points pairwise separated by at least ηBzη B_z, all lying in the radius-BzB_z ball of ℝrR^r. The open balls of radius ηBz/2η B_z/2 around these points are therefore disjoint and contained in the ball of radius Bz(1+η/2)B_z(1+η/2). Comparing volumes, |ℱ|⋅(ηBz2)r≤(Bz(1+η2))r,|F|· ( η B_z2 )^r\;≤\; (B_z (1+ η2 ) )^r, which rearranges to |ℱ|≤(1+2/η)r|F|≤(1+2/η)^r. ∎ Proof of Theorem 2. Apply Theorem 3 with Bu=Bz=1B_u=B_z=1. ∎ We next prove the upper bound. The construction first realizes the desired assignments in dimension M using the standard basis, then applies a Johnson–Lindenstrauss projection (Johnson and Lindenstrauss, 1984) to reduce the dimension while preserving the inner products that determine the top-k ordering. Theorem 4 (Upper bound on required rank, general norm bounds). Fix 1≤k≤M−11≤ k≤ M-1, ℱ⊆([M]k)F [M]k, and norm bounds Bu,Bz>0B_u,B_z>0. Set ρ:=1/k+1/(M−k)>1/kρ:= 1/k+1/(M-k)>1/ k. For any target normalized margin η∈(0,ρ)η∈(0,ρ), h≥rh≥ r, and r≥Clog(M+|ℱ|)(ρ−η)2r\;≥\; C (M+|F|)(ρ-η)^2 with universal C>0C>0, there is a rank-r router with ‖ui‖2≤Bu\|u_i\|_2≤ B_u realizing every S∈ℱS at some input with ‖R2xS‖2≤Bz\|R_2x_S\|_2≤ B_z and margin at least ηBuBzη B_uB_z. Proof. We first solve the problem in the uncompressed regime r=Mr=M, then use a Johnson–Lindenstrauss projection to compress down to r=O(log(M+|ℱ|)/(ρ−η)2)r=O( (M+|F|)/(ρ-η)^2) while preserving the inner products that determine routing. To begin, let e1,…,eMe_1,…,e_M be the standard basis of ℝMR^M, which serve as expert vectors. For each S∈ℱS , define zS:=1ρ(Sk−[M]∖SM−k).z_S\;:=\; 1ρ\! ( 1_Sk- 1_[M] SM-k ). A direct computation gives ‖zS‖22=1ρ2(k⋅1k2+(M−k)⋅1(M−k)2)=1ρ2(1k+1M−k)= 1.\|z_S\|_2^2\;=\; 1ρ^2\! (k· 1k^2+(M-k)· 1(M-k)^2 )\;=\; 1ρ^2\! ( 1k+ 1M-k )\;=\;1. For i∈Si∈ S and j∉Sj∉ S, ⟨ei,zS⟩=1kρ,⟨ej,zS⟩=−1(M−k)ρ,⟨ei,zS⟩−⟨ej,zS⟩=ρ. e_i,z_S \;=\; 1kρ, e_j,z_S \;=\;- 1(M-k)ρ, e_i,z_S - e_j,z_S \;=\;ρ. So in dimension M, the trivial choice of expert vectors eie_i realizes every S∈ℱS with margin exactly ρ, which is strictly larger than the target η. We now compress this construction via Johnson–Lindenstrauss. Let =e1,…,eM∪zS:S∈ℱV=\e_1,…,e_M\∪\z_S:S \, so ||≤M+|ℱ||V|≤ M+|F|. By the Johnson–Lindenstrauss lemma, for r=O(log||/ε2)r=O( |V|/ ^2) there is a linear map Φ:ℝM→ℝr :R^M ^r preserving squared Euclidean distances on ∪0V∪\0\ multiplicatively to within (1±ε)(1± ), i.e. (1−ε)‖u−v‖22≤‖Φu−Φv‖22≤(1+ε)‖u−v‖22(1- )\|u-v\|_2^2≤\| u- v\|_2^2≤(1+ )\|u-v\|_2^2 for all u,v∈∪0u,v ∪\0\. Since every u∈u has unit norm, ‖u−v‖22≤4\|u-v\|_2^2≤ 4 for all such pairs, so this is equivalent to additive preservation up to O(ε)O( ); absorbing the constant, we assume preservation up to additive error ε in what follows. The polarization identity ⟨u,v⟩=12(‖u‖22+‖v‖22−‖u−v‖22) u,v \;=\; 12\! (\|u\|_2^2+\|v\|_2^2-\|u-v\|_2^2 ) then preserves all pairwise inner products of vectors in V up to additive error O(ε)O( ), and absorbing the constant we may assume preservation up to additive ε . Choose ε:=c0(ρ−η) \;:=\;c_0\,(ρ-η) for a sufficiently small universal constant c0>0c_0>0, which gives r=O(log(M+|ℱ|)/(ρ−η)2)r=O\! ( (M+|F|)/(ρ-η)^2 ). To extract the rank-r router from Φ , define normalized compressed vectors u~i:=Φei‖Φei‖2,z~S:=ΦzS‖ΦzS‖2, u_i\;:=\; e_i\| e_i\|_2, z_S\;:=\; z_S\| z_S\|_2, and set ui=Buu~iu_i=B_u u_i, z¯S=Bzz~S z_S=B_z z_S. The norm bounds are immediate. Since normalization changes the preserved inner products by at most another constant multiple of ε , for every i∈Si∈ S and j∉Sj∉ S, ⟨ui,z¯S⟩−⟨uj,z¯S⟩=BuBz(⟨u~i,z~S⟩−⟨u~j,z~S⟩)≥BuBz(ρ−C′ε)≥ηBuBz u_i, z_S - u_j, z_S =B_uB_z\! ( u_i, z_S - u_j, z_S )≥ B_uB_z(ρ-C )≥η B_uB_z for a universal constant C′>0C >0, by the choice of c0c_0. Hence every S∈ℱS is realized with margin at least ηBuBzη B_uB_z. Finally, we lift to input space: since h≥rh≥ r, set R2=[Ir∣0]∈ℝr×hR_2=[I_r 0] ^r× h and choose xSx_S whose first r coordinates are z¯S z_S, so that R2xS=z¯SR_2x_S= z_S. This realizes the desired router with the claimed parameters. ∎ Proof of Theorem 1. Apply Theorem 4 with Bu=Bz=1B_u=B_z=1 and target margin η=c/kη=c/ k for a sufficiently small universal constant c>0c>0. Since ρ>1/kρ>1/ k, we have (ρ−η)2=Ω(1/k)(ρ-η)^2= (1/k), and therefore r=O(klog(M+|ℱ|))r=O(k (M+|F|)). The constructed margin is η=Ω(1/k)η= (1/ k), as claimed. ∎ Proof of Corollary 1. For ℱ=([M]k)F= [M]k, Stirling’s bounds (M/k)k≤|ℱ|≤(eM/k)k(M/k)^k≤|F|≤(eM/k)^k give log|ℱ|=Θ(klog(M/k)). |F|\;=\; (k (M/k)). Since M≤|ℱ|M≤|F| for 1≤k≤M−11≤ k≤ M-1, we also have log(M+|ℱ|)=Θ(log|ℱ|)=Θ(klog(M/k)) (M+|F|)= ( |F|)= (k (M/k)). Upper bound. By Theorem 1, any rank r≥Cklog(M+|ℱ|)r≥ Ck (M+|F|) admits the desired router at margin Ω(1/k) (1/ k). Substituting log(M+|ℱ|)=Θ(klog(M/k)) (M+|F|)= (k (M/k)) yields rank r=O(k2log(M/k))r=O(k^2 (M/k)) as a sufficient condition. This margin is resolvable under p=Ω(logk)p= ( k) precision. Lower bound. At p-bit top-k precision, η=Ω(2−p)η= (2^-p) implies log(1+2/η)=O(p) (1+2/η)=O(p). Theorem 2 then gives r≥log|ℱ|log(1+2/η)=Ω(klog(M/k)p).r\;≥\; |F| (1+2/η)\;=\; \! ( k (M/k)p ). ∎ Proposition 1 (restated). Let x1,…,xn∼(0,Ih)x_1,…,x_n (0,I_h) be independent, and suppose M=2rM=2^r (r=Θ(log2M)r= ( _2M)) with h≥rh≥ r. For any 1≤k≤M1≤ k≤ M, there exists a rank-r top-k router with M experts that selects every expert with probability exactly k/Mk/M on each example. Consequently, if NaN_a is the number of examples routed to expert a, then writing L:=log2MδL:= 2Mδ, for any δ∈(0,1)δ∈(0,1) the following holds with probability at least 1−δ1-δ: maxa∈[M]|Na−nkM|≤2nkLM+23L. _a∈[M] |N_a- nkM |≤ 2nkLM\;+\; 23L. Construction. The construction follows the routing step in the Gaussian memorization argument of Jelassi et al. (2025). Let R2∈ℝr×hR_2 ^r× h select the first r coordinates, so zi:=R2xi∼(0,Ir)z_i:=R_2x_i (0,I_r) independently for i∈[n]i∈[n]. Index the M=2rM=2^r experts by sign vectors a∈±1ra∈\± 1\^r and set the corresponding expert vector to ua=a/ru_a=a/ r. At k=1k=1 the winner is available in closed form: the score ⟨ua,z⟩=1r∑ℓaℓzℓ u_a,z = 1 r _ a_ z_ is maximized coordinate-wise by aℓ=sign(zℓ)a_ =sign(z_ ), so the selected expert is the sign pattern of z, which is uniform on ±1r\± 1\^r. At k>1k>1 no such closed form is available, since the identity of the runner-up experts depends on the magnitudes |zℓ||z_ | and not only on their signs. We therefore argue from a symmetry of the construction rather than from an explicit description of the selected set. The symmetry specializes to the sign-pattern argument at k=1k=1, so the proof below also covers the top-11 case. Definition 2 (Score vector). The score vector at a projected input z∈ℝrz ^r is the M-dimensional vector s(z)∈ℝMs(z) ^M that collects the routing scores of all M experts, with entries indexed by the experts themselves: s(z)a:=⟨ua,z⟩=1r∑ℓ=1raℓzℓ,a∈±1r.s(z)_a\;:=\; u_a,z \;=\; 1 r _ =1^ra_ z_ , a∈\± 1\^r. (13) The router selects the k experts whose entries in s(z)s(z) are largest. The key point is that under this construction the standard Gaussian has no special preference for one expert over another in expectation. Each expert is a vertex of the hypercube ±1r\± 1\^r, and negating a subset of the input coordinates is the same as negating the corresponding entries of the expert index, because the two negations cancel in the inner product ⟨a,z⟩ a,z . Negating input coordinates therefore only relabels the experts. We capture this intuition using the reflection action. Definition 3 (Reflection action). Let :=±1rG:=\± 1\^r, a group under coordinatewise multiplication ⊙ . It acts on ℝrR^r by σ⋅z:=σ⊙zσ· z:=σ z and on the expert set ±1r\± 1\^r by πσ(a):=σ⊙a _σ(a):=σ a. For a score vector s∈ℝMs ^M, write s∘πσs _σ for the relabelled vector with entries (s∘πσ)a=sσ⊙a(s _σ)_a=s_σ a. The lemma below states the three properties of this action that the proof uses. Lemma 2 (Symmetry properties). For every σ∈σ : (G1) (Invariance of the input distribution.) If z∼(0,Ir)z (0,I_r) then σ⋅zσ· z has the same distribution as z. (G2) (Equivariance of the router.) s(σ⋅z)=s(z)∘πσs(σ· z)=s(z) _σ for all z∈ℝrz ^r. Reflecting the input relabels the score vector by πσ _σ and changes nothing else. (G3) (Transitivity.) The map σ↦πσ _σ is a group action of G on the expert set which is simply transitive: for all a,b∈±1ra,b∈\± 1\^r there is exactly one σ∈σ with πσ(a)=b _σ(a)=b, namely σ=a⊙bσ=a b. Moreover πσ−1=πσ _σ^-1= _σ. Proof. (G1) The coordinates of z are independent and symmetric about 00, so negating any subset of them preserves the joint distribution. (G2) By (13), s(σ⋅z)a=1r∑ℓaℓσℓzℓ=1r∑ℓ(σ⊙a)ℓzℓ=s(z)σ⊙a=(s(z)∘πσ)as(σ· z)_a= 1 r _ a_ _ z_ = 1 r _ (σ a)_ z_ =s(z)_σ a=(s(z) _σ)_a. (G3) Since πσπτ=πσ⊙τ _σ _τ= _σ τ and π=id _1=id, the map is a group action. If πσ(a)=b _σ(a)=b then σ=σ⊙a⊙a=b⊙aσ=σ a a=b a, which is unique, and πb⊙a(a)=b⊙a⊙a=b _b a(a)=b a a=b confirms that it works. Finally πσ(πσ(a))=σ⊙σ⊙a=a _σ( _σ(a))=σ σ a=a. ∎ Top-k selection reads the score vector only through the ordering of its entries, never through which expert happens to carry which entry. Relabelling the experts therefore relabels the selected set in the same way. Lemma 3 (Top-k selection is equivariant and almost surely well defined). Let s∈ℝMs ^M be a score vector whose M entries are pairwise distinct, and let TopK(s)TopK(s) be the set of the k experts whose entries in s are largest. Relabelling does not create equal entries, so TopK(s∘πσ)TopK(s _σ) is defined for every σ∈σ , and the selection is G-equivariant: TopK(s∘πσ)=πσ(TopK(s)).TopK(s _σ)\;=\; _σ (TopK(s) ). Moreover, for z∼(0,Ir)z (0,I_r) the entries of s(z)s(z) are almost surely pairwise distinct, so TopK(s(z))TopK(s(z)) is almost surely well defined. Proof. Relabelling the entries of a vector relabels its order statistics. The entry of s∘πs π at index a is the entry of s at index π(a)π(a), so the indices carrying the k largest entries of s∘πs π are exactly π−1π^-1 applied to those for s. With π=πσπ= _σ and πσ−1=πσ _σ^-1= _σ from (G3), this is the stated equivariance. For distinctness, if a≠ba≠ b then ua−ub=(a−b)/r≠0u_a-u_b=(a-b)/ r≠ 0, so z:s(z)a=s(z)b\z:s(z)_a=s(z)_b\ is a hyperplane through the origin and therefore has Lebesgue measure zero. A union over the finitely many pairs is still null. ∎ Lemma 4 (Uniform marginals). Let z∼(0,Ir)z (0,I_r) and let 1≤k≤M1≤ k≤ M. Then Pr[a∈TopK(s(z))]=kMfor every expert a∈±1r. [\,a (s(z))\, ]\;=\; kM every expert a∈\± 1\^r. Proof. By Lemma 3 we may work on the almost sure event that the entries of s(z)s(z) are pairwise distinct. Fix σ∈σ and an expert a. Using (G1) in the first step, (G2) in the second, and equivariance in the third, Pr[a∈TopK(s(z))] [\,a (s(z))\,] =Pr[a∈TopK(s(σ⋅z))]=Pr[a∈TopK(s(z)∘πσ)] = [\,a (s(σ· z))\,]= [\,a (s(z) _σ)\,] =Pr[a∈πσ(TopK(s(z)))]=Pr[πσ(a)∈TopK(s(z))], = [\,a∈ _σ(TopK(s(z)))\,]= [\, _σ(a) (s(z))\,], where the last step uses πσ−1=πσ _σ^-1= _σ from (G3). By transitivity (G3), for any two experts a and b there is a σ with πσ(a)=b _σ(a)=b, so the function a↦Pr[a∈TopK(s(z))]a [a (s(z))] takes the same value at every expert. Since |TopK(s(z))|=k|TopK(s(z))|=k always, we have ∑a=1Ma∈TopK(s(z))=k _a=1^M1\a (s(z))\=k, and taking expectations gives ∑a=1MPr[a∈TopK(s(z))]=k _a=1^M [a (s(z))]=k. A constant function on M points summing to k equals k/Mk/M at each point. ∎ Proof of Proposition 1. Lemma 4 gives selection probability k/Mk/M for every expert. Fix an expert a. The indicators a selected at xi1\a selected at x_i\, i∈[n]i∈[n], are independent because the xix_i are independent, and each is Bernoulli(k/M)(k/M). Hence Na∼Binomial(n,k/M)N_a (n,k/M) with mean nk/Mnk/M and variance v:=nkM(1−kM)v:=n kM(1- kM). Bernstein’s inequality with increments bounded by 11 gives Pr[|Na−nkM|≥t]≤ 2exp(−t22v+23t). \! [ |N_a- nkM |≥ t ]\;≤\;2 \! (- t^22v+ 23t ). Setting the right-hand side to δ/Mδ/M and solving t2−23Lt−2vL=0t^2- 23Lt-2vL=0 yields t=L3+L29+2vL≤2vL+23Lt= L3+ L^29+2vL≤ 2vL+ 23L, using x+y≤x+y x+y≤ x+ y. Bounding v=nkM(1−kM)≤nkMv=n kM(1- kM)≤ nkM and taking a union bound over the M experts gives the claimed bound. This proposition isolates only the routing step. The dependence on n enters through the number of examples assigned to each expert, approximately nk/Mnk/M. The router rank only needs to distinguish the M expert labels, which requires r=log2Mr= _2M in this construction, independently of k. The expert MLPs then memorize their assigned examples, as in the memorization upper bound of Jelassi et al. (2025). ∎