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

arXiv:2610.02701 · cs.LG, hep-ex · Submitted 2026-10-02 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Conditional Capacity and Routing in Mixture-of-Experts Particle Transformers".

Jane: The gist: Sparse conditional capacity improves over dense at nearly fixed compute, while activating multiple experts per token improves performance while increasing active computation.

Tom: First, who's behind it and why it matters.

Title and authors: Tom: We just talked about the title and authors of "Conditional Capacity and Routing in Mixture-of-Experts Particle Transformers" and how it sets up this comparison between dense and sparse models. It’s all about those two concepts interacting.

Jane: So, what does that actually mean for us when we think about building something? We're looking at jet classification on JetClass-II data, right? How does this comparison translate to practical application in a real AI system?

Lu: The paper is using Particle Transformers and comparing the dense version against MoEParT, which replaces the dense feed-forward networks with a sparsely routed mixture of experts <ref:2610.02701#pg3>. This is the core architectural change they are testing.

Meng: So, what’s on that page about how they set up both architectures? Is it just replacing the dense layers with experts?

Tom: It's more than just replacement. Both models start off with the same things: embedding seventeen particle features with widths of one hundred twenty-eight five hundred twelve and one hundred twenty-eight and they use eight heads for attention blocks and two class-attention blocks <ref:2610.02701#pg3>.

Jane: That’s a lot of detail about the base architecture. So when they talk about routing, what exactly is that learned linear router doing? Is it just picking one expert randomly?

Lalam: No, it's a learned linear router that produces a probability distribution p(x) through softmax of W x + b, and then it selects the K largest probabilities for active experts per token <ref:2610.02701#pg3>. That’s the mechanism for sparsity.

Tom: Right, so it's a learned decision process based on the input particle features, not just random selection <ref:2610.02701#pg3>. It’s tailored to the data at hand.

Lu: And they then look at how they define the capacity constraint for assignments using that router, which is C = c TK/E, where E is the number of experts and c is the capacity factor <ref:2610.02701#pg3>. That’s where it gets technical about limits.

Meng: That formula shows they are thinking about how to manage that load at scale, which I need to see because if E grows really big, that capacity constraint might become the real bottleneck <ref:2610.02701#pg3>.

Jane: So the paper is showing us a system where we have storage capacity but performance depends on how well we manage the routing and that defined capacity formula <ref:2610.02701#pg3>. It’s not just about having more experts in the system.

Tom: Exactly. It's about the interplay between those fixed features, the attention blocks, and this sparse mixture-of-experts routing mechanism <ref:2610.02701#pg3>. It’s a tightly coupled system for jet prediction.

The paper's summary: Jane: Now that we know the setup, let’s get to the actual summary of what they found in "Conditional Capacity and Routing in Mixture-of-Experts Particle Transformers." What is the big picture result they are reporting here?

Tom: They summarize it by saying that MoE models can increase parameter capacity without proportionally increasing active computation, but they aren't clear on how this trade-off behaves specifically in particle physics transformers <ref:2610.02701#pg1>.

Lu: They found that when token dropping is avoided, top-one MoE models perform better than the dense baseline at nearly unchanged nominal forward compute <ref:2610.02701#pg1>. This is a key finding because it’s efficiency over raw parameter count.

Meng: So they are saying you can get better accuracy on the same hardware if you switch to this sparse setup, provided you don't lose any tokens during routing <ref:2610.02701#pg1>. That sounds like something we need to think about for our training pipeline.

Lalam: They also noted that further increasing the number of stored experts doesn't bring more accuracy gain once you’re already using a top-one approach <ref:2610.02701#pg1>. So there’s a point where adding storage stops paying off on its own.

Tom: And they emphasize that activating multiple experts per token does yield additional predictive improvements, but only at a higher computational cost <ref:2610.02701#pg1>. It’s an explicit trade-off they are showing.

Jane: So the summary is this: sparse conditional capacity can improve over dense at nearly fixed nominal forward compute, but simply adding more experts is not sufficient for that improvement <ref:2610.02701#pg1>.

Lu: They also concluded that performance and stronger routing association are not a monotonic indicator of better classification <ref:2610.02701#pg2>. This means the way representations are partitioned across experts isn't automatically a sign of better results <ref:2610.02701#pg2>.

Meng: That’s important for interpretation because we shouldn't just look at the routing structure and assume it’s telling us which expert is doing the best job on a particle level <ref:2610.02701#pg4>.

Tom: So they are separating the effects of stored capacity, active computation, and routing structure in MoE-based jet classifiers using this study <ref:2610.02701#pg2>. That’s a very structured way to look at it.

Jane: It seems like the main message is that we need to consider all three of those factors together for a good sparse setup <ref:2610.02701#pg3>.

The paper's improvements: Tom: Let’s look at the specific improvements they suggest. They found that avoiding token dropping in top-one MoE models leads to better performance compared to the dense baseline <ref:2610.02701#pg1>.

Jane: So if we avoid those drops, we get a gain even if the nominal forward compute stays almost the same, which is a pretty strong efficiency claim for this paper <ref:2610.02701#pg3>.

