A theoretical model for task routing in mixture-of-expert transformers

• URL: http://arxiv.org/abs/2606.14398v2
• Date: Mon, 15 Jun 2026 02:01:46 GMT
• Title: A theoretical model for task routing in mixture-of-expert transformers
• Abstract summary: We prove formally that a single-layer MoE transformer can encode knowledge by using experts that specialize in the corresponding tasks.<n>Our construction provides a theoretical support for empirical results on localized knowledge circuits in MoE models.
• Score: 61.91842664811864
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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 \textit{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.

Related papers

• Mixture-of-Transformers Learn Faster: A Theoretical Study on Classification Problems [59.94955550958074]
  We study a tractable theoretical framework in which each transformer block acts as an expert governed by a continuously trained gating network.<n>We show that expert specialization reduces gradient conflicts and makes each subtask strongly convex.<n>We prove that the training drives the expected prediction loss to near zero in $O(log(epsilon-1)$ steps, significantly improving over the $O(epsilon-1)$ rate for a single transformer.
  arXiv  Detail & Related papers  (2025-10-30T21:07:36Z)
• MoLAE: Mixture of Latent Experts for Parameter-Efficient Language Models [10.623996218106564]
  Mixture of Experts (MoE) has become a key architectural paradigm for efficiently scaling Large Language Models (LLMs)<n>We introduce MoLAE, a novel parameterization that reformulating expert operations through a shared projection into a lower-dimensional latent space, followed by expert-specific transformations.<n>We show that MoLAE significantly improves efficiency across multiple dimensions while preserving model capabilities.
  arXiv  Detail & Related papers  (2025-03-29T14:35:34Z)
• Convergence Rates for Softmax Gating Mixture of Experts [78.3687645289918]
  Mixture of experts (MoE) has emerged as an effective framework to advance the efficiency and scalability of machine learning models.<n>Central to the success of MoE is an adaptive softmax gating mechanism which takes responsibility for determining the relevance of each expert to a given input and then dynamically assigning experts their respective weights.<n>We perform a convergence analysis of parameter estimation and expert estimation under the MoE equipped with the standard softmax gating or its variants, including a dense-to-sparse gating and a hierarchical softmax gating.
  arXiv  Detail & Related papers  (2025-03-05T06:11:24Z)
• Revisiting SMoE Language Models by Evaluating Inefficiencies with Task Specific Expert Pruning [78.72226641279863]
  Sparse Mixture of Expert (SMoE) models have emerged as a scalable alternative to dense models in language modeling. Our research explores task-specific model pruning to inform decisions about designing SMoE architectures. We introduce an adaptive task-aware pruning technique UNCURL to reduce the number of experts per MoE layer in an offline manner post-training.
  arXiv  Detail & Related papers  (2024-09-02T22:35:03Z)
• The Extrapolation Power of Implicit Models [2.3526338188342653]
  Implicit models are put to the test across various extrapolation scenarios: out-of-distribution, geographical, and temporal shifts. Our experiments consistently demonstrate significant performance advantage with implicit models.
  arXiv  Detail & Related papers  (2024-07-19T16:01:37Z)
• Diversifying the Expert Knowledge for Task-Agnostic Pruning in Sparse Mixture-of-Experts [75.85448576746373]
  We propose a method of grouping and pruning similar experts to improve the model's parameter efficiency.<n>We validate the effectiveness of our method by pruning three state-of-the-art MoE architectures.<n>The evaluation shows that our method outperforms other model pruning methods on a range of natural language tasks.
  arXiv  Detail & Related papers  (2024-07-12T17:25:02Z)
• Estimating Knowledge in Large Language Models Without Generating a Single Token [12.913172023910203]
  Current methods to evaluate knowledge in large language models (LLMs) query the model and then evaluate its generated responses. In this work, we ask whether evaluation can be done before the model has generated any text. Experiments with a variety of LLMs show that KEEN, a simple probe trained over internal subject representations, succeeds at both tasks.
  arXiv  Detail & Related papers  (2024-06-18T14:45:50Z)
• On Least Square Estimation in Softmax Gating Mixture of Experts [78.3687645289918]
  We investigate the performance of the least squares estimators (LSE) under a deterministic MoE model. We establish a condition called strong identifiability to characterize the convergence behavior of various types of expert functions. Our findings have important practical implications for expert selection.
  arXiv  Detail & Related papers  (2024-02-05T12:31:18Z)
• Transformers are uninterpretable with myopic methods: a case study with bounded Dyck grammars [36.780346257061495]
  Interpretability methods aim to understand the algorithm implemented by a trained model. We take a critical view of methods that exclusively focus on individual parts of the model.
  arXiv  Detail & Related papers  (2023-12-03T15:34:46Z)
• Interpretable Mixture of Experts [71.55701784196253]
  Interpretable Mixture of Experts (IME) is an inherently-interpretable modeling framework. IME is demonstrated to be more accurate than single interpretable models and perform comparably with existing state-of-the-art Deep Neural Networks (DNNs) in accuracy. IME's explanations are compared to commonly-used post-hoc explanations methods through a user study.
  arXiv  Detail & Related papers  (2022-06-05T06:40:15Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.