Mixture-of-Experts Graph Transformers for Interpretable Particle Collision Detection

Technical Summary: Mixture-of-Experts Graph Transformers for Interpretable Particle Collision Detection

Problem Statement
The Large Hadron Collider (LHC) at CERN generates vast volumes of complex data from high-energy particle collisions, necessitating sophisticated analytical techniques for event classification and object identification. While Graph Neural Networks (GNNs) have demonstrated state-of-the-art performance in representing collisions as graphs, their "black box" nature limits their explainability. In high-energy physics (HEP), where establishing confidence in novel discoveries is paramount, the inability to interpret decision-making processes is a critical limitation. Existing post-hoc explainability methods (e.g., PGExplainer, GNNExplainer) often lack robustness and fail to provide a comprehensive understanding of internal model workings. This work addresses the challenge of developing machine learning models that achieve high predictive accuracy while embedding intrinsic explainability directly into the architecture.

Methodology
The authors propose a novel architecture, the Mixture-of-Experts Graph Transformer (MGT), designed to classify simulated collision events from the ATLAS experiment. The task involves distinguishing rare Supersymmetric (SUSY) signal events from dominant Standard Model (SM) backgrounds (specifically $t\bar{t}$ and $Wt$ production).

1. Data Representation: Each collision event is modeled as a fully connected graph $G$ with 6 or 7 nodes. Nodes represent reconstructed particles (jets, b-tagged jets, leptons, and missing transverse energy) and contain up to 6 kinematic features (e.g., transverse momentum $p_T$, pseudorapidity $\eta$, azimuthal angle $\phi$, b-tagging quantile, and mass).
2. Architecture: The model builds upon the Graph Transformer (GT) framework.
   • Graph Transformer: Utilizes multi-head self-attention mechanisms to update node representations, allowing the model to visualize attention maps and quantify relationships between particles.
   • Mixture of Experts (MoE): The standard feed-forward network (FFN) within the Transformer layers is replaced by an MoE block. This block consists of multiple specialized subnetworks ("experts") and a gating mechanism. The router dynamically assigns each node to a subset of experts based on learned patterns.
   • Training: The model employs a noisy top-K gating mechanism with a load-balancing loss to ensure even expert utilization and prevent the training of only a small subset of experts.
3. Explainability Mechanism: The architecture provides intrinsic explainability through two primary channels:
   • Attention Maps: Visualizing attention weights reveals which nodes (particles) the model focuses on when making predictions.
   • Expert Specialization: Analyzing the gating network's output identifies which experts are activated for specific node types, revealing how the model decomposes the problem into specialized sub-tasks.

Key Contributions

• Intrinsic Explainability: Unlike post-hoc methods, this approach integrates explainability into the model's core architecture, allowing for the direct tracing of predictions to specific attention mechanisms and expert activations.
• Physics-Aligned Interpretation: The study demonstrates that the model's internal mechanisms align with known physical principles. For instance, attention maps show high focus on b-jets ($b_1, b_2$) and missing energy for signal events, consistent with the decay of Higgs bosons and dark matter candidates.
• Expert Specialization Patterns: The analysis reveals that MoE experts naturally specialize in distinct particle types (e.g., specific experts focusing on lepton/energy pairs or b-jets), offering a granular view of the model's decision-making process.

Results
The model was evaluated on a dataset of 800,000 Monte Carlo simulated events (400,000 signal, 400,000 background) using ten runs with different random seeds.

• Performance: The MGT model achieved superior performance compared to baseline architectures (MLP, GCN, and standard GT). It attained an accuracy of 0.852 ± 0.0005 and an AUC of 0.929 ± 0.0039, outperforming the standard GT (0.849 accuracy, 0.928 AUC) and significantly exceeding the GCN (0.750 accuracy, 0.832 AUC).
• Interpretability Findings:
  • Attention Analysis: In correctly classified signal events, attention strongly concentrated on $b_1$, $b_2$, and energy nodes, reflecting the physics of Higgs decay and missing energy from dark matter. In background events, attention was more evenly distributed.
  • Expert Analysis: Experts in the first layer showed distinct specialization (e.g., Experts 2 and 4 focused on lepton and energy, while Experts 1 and 5 focused on b-jets). Deeper layers showed more overlap, suggesting a higher level of abstraction in combining features.
  • Ablation Studies: The authors found that load balancing loss is critical for expert specialization; without it, experts fail to specialize. The optimal configuration involved 6 experts with top-2 routing.

Significance and Claims
The paper claims that the MGT model successfully balances high predictive performance with robust explainability, offering a transparent tool for HEP data analysis. By linking predictions to physics-informed features through attention maps and expert specialization, the model fosters trust in AI-driven discoveries. The authors emphasize that explainability is as crucial as accuracy when exploring potential new physics. They conclude that while the current study focuses on a specific SUSY search, the framework demonstrates the viability of intrinsic explainable GNNs for complex scientific data, paving the way for more reliable AI methodologies in particle physics. The authors remain modest regarding scalability, noting that generalization to larger datasets and broader physics tasks requires further validation.