⚛️ high-energy experiments

Mind the Gap: Navigating Inference with Optimal Transport Maps

This paper introduces an optimal transport-based model calibration framework that resolves simulation-data discrepancies in high-dimensional particle physics simulations, enabling the unbiased application of powerful foundation models for tasks like jet tagging at the Large Hadron Collider.

Original authors: Malte Algren, Tobias Golling, Francesco Armando Di Bello, Christopher Pollard

Published 2026-09-09

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Malte Algren, Tobias Golling, Francesco Armando Di Bello, Christopher Pollard

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the vast, high-energy collisions that take place inside particle accelerators, scientists are trying to answer some of the most fundamental questions about the universe. To do this, they rely on a delicate partnership between real-world experiments and computer simulations. When particles smash together, they create a shower of debris that detectors record, but these raw signals are often too complex to interpret directly. Scientists use sophisticated computer programs to simulate what should happen during these collisions, creating a theoretical map of the data. By comparing the real data to these simulations, they can spot new particles or measure known ones with incredible precision. However, this partnership has a flaw: the simulations are never perfect. They are based on our current understanding of physics, but they often miss subtle details of how the detectors actually work or how particles interact. When the simulation does not match the real data, the results of the analysis can become biased, leading scientists to draw the wrong conclusions about the nature of the universe.

A team of researchers has developed a new way to fix this mismatch, specifically for the complex, high-dimensional data generated by modern particle physics experiments. Instead of trying to tweak the simulation after the fact or ignoring the differences, they used a mathematical framework called optimal transport to reshape the simulated data so that it aligns perfectly with the real data. They tested this method on a specific type of particle spray known as a "jet," which is created when quarks or gluons are produced in a collision. The researchers trained a powerful artificial intelligence system to recognize the origin of these jets, creating a rich, 128-dimensional internal representation of the data. They then applied their new calibration technique to this internal representation, effectively translating the imperfect simulation into a form that matches the real-world observations. The result was a calibrated dataset where the simulated jets looked and behaved just like the real ones, allowing for much more accurate and unbiased scientific inference.

The challenge of aligning simulations with reality has long been a bottleneck in particle physics. For decades, scientists have used a method known as "vertical calibration" to correct these discrepancies. This approach involves calculating a weight for each simulated event, essentially telling the computer to count some events more often and others less often to make the simulation match the data. While this works well for simple, low-dimensional data, it breaks down when the data becomes complex and high-dimensional. In these cases, the differences between the simulation and reality are often so large that the required weights become impossibly high, or the simulation simply does not cover the range of possibilities seen in the real data. The researchers found that for the complex jet data they were studying, this traditional method was not just inefficient; it was mathematically impossible to apply without losing all statistical power. The simulated data and the real data were so different that they barely overlapped, making it impossible to simply reweight one to look like the other.

To solve this, the team turned to a different strategy called "horizontal calibration," which relies on the mathematics of moving mass from one distribution to another with the least amount of effort. Imagine you have a pile of sand representing your simulation and a different pile representing your real data. Instead of trying to change the height of individual grains of sand to match the second pile, you physically move the grains from the first pile to new positions until the shape of the pile matches the second one. This is the essence of the optimal transport map the researchers developed. They built a neural network that learned the most efficient way to transform the simulated jet data into the real jet data. This transformation was applied directly to the 128-dimensional internal "latent" representation of the jets, a space where the artificial intelligence had already summarized all the complex features of the particle spray. By calibrating this internal space, they ensured that the fundamental structure of the data was preserved while correcting the mismodeling.

The researchers tested their method using a dataset inspired by the Compact Muon Solenoid experiment at the Large Hadron Collider. They created two sets of data: a source set that represented the standard simulation, and a target set that mimicked real experimental data by introducing known errors and variations, such as changes in how the detector measures particle positions. They then trained their artificial intelligence classifier to distinguish between the different types of jets, such as those coming from Higgs bosons versus those coming from ordinary quarks. Before calibration, the classifier could easily tell the difference between the simulated data and the target data, indicating a significant mismatch. After applying the optimal transport map, the classifier could no longer distinguish between the two; the simulated data had been successfully transformed to look like the real data. The researchers verified this by checking various physical quantities derived from the data, finding that the calibrated simulation now matched the target distribution across a wide range of measurements.

