Towards replacing detector simulation with heterogeneous GNNs in flavour physics analyses

This paper presents a novel fast simulation tool for the LHCb experiment that utilizes heterogeneous graph neural networks to emulate detector responses for arbitrary multibody decay topologies, enabling generalizable interpolation to unseen channels and offering a scalable solution to the growing computational demands of particle physics simulations.

The Gist

Imagine you are trying to predict exactly how a complex machine, like a car engine, will behave when you turn the key. In the world of particle physics, the "machine" is the LHCb detector at the Large Hadron Collider, and the "turning of the key" is a particle collision.

To understand what happens after a collision, scientists usually run a massive, incredibly detailed computer simulation. It's like running a full-scale, hour-long movie of every single atom in the detector reacting to the crash. The problem is, the LHCb experiment is recording data so fast that they would need to run these "movies" for millions of hours every year. They simply don't have enough computer power or storage space to keep up.

Enter "Rex": The Fast-Forward Simulator

This paper introduces a new tool called Rex. Think of Rex not as a movie camera, but as a highly skilled artist who has memorized the style of the original movies.

Instead of simulating every tiny atom and every second of interaction (which takes forever), Rex looks at the "blueprint" of a particle decay (what particles were created) and instantly paints a picture of what the detector would have seen. It doesn't re-enact the physics step-by-step; it learns the patterns of the detector's response and generates the final result directly.

How Does Rex Learn? (The "Graph" Analogy)

The paper explains that Rex uses a special type of AI called a Heterogeneous Graph Neural Network. Here is a simple way to visualize that:

• The Graph: Imagine a party where guests are particles. Some guests are electrons, some are pions, some are muons. In a normal simulation, you might treat everyone the same. But in Rex's "party," the AI knows that an electron behaves differently than a muon.
• The Nodes and Edges: Each guest is a "node." The connections between them (who is talking to whom) are "edges."
• Heterogeneous: This just means the AI knows there are different types of guests and different types of conversations. It understands that a "kaon-to-electron" conversation is different from a "muon-to-pion" one.
• The Magic: By studying millions of real detector "movies," Rex learns the rules of these conversations. It learns that if two particles come in very close together, the detector gets confused (a "smearing" effect). If a particle is an electron, it tends to lose energy in a specific way.

What Rex Can Do

The paper claims Rex is a "generalist." It doesn't just memorize one specific decay (like a specific car crash). Instead, it learns the principles of how the detector works.

• The "Interpolation" Trick: If you show Rex a decay it has never seen before (a new type of particle combination), it can still guess the outcome accurately because it understands the underlying rules, just like an artist who can draw a new type of car because they understand how wheels and engines work, even if they've never seen that specific model.
• Speed: The paper states that generating data for 10 million events takes about one hour on a standard computer. Doing the same thing with the old, full simulation would take roughly 100,000 times longer (about 100,000 hours). It's the difference between watching a movie in real-time versus watching a 100,000-hour marathon.

Does It Work? (The "Taste Test")

The researchers tested Rex by running a "blind taste test." They took real physics analyses (looking for specific rare particle decays) and replaced the slow, full simulation data with Rex's fast data.

• The Results: The paper shows that the "taste" (the statistical distributions of the data) was almost identical. Rex correctly predicted how often particles would be detected, how their paths would curve, and how well they could be identified.
• The "J/ψ" Test: They even tested a specific ratio called $R_K$, which is a famous measurement in particle physics. When they swapped in Rex's data, the result only shifted by a tiny amount (0.5%), which is considered a very small error in this field.

Limitations and Future Plans

The paper is honest about what Rex can't do yet:

• The "Guest List": Currently, Rex is great at handling charged particles (like pions, kaons, electrons, and muons) but doesn't handle protons or neutral particles yet.
• The "Room Layout": It approximates the detector's physical boundaries (geometric acceptance) rather than simulating them perfectly.
• The "Training": The AI is still learning. Sometimes it gets a little "jittery" during training, which can lead to small inaccuracies in very specific, rare scenarios.

The Bottom Line

This paper presents a tool that acts as a fast-forward button for particle physics. By using a smart, pattern-recognizing AI (the Graph Neural Network), Rex can generate the data scientists need for their analyses in a fraction of the time and storage space required by traditional methods. It allows physicists to run more experiments, search for more background noise, and potentially discover new physics without being bottlenecked by slow computers.

Technical Summary

Technical Summary: Towards replacing detector simulation with heterogeneous GNNs in flavour physics analyses

Problem Statement

Experiments at the Large Hadron Collider (LHC), particularly LHCb, face an unprecedented increase in data recording rates. To maintain the precision of physics analyses, the volume of simulated data required to model detector efficiencies and backgrounds must grow proportionally. However, the current standard workflow relies on computationally intensive full simulation (using Pythia, Geant4, and reconstruction algorithms), which consumes a significant fraction of the computing budget and limits the size of available simulation samples. This limitation introduces major sources of uncertainty in high-profile analyses. Existing fast simulation tools either rely on simplifications of full frameworks (offering modest speed-ups), use fully parametric approximations (lacking necessary detail for precision studies), or attempt to model specific detector bottlenecks. There is a critical need for a tool that can emulate the full detector response with the speed of parametric tools but the fidelity of full simulation, capable of generalizing to arbitrary decay topologies without memorizing specific channels.

Methodology

The paper introduces Rex, a fast simulation tool based on Heterogeneous Graph Neural Networks (HGNNs) and Conditional Generative Adversarial Networks (cGANs). Unlike traditional approaches that simulate low-level energy deposits and pass them through reconstruction algorithms, Rex learns a stochastic mapping to directly generate high-level, analysis-ready variables conditioned on the kinematics of the decay.

Core Architecture and Training

• Graph Representation: The detector response is modeled using HGNNs where nodes represent particle tracks and edges represent interactions. The graph structure is heterogeneous, allowing for distinct node types (e.g., electron, muon, pion, kaon tracks) and edge types (e.g., kaon-to-electron interactions). This structure embeds physical characteristics directly into the network, removing the need for explicit particle species labels in the conditional input.
• Conditional Generation: The networks are conditioned on physical properties derived from event generators (e.g., EvtGen, RapidSim), including true particle identification, momenta, pseudorapidity, and geometric relationships between tracks.
• Modular Decomposition: The detector response is split into three interdependent but independently trained components:
  1. Primary Vertex (PV) Smearing: A simple GAN generates realistic PV locations conditioned on the decaying meson's momentum.
  2. Momentum Smearing: An HGNN smears track momenta. It utilizes self-loop layers for species-specific processing and Graph Attention Convolution (GAT) layers for message passing between tracks to capture inter-track correlations (e.g., resolution degradation when tracks are close in angle).
  3. Particle Identification (PID) and Trigger Response: A similar HGNN architecture generates PID variables (e.g., PIDi, ProbNNi) and trigger responses, conditioned on the generated reconstructed momenta to ensure correct correlations.
  4. Vertexing: A hierarchical HGNN reconstructs decay trees. It employs a two-pass message passing mechanism: a downward pass (parent to tracks) to correlate track features with object properties, and an upward pass (tracks to parent) mimicking the LHCb reconstruction algorithm's reliance on track momenta to build candidates. This network handles complex topologies, including partially reconstructed decays with missing particles (e.g., neutrinos) and misidentified tracks.

Data Preparation

Training data is extracted from the LHCb archive of fully simulated events. The process involves truth-matching tracks to ensure they originate from heavy meson decays and systematically combining them to form candidates representing all possible decay topologies (up to four tracks in the initial implementation). Variables are preprocessed using quantile transformations to map distributions to a truncated normal distribution, aligning with the network's activation functions and improving training stability.

Key Contributions

1. Generalized Topology Modeling: Rex is trained to be agnostic to specific final states. It learns generalizable patterns of detector response, allowing it to interpolate to decay channels not present in the training data, including fully reconstructed, partially reconstructed, and modes with misidentified particles.
2. Heterogeneous Graph Architecture: The use of HGNNs allows the model to naturally handle varying particle multiplicities and species-specific resolution effects (e.g., the distinct smearing behavior of electrons due to bremsstrahlung) without architectural changes.
3. End-to-End High-Level Generation: The tool bypasses the internal steps of LHCb reconstruction algorithms, directly outputting analysis-level variables (kinematics, PID, vertex quality) conditioned on truth-level kinematics.
4. Integration and Speed: Rex is implemented as a lightweight Python package compatible with existing LHCb workflows (e.g., RapidSim). It achieves a speed-up of approximately $O(10^5)$ compared to full simulation, enabling on-demand generation of analysis-ready samples.

Results

The paper validates Rex across a range of decay topologies, including $B^+ \to K^+ e^+ e^-$, $B^+ \to K^+ \mu^+ \mu^-$, and partially reconstructed decays like $B^+ \to \bar{D}^0 (\to K^+ e^- \bar{\nu}_e) \pi^+$.

• Momentum and PID: The generated distributions for momentum smearing ($\Delta E/E$) and PID efficiencies closely match full simulation, correctly capturing species-specific behaviors and correlations between tracks (e.g., resolution degradation for small opening angles).
• Vertexing and Reconstruction: The tool successfully models complex topological features, such as the broadening of vertex fit quality ($\chi^2_{DV/ndof}$) and impact parameter significance ($\chi^2_{IP}$) in decays with long-lived intermediates or missing particles. It correctly reproduces the effects of different reconstruction hypotheses (e.g., swapping track combinations).
• Analysis Performance: In a case study involving the $B^+ \to K^+ e^+ e^-$ analysis, Rex-generated samples passed standard selection cuts and Boosted Decision Tree (BDT) classifiers with high fidelity. The selection efficiencies and invariant mass distributions were well-reproduced.
• Systematic Uncertainties: Replacing full simulation with Rex in the measurement of $R_K$ results in a systematic shift of approximately 0.5%, which is subleading. However, for the single-ratio measurement $r_{J/\psi}$, which is more sensitive to mismodeling, the shift is around 2%, comparable to existing systematic uncertainties. The authors note that applying standard data-driven weighting procedures (used to correct full simulation) to Rex samples would further reduce these shifts.

Significance and Claims

The paper claims that Rex represents a significant step toward replacing computationally expensive detector simulations in flavour physics. By demonstrating that heterogeneous GNNs can learn the stochastic mapping of detector response across diverse decay topologies, the tool offers a pathway to:

• Decouple analysis from simulation bottlenecks: Enabling on-demand generation of large samples without the need for long-term storage of full detector readout.
• Facilitate novel studies: Allowing exhaustive grid searches for background processes that are currently computationally prohibitive.
• Generalizability: The architectures and methods are presented as generic, suggesting they could be adapted to emulate workflows at other particle physics experiments facing similar simulation demands.

The authors remain modest regarding current limitations, acknowledging that the tool is in its first iteration. They note that training stability, convergence criteria, and the modeling of specific variables (like trigger decisions and calorimetric variables) require further development. Future work aims to extend support to protons and higher multiplicity final states, improve the modeling of geometric acceptance, and potentially transition from GANs to diffusion models to enhance training stability.