JetPrism: diagnosing convergence for generative simulation and inverse problems in nuclear physics

The paper introduces JetPrism, a configurable Conditional Flow Matching framework that addresses the misleading nature of standard training losses in nuclear physics simulations by establishing a multi-metric evaluation protocol to ensure generative models achieve true physical fidelity and convergence beyond generic loss indicators.

The Gist

The Big Picture: The "Blurry Photo" Problem

Imagine you are a detective trying to solve a crime. You have a high-definition photo of the crime scene (the Truth), but the only evidence you have is a photo taken through a foggy, scratched-up window (the Detector Data). The photo is blurry, distorted, and missing details.

In nuclear physics, scientists face this exact problem. They want to know exactly what happened when particles smashed together (the Truth), but their detectors only give them a "smeared" version of the event. To figure out the truth, they usually have to run massive, slow computer simulations to guess how the blur happened and then reverse it. This takes forever and costs a lot of computing power.

The Solution: A New "AI Detective" (JetPrism)

The authors of this paper built a new AI tool called JetPrism. Think of it as a super-smart AI detective that learns to look at the blurry photo and instantly reconstruct the high-definition crime scene without needing to run the slow simulations every time.

They used a specific type of AI called Conditional Flow Matching (CFM).

• The Metaphor: Imagine you have a cup of clear water (simple noise) and a cup of muddy water (complex data). CFM is like a machine that learns the exact path to turn the clear water into the muddy water, and then learns how to run that process in reverse to turn the muddy water back into clear water.

The Trap: The "False Alarm" in the Dashboard

Here is the most important discovery in the paper.

When you train an AI, it usually has a "loss function"—a dashboard gauge that tells you how well it's doing. Usually, when the gauge stops moving and flattens out, you think, "Great! The AI is finished learning."

The authors found that this gauge is lying to you.

• The Analogy: Imagine you are baking a cake. The "loss gauge" is like a timer that stops when the batter is mixed. But just because the batter is mixed doesn't mean the cake is baked! If you take it out of the oven too early, it's still raw in the middle.
• The Reality: In this physics AI, the "loss gauge" stops moving (plateaus) very quickly, making scientists think the AI is done. But the AI is actually still "raw." It hasn't learned the subtle, complex physics rules yet. If you stop training when the gauge says "Done," your AI will produce fake, inaccurate physics data.

The Fix: The "Physics Report Card"

To fix this, the authors created JetPrism, which doesn't just look at the "loss gauge." Instead, it uses a multi-metric report card to check if the AI is actually telling the truth.

They check four specific things:

1. The Shape Check ($\chi^2$): Does the curve of the data look like the real thing?
2. The Distance Check ($W_1$): How far away is the fake data from the real data?
3. The Relationship Check ($D_{corr}$): Do the variables still talk to each other correctly? (e.g., If one particle goes fast, does the other slow down, just like in real life?)
4. The Memory Check ($R_{NN}$): Is the AI just memorizing the training photos and copying them, or is it actually learning the rules?

The Result: They found that even after the "loss gauge" stopped moving, these "Physics Report Card" metrics kept improving for hundreds more training sessions. They had to keep training the AI long after it looked finished to get a truly accurate result.

Why This Matters

1. Speed: JetPrism can generate these complex particle simulations thousands of times faster than traditional methods. It's like switching from hand-drawing a map to using GPS.
2. Reliability: By proving that the standard "loss gauge" is unreliable, they prevent scientists from making mistakes based on "fake" finished models.
3. Versatility: While they tested this on nuclear physics (specifically for the future Electron-Ion Collider), this method works for anything where you need to reverse-engineer a blurry signal.
   • Medical Imaging: Turning a noisy MRI scan into a crystal-clear diagnosis.
   • Astrophysics: Cleaning up blurry telescope images to see distant stars.
   • Finance: Reconstructing market trends from noisy data.

Summary

JetPrism is a new AI framework that acts as a fast, reliable simulator for nuclear physics. Its biggest contribution isn't just the speed, but the diagnostic tool: it teaches us that in complex physics problems, you cannot trust the standard "stop training" signal. You must keep training until the specific physics rules are perfectly satisfied, ensuring the AI isn't just guessing, but truly understanding the universe.

Technical Summary

1. Problem Statement

High-fidelity Monte Carlo (MC) simulations are computationally expensive and essential for analyzing data in nuclear and particle physics, particularly for tasks like detector unfolding (recovering true particle-level observables from smeared detector measurements). While deep generative models, specifically Conditional Flow Matching (CFM), offer a promising path to accelerate these tasks by learning simulation-free mappings, a critical diagnostic gap exists:

• Misleading Convergence: The standard CFM training loss (velocity regression loss) often plateaus prematurely.
• The Disconnect: A converged loss does not guarantee that the model has learned the correct physical kinematics or that the generated distributions match the ground truth.
• Risk: Relying solely on the standard loss can lead to models that appear trained but fail to capture complex physical constraints, correlations, or sharp kinematic boundaries, or conversely, models that continue training unnecessarily long after the loss stabilizes.

2. Methodology

