Stress testing of fast reconstruction pipelines using machine learning

This paper introduces a domain-agnostic, context-aware stress testing pipeline to evaluate the robustness of fast reconstruction models in fields like High Energy Physics, demonstrating that while local assumptions can fail to capture global dependencies (as seen in HL-LHC $Z \rightarrow \ell\ell$ decays), an unsupervised regime-mapping framework can effectively restore reconstruction accuracy and resolution.

The Gist

Imagine you are trying to take a perfect photo of a moving car. In the world of high-energy physics (like at the Large Hadron Collider), scientists do something similar: they try to "reconstruct" what happened when particles smash into each other. They use computer programs called pipelines to figure out the speed, direction, and identity of these tiny particles.

For a long time, these programs have operated on a simple rule: "What happens here depends only on what's right next to it."

Think of it like a weather app that only looks at the temperature in your specific backyard to predict the weather for the whole city. It assumes the rest of the city is just like your backyard. This works fine most of the time, but it fails miserably if a massive storm is brewing just a few miles away.

The Problem: The "Local" Rule Breaks Down

This paper tests that "local rule" using a specific event: the decay of a Z-boson (a heavy particle) into two lighter particles (leptons). The researchers wanted to see if their reconstruction program stayed accurate when the "storm" (the overall energy of the collision) got too big.

They found that the program works great when the energy is low and the particles are close together. But, as the energy gets higher (like a storm rolling in), the program starts to get confused. It begins to make mistakes, misidentifying the mass of the particles. The "local" view wasn't enough; the program needed to know about the global context—the big picture of the entire collision.

The Solution: A "Context-Aware" Map

To fix this, the researchers used a special type of Machine Learning (called unsupervised learning).

Imagine you have a giant box of mixed-up LEGO bricks. Instead of sorting them by color (which requires you to know the colors beforehand), you let a robot sort them by how they naturally fit together. The robot notices that some bricks always stick together in certain patterns, while others don't.

The researchers did this with their data. They let the computer group the particle collisions into different "regimes" or "neighborhoods" based on two main things:

1. Total Energy ($H_T$): How much "oomph" the collision had.
2. Separation ($|\Delta\eta|$): How far apart the particles were flying.

The computer discovered four distinct neighborhoods:

• The Safe Zones (Regimes 0 & 1): Where the energy is moderate to high, and the particles are close together. Here, the reconstruction is usually good.
• The Danger Zones (Regimes 2 & 3): Where the energy is low or the particles are flying wildly apart. Here, the reconstruction gets messy.

The "Stress Test"

The researchers then put their reconstruction pipeline through a "stress test" in these different neighborhoods.

• In the Safe Zones: The pipeline was a hero. It reconstructed the particle masses perfectly.
• In the High-Energy Danger Zones: The pipeline stumbled. It started to drift, creating a "blur" in the results. The higher the energy, the worse the blur became. It was like trying to read a sign while driving 100 mph; the local details (the letters) became impossible to see clearly because the speed (global context) was overwhelming the system.

The Fix: A Smart Correction Layer

Once they identified where the pipeline was failing (the high-energy "hot spots"), they added a simple "correction layer."

Think of this like a GPS that realizes you are in a heavy traffic jam (a specific regime). Instead of giving you the standard "average speed" route, it adjusts its instructions specifically for that traffic jam. By telling the computer, "Hey, we are in the high-energy zone, adjust your calculations," they were able to restore the accuracy. The "blur" disappeared, and the true picture of the particles came back into focus.

What This Means

The paper concludes that we can no longer treat all particle collisions the same way. We can't just use one "average" rule for everything. Instead, we need to map out the different "neighborhoods" of the data and tune our tools specifically for each one.

Note on Medical Imaging:
The paper mentions that this same logic—checking if a tool works well in every "neighborhood" of a patient's scan—could be useful for medical imaging (like MRI or CT scans) in the future. However, the paper only tested this on particle physics data. It did not actually apply this to medical scans yet; it simply suggests that the strategy could be useful there later.

In short: The old way of looking at particle collisions was like using a single map for the whole world. This paper says, "That doesn't work when the terrain gets rough." They built a new system that recognizes different terrains and adjusts the map accordingly, ensuring the scientists get a clear picture even in the most chaotic, high-energy crashes.

Technical Summary

Technical Summary: Stress Testing of Fast Reconstruction Pipelines Using Machine Learning

