Source-linked AI summary
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, Jennifer Ngadiuba, Richard Cavanaugh, Javier Duarte
TL;DR
Sparse MoE layers can add stored parameter capacity without proportional active computation, but how this trade-off behaves in particle-physics transformers is unclear. This study compares dense and MoE Particle Transformers on JetClass-II and finds that stored capacity, active computation, and routing organization have distinct effects on prediction.
Problem
It is unclear how MoE’s trade-off between increased stored parameter capacity and active computation behaves in particle-physics transformers, including whether structured routing accompanies better classification.
Source: §1 · p. 1Method
Results
Top-1 expert scaling expands stored capacity with quickly saturating predictive gains, while activating multiple experts per token improves performance at higher computation.
Source: §4 · pp. 5–6, §5 · p. 6Takeaways & Limitations
MoEParT’s effectiveness depends on stored capacity, active computation, and routing capacity, while routing organization is informative but not a direct proxy for predictive quality.
Source: §5 · pp. 6–7The study uses one Particle Transformer backbone, one JetClass-II setup, single checkpoints, and a restricted training subset, so trade-offs may not transfer unchanged or generalize to higher-data regimes.
Source: §6 · p. 7
Abstract
from arXiv · showhide
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.
1 Introduction
Sparse MoE layers offer a way to expand stored capacity without proportionally increasing active computation, but their predictive and routing effects in particle-physics transformers are uncertain. This study compares dense and MoE Particle Transformers across routing configurations and finds that stored expert count, active computation, and routing structure have distinct effects.
1 Introduction
- Sparse expert routing may drop tokens under restricted capacity, and activating more experts changes both computation and expert-output combinations.Routing associations with identifiable particle properties do not necessarily improve whole-jet classification.Source: §1 · p. 1
- No-drop top-1 MoE scaling changes routing without improving accuracy, while top-2 improves prediction at higher nominal compute.The study also finds that capacity constraints accompany degraded performance and stronger routing association is not a monotonic indicator of better classification.Source: §1 · p. 1
- The results separate the effects of stored capacity, active computation, and routing structure in MoE-based jet classifiers.The study compares dense ParT and MoEParT across expert count, capacity, top-K, and auxiliary loss.Source: §1 · pp. 1–2
2 Related Work
Prior work established particle-based jet tagging, large-scale JetClass-II classification, and sparse expert routing in transformers. This study uses JetClass-II as a controlled setting to examine sparse expert routing within a Particle Transformer.
Deep learning for jet tagging
- ParticleNet uses dynamic graph convolution, while Particle Transformer incorporates pairwise particle interactions into attention and improved performance on JetClass.This work retains ParT’s representation and attention structure while modifying its feed-forward sublayers.Source: §2 · p. 2
Large-scale jet pre-training and JetClass-II
- Sophon uses a ParT pretrained on 188-class JetClass-II for heavy-resonance searches and transfer learning.Here, JetClass-II is instead the controlled setting for comparisons of conditional capacity, active computation, and routing organization.Source: §2 · p. 2
Sparse mixture-of-experts models
- Sparse MoE research spans conditional capacity, distributed Transformer routing, single-expert routing, capacity and load balancing, and training stability [Shazeer et al., 2017, Fedus et al., 2021].These studies motivate separating stored expert capacity, active computation, and routing capacity.Source: §2 · p. 2
Mixture-of-experts models in high-energy physics
- MoE has been applied to event-level signal-background classification, but this study examines constituent-level routing within a Particle Transformer.The authors state that prior work had not systematically studied sparse MoE routing within ParT for large-scale jet classification.Source: §2 · p. 2
3 Model, data, and evaluation
The study replaces Particle Transformer feed-forward layers with sparsely routed experts and evaluates the resulting models on 188-class JetClass-II classification. It varies routing choices and training settings while assessing prediction, computation, and routing organization.
Architecture and routing
- A learned linear router selects each token’s top-K experts, with per-expert assignment capacity set by expert count and capacity factor; excess assignments are dropped.For K > 1, selected routing weights are renormalized before expert outputs are combined, and an auxiliary load-balancing loss discourages imbalanced routing.Source: §3 · pp. 2–3
- MoEParT replaces the feed-forward networks in all eight particle-attention and two class-attention blocks with expert networks, while retaining the shared embedding and classifier architecture.Both models produce a single 188-class jet prediction; the models use eight heads and particle-attention blocks, two class-attention blocks, and feed-forward widths of 128→512→128.Source: §3 · p. 2, Figure 2 · p. 3, Figure 1 · p. 3
Training setup
- Dense and MoE configurations use matched data exposure and optimization budgets, but each configuration is represented by one training run without random-seed variation.Models train for 20 epochs with 1,000,000 sampled training examples per epoch and checkpoints selected by validation accuracy; final evaluation uses a held-out subset of 456,822 selected jets.Source: §3 · p. 4
Experimental configurations and task
- The study compares capacity-limited and no-drop top-1 models, top-K variants, and an auxiliary-loss sweep on Pythia JetClass-II’s 188 jet classes.Jets contain up to 128 particles; the reported signal score sums probabilities for classes 0–160, with classes 161–187 treated as QCD background, and AUROC and QCD rejection are also reported.Source: §3 · pp. 3–4
Computation and routing diagnostics
- Forward-pass FLOPs are profiled for one jet with 128 particles as arithmetic cost, not as wall-clock latency or throughput.Fused scaled-dot-product attention is disabled during profiling to expose relevant operations.Source: §3 · p. 4
- Routing structure is measured with normalized mutual information between the primary selected expert and particle properties such as particle type, relative pT, and pT rank.The analysis compares aggregate associations rather than matching expert identities across runs; padded tokens are excluded.Source: §3 · p. 4
4 Results
No-drop top-1 MoE improves test accuracy over dense ParT at nearly unchanged forward compute, while added experts bring little further accuracy gain. Capacity constraints and higher top-K alter performance, and routing associations do not track classification quality monotonically.
4 Results
- 42.73–42.82% test accuracy for no-drop top-1 models with E = 2, 4, 8 compares with 41.69% for dense ParT at nearly identical standardized forward FLOPs.This comparison separates increased stored parameter capacity from a comparable increase in active computation.Source: §4 · p. 4, Figure 3 · p. 5
- At c = 1.5, E = 8 drops 21.1% of valid assignments and reaches 39.90% accuracy versus 42.73% with sufficient capacity; overflow is associated with degradation, not proven causal.The capacity-limited and no-drop models were trained separately.Source: Figure 3 · p. 5, §4 · p. 5
- Moving from K = 1 to K = 2 improves accuracy for both E4 and E8, with roughly 51% more nominal forward FLOPs.The gains occur in a higher-compute regime, rather than from stored capacity alone.Source: Figure 3 · p. 5, §4 · p. 5
- As E rises from 2 to 8, parameters grow from 3.64M to 11.62M while accuracy stays nearly flat; R50 nevertheless rises from 83.03 to 88.71.The accuracy gain from stored experts appears to saturate quickly, while background rejection continues to change.Source: §4 · p. 5, Figure 3 · p. 5
- In the no-drop K = 1 sweep, first-block particle-type NMI rises from 0.118 at E = 2 to 0.272 at E = 8 as accuracy remains nearly unchanged.These associations do not establish semantic expert specialization.Source: Figure 4 · p. 6, §4 · p. 6
- At E = 4, K = 1 to K = 2 strengthens late-layer relative-pT routing, yet E8/K4 improves accuracy over E8/K2 while first-block particle-type NMI falls from 0.378 to 0.205.Routing association is not a monotonic predictor of performance.Source: Figure 4 · p. 6, §4 · p. 6
- For E8, K4 slightly improves accuracy and R80 relative to K2, while K2 retains the best R50; neither configuration therefore dominates across all evaluated operating points.Source: §4 · p. 6
- At fixed E4/K1, α = 0.05 improves rejection while lowering first-block particle-type NMI, and varying α produces no monotonic accuracy trend.Stronger load balancing does not consistently improve accuracy or routing structure.Source: §4 · p. 6
5 Discussion and conclusion
MoEParT’s effectiveness depends on more than stored expert count: active computation, routing capacity, and routing organization also matter. The results distinguish these factors while showing that routing structure is not a direct measure of predictive quality.
5 Discussion and conclusion
- Top-1 expert scaling adds stored capacity with nearly unchanged active computation, but gains saturate; activating multiple experts improves performance at higher computation.Sparse computation can improve over a dense Particle Transformer at comparable nominal forward compute, but adding experts alone is insufficient.Source: §5 · pp. 6–7
- Insufficient routing capacity can drop more tokens as expert count grows, confounding added model capacity with reduced effective computation.Meaningful comparisons must account for whether the router can successfully assign tokens to the available experts.Source: §5 · p. 6
- Expert assignments associate with particle identity and kinematics, but stronger routing structure does not correspond monotonically to higher classification performance.Routing organization describes how representations are allocated, not a direct proxy for predictive quality or semantic specialization.Source: §5 · p. 6
6 Limitations and Future Work
The study’s findings are bounded by its single-backbone, JetClass-II setup, one checkpoint per configuration, and restricted training subset. Future work proposes testing larger samples and routing dynamics; code and experiment configurations are available.
6 Limitations and Future Work
- The study uses one Particle Transformer backbone, one JetClass-II training setup, a restricted training subset, and a single checkpoint per configuration.Consequently, transfer to other architectures or tasks, variation across random seeds, and behavior in a higher-data regime remain unmeasured.Source: §6 · p. 7
- Future work proposes comparing dense and MoE models across larger training samples and tracking routing dynamics throughout training.These studies would assess whether capacity value changes with data availability and when routing specialization emerges or reorganizes.Source: §6 · p. 7
- Code, model configurations, training scripts, and experiment settings are available in the study repository.The repository revision is df27f8e912a08ea887f77c7940e3932e2a47fc11.Source: §6 · p. 7