One of the most significant findings of this work is that the calibration performed on the internal 128-dimensional representation worked for all the downstream tasks that relied on it. Even though the calibration was applied deep inside the neural network, the improvements carried through to the final output, such as the probability scores used to identify specific particles. The researchers showed that the calibrated data could be used to construct statistical tests that were free from the bias introduced by the simulation errors. This is a crucial step toward the use of "foundation models" in particle physics. These are large, pre-trained artificial intelligence models that can be adapted for many different tasks. If these models are not properly calibrated, their predictions will be biased, limiting their usefulness. By demonstrating that optimal transport can calibrate these high-dimensional internal representations, the researchers have provided a path forward for using these powerful tools in a way that is both accurate and unbiased.

The implications of this work extend beyond just identifying jets. The method offers a general framework for correcting high-dimensional simulations across the sciences, wherever complex models are used to predict real-world phenomena. The researchers noted that while their specific application was in particle physics, the underlying mathematics is universal. They emphasized that this approach does not require the simulation to be perfect; it only requires that there is a way to map the imperfect simulation to the real data. By moving the focus from correcting the final output of a model to correcting its internal representations, scientists can now handle much more complex data than was previously possible. This opens the door to more precise measurements of fundamental particles and potentially the discovery of new physics that was previously hidden by the limitations of simulation.

The study concludes that this calibration strategy is a viable path for integrating advanced machine learning into particle physics experiments. It suggests that the field can move away from relying on numerous, ad-hoc corrections for every single analysis and instead adopt a unified approach to calibrating the internal representations of their models. The researchers are careful to note that while their simulations showed excellent results, further work is needed to ensure that the calibration does not introduce new, subtle discrepancies. They also suggest that future research could explore how to fine-tune models using these calibrated representations to improve performance even further. For now, the work stands as a proof of concept that the gap between simulation and reality can be bridged, not by ignoring the differences, but by mathematically transforming the simulation to fit the world as it truly is.

Technical Summary: Mind the Gap: Navigating Inference with Optimal Transport Maps

Problem Statement

Machine learning (ML) has revolutionized sensitivity to new phenomena in particle physics, particularly through "simulation-based inference" (SBI). However, these methods rely heavily on high-quality training samples generated by simulations. A critical limitation arises from model misspecification: discrepancies between the simulated data ( $p_{sim}$ ) and real experimental data ( $p_{data}$ ). When the statistical model used for inference does not cover the true data generation process, constraints on model parameters become non-robust.

While previous work has addressed this via "vertical calibration" (reweighting samples based on density ratios), this approach faces severe limitations in high-dimensional spaces. Specifically, density-ratio corrections require significant overlap between the simulated and data distributions. In high-dimensional scenarios, even slight mismodeling can lead to a "dilution factor" where the effective number of samples becomes negligible, or the solution may not exist at all if the support of the simulation does not fully cover the data. This is particularly problematic for foundation models in particle physics, which generate high-dimensional latent representations (e.g., 128-dimensional vectors) that are difficult to calibrate using traditional low-dimensional binning or reweighting techniques.

Methodology

The authors propose a horizontal calibration strategy based on Optimal Transport (OT) theory. Instead of reweighting samples, this approach learns a transport map that transforms the simulated distribution to match the data distribution while minimally altering the underlying structure.