Lu: They also showed that moving from K equals one to K equals two at E equals four substantially strengthens the late-layer relative pT routing association <ref:2610.02701#pg4>. This means increasing the number of active experts per token can actually help refine how the model makes decisions on particle properties.

Meng: So we get better predictive performance by paying that higher computational cost when we know the routing mechanism is robust enough to handle it <ref:2610.02701#pg1>. That’s a clear signal for where we might deploy this kind of setup.

Tom: But they also showed that capacity constraints accompany degraded performance, and increasing expert count without sufficient assignment capacity causes a growing fraction of routed tokens to be dropped <ref:2610.02701#pg3>. This is the main caveat we have to watch for in implementation.

Jane: That's the operational risk, Tom. If your capacity factor c isn't high enough relative to your expert count E, you’re going to end up dropping valid assignments, and that completely invalidates the intended benefit <ref:2610.02701#pg3>.

Lu: They also tested how auxiliary load-balancing loss influences things, and they found that varying alpha changes QCD rejection and routing association without a monotonic accuracy trend at fixed E equals four/K equals one <ref:2610.02701#pg3>.

Meng: So even tuning the loss function isn't a simple way to guarantee better accuracy; it’s more about controlling the balance between different aspects of the model <ref:2610.02701#pg3>.

Tom: The overall improvement they highlight is that MoEParT allows us to separate these three effects—stored capacity, active computation, and routing structure—so we can see which one is driving the performance gains <ref:2610.02701#pg2>.

Jane: So it’s not just about picking one thing to optimize; it’s about understanding how all those levers interact in this specific architecture <ref:2610.02701#pg3>.

Conclusion: Tom: So we're wrapping up with the conclusions from "Conditional Capacity and Routing in Mixture-of-Experts Particle Transformers." The main idea is that sparse conditional computation can improve over a dense Particle Transformer at comparable nominal forward compute <ref:2610.02701#pg1>.

Jane: But we have to be careful because simply adding more experts isn't the only way to get those improvements, and meaningful comparisons require looking at how the routing mechanism can actually utilize those experts <ref:2610.02701#pg3>.

Lu: They stress that the effectiveness of MoEParT depends on the interaction between stored parameter capacity, active computation, and routing capacity <ref:2610.02701#pg3>. That’s the core principle here.

Meng: So to deploy this effectively, you have to account for all those factors—the storage, the compute cost per token, and how well your router is distributing the load <ref:2610.02701#pg3>.

Lalam: I’m confident that focusing on that interaction will be key for future work <ref:2610.02701#pg3>. It moves beyond just counting parameters or computation <ref:2610.02701#pg3>.

Tom: Exactly. They also noted that the study is focused on a single Particle Transformer backbone and a JetClass-II training setup, so these observed trade-offs might not transfer unchanged to other architectures or jet classification tasks <ref:2610.02701#pg5>.

Jane: That’s a fair limitation. The authors are warning us that we can't just take this result and assume it works everywhere without testing it on different models or different tasks <ref:2610.02701#pg5>.

Lu: Future work will need to evaluate the principal dense and MoE configurations across multiple seeds and larger training subsets to determine the stability of the observed performance and routing trends <ref:2610.02701#pg5>.

Meng: And I’m hoping they do that. Stability under different conditions is what we actually need when we move from a research paper to something that runs reliably in production <ref:2610.02701#pg5>.

Tom: So, the final word on "Conditional Capacity and Routing in Mixture-of-Experts Particle Transformers" is that sparse conditional computation can improve over a dense Particle Transformer at comparable nominal forward compute, but simply adding more experts is not sufficient <ref:2610.02701#pg5>.

Jane: It’s a nuanced result. It’s about smart architectural design involving capacity, computation, and routing working together <ref:2610.02701#pg3>.

Kaushik Pendiyala, *, Haris Zia, *, Trevin Lee, *, Timothy Legge, *, Alejandro J. De Leon 2 Zihan Zhao2 Aaron Wang3 Abhijith Gandrakota4 Jennifer Ngadiuba4 Richard Cavanaugh4 Javier Duarte2

University of California Davis University of California San Diego University of Illinois at Chicago Fermi National Accelerator Laboratory

cs.LG, hep-ex

Submitted: 2026-10-02

Updated: 2026-10-07

Comments: 9 pages, 4 figures. Submitted to the ML4PS 2026 workshop

Code: https://github.com/kpendiyala/MPT

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 81/100

The gist: The gist: Sparse conditional capacity improves over dense at nearly fixed compute, while activating multiple experts per token improves performance while increasing active computation.

Key concepts

Mixture-of-Experts (MoE)
MoE models use several specialized 'experts' to process input tokens. Instead of every part of the network doing all the work, a learned router decides which few experts are needed for each token. This increases model capacity without proportionally increasing the total computation required.
Dense vs. MoE Particle Transformer
The study compares a standard 'dense' transformer where all processing happens in one large network against a 'MoE Particle Transformer' where the feed-forward layers are replaced by experts. The goal is to see if using sparse, expert-based computation can outperform the dense version under similar computational constraints.
Routing Capacity
This refers to how effectively the model can assign tokens to available experts. If there aren't enough slots or assignments for a given number of experts, some tokens might be dropped. This capacity constraint limits how much extra model capacity can actually be utilized during inference.
Routing Organization
This examines *how* the model chooses which expert to use for a token. The paper finds that while expert assignments become more tied to particle properties (like identity), this specific organizational structure does not always lead to better classification accuracy.

Terminology

Summary

The gist: Sparse conditional capacity improves over dense at nearly fixed compute, while activating multiple experts per token improves performance while increasing active computation.

Model Architecture and Comparison

The study compares dense Particle Transformer (ParT) with Mixture-of-Experts ParT (MoEParT) across various configurations including expert count, routing capacity, top-K settings, and auxiliary loss Both architectures embed 17 particle features with widths (128, 512, 128) and use eight heads, eight particle-attention blocks, and two class-attention blocks. MoEParT replaces the dense feed-forward networks in all ten blocks with E experts For a token x, a learned linear router produces p(x) = softmax(W x + b) and selects the K largest probabilities, where K is the number of active experts per token

Scaling Effects on Capacity and Computation

The research investigates how stored parameter capacity, active computation, and routing structure interact When using top-1 MoE models without dropping tokens, top-1 MoE models improve over the dense baseline at nearly unchanged nominal forward compute However, further increasing the number of stored experts produces little additional accuracy gain Conversely, Activating multiple experts per token yields additional predictive improvements at higher computational cost This suggests that sparse conditional capacity can improve over dense at nearly fixed compute, but simply adding more experts is not sufficient

Routing Capacity and Token Dropping

Routing capacity plays a critical role in the performance trade-off Increasing expert count without sufficient assignment capacity can cause a growing fraction of routed tokens to be dropped, confounding the effect of additional model capacity with reduced effective computation The per-expert assignment capacity is defined by C = ⌊c ⌈TK/E⌉⌋ Furthermore, Capacity constraints accompany degraded performance For instance, with a fixed capacity factor c = 1.5, the fraction of dropped valid assignments grows rapidly with expert count, reaching 21.1% for E = 8

Routing Structure and Predictive Performance

The analysis shows that routing structure does not monotonically correspond to better classification performance Routing organization is therefore informative about how the model allocates representations across experts, but should not be treated as a direct proxy for predictive quality or semantic expert specialization While expert assignments develop measurable associations with particle identity and kinematic properties, stronger routing association is not a monotonic indicator of better classification For example, at E = 4, moving from K = 1 to K = 2 substantially strengthens late-layer relative-pT routing

Auxiliary Loss and Metric Dependence

The auxiliary load-balancing loss influences the results depending on the metric being optimized At fixed E4/K1, varying α changes QCD rejection and routing association without a monotonic accuracy trend Specifically, α = 0.05 improves rejection while producing lower first-block particle-type NMI, indicating that stronger load balancing does not consistently improve either accuracy or routing structure Overall, the effectiveness of MoEParT depends on the interaction between stored parameter capacity, active computation, and routing capacity

Conclusion

The study concludes that sparse conditional computation can improve over a dense Particle Transformer at comparable nominal forward compute, but simply adding more experts is not sufficient Meaningful comparisons between sparse expert models therefore require considering not only how many experts are available, but also whether the routing mechanism can make use of them The effectiveness of MoEParT depends on the interaction between stored parameter capacity, active computation, and routing capacity Sparse conditional computation can improve over a dense Particle Transformer at comparable nominal forward compute, but simply adding more experts is not sufficient The study focuses on a single Particle Transformer backbone and JetClass-II training setup, so the observed trade-offs may not transfer unchanged to other architectures or jet-classification tasks Future work will therefore evaluate the principal dense and MoE configurations across multiple seeds and larger training subsets to determine the stability of the observed performance and routing trends The code, model configurations, and training scripts used in this study are available at https://github.com/kpendiyala/MPT The corresponding repository revision is df27f8e912a08ea887f77c7940e3932e2a47fc11 The repository includes the dense and MoE Particle Transformer implementations, JetClass-II training configuration, and experiment settings used for the reported results The paper is 2610.02701

--- Page 5 ---

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) R50 versus active experts

--- Page 6 ---

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

--- Page 7 ---

References

--- Page 8 ---

References

--- Page 9 ---

References

--- Page 10 ---

(No text provided for page 10) (End of extraction based on provided text)

Improvements for AI systems

  1. System can achieve improved predictive performance at nearly unchanged nominal forward compute by utilizing top-1 MoE models when token dropping is avoided, as stated in Figure 3(b). This suggests a more efficient way to scale parameter capacity without proportionally increasing active computation.

  2. The system can utilize multiple experts per token to yield additional predictive improvements at higher computational cost, indicating that increased active computation can be leveraged for better performance when routing capacity is sufficient.

  3. The system can implement routing analyses that become more strongly associated with particle identity and kinematics in some configurations, providing a structured view of how the model partitions particle representations, even if this structure does not increase monotonically with classification performance.

  4. The system can distinguish between stored parameter capacity, active computation, and routing organization when evaluating sparse expert models for jet classification, addressing the need to separate these distinct scaling effects.

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.

Sources

Related papers