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

2026-06-12

AI Summary

This work addresses the lack of a formal theoretical explanation for task-expert specialization in Mixture-of-Experts (MoE) Transformers. It proposes, for the first time, a discrete linguistic framework that formalizes this phenomenon by introducing syntactic templates and a finite key-value lexicon. Within this model, the authors rigorously prove that a single-layer MoE Transformer can achieve task-specific expert routing—where each query is uniquely mapped to a dedicated expert based solely on its task—and that the required number of experts scales only with the intrinsic complexity of the tasks. This theory not only elucidates the mechanism underlying task-driven routing but also provides a formal foundation for the empirically observed localized knowledge circuits. Experimental validation across multiple MoE loss functions further corroborates the theoretical findings.

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.

Problem

• mixture-of-experts
• task-expert specialization
• discrete language models
• transformer
• theoretical explanation

Innovation

• mixture-of-experts
• task-expert specialization
• discrete language model
• syntactic templates
• knowledge routing