Problem Statement
In scientific domains such as High Energy Physics (HEP) and Medical Imaging, fast reconstruction and detector-simulation pipelines are essential for handling computationally challenging tasks where full device-level simulation is impractical. These pipelines typically rely on simplified response models that assume reconstruction uncertainty depends solely on local input parameters, ignoring the global data context. However, this "local assumption" is frequently challenged in non-bulk regions of the phase space. Specifically, in HEP, the reconstruction accuracy of particle masses and decay kinematics degrades in high-energy regimes or distribution tails where data is sparse and particles are boosted. The standard approach of relying on a uniform distribution of data across the entire phase space fails to account for these context-dependent variations, leading to significant reconstruction bias and resolution degradation.

Methodology
To probe the robustness of the local assumption, the authors introduce a domain-agnostic, context-aware stress testing pipeline utilizing an unsupervised Machine Learning (ML) framework. The study focuses on the $Z \to \ell\ell$ decay channel at the High-Luminosity LHC (HL-LHC) as a benchmark.

• Data Generation: The analysis uses a dataset of 100,000 signal events generated by MadGraph5 aMC@NLO at $\sqrt{s} = 14$ TeV. To isolate the intrinsic behavior of the reconstruction pipeline, the study performs parton-level analysis, avoiding uncertainties from parton showering and detector smearing (e.g., Pythia8 and Delphes3 are excluded).
• Feature Space: Events are characterized by local parameters (individual four-momenta $p^\mu$ of leptons) and global context parameters: the scalar sum of lepton transverse momenta ($H_T = p_{T1} + p_{T2}$) and the absolute pseudorapidity separation ($|\Delta\eta|$).
• Unsupervised Regime Mapping: Instead of using supervised learning which requires pre-labeled data and introduces human bias, the framework employs K-Means clustering. This algorithm partitions the continuous phase space of $(H_T, |\Delta\eta|)$ into $k$ distinct context regimes without prior knowledge of the underlying physics. The dataset is split into 70% for training the regime map, 15% for validation, and 15% for stress testing.
• Stress Testing Metric: Reconstruction robustness is evaluated using the relative reconstruction error $\delta = (m_{reco} - m_{truth}) / m_{truth}$. A constant $\delta$ across varying global parameters indicates a robust pipeline, while a varying $\delta$ indicates bias.
• Correction Mechanism: A lightweight correction layer is implemented to restore truth-level resolution, taking the identified regime ID as an input feature.

Key Results
The stress testing reveals that the local assumption is violated in specific kinematic contexts:

• Context-Dependent Bias: While the reconstruction pipeline performs robustly in the "bulk" regions (low to moderate $H_T$), it exhibits significant divergence in performance in high-energy regimes.
• Regime Analysis: The unsupervised clustering identified four distinct regimes. Regimes 0 and 1, representing moderate-to-high energy events with varying separation angles, were found to be the primary signal regions for Z-boson decays.
  • Regime 1 (High Energy, Low Separation): Shows a narrow, tall distribution of reconstruction error centered at zero, indicating robust performance.
  • Regime 0 (Moderate Energy, Medium Separation): Shows a broad distribution of error, indicating non-robustness where the algorithm struggles to reconstruct mass accurately despite the mean error being near zero.
• Global Context Impact: Heatmap analysis demonstrates a clear correlation between increasing $H_T$ and the mean absolute bias. As the system moves from low to high energy regions, the bias increases, creating a "danger zone" above 400 GeV where the reconstruction pipeline becomes erratic. The $Z$-peak position shifts due to influences outside the local lepton parameters, confirming the necessity of global context awareness.
• Restoration: By applying the unsupervised regime-mapping framework, the study successfully restores peak stability and recovers truth-level resolution, effectively mitigating the bias observed in the uncorrected pipeline.

Significance and Claims
The paper claims that this work provides a concrete roadmap for quantifying performance in targeted, context-sensitive regimes rather than relying on uniform global averages. The primary significance lies in:

1. Diagnostic Capability: The framework serves as a robust diagnostic tool for next-generation fast simulation pipelines, capable of identifying "hot spots" of systematic degradation.
2. HL-LHC Readiness: The approach offers a path toward improving the overall framework robustness and precision required for future High-Luminosity LHC physics, where non-bulk regions are critical for new particle searches.
3. Methodological Shift: It advocates for a shift from global average performance metrics to targeted calibration within non-robust kinematic regions.
4. Future Utility: While the current work is limited to HEP, the authors note that the logic of "context-aware stress testing" is directly applicable to Medical Imaging (e.g., MRI or CT), where variables like Signal-to-Noise Ratio (SNR) and tissue density could replace $H_T$ and invariant mass to minimize diagnostic misclassification. This extension is identified as a direction for future work.