A theoretical model for task routing in mixture-of-expert transformers Vinoth Nandakumar* University of Sydney Yunzhi Yao Zhejiang University
Yongli Xiang University of Sydney Peike Li Google Research
arXiv:2606.14398v1 [cs.LG] 12 Jun 2026
Abstract
duced inference cost (Shazeer et al., 2017; Fedus et al., 2022). Empirical studies have observed task–expert specialization in MoE transformers (Li et al., 2025b), where experts learn to specialize in fine-grained semantic tasks. Ideally, each task would be routed to a well-defined sparse subset of experts, whose size scales with the intrinsic complexity of that task, leading to inference-time speedups (Kudugunta et al., 2021). Mapping specific tasks to distinct sub-networks could also enhance interpretability by allowing practitioners to trace decision-making using these functional circuits (Ostapenko et al., 2024), thereby leading to safer and more reliable systems. An important open question is whether we can establish a rigorous theoretical foundation for explicit task-expert specialization in MoE language models (Chen et al., 2023). From a theoretical perspective, recent studies have analyzed task-expert specialization in mixtureof-experts architectures (Chen et al., 2022; Dikkala et al., 2023), and establish theoretical guarantees by using continuous data distributions, such as Gaussian mixture models. However, these frameworks analyze feedforward-only MoE models lacking multi-head attention. They don’t extend to discrete data distributions, such as n-gram templates (Svete and Cotterell, 2024; Nandakumar et al., 2025), which use a predefined pattern with fixed tokens and wildcard symbols to generate sentences. To address these gaps in the theoretical foundations of MoE transformer models, we theoretically explain how they can utilize attention heads and task-specialized experts to process and store discrete structures that can be used to model language. Our contributions. To study this question theoretically, we introduce a simplified discrete model of structured knowledge based on syntactic templates and finite key–value dictionaries. Within this framework, we establish a formal result show-
Mixture-of-experts (MoE) layers enable the scaling of transformer models while keeping the inference compute fixed. While task-expert specialization has been observed in empirical studies of frontier MoE transformer models, existing theoretical work analyzes this using continuous mixture models that cannot be used to model natural language effectively. An important open question is to theoretically explain task-expert specialization in transformer MoE models using discrete models of language. To address this, we represent structured knowledge via syntactic templates and finite key-value dictionaries, and prove formally that a single-layer MoE transformer can encode knowledge by using experts that specialize in the corresponding tasks. Our construction shows how queries are routed to unique, task-specific experts whose size depends solely on the intrinsic complexity of the given task (i.e. the combined size of its syntactic templates and factual dictionary). Our construction provides a theoretical support for empirical results on localized knowledge circuits in MoE models. We support our theoretical findings with experiments evaluating model performance under varying MoE loss functions.
1
Introduction
Modern transformer models achieve state-of-theart accuracy by leveraging hundreds of attention heads and billions of parameters, yet mounting evidence shows that only a small fraction of these computations is required for any single input (Liu et al., 2023; Gao et al., 2025). Frontier large language models increasingly adopt mixture-of-experts transformers (Dai et al., 2024; Yang et al., 2024), which dynamically route tokens through a sparse collection of expert modules at each layer, achieving state-of-the-art performance with substantially re*
Tongliang Liu University of Sydney
Correspondence to: [email protected]
1
Input Tokens
Attention Block
task 2
Attention
task 1 MoE Router
task 3
Output token expert 1 expert 2 expert 3
Figure 1: An architectural overview of task routing in mixture-of-expert transformers. The attention block separates dictionary tokens from templates, allowing the router to use template structure to select a task-specific expert for prediction. Expert size scales additively with task complexity.
ing that a single–layer mixture-of-experts transformer can represent structured knowledge using task-specialized experts. We theoretically construct MoE transformers in which the attention mechanism disentangles template structure from the factual subject, while the routing mechanism maps inputs to task-specific experts. Each expert then performs task-specific associative retrieval over the corresponding knowledge dictionary. In particular, our key contributions are as follows.
formers with task-specific experts (Theorem 4.2). Section 5 gives an overview of the proofs for these core theorems, and in Section 6 we conduct experiments on synthetic data that support our key theoretical findings.
2
Related work
2.1
Theoretical models for mixture-of-expert transformers
To bridge the gap between empirical success and mathematical understanding, recent theoretical studies have sought to formally characterize the optimization and generalization properties of mixtureof-experts architectures under simplified settings. For instance, Chen et al. (2022) provides a foundational analysis of the MoE model trained with gradient descent where each expert is a two-layer CNN, using a data distribution with cluster structure. Building upon this, Dikkala et al. (2023) uses synthetic data generated from a Gaussian mixture model, and shows that the router in a MoE model with a feedforward network learns the latent cluster structures. Further investigations have expanded on these properties by establishing formal generalization error bounds for sparse MoEs (Zhao et al., 2024) and demonstrating their theoretical advantages in mitigating catastrophic forgetting during continual learning (Li et al., 2025a). However, a significant limitation of these existing frameworks is that they do not use MoE layers with attention mechanisms, nor do they model the discrete, sequential data typical of natural language processing tasks.
• We identify a gap in existing MoE theory, which primarily explains routing using mixture models that cannot model natural language effectively. To address this, we propose a theoretical framework for sparse routing in MoE transformers based on syntactic templates and finite key–value dictionaries. • Within this framework, we show that attention blocks can separate structural templates from factual information, enabling routing based on relational structure while sparse experts perform task-specific associative retrieval, and the size of the expert scales with task complexity. • Using synthetic knowledge datasets, we empirically study routing behavior under different training objectives and observe that explicit taskaware routing objectives produce substantially clearer task–expert alignment than standard objectives alone. The remainder of this paper is organized as follows. Section 2 reviews related work on MoE transformers and knowledge circuits. Section 3 introduces a synthetic model for knowledge data, and transformer architectures that are used in subsequent sections. Section 4 presents our key theoretical results, explicitly constructing dense transformers with sparse circuits (Theorem 4.1) and MoE trans-
2.2
Task-specific experts for language models
In Gururangan et al. (2022), the authors introduce DeMix Layers, a modular mixture-of-experts (MoE) architecture that embeds domain specializa2
tion directly within the transformer network. Each layer contains a mixture of feed-forward experts, where every expert corresponds to a specific domain, and attention layers are shared across domains. In (Gururangan et al., 2023; Li et al., 2022), the authors propose an expert-style training recipe that discovers domains without supervision, clusters the corpus accordingly, trains a separate expert LM per cluster, and combines them as a sparse ensemble at inference (so only a few experts are active per input). While these studies highlight the practical success of task-specific routing, they primarily offer empirical observations without formal guarantees. Our results bridge this gap by providing a mathematically rigorous foundation for these architectures. 2.3
with words drawn from a task-specific knowledge table. This framework allows us to separate syntactic structure (captured by templates) from factual content (captured by key–value pairs). Definition 3.1 (Templates). Let Σ be a finite alphabet containing every token (word or punctuation mark) that may occur in the corpus. We introduce two distinguished wildcard symbols ⋆1 , ⋆2 ∈ / Σ. A template is a string τ = w1 w2 . . . wL ∈ (Σ ∪ {⋆1 , ⋆2 })∗ , that contains each wildcard ⋆1 and ⋆2 exactly once, with ⋆1 occurring before ⋆2 . For any template τ and any pair (k, v) ∈ Σ, define the instantiation τ [k, v] to be the string obtained by simultaneously replacing ⋆1 7→ k and ⋆2 7→ v in τ .
Knowledge circuits
Recent investigations into how transformers store and retrieve facts have identified ‘knowledge circuits’ (Yao et al., 2024; Wang et al., 2025), which are localized sub-networks within pretrained models dedicated to specific factual associations (Geva et al., 2023). Building on this with mechanistic interpretability, Ou et al. (2025) utilizes circuits to analyze how large language models acquire and update information during continual pre-training. Our results complement recent theoretical results understanding factual recall with dense transformers (Nichani et al., 2025), by showing how mixtureof-expert models can store key-value information using sparse, task-specific circuits. Our theoretical analysis further supports the empirical findings on knowledge circuits in Section 5 of Li et al. (2025b), which show how experts in Qwen 1.5-MoE specialize in structured key-value associations like name-birthplace and country-capital pairs (see also Bandarkar et al. (2026) for a related analysis with multilingual datasets).
3
Preliminaries
3.1
Data model for knowledge
Definition 3.2 (Knowledge tasks). A knowledge task k consists of a knowledge table ∆k , and a set of templates T (k). Here a knowledge table is a finite subset ∆k ⊆ Σ∗ × Σ∗ , whose first component is called the subject and whose second component is the associated object. Define the sentences D(k) generated by the collection as follows. D(k) =
τ [k, v] : τ ∈ T (k), (k, v) ∈ ∆k
Below we define the task complexity c(k) of a given task k as the combined size of its syntactic patterns and semantic facts. Here |τ | is the length of the template and |∆k | is the number of key-value pairs in the knowledge table. c(k) =
X
|τ | + |∆k |
τ ∈T (k)
Definition 3.3 (Knowledge dataset). Define a knowledge dataset to be a pair (K, T ), where K is a set of knowledge tasks and T is a set of templates. Here for each knowledge task k ∈ K, the set of templates T (k) is a subset of T . Define the sentence set generated by the collection as follows.
We start with a simplified theoretical model for structured knowledge datasets, which are defined using a collection of syntactic patterns in conjunction with a finite set of subject–object associations (see (Geva et al., 2023), and Section 4 of (Nichani et al., 2025), for similar models of synthetic factual recall tasks). The patterns are represented by templates containing designated wildcard positions, which are instantiated by replacing the wildcards
D(K, T ) =
[
D(k).
k∈K
That is, we instantiate every template in every family with every key–value pair drawn from the corresponding dictionary and then take the union over all knowledge tasks. 3
Example. Consider the two templates τ1 = “ ⋆1 is in ⋆2 .”, τ2 = “ ⋆1 is spoken widely in ⋆2 .”, with dictionaries ∆1 = {(Paris, France), (Madrid, Spain)}, ∆2 = {(English, Canada), (Hindi, India)}. Applying Definition 3.2 gives the following, using the knowledge tasks k1 = (∆1 , {τ1 }) and k2 = (∆2 , {τ2 }).
The top-1 MoE feedforward layer is then given as follows, and is applied independently at each position. For an input vector x ∈ Rd , here e∗ (x) denotes the index of the expert. e∗ (x) = arg max G(x) e e∈{1,...,E}
MoE(x) = FFe∗ (x) (x), For an input sequence A ∈ Rn×d , the MoE transformer layer with H attention heads is defined as
D(k1 ) = {“Paris is in France.”, “Madrid is in Spain.”},
A′ = MHA(A) + A, D(k2 ) = {“English is spoken widely in Canada.”,
A mixture-of-experts transformer model is obtained by stacking such layers and adding a linear output head, as in the dense case.
“Hindi is spoken widely in India.”}. 3.2
A′′ = MoE(A′ ) + A′ .
Transformer models
Dense transformers. We assume familiarity with transformers (Vaswani et al., 2017); see Appendix A.1 for the definitions. A transformer block first applies multi-head self-attention, then a position-wise feedforward neural network, wrapping each sub-block in residual connections. The multi-head attention mechanism (MHA) processes an input sequence by projecting it into queries, keys, and values using learned weight matrices WQ , WK , and WV . For each attention head, the output is formed by computing attention weights via a softmax using scaled dot products of the queries and keys. The outputs from multiple heads are concatenated via an output matrix WO . A transformer model M is obtained by composing transformer blocks followed by a linear output head. We define a circuit C as a sparse sub-network within the transformer M, consisting of a designated subset of attention heads and feedforward neurons (Yao et al., 2024). We denote by MC the restricted model where all components outside of C are set to zero.
4
Theoretical results.
4.1
Preliminaries: knowledge circuits in transformers.
We first investigate how structured knowledge is represented within dense transformers. A core premise of our framework is that they have the capacity to partition learned facts into localized, decoupled sub-networks, commonly referred to as knowledge circuits. The following theorem formalizes this intuition using our synthetic data model, and shows that a single-layer transformer can encode knowledge, using a sparse circuits to perform the computation for any individual task. Because the size of this task-specific circuit depends only on the task’s intrinsic complexity, this result provides the fundamental theoretical justification for task routing in MoE architectures, which can use a router to eliminate redundant computation. Theorem 4.1 (Sparse circuits in a single–layer transformer). Let Σ be a finite alphabet, let T ⊆ (Σ ∪ {⋆1 , ⋆2 })∗ be any finite set of wildcard templates, each with length at most L. Let K be a finite set of knowledge tasks. For each task k ∈ K, let ∆k ⊆ Σ∗ × Σ∗ be the corresponding dictionary with size |∆k | < N , and let T (k) ⊆ T be the corresponding set of templates. For any ϵ > 0, there exists a transformer M with the following properties.
Mixture of experts transformers. A mixtureof-experts (MoE) transformer modifies the feedforward component of a transformer layer by replacing a single feedforward network with a collection of expert feedforward networks and a gating mechanism that selects experts per position (Shazeer et al., 2017). For simplicity, we restrict our formulation to top-1 expert routing, activating only a single expert per token (Fedus et al., 2022). Let FF1 , . . . , FFE be E position-wise feedforward networks, each with input/output width d. The gating function G : Rd → RE , is defined by a linear map followed by a softmax.
• The model M has one transformer layer with dimension 2L|T |. It consists of a multi-head attention layer with 2 attention heads, followed by a feedforward network with at most L|T |+N |K| neurons. This transformer block is then followed by a linear output layer. 4
• For each task k ∈ K, there exists a sparse circuit C(k) with at most L|T (k)| + N neurons in the hidden layer of the feedforward network that solves this task. MC(k) (x) = M(x)
dictionary with size |∆k | < N , and let T (k) ⊆ T be the corresponding set of templates. For any ϵ > 0, there exists a transformer M with the following properties. • The model M has one transformer layer with dimension 2L|T |. It consists of a multi-head attention layer with 2 attention heads, followed by a mixture-of-experts layer with experts indexed by K. The expert corresponding to k ∈ K has at most L|T (k)| + N neurons. This transformer layer is then followed by a linear output layer.
for every x ∈ D(k)
• The model can achieve an error of less than ϵ: for any prefix w, the model’s output probability distribution M(w) = P satisfies the following. Here V (w) ⊆ Σ denotes the set of all valid next tokens from the distribution D(K, T ). X P (w) > 1 − ϵ
• Given an input sequence w consisting of the first i words generated by a template τ ∈ T (k), the transformer routes the sequence to the expert corresponding to k.
w∈V (w)
This theoretical construction provides a formal mathematical foundation for the empirical phenomena recently observed by (Yao et al., 2024), which demonstrate that large, pretrained language models naturally encode specific relational facts within highly localized, sparse sub-networks, which they term “knowledge circuits.” While their findings rely on empirical techniques like causal tracing to locate these circuits, Theorem 4.1 proves that transformers have the structural capacity to modularize knowledge using sparse knowledge circuits. 4.2
• The model can achieve an error of less than ϵ in predicting the valid next words. Specifically, for any prefix w, the model’s output probability distribution P satisfies the following. Here V (w) ⊆ Σ denotes the set of all valid next tokens from the distribution D(K, T ). X P (w | w) > 1 − ϵ w∈V (w)
While recent theoretical work has analyzed the dynamics of optimization in MoE models (Chen et al., 2022; Li et al., 2025a), our result focuses on expressiveness and diverges from these frameworks in two critical aspects. First, these studies analyze isolated MoE layers, consisting solely of a router and feedforward experts, without using attention layers. Second, they define tasks using clustering approaches, such as Gaussian mixtures, which do not effectively represent natural language processing tasks. In contrast, Theorem 4.2 explicitly integrates the multi-head attention mechanism with the MoE layer to process data, using a symbolic data model built on structural templates and relational dictionaries.
Main results: task routing in MoE transformers.
Recent empirical studies in mechanistic interpretability have demonstrated that MoE transformers naturally develop specialized experts that handle distinct knowledge tasks (Li et al., 2025b). These findings lack a formal guarantee that such MoE models can perfectly isolate discrete factual associations, and the following theorem bridges this gap by providing a mathematically rigorous foundation for these empirical observations. We establish that a single-layer MoE transformer can explicitly partition symbolic knowledge into distinct, task-specific experts without entanglement. Crucially, our construction proves that the required capacity of each expert scales strictly with the complexity of its assigned templates and dictionary.
5
Proofs of key results
In this section, we analyze the expressiveness of MoE transformer models and establish a rigorous mathematical foundation for explicit task-expert specialization. We prove that while dense transformers have the structural capacity to encode facts using sparse knowledge circuits, mixture-ofexperts architectures can explicitly modularize this
Theorem 4.2 (Task routing in a single–layer mixture-of-experts transformer). Let Σ be a finite alphabet, let T ⊆ (Σ ∪ {⋆1 , ⋆2 })∗ be any finite set of wildcard templates, each with length at most L. Let K be a finite set of knowledge tasks. For each task k ∈ K, let ∆k ⊆ Σ∗ × Σ∗ be the corresponding 5
computation by routing discrete tasks to specialized experts whose capacity scales strictly with task complexity. 5.1
rons to map all unique template prefixes to vectors that encode next-word tokens, while the semantic block utilizes |Ik′ | ≤ N neurons to map keys to vectors that encode their corresponding values. Construction of the output layer.
Proof of Theorem 4.1.
In this section, we outline the proof of Theorem 4.1. Given an input sequence w = (wj )1≤j≤i , and an index c ∈ {1, 2}, we denote by ac (w) ∈ R the image of w in the corresponding attention head. We denote by z ′ (w) ∈ R2L|T | the image after concatenating and passing through the WO matrix, and z(w) ∈ R2L|T | the image after passing through the feedforward network.
Lemma 5.1. For every integer n ≥ 3, there exist vectors v1 , . . . , vn ∈ Rn spanning a subspace of dimension 2 such that, for each i ∈ {1, . . . , n}, the i-th coordinate of vi is the unique maximal coordinate of vi . To construct the output logits, we use the above Lemma, which implies that for each token in the vocabulary, we can construct a vector within a 2dimensional subspace of the output logits, such that the maximal coordinate of that vector is indexed by the corresponding token. We construct the matrix WU so that its output is this 2-dimensional subspace, and choose the target outputs of the feedforward network above to be the pre-image of these vectors under the mapping induced by WU . By applying a sufficiently large scaling factor to the WU , we can ensure that the resulting softmax probabilities for the target token are larger than > 1 − ϵ for a given ϵ, concluding the proof.
Proof. Construction of embedding layer. We partition the embedding dimension into two blocks: one dedicated to encoding the tokens that appear as keys or values in the dictionaries, and the other for the syntactic tokens forming the templates. Construction of attention heads. Let the input sequence w consist of the first i words generated by a template τ ∈ T (k) and a subject-object pair (k, v) ∈ ∆k , for a given task k ∈ K. We choose the query, key and value matrices of the first attention head so that a1 (w) = vi (τ ) ∈ RL|T | is a vector that depends on the template τ and the index i, and not the subject k. We also choose the weights so that the vectors {vi (τ )} are linearly independent in R2L|T | . We choose the query, key and value matrices of the second attention head so that a2 (w) is a vector that depends only on the subject k. We also set WO to be the identity matrix, so that z ′ (w) is obtained by concatenating a1 (w) and a2 (w). Construction of feedforward network. We decompose the neurons I in the hidden layer into 2|K| blocks as follows. G G I= Ik ∪ Ik′ k∈K
5.2
Proof of Theorem 4.2.
In this section, we outline the proof of Theorem 4.2. Let w = (wj )1≤j≤i be an input sequence corresponding to a task k ∈ K, that consists of the first i words generated by a template τ ∈ T (k) and a subject-object pair (k, v) ∈ ∆k . Given an index c ∈ {1, 2}, we denote by ac (w) ∈ R the image of w in the corresponding attention head. We denote by z ′ (w) ∈ R2L|T | the image after concatenating and passing through the WO matrix, and z(w) ∈ R2L|T | the image after passing through the feedforward network. We partition the vectors z(w) and z ′ (w) into two block components, each of dimension L|T |.
k∈K
For each k ∈ K, the sparse circuit C(k) is obtained by selecting neurons in Ik , which are used for syntactic tokens, and Ik′ which are used for factual recall. To enforce strict circuit isolation, large negative biases are applied to all off-task neurons, ensuring that the ReLU activations are zero for any task k ′ ̸= k. The syntactic block Ik connects exclusively to the output of the first attention head, while the semantic block Ik′ connects exclusively to the output of the second attention head. We use Lemma A.4 to construct the weights of these blocks, which memorize the respective mappings. The syntactic block has |Ik | = L|T (k)| neu-
Proof. Construction of attention heads and embedding layer. We follow the approach used in the proof of Theorem 4.1 to construct the embedding layer, and choose the query, key and value matrices so that a1 (w) = vi (τ ) ∈ RL|T | is a vector that depends only on the template τ and the index i (and not the subject k), while a2 (w) depends only on the subject k. Our choice of these matrices also ensures that the vectors {vi (τ )} are linearly independent in RL|T | . 6
0.11 0.13 0.11 0.07 0.02 0.04 0.17 0.05 0.03 0.04 0.08 0.14 0.00 0.09 0.12 0.11 0.07 0.04 0.11 0.08 1
0.09 0.05 0.01 0.19 0.18 0.12 0.03 0.09 0.03 0.13 0.13 0.14 0.18 0.02 0.01 0.03 0.14 0.10 0.01 0.02 2
0.11 0.09 0.06 0.12 0.19 0.05 0.13 0.12 0.20 0.15 0.12 0.09 0.06 0.16 0.08 0.11 0.11 0.18 0.15 0.13 3 Expert ID
0.14 0.22 0.21 0.11 0.08 0.11 0.17 0.19 0.16 0.05 0.12 0.07 0.03 0.12 0.20 0.18 0.12 0.16 0.18 0.22 4
0.24 0.17 0.40 0.25 0.22 0.15 0.22 0.30 0.15 0.20 0.22 0.13 0.15 0.18 0.36 0.21 0.10 0.29 0.31 0.17 5
(a) CE loss only
0.20 0.25 0.11 0.25 0.27 0.44 0.14 0.15 0.34 0.34 0.32 0.29 0.45 0.32 0.18 0.31 0.36 0.21 0.19 0.35 6
0.40 0.35 0.30 0.25 0.20 0.15 0.10 0.05
P86 P407 P137 P131 P361 P39 P735 P50 P22 P19 P495 P171 P607 P161 P105 P136 P463 P17 P108 P27
0.16 0.24 0.13 0.03 0.04 0.07 0.03 0.07 0.05 0.05 0.02 0.13 0.16 0.14 0.15 0.10 0.14 0.05 0.12 0.07 0
0.13 0.09 0.08 0.17 0.17 0.07 0.18 0.10 0.08 0.09 0.10 0.13 0.07 0.13 0.08 0.10 0.13 0.10 0.15 0.13 1
0.17 0.13 0.22 0.19 0.17 0.25 0.14 0.19 0.16 0.19 0.18 0.15 0.16 0.14 0.17 0.20 0.12 0.16 0.16 0.17 2
0.12 0.11 0.11 0.15 0.21 0.09 0.25 0.13 0.23 0.16 0.14 0.07 0.13 0.14 0.07 0.10 0.11 0.17 0.17 0.15 3 Expert ID
0.05 0.23 0.06 0.04 0.01 0.12 0.04 0.11 0.04 0.13 0.16 0.11 0.07 0.03 0.12 0.10 0.19 0.08 0.00 0.11 4
0.28 0.12 0.30 0.27 0.28 0.18 0.17 0.31 0.26 0.23 0.27 0.19 0.19 0.20 0.20 0.20 0.14 0.27 0.31 0.27 5
0.10 0.08 0.09 0.15 0.12 0.23 0.20 0.09 0.17 0.16 0.13 0.22 0.23 0.23 0.21 0.18 0.18 0.17 0.08 0.10 6
0.30
0.25
0.20
0.15
0.10
0.05
Relation ID
0.11 0.09 0.10 0.01 0.05 0.08 0.15 0.09 0.09 0.07 0.02 0.14 0.13 0.10 0.06 0.04 0.10 0.03 0.06 0.04 0
Relation ID
Relation ID
P86 P407 P137 P131 P361 P39 P735 P50 P22 P19 P495 P171 P607 P161 P105 P136 P463 P17 P108 P27
P86 P407 P137 P131 P361 P39 P735 P50 P22 P19 P495 P171 P607 P161 P105 P136 P463 P17 P108 P27
(b) CE + Load Balancing loss
0.05 0.02 0.02 0.00 0.05 0.03 0.02 0.02 0.02 0.95 0.03 0.90 0.93 0.01 0.00 0.00 0.00 0.00 0.00 0.04 0
0.09 0.04 0.05 0.90 0.28 0.00 0.01 0.88 0.00 0.01 0.78 0.02 0.00 0.02 0.01 0.04 0.02 0.09 0.08 0.03 1
0.03 0.00 0.00 0.01 0.00 0.91 0.02 0.02 0.01 0.00 0.00 0.00 0.04 0.08 0.01 0.02 0.97 0.00 0.73 0.03 2
0.03 0.00 0.01 0.04 0.22 0.02 0.01 0.04 0.04 0.01 0.05 0.03 0.01 0.84 0.02 0.86 0.00 0.88 0.12 0.01 3 Expert ID
0.77 0.04 0.05 0.04 0.43 0.03 0.05 0.03 0.07 0.01 0.09 0.03 0.01 0.02 0.93 0.06 0.00 0.02 0.06 0.02 4
0.02 0.00 0.86 0.01 0.00 0.01 0.00 0.00 0.84 0.01 0.05 0.00 0.00 0.00 0.01 0.00 0.00 0.00 0.00 0.86 5
0.02 0.89 0.00 0.00 0.02 0.00 0.88 0.01 0.02 0.01 0.00 0.01 0.00 0.02 0.03 0.01 0.00 0.00 0.00 0.00 6
0.8
0.6
0.4
0.2
0.0
(c) CE + Router loss
Figure 2: Task-Expert routing distributions under different training objectives with 20 tasks and 7 experts.
Construction of the router. We construct a router G with the property that for the input sequence w, the chosen expert corresponds to the task k. To do this, we choose the gating function G : R2L|T | → R|K| so that it only depends on the first L|T | coordinates of the input vector. Since the vectors {vi (τ )} are linearly independent, we can choose the linear map so that for each i and template τ ∈ T (k), G(vi (τ )) is a one-hot vector in R|K| indexed by k.
6
Experiments.
6.1
Dataset
We construct a synthetic knowledge dataset from the setup in Section 3.1, using entity pairs and natural language templates from WikiData5M (Wang et al., 2021), which is publicly available under an open-source license. Our use of WikiData5M is consistent with its intended use for academic research, and the synthetic dataset is similarly intended solely for research purposes. We sample 50 knowledge tasks from WikiData5M, where each task k is indexed by a relation IDs. The templates T (k) are generated by sampling 5 sentences from the relation ID using Gemini-2.5. The corresponding knowledge table D(k) is obtained by sampling 200 entity pairs for that relation. We tokenize each sentence using a custom tokenizer that preserves quoted entity names as atomic tokens. The dataset is then split into training and test subsets, with an 80% − 20% train-test split.
Construction of the experts. To construct the expert for each task k ∈ K, we separate its neurons into two blocks, and follow the approach used in the proof of Theorem 4.1. The first block is responsible for identifying the next word in the template τ , and maps the image of current template vi (τ ) ∈ RL|T | to an embedding of the next word. The second block functions as an associative memory responsible for retrieving factual associations from the knowledge table ∆k using the representation a2 (w) of the key k for every dictionary entry (k, v) ∈ ∆k .
6.2
Output logits and probability bound. Finally, we map the expert’s output z(w) to a probability distribution over the vocabulary Σ using a linear unembedding matrix WU , following the proof of Theorem 4.1. When using Lemma A.4 above, for each of the two blocks we choose the outputs so that they lie in a 2-dimensional subspace of RL|T | . We then construct a linear mapping WU between these two-dimensional subspaces, so that the logits l = WU z ′ (w) satisfy the following property: the logit for the valid next token w ∈ V (w) is strictly greater than the logit for any invalid token y ∈ / V (w). To satisfy the theorem’s error bound, we apply a sufficiently large scaling factor γ > 0 to the weight matrix WU , so that the model predicts the correct next tokens with a total probability greater than 1 − ϵ.
Models
We use a one-layer mixture-of-experts transformer model, following the setup in Section 3.2. We train the model on the synthetic dataset obtained from WikiData5M using stochastic gradient descent. We evaluate the model by autoregressively generating sequences from randomly truncated sentences from the test set, until the end-of-sentence token is produced or a maximum length is reached. The model’s accuracy is the proportion of generated sequences that exactly match sentences in the dataset. We refer the reader to Appendix B for more details about the experimental setup. We optimize expert utilization using two distinct auxiliary objectives. The load balancing loss (Shazeer et al., 2017; Fedus et al., 2022) encourages an even distribution of tokens across all available 7
(a) Number of Tasks = 10 CE
90 88
1
3
5
7
Number of Experts
CE
CE+LB
Accuracy (%)
90 88
1
3
5
7
Number of Experts
CE+Router
88
1
3
5
Number of Experts
7
9
7
9
(d) Number of Tasks = 40 CE
CE+LB
CE+Router
92 90 88 86
9
CE+LB
90
94
CE+Router
92
86
CE
92
86
9
(c) Number of Tasks = 30
94
(b) Number of Tasks = 20
94
CE+Router
92
86
Accuracy (%)
CE+LB
Accuracy (%)
Accuracy (%)
94
1
3
5
Number of Experts
Figure 3: Accuracy under different training objectives across varying numbers of tasks and experts. CE+Router achieves comparable performance in most settings.
experts to prevent routing bottlenecks. Our task router loss explicitly enforces task-specific routing by training the model to route each token directly to the expert associated with that task. For a batch of N tokens and E experts, let px,e and zx,e denote the routing probability and one-hot assignment, respectively, for token x ∈ {1, . . . , N } at expert e ∈ {1, . . . , E}. Let ex ∈ {1, . . . , E} be the target task expert index for token x. The loss functions are defined as follows: ! ! N N E X 1 X 1 X LLB = E · px,e zx,e N N e=1
x=1
substantially clearer correspondence, where each task is predominantly routed to a specific expert. 6.4
We compare the overall model performance under the three different training objectives in Figure 3. As we vary the number of tasks and experts, the three variants achieve comparable accuracy, indicating that enforcing task-specific routing does not degrade predictive performance. Notably, the model trained with task-based routing typically achieves a slightly higher accuracy when more experts are present. These results suggest that explicit taskspecific routing can also reduce the active parameter count while maintaining model performance (see Appendix B for more details).
x=1
N
Lrouter = −
1 X log px,ex N x=1
The total training objective is L = LCE + λLaux , where LCE is the cross-entropy loss, Laux ∈ {LLB , Lrouter } is the chosen auxiliary objective, and λ is a weighting coefficient (e.g., λ = 0.2). 6.3
Model Performance
7
Conclusion.
In this work, we establish a rigorous theoretical foundation for task-expert specialization in MoE transformers. By using a synthetic framework based on templates, we prove that a single-layer MoE transformer can solve this language modelling task using task-specific experts whose size scales with task complexity. Our experiments show that, in simplified settings with synthetic data, task-level routing encourages specialized experts to store knowledge. Our framework provides a mathematical blueprint for designing interpretable mixture-of-expert architectures with taskspecialized experts, that could yield inference-time speedups by reducing the active parameter count.
Routing specialization in trained mixture-of-experts models
We examine the relationship between tasks and experts by visualizing token-level routing distributions across different objectives in Figure 2. While training with cross-entropy alone (Figure 2a) or with an additional load balancing loss (Figure 2b) results in diffuse routing patterns that lack clear task–expert alignment, the task routing loss yields more uniform expert utilization. In contrast, incorporating task-based routing (Figure 2c) leads to 8
8
Limitations.
8.1
Theoretical models for deep mixture-of-experts transformers.
circuits. We utilized an AI assistant, Gemini 3 Pro, to iteratively draft and refine prose throughout the manuscript. We carefully reviewed and edited all generated text.
Our construction could be generalized to deeper mixture-of-experts (MoE) transformers by replacing templates with sentences generated using context-free grammars that can model the hierarchical structure of language (Zhao et al., 2023). In this setting, sentences are generated by derivation trees rather than single-step wildcard instantiations, with non-terminals corresponding to phrases in the sentence (Chomsky, 1957). In a deep MoE transformer, earlier attention layers can be used to parse grammar, and later expert layers route tokens to experts that encode knowledge (Li et al., 2025b). Extending this hierarchical formulation could provide a theoretical mechanism for modeling multistep reasoning in transformer models, where deeper MoE layers sequentially route intermediate tokens through specialized experts to obtain a reasoning chain (Merrill and Sabharwal, 2024). 8.2
References Lucas Bandarkar, Alan Ansell, and Trevor Cohn. 2026. Knowledge localization in mixture-of-experts llms using cross-lingual inconsistency. Preprint, arXiv:2603.17102. Zitian Chen, Yikang Shen, Mingyu Ding, Zhenfang Chen, Hengshuang Zhao, Erik Learned-Miller, and Chuang Gan. 2023. Mod-squad: Designing mixtures of experts as modular multi-task learners. 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 11828–11837. Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. 2022. Towards understanding the mixture-of-experts layer in deep learning. In Advances in Neural Information Processing Systems, volume 35, pages 23049–23062. Noam Chomsky. 1957. Syntactic Structures. Mouton.
Sequence-level routing in mixture-of-experts models.
Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. 2024. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models. arXiv preprint arXiv:2401.06066.
Our theoretical framework supports the thesis in (Kudugunta et al., 2021), which shows that routing entire inputs in mixture-of-expert models based on predefined task labels explicitly isolates domain knowledge into specialized sub-networks. While this approach relies on task labels for input data, in (Fan et al., 2024) the authors show that sequencelevel routing in a self-supervised setting is a viable alternative. It would be interesting to empirically compare our simplified experiments from Section 6.2 using token-level task routing in mixture-ofexperts with the sequence-level routing approach from (Fan et al., 2024). In this paradigm, the gating network computes a single, global routing decision for the sequence, typically by pooling token representations, thereby forcing all tokens within that sequence to be processed by the same chosen expert without requiring task labels. In (Fan et al., 2024), it is observed that while token-level routing often causes experts to specialize in shallow syntactic features, sequence-level routing encourages experts to capture higher-level, domain-specific semantic concepts.
Nishanth Dikkala, Nikhil Ghosh, Raghu Meka, Rina Panigrahy, Nikhil Vyas, and Xin Wang. 2023. On the benefits of learning to route in mixture-of-experts models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 9376–9396. Dongyang Fan, Bettina Messmer, and Martin Jaggi. 2024. Towards an empirical understanding of mixture of experts design choices. In ICLR 2024 Workshop on Understanding of Foundation Models (MEFoMo). William Fedus, Barret Zoph, and Noam Shazeer. 2022. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. In Advances in Neural Information Processing Systems (NeurIPS), volume 35, pages 24642–24655. Leo Gao, Achyuta Rajaram, Jacob Coxon, Soham V. Govande, Bowen Baker, and Dan Mossing. 2025. Weight-sparse transformers have interpretable circuits. arXiv:2511.13653. Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson. 2023. Dissecting recall of factual associations in auto-regressive language models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing.
Acknowledgments We would like to thank Goncalo Paulo and Catherine Arnett for insightful discussions on transformer 9
Suchin Gururangan, Mike Lewis, Ari Holtzman, Noah A. Smith, and Luke Zettlemoyer. 2022. Demix layers: Disentangling domains for modular language modeling. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies.
Yixin Ou, Yunzhi Yao, Ningyu Zhang, Hui Jin, Jiacheng Sun, Shumin Deng, Zhenguo Li, and Huajun Chen. 2025. How do llms acquire new knowledge? a knowledge circuits perspective on continual pre-training. In Findings of the Association for Computational Linguistics: ACL 2025, pages 14612–14631. Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In Proceedings of the International Conference on Learning Representations (ICLR).
Suchin Gururangan, Margaret Li, Mike Lewis, Weijia Shi, Tim Althoff, Noah A. Smith, and Luke Zettlemoyer. 2023. Scaling expert language models with unsupervised domain discovery. arXiv preprint arXiv:2303.14177. Sneha Kudugunta, Yanping Huang, Ankur Bapna, Maxim Krikun, Dmitry Lepikhin, Minh-Thang Luong, and Orhan Firat. 2021. Beyond distillation: task-level mixture-of-experts for efficient inference. In Proceedings of the EMNLP Findings, pages 3577– 3599.
Anej Svete and Ryan Cotterell. 2024. Transformers can represent n-gram language models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 6845–6881. Association for Computational Linguistics.
Hongbo Li, Sen Lin, Lingjie Duan, Yingbin Liang, and Ness B. Shroff. 2025a. Theory on mixture-of-experts in continual learning. In International Conference on Learning Representations.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Processing Systems.
Junzhuo Li, Bo Wang, Xiuze Zhou, Peijie Jiang, Jia Liu, and Xuming Hu. 2025b. Decoding knowledge attribution in mixture-of-experts: A framework of basic-refinement collaboration and efficiency analysis. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 22431–22446.
Xiaozhi Wang, Tianyu Gao, Zhaocheng Zhu, Zhengyan Zhang, Zhiyuan Liu, Juanzi Li, and Jian Tang. 2021. Kepler: A unified model for knowledge embedding and pre-trained language representation. Transactions of the Association for Computational Linguistics, 9:176–194.
Margaret Li, Suchin Gururangan, Tim Dettmers, Mike Lewis, Tim Althoff, Noah A Smith, and Luke Zettlemoyer. 2022. Branch-train-merge: Embarrassingly parallel training of expert language models. arXiv preprint arXiv:2208.03306.
Xuehao Wang, Liyuan Wang, Binghuai Lin, and Yu Zhang. 2025. Headmap: Locating and enhancing knowledge circuits in llms. In The Thirteenth International Conference on Learning Representations.
Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, and 1 others. 2023. Deja vu: Contextual sparsity for efficient llms at inference time. In International Conference on Machine Learning, pages 22137–22176. PMLR.
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, and 1 others. 2024. Qwen2 technical report. arXiv preprint arXiv:2407.10671.
William Merrill and Ashish Sabharwal. 2024. The expressive power of transformers with chain of thought. In The Twelfth International Conference on Learning Representations.
Yunzhi Yao, Ningyu Zhang, Zekun Xi, Mengru Wang, Ziwen Xu, Shumin Deng, and Huajun Chen. 2024. Knowledge circuits in pretrained transformers. In Advances in Neural Information Processing Systems, volume 37, pages 33215–33234.
Vinoth Nandakumar, Qiang Qu, Peng Mi, and Tongliang Liu. 2025. State space models can express n-gram languages. In Transactions on Machine Learning Research.
Chulhee Yun, Suvrit Sra, and Ali Jadbabaie. 2019. Small relu networks are powerful memorizers: a tight analysis of memorization capacity. In Advances in Neural Information Processing Systems.
Eshaan Nichani, Jason D Lee, and Alberto Bietti. 2025. Understanding factual recall in transformers via associative memories. In International Conference on Learning Representations.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. 2017. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations.
Oleksiy Ostapenko, Zhan Su, Edoardo Maria Ponti, Laurent Charlin, Nicolas Le Roux, Lucas Caccia, and Alessandro Sordoni. 2024. Towards modular llms by building and reusing a library of loras. In Proceedings of the 41st International Conference on Machine Learning, volume 235. PMLR.
Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. 2023. Do transformers parse while predicting the masked word? In Proceedings of the Conference on Empirical Methods in Natural Language Processing.
10
Jinze Zhao, Peihao Wang, and Zhangyang Wang. 2024. Generalization error analysis for sparse mixture-ofexperts: A preliminary study. In ICLR 2024 Workshop on Understanding of Foundation Models (MEFoMo).
We also recall a result from (Zhang et al., 2017) that analyzes the memorization capacity of feedforward neural networks and shows that they can fit an arbitrary finite set of input-output vectors (see also (Yun et al., 2019)).
A
Appendix: Proofs
A.1
Definitions
Lemma A.4. Let {(xi , yi )}K i=1 be a finite dataset where each input xi ∈ Rn and each output yi ∈ Rq . Assume that all inputs are distinct, i.e., xi ̸= xj for i ̸= j. Consider a feedforward neural network fθ with one hidden layer of width K, ReLU activations, input dimension n, and output dimension q. Then there exists a choice of parameters θ such that fθ (xi ) = yi for all i ∈ [K].
Definition A.1. A position-wise feedforward network (FFN) with input/output width d and hidden width dff is the function FF : Rd → Rd , x 7→ W (2) σ(W (1) x+b(1) )+b(2) , with W (1) ∈ Rdff ×d , b(1) ∈ Rdff , W (2) ∈ Rd×dff , b(2) ∈ Rd , and σ is the ReLU applied elementwise.
A.3
Definition A.2. A transformer layer with H heads and width d is a function defined as follows.
Proof. Construction of attention heads. Recall that the embedding dimension is partitioned into two components: the first for template tokens and the second for dictionary tokens. We choose the query and key matrices of the first attention head so that they operate exclusively on coordinates within the first component. This yields an output a1 (w) = vi (τ ) ∈ RL|T | that depends exclusively on the underlying template τ ∈ T and the current sequence index i. Similarly, we choose the query and key matrices of the second attention head so that they operate exclusively on coordinates within the second component. The output vector a2 (w) depends only on the dictionary tokens present in the input sequence w. Because the total number of unique template prefixes is bounded by L|T |, which does not exceed the dimension of the target space RL|T | , we can rely on a generic choice of the matrices WQ , WK and WV to ensure the linear independence of the vectors {vi (τ )}. Using the standard attention mechanism, each output vector vi (τ ) can be expressed as a linear combination below, where αj are the attention weights, and xj denotes the input embedding (comprising both the token embedding and its positional encoding) for the j-th token in the sequence.
Layer : Rn×d → Rn×d , Layer(A) = A′′ A′ = MHA(A) + A A′′ = FF(A′ ) + A′ Here MHA is the multi-head attention mechanism, and FF is a position-wise feedforward network (FFN). The terms +A and +A′ are residual connections. ■ A.2
Proof of Theorem 4.1.
Lemmas
We start with a proof of Lemma 5.1 above. Lemma A.3. For every integer n ≥ 3, there exist vectors v1 , . . . , vn ∈ Rn spanning a subspace of dimension 2 such that, for each i ∈ {1, . . . , n}, the i-th coordinate of vi is the unique maximal coordinate of vi . Proof. Let θk = 2π(k−1) for k = 1, . . . , n, and n define the vectors below. u = (cos θ1 , cos θ2 , . . . , cos θn )⊤ , w = (sin θ1 , sin θ2 , . . . , sin θn )⊤ . vk = cos θk u + sin θk (vk )j = cos θk cos θj + sin θk sin θj = cos(θk − θj ).
vi (τ ) =
Hence (vk )k = cos 0 = 1, while for j ̸= k we have cos(θk − θj ) < 1, so the maximum coordinate of vk is attained uniquely at position k. All vk lie in span{u, w}, which has dimension 2 because u and w are linearly independent when n ≥ 3. This gives the desired family of n vectors in a 2-dimensional subspace of Rn .
i X
αj (xj WV )
j=1
Because the query and key matrices only operate on the first component, the attention weights αj are uniquely determined by the first i tokens in the template τ . With the standard positional encodings, each of these prefixes yields a distinct 11
logits l, we use Lemma 5.1. Letting n = |Σ| denote the vocabulary size, by Lemma 5.1, there exist n vectors {vw }w∈Σ ⊂ Rn spanning a 2-dimensional subspace, such that the w-th coordinate of vw is its unique maximal coordinate. If {b1 , b2 } is a basis for this subspace, each target vector vw corresponding to a valid token w can be uniquely expressed as a linear combination of these basis vectors: vw = αw b1 + βw b2 for some scalars αw , βw ∈ R. When applying Lemma A.4 to construct the weights of the feedforward network, we choose the output vectors so that they lie within a 2dimensional subspace of RL|T | . The linear unembedding matrix WU ∈ Rn×2L|T | is used to construct a mapping between these 2-dimensional subspaces, so that that the output is a linear combination of vw for w ∈ V (w). Because the w-th coordinate of vw is uniquely maximal by Lemma 5.1, it guarantees that the logit lw for the valid next token is strictly greater than the logit ly for any invalid token y ∈ / V (w). To ensure that the model has error at most ϵ, we scale the weight matrix WU by a sufficiently large constant γ > 0. Recalling that the final probability P (w) for any token w is given by the softmax function below. exp(γlw ) P (w) = P y∈Σ exp(γly )
combination of token representations and attention weights. Since there are at most L|T | distinct vectors {vi (τ )} in a space of dimension L|T |, for a generic choice of weight matrices WQ , WK and WV , the vectors will be linearly independent. Construction of feedforward network. We denote by h(w) the activations of the hidden layer of the feedforward network. We decompose the neurons I in the hidden layer into 2|K| blocks as follows. G G I= Ik ∪ Ik′ k∈K
k∈K
Here |Ik | = L|T (k)| and |Ik′ | ≤ N neurons. We choose the weights of the feedforward network such that the following properties hold. • Sparse activation (Circuit isolation): By adding a large negative bias to neurons not associated with the current task, we ensure that if k ′ ̸= k, then the ReLU activation yields hj (w) = 0 for all neurons j ∈ Ik′ ∪ Ik′ ′ . This isolates the computation to the sparse circuit C(k) = Ik ∪ Ik′ . • Decoupled processing: The block Ik connects exclusively to the first half of z ′ (w) (the template representation a1 (w)), while the block Ik′ connects exclusively to the second half (the subject representation a2 (w)). The final output of the network, z(w), consists of two blocks: the first predicts the next structural word of the template (computed solely by Ik ), and the second predicts the factual object from the dictionary (computed solely by Ik′ ).
Because the logits of the valid tokens are strictly greater than those of the invalid tokens, taking the limit as γ → ∞ forces the probability mass to concentrate entirely on the set V (w) of valid nextword tokens. Therefore, for any strictly positive ϵ > 0, there exists a scaling factor P γ such that the probability distribution satisfies w∈V (w) P (w) > 1 − ϵ. This concludes the proof of the theorem.
We use Lemma A.4 above to construct the weights for these blocks. The template block Ik , maps the representation a1 (w) to a vector representation of the next syntactic token. Because there are at most L|T (k)| unique linearly independent prefixes, Lemma A.4 guarantees this exact mapping can be memorized using |Ik | neurons. Concurrently, the dictionary block Ik′ must map the subject representations a2 (w) to their corresponding object representations. Since there are at most N entries in ∆k , Lemma A.4 guarantees this can be memorized using |Ik′ | ≤ N neurons. Construction of the output layer. Finally, we map the output z(w) of the transformer block to a probability distribution over the vocabulary Σ. To construct the weights WU of the linear unembedding layer, which maps z(w) to the
A.4
Proof of Theorem 4.2.
Proof. Construction of attention heads. We follow the approach used in the proof of Theorem 4.1, and choose the query, key and value matrices so that a1 (w) = vi (τ ) ∈ RL|T | is a vector that depends only on the template τ and the index i (and not the subject k), while a2 (w) depends only on the subject k. Our choice of these matrices also ensures that the vectors {vi (τ )} are linearly independent in RL|T | . Construction of the router. We construct a linear router G : R2L|T | → R|K| whose input is the concatenated representation z ′ (w) = 12
Table 1: Model performance under different MoE loss formulations across varying numbers of tasks and experts. We report accuracy (%) averaged over three runs to ensure robustness. num_tasks 10
20
30
40
num_experts 1 3 5 7 9 1 3 5 7 9 1 3 5 7 9 1 3 5 7 9
acc (ce loss) 88.8% 89.7% 89.7% 87.5% 88.7% 87.6% 89.4% 88.7% 89.6% 86.9% 86.3% 88.5% 88.2% 88.6% 88.8% 88.8% 89.1% 90.8% 90.1% 90.9%
acc (ce + router loss) 88.8% 90.7% 90.0% 88.7% 90.1% 87.6% 90.2% 90.0% 91.2% 90.7% 86.3% 88.4% 88.6% 89.9% 90.8% 88.8% 89.6% 91.3% 89.7% 92.4%
acc (ce + load balancing loss) 88.8% 89.4% 88.6% 88.4% 88.1% 87.6% 89.7% 88.5% 90.1% 87.5% 86.3% 87.0% 87.9% 88.9% 89.5% 88.8% 88.3% 90.2% 89.5% 90.9%
Figure 4: Example of template-based generation from WikiData5M. Given templates T and knowledge pairs D, sentences are generated by filling subject-object pairs into templates.
[a1 (w); a2 (w)]. We set the weights corresponding to the second component a2 (w) to strictly zero, so that the routing decision depends solely on the syntactic template: a1 (w) = vi (τ ). Since the vectors {vi (τ )} are linearly independent, we can choose the weight matrix of G that maps each vi (τ ) (for τ ∈ T (k)) exactly to the one-hot basis vector ek ∈ R|K| . This ensures that all tokens generated by templates in T (k) are routed to the k-th expert. Construction of the experts. For each k ∈ K, we separate the neurons of the corresponding expert into two block components. We choose the weights such that the first and second blocks of the input z(w) are routed strictly through their respective expert components and mapped to the corresponding blocks of the output z ′ (w), setting all other weights to zero. We use Lemma A.4, about memorizing input-output vectors with a neural network with a single hidden layer, to construct the weights. The first block of the expert is responsible for identifying the next word in the template τ . Its set of input-output pairs consists of (vi (τ ), ei (τ )), where vi (τ ) ∈ RL|T | is the representation of the current template τ ∈ T (k) at position i, and ei (τ ) is the embedding of the next token in the template. Since there are at most L positions across |T (k)| templates, from Lemma A.4, this first block requires at most L|T (k)| neurons to memorize these input-output pairs. The second block is responsible for retrieving dictionary values. We use Lemma A.4 again, where the inputs are the subject representation from
Knowledge Task This knowledge describes the geographic relationship between an entity and its country. Templates T • {subject} is located in {object}. • {subject} can be found in the nation of {object}. • {subject} is a landmark situated in {object}. • The country associated with {subject} is {object}. • {subject} is a site within {object}. Knowledge Pairs D • (‘bent county high school’, ‘united stated’) • (‘Opachychi’, ‘ucrania’) Generated Sentences • "bent county high school" is located in "united stated". • "Opachychi" is located in "ucrania". • "bent county high school" can be found in the nation of "united stated". • "Opachychi" can be found in the nation of "ucrania". • "bent county high school" is a landmark situated in "united stated". • "Opachychi" is a landmark situated in "ucrania". • The country associated with "bent county high school" is "united stated". • The country associated with "Opachychi" is "ucrania". • "bent county high school" is a site within "united stated". • "Opachychi" is a site within "ucrania".
a2 (w), and the target outputs are vector representation of the corresponding objects. Since there are at most N entries in the dictionary ∆k , Lemma 13
A.4 guarantees this mapping can be memorized using N neurons. By using separate blocks for template processing and dictionary retrieval, the total number of neurons required for the k-th expert is bounded by L|T (k)| + N . Output logits and probability bound. Finally, we map the output z ′ (w) of the k-th expert to a probability distribution over the vocabulary Σ. To construct the weights WU of the linear unembedding layer, which maps z ′ (w) to the logits l, we follow the approach used in the proof of Theorem 4.1 above, using Lemma 5.1. When applying Lemma A.4 to construct the experts, we choose the weights of the k-th expert so that their outputs lie entirely within a 2dimensional subspace of RL|T | . Using the argument from the proof of Theorem 4.1 above, the linear unembedding matrix WU ∈ Rn×2L|T | yields a mapping between these 2-dimensional subspaces, so that the logit lw for the valid next token w ∈ V (w) is strictly greater than the logit ly for any invalid token y ∈ / V (w). To ensure that the model has error at most ϵ, we scale the weight matrix WU by a sufficiently large constant γ > 0. Using the argument from the proof of Theorem 4.1 above, for any ϵ > 0, there exists a scaling factor P γ such that the probability distribution satisfies w∈V (w) P (w) > 1 −ϵ. This concludes the proof of the theorem.
B
Appendix: Experiment Details
We conduct all experiments on PyTorch using a single NVIDIA GeForce RTX 4090 GPU with 24GB memory, with the total training time taking less than 10 hours, and report the average results over three random seeds. For data generation, we construct 200 sentences per template (see Figure 4 for an illustration). For models trained with router supervision or load balancing, we set the loss weight to λ = 0.2. The transformer backbone uses a model dimension of 30, feedforward dimension 10, and 10 attention heads, with a single transformer layer in all configurations. Models are trained with a batch size of 64 for 50 epochs using the Adam optimizer with an initial learning rate of 0.006 and exponential decay factor 0.98 per epoch. Full numerical results are provided in Table 1.
14