Hierarchical Model Framework: The problem is framed within hierarchical models where an observable $x$ depends on parameters $\theta$ through a latent variable $\omega$ . The goal is to correct the conditional density $p_{sim}(x|\omega)$ to match $p_{data}(x|\omega)$ .
Optimal Transport Map: The authors seek a transport map $T_\omega$ such that the pushforward of the simulated distribution equals the data distribution: $(T_\omega)_\# p_{sim}(x|\omega) = p_{data}(x|\omega)$ . The optimal map is defined as the one that minimizes the Euclidean transport cost, ensuring the transformation is the "least invasive" possible.
Neural Implementation: The map is learned using Input-Convex Neural Networks (ICNNs). The optimization minimizes the squared Wasserstein distance ( $W_2^2$ ) between the simulated and target distributions. The network architecture involves two ICNNs, $f$ and $g$ , where the optimal map is derived as the gradient of the convex potential: $\hat{T}_\omega(z|\omega) = \nabla_z g(z|\omega)$ .
Application to Jet Tagging:
- Dataset: The study utilizes a modified version of the JetClass benchmark, simulating QCD multijet events. A "source" dataset (standard simulation) is compared against a "target" dataset where detector smearing and QCD scales are artificially altered to mimic real-world discrepancies.
- Latent Space Calibration: A transformer-based classifier is trained to identify jet origins (10 classes). The calibration is applied directly to the 128-dimensional latent representation ( $z_{128}$ ) of the classifier, conditioned on the jet flavor label ( $q/g$ ).
- Downstream Evaluation: The calibrated latent vectors are passed through the remaining layers of the classifier (MLP Head) to evaluate the impact on the final 10-dimensional output space and physics-motivated discriminants.

Key Contributions

First High-Dimensional OT Calibration: This work presents the first application of OT-based calibration to high-dimensional latent spaces (up to 128 dimensions) in particle physics. Previous applications were limited to low-dimensional observables or classifier scores.
Demonstration of Horizontal vs. Vertical Superiority: The paper provides empirical evidence that vertical calibration (reweighting) fails in high-dimensional latent spaces due to sparse distributional overlap, leading to infinite or extremely large dilution factors. In contrast, the horizontal OT approach successfully aligns distributions even when they are nearly disjoint.
Foundation Model Calibration: The study demonstrates that calibrating the internal "latent" representation of a general-purpose foundation model allows for unbiased inference across a wide variety of downstream tasks without requiring individual calibration for each specific analysis.

Results

Latent Space Alignment: Principal Component Analysis (PCA) of the 128-dimensional latent space shows that the OT map successfully aligns the source (simulated) distribution with the target (synthetic data) distribution, significantly reducing discrepancies.
Discriminator Performance:
- Before calibration, a discriminator could easily distinguish between source and target distributions in the latent space (AUC $\approx$ 0.98).
- After calibration, the discriminator's ability to distinguish the calibrated source from the target drops to random chance (AUC $\approx$ 0.5) in the deeper layers ( $z'_{128}$ and $z'_{10}$ ), indicating successful closure.
- A minor residual discrepancy remains in the initial $z_{128}$ layer (AUC $\approx$ 0.52), but the authors note this does not propagate to the final output layers, rendering it inconsequential for the classification task.
Physics Observables: The calibration significantly improves the agreement between simulated and target distributions for physics-motivated discriminants (e.g., log-likelihood ratios for Higgs boson decays vs. background), validating the utility of the calibrated latent space for downstream inference.

Significance and Claims

The paper claims that this work is a key step toward enabling the unbiased use of "foundation models" in particle physics. By calibrating the high-dimensional internal representations rather than just the final outputs, the method:

Mitigates Misspecification: It addresses the "misspecification gap" in a way that is robust to the high dimensionality of modern ML features, where traditional reweighting fails.
Generalizes Across Tasks: A single calibration of a rich, informative latent representation benefits hundreds of downstream analyses (e.g., Higgs, Top quark, and new physics searches) without the need for individual, ad-hoc calibration procedures for each specific measurement.
Broad Applicability: While demonstrated on jet tagging, the authors assert the framework is broadly applicable to correcting high-dimensional simulations across the sciences, offering a path to more reliable, simulation-based inference.

The authors remain modest, noting that while the calibration is effective, residual discrepancies in the latent space warrant further investigation, and the method does not yet introduce new discrepancies in the final output space. The work lays the groundwork for properly calibrated foundation models in scenarios with suitable calibration data.