The authors introduce JetPrism, a configurable CFM framework designed to address these convergence issues through rigorous, physics-informed validation.

A. Data and Feature Representation

• Dataset: The primary testbed is the MC-POM dataset (8 million events) modeling exclusive photoproduction ($\gamma p \to \rho^0 p \to \pi^+ \pi^- p$) relevant to the Jefferson Lab (JLab) and the future Electron-Ion Collider (EIC).
• Preprocessing: Events are reduced from 24 dimensions to a 10-dimensional phase space by removing redundant variables (e.g., recoil proton momentum inferred via conservation laws).
• Scaling: A deliberate scaling factor of 5.0 is applied to features to separate source and target supports, amplifying the deterministic velocity signal and improving training stability.
• Unfolding Simulation: Detector effects are simulated by applying Gaussian momentum smearing to pion tracks, breaking energy-momentum conservation. The model must learn to map this "broken" conditional data back to the exact ground-truth manifold.

B. Network Architecture

JetPrism utilizes two variants of a residual MLP backbone with SiLU activations:

1. Unconditional Generation Network: Maps pure Gaussian noise to the target physics distribution.
2. Conditional Unfolding Network: Takes a detector-level measurement ($c$) as a condition (processed via an MLP) to deterministically recover particle-level kinematics.

• Training: Both use the standard CFM objective, regressing a neural network to predict the velocity field transporting noise to data.
• Inference: Uses an adaptive Dormand-Prince ODE solver. Strict tolerances ($10^{-7}$) are used for generation, while relaxed tolerances ($10^{-3}$) are used for unfolding to maximize throughput.

C. Physics-Informed Evaluation Protocol

The core innovation is the replacement of loss-based stopping criteria with a multi-metric evaluation protocol:

1. Marginal Metrics: $\chi^2$ statistic and Wasserstein-1 ($W_1$) distance for 1D feature distributions.
2. Pairwise Joint Metrics: 2D binned $\chi^2$ to verify bivariate dependencies.
3. Global Correlation: Correlation matrix distance ($D_{corr}$) to measure the reproduction of linear dependencies across all channels.
4. Memorization Check: Nearest-neighbor distance ratio ($R_{NN}$). A ratio $\approx 1$ indicates generalization; $R_{NN} \ll 1$ indicates overfitting/memorization.

3. Key Contributions

1. Convergence Diagnostics: The paper demonstrates that standard CFM loss is an unreliable indicator of physical fidelity. It shows that while loss plateaus early (25 epochs), physics metrics ($W_1$, $D_{corr}$) continue to improve significantly until much later (500 epochs).
2. JetPrism Framework: A modular, PyTorch-based tool for generative simulation and detector unfolding, validated on realistic JLab data and synthetic benchmarks.
3. Synthetic Stress Tests: A suite of 1D synthetic benchmarks (gaussian, multi-peak, high-frequency, delta functions) used to isolate topological failure modes (e.g., mode collapse, boundary smearing) before deployment on real data.
4. Multi-Metric Protocol: Establishes a rigorous validation suite ($\chi^2, W_1, D_{corr}, R_{NN}$) that must supersede generic loss metrics to ensure generative reliability.

4. Results

• Convergence Dynamics: On the MC-POM dataset, the CFM loss stabilized after ~25 epochs, but the Number of Function Evaluations (NFE) for the ODE solver continued to decrease until epoch 600, and physical metrics improved until epoch 500. This confirms that extended training is required to find stable, straight probability paths, even when the loss suggests convergence.
• Generation Performance: The model achieved high-fidelity sampling across the 10D phase space with $R_{NN} \approx 1.00$, confirming excellent generalization without memorization. Visual comparisons showed strong agreement with ground truth, with only minor deviations near hard kinematic cutoffs (a known limitation of continuous flows).
• Unfolding Performance: The model successfully mapped smeared detector data back to particle-level truth across various smearing scales ($\sigma_{smear} \in \{0.5, 1.0, 2.0\}$). The metrics remained comparable across scales, indicating robustness.
• Computational Efficiency:
  • Training: Achieved ~880K samples/s on an NVIDIA A100 GPU.
  • Inference: Conditional unfolding was significantly faster (101K events/s on A100) than unconditional generation (3.7K events/s) due to relaxed ODE solver tolerances, offering a ~27x speedup.

5. Significance and Impact

• Reliability in HEP: JetPrism provides a dependable generative surrogate that ensures precise statistical agreement with ground-truth data, addressing a critical gap in the application of AI to High-Energy Physics (HEP).
• EIC Readiness: The framework is directly applicable to the upcoming Electron-Ion Collider (EIC), enabling rapid systematic iterations of detector configurations without the computational cost of repeated GEANT simulations.
• Generalizability: While demonstrated in nuclear physics, the diagnostic framework (multi-metric validation over loss minimization) is extensible to other fields requiring high-fidelity simulation and inverse problem solving, such as medical imaging, astrophysics, semiconductor discovery, and quantitative finance.
• Paradigm Shift: The work advocates for a shift from "loss-driven" training to "physics-driven" validation, ensuring that generative models satisfy strict physical constraints rather than just minimizing probability divergences.