Conditional Capacity and Routing in Mixture-of-Experts Particle Transformers

Kaushik Pendiyala; Haris Zia; Trevin Lee; Timothy Legge; Alejandro J. De Leon; Zihan Zhao; Aaron Wang; Abhijith Gandrakota

cs.LGOct 2, 2026

Conditional Capacity and Routing in Mixture-of-Experts Particle Transformers

Authors: Kaushik Pendiyala, Haris Zia, Trevin Lee, Timothy Legge, Alejandro J. De Leon, Zihan Zhao, Aaron Wang, Abhijith Gandrakota, +3 more

Organizations: University of California Davis, Davis, CA, USA · University of California San Diego, San Diego, CA, USA · University of Illinois at Chicago, Chicago, IL, USA · Fermi National Accelerator Laboratory, Batavia, IL, USA

Abstract

Mixture-of-Experts (MoE) models can increase parameter capacity without proportionally increasing active computation, but it is unclear how this trade-off behaves in particle-physics transformers. We study dense and MoE Particle Transformers on 188-class JetClass-II, varying expert count, routing capacity, top-K, and auxiliary loss. We find that, when token dropping is avoided, top-1 MoE models improve over the dense baseline at nearly unchanged nominal forward compute, while further increasing the number of stored experts produces little additional accuracy gain. Activating multiple experts per token yields additional predictive improvements at higher computational cost. Routing analyses show that expert assignments become more strongly associated with particle identity and kinematics in some configurations, but this structure does not increase monotonically with classification performance. These results highlight the need to distinguish stored parameter capacity, active computation, routing capacity, and routing organization when evaluating sparse expert models for jet classification. Code and experiment configurations are available at https://github.com/kpendiyala/MPT.

Figures & tables

Figure 1 : Overview of the dense Particle Transformer and MoEParT architectures. Dense ParT and MoEParT share the embedding and classifier architecture; MoEParT differs by replacing the dense feed-forward network with a sparsely routed mixture of experts. Both models produce a single 188-class jet prediction.

Panel of Figure 2 — Figure 2 : Dense and mixture-of-experts transformer blocks. Top: dense ParT uses a feed-forward network after attention. Bottom: MoEParT replaces this FFN with $E$ expert networks of the same architecture and routes each token to the top- $K$ selected experts. The replacement is applied in all eight particle-attention and two class-attention blocks.

Model	$E$	$K$	$c$	$\alpha$	$P$	$G$	Drop (%)	Acc. (%)	$R_{50}$	$R_{80}$
Dense	–	–	–	–	2.30	0.524	–	41.694	75.74	7.99
Limited	2	1	1.5	.01	3.64	0.524	0.465	41.396	74.76	8.00
Limited	4	1	1.5	.01	6.30	0.525	9.370	41.276	75.91	8.00
Limited	8	1	1.5	.01	11.62	0.526	21.126	39.902	54.53	6.31
No drop	2	1	2	.01	3.64	0.524	0	42.824	83.03	8.28
No drop	4	1	4	.01	6.30	0.525	0	42.819	87.88	8.50

Table 1: Model configurations and held-out test performance.

E

: number of experts;

K

: active experts per token;

c

: capacity factor;

\alpha

: auxiliary-loss coefficient;

P

: total parameters in millions;

G

: standardized forward GFLOPs per jet at 128 particles; Drop: validation valid-assignment drop rate.

R_{50}

and

R_{80}

denote QCD rejection at 50% and 80% signal efficiency, respectively.

Figure 3 : Performance under changes in routing capacity and active computation. (a) Capacity-limited and no-drop top-1 expert scaling; annotations give valid-assignment drop rates. (b) Test accuracy versus standardized forward GFLOPs. (c) $R_{50}$ versus active experts.

Figure 4 : Layer-wise routing NMI on validation data. (a) Particle-type association for no-drop top-1 models. (b) Relative- $p_{T}$ association for E4 K1/K2. (c) Particle-type association for E8 K2/K4.

Explore similar work

A theoretical model for task routing in mixture-of-expert transformers
Jun 12, 2026Vinoth Nandakumar, Yongli Xiang, Yunzhi Yao +2Mixture Models Theory
University of Sydney · Zhejiang University · Google Research
Distributionally Robust Mixture-of-Experts Training
Oct 5, 2026Xin Teng, Muxiao Li, Hongyi WenBalanced Learning
New York University Center for Data Science, NYU Shanghai
ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts
Jun 1, 2026Heng Zhao, Zilei Shao, Guy Van den Broeck +1Large Language Model Routing Probabilistic Inference
Department of Computer Science, University of Virginia, Charlottesville, USA · Department of Computer Science, University of California, Los Angeles, Los Angeles, USA

Jun 12, 2026, cs.LGJ/K move · Enter open · S save

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

Vinoth Nandakumar, Yongli Xiang, Yunzhi Yao, Peike Li+1

University of Sydney · Zhejiang University · Google Research

Mixture Models Theory Natural Language

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

Distributionally Robust Mixture-of-Experts Training

ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts