Improving Generative Model-based Unfolding with Schrödinger Bridges

This paper introduces SBUnfold, a novel generative unfolding method leveraging Schrödinger Bridges and diffusion models to combine the strengths of discriminative and generative approaches, enabling efficient, high-dimensional cross-section measurements without relying on known probability densities.

The Gist

Imagine you are a detective trying to solve a crime, but the only evidence you have is a blurry, distorted photo taken by a faulty security camera. Your goal is to reconstruct the original, crystal-clear image of what actually happened. In the world of particle physics, this is called "unfolding."

Scientists smash particles together to see what's inside them, but the detectors they use are like that faulty camera. They add "noise," blur the edges, and sometimes miss details entirely. To understand the true laws of nature, physicists need to reverse-engineer these messy detector readings back into the clean, original particle events.

This paper introduces a new, super-smart detective tool called SBUnfold (Schrödinger Bridge Unfolding). Here is how it works, explained without the heavy math.

The Problem: Two Old Ways of Solving the Puzzle

Before this new method, scientists had two main ways to fix these blurry photos, and both had flaws:

1. The "Correction" Team (Discriminative Models):

   • How it works: They start with a perfect simulation of what should happen and ask a computer, "How much do we need to tweak this to look like the blurry photo?"
   • The Flaw: It's great if the simulation is already very close to reality. But if the real world is weird and the simulation is way off, this team gets confused. Also, if there are very few blurry photos (rare events), they run out of clues to work with.

2. The "Reconstruction" Team (Generative Models):

   • How it works: They start with a blank canvas (pure random noise) and try to paint a picture that looks exactly like the blurry photo.
   • The Flaw: It's very flexible and can handle weird, rare events well. But starting from "random noise" is a huge leap. It's like trying to paint a portrait of a specific person starting from a pile of random paint splatters. It's hard to get the small details right, and the process can be unstable.

The New Solution: The "Schrödinger Bridge"

The authors of this paper realized: Why not combine the best of both worlds?

They created SBUnfold, which uses a concept called a Schrödinger Bridge.

The Analogy: The River Crossing
Imagine you have a river. On one bank, you have a messy, distorted crowd of people (the Detector Data). On the other bank, you have the clean, original crowd (the True Physics).

• Old Method 1 tried to push the messy crowd slightly to make them look like the clean crowd.
• Old Method 2 tried to summon the clean crowd out of thin air (random noise) and hoped they would match the messy crowd.
• SBUnfold builds a bridge directly between the two banks.

This bridge is special because it doesn't need to know the "map" (the exact probability) of the clean crowd beforehand. It just learns how to walk from the messy side to the clean side by looking at pairs of "before and after" photos from a simulation.

How SBUnfold Wins

1. It Starts Where It Matters: Instead of starting from random noise (like the Reconstruction Team), SBUnfold starts with the actual messy data. It treats the detector reading as the "starting point" of a journey.
2. It Takes Small Steps: Using a technique called a Diffusion Model, it slowly "denoises" the data. Imagine taking a blurry photo and slowly sharpening it, pixel by pixel, until the image is clear.
3. It Handles Rare Events: Because it starts with the real data, it doesn't get lost when dealing with rare, weird events that the simulation might have missed. It learns the "small corrections" needed to fix the specific event in front of it.

The Results: A Better Detective

The team tested SBUnfold on a synthetic dataset involving Z-bosons (a type of particle) and jets (sprays of particles).

• Accuracy: When they compared the "unblurred" results to the true reality, SBUnfold was more accurate than the previous best methods.
• Stability: When they reduced the amount of data available (simulating a scenario where the experiment didn't catch many events), SBUnfold stayed calm and accurate. The old "Correction" team struggled and became very unstable with less data.
• Sharp Details: SBUnfold was particularly good at preserving sharp, distinct features in the data, which is often where other methods fail.

The Bottom Line

Think of SBUnfold as the ultimate photo-editing AI. It doesn't just guess what the photo should look like from scratch, nor does it just try to nudge the existing photo slightly. Instead, it learns the exact path of transformation between a "bad photo" and a "good photo," allowing it to restore the image with incredible precision, even when the original photo is very damaged or rare.

This is a big step forward for particle physics, allowing scientists to see the universe more clearly, even when their detectors are imperfect.

Technical Summary

1. Problem Statement

In particle physics, nuclear physics, and astrophysics, unfolding (or deconvolution) is the statistical process of correcting detector effects to infer the true "particle-level" distribution ($Z$) from observed "detector-level" data ($X$).

• Classical Limitations: Traditional methods rely on binned histograms and regularized matrix inversion (e.g., Lucy-Richardson), which struggle with high-dimensional, unbinned data.
• Machine Learning Approaches: Two main ML-based approaches have emerged:
  1. Discriminative Models (e.g., OmniFold): Uses classifiers to iteratively reweight simulated events.
     • Strength: Learns small corrections if the simulation is close to reality.
     • Weakness: Performance degrades significantly when real data samples are scarce because it trains directly on data.
  2. Generative Models (e.g., cINN): Uses models like Normalizing Flows to map a known prior (e.g., Gaussian) to the data distribution.
     • Strength: Can handle sparse data because the "Expectation" (E) step trains on simulation, not real data.
     • Weakness: Requires mapping a simple prior to complex data, which can be difficult if the distributions differ significantly. It struggles to learn "small corrections" efficiently.

The core challenge is to combine the data efficiency of generative models with the correction capability of discriminative models, specifically avoiding the need for a tractable prior distribution while maintaining stability with low data statistics.

2. Methodology: SBUnfold

The authors propose SBUnfold, a new unfolding framework that integrates Schrödinger Bridges (SB) with Diffusion Models.

Core Concept: Schrödinger Bridges

Unlike standard diffusion models or normalizing flows that map a fixed prior (like a Gaussian) to data, a Schrödinger Bridge learns to transport mass between two arbitrary distributions ($p_A$ and $p_B$) without requiring knowledge of the probability density of one of them.

• In the context of unfolding:
  • $p_A$: The detector-level (reconstructed) distribution.
  • $p_B$: The particle-level (truth) distribution.
• The SB learns a stochastic path connecting these two states, effectively learning a "small correction" from detector-level to particle-level, similar to discriminative models, but using the generative framework of diffusion models.

Technical Implementation

1. Forward and Backward SDEs: The method utilizes Stochastic Differential Equations (SDEs) to define a diffusion process.
   • The forward process perturbs data toward a noise distribution.
   • The reverse process (the generative model) denoises the data.
2. I2SB (Interpolating Schrödinger Bridge): The authors adopt the tractable solution proposed in Ref. [31] (I2SB). This assumes pairs of observations $(x_a, x_b)$ exist in the simulation (particle-level and detector-level pairs).
   • The posterior distribution $q(x|xa, xb)$ is derived analytically as a Gaussian.
   • The model is trained to predict the noise term required to reverse the diffusion process between the paired datasets.
3. Workflow:
   • Training: The model is trained exclusively on simulation pairs (Pythia particle-level + Delphes detector-level). It does not see real data during training.
   • Inference: Given a real detector-level event, the model uses the learned SB to generate the corresponding particle-level event by reversing the diffusion process starting from the observed event (acting as the prior).
4. Architecture: A fully connected neural network with ResNet blocks and skip connections is used to approximate the score function (noise prediction).

3. Key Contributions

• Novel Unfolding Algorithm: Introduction of SBUnfold, which uniquely combines the strengths of discriminative (small correction learning) and generative (simulation-based training) approaches.
• Prior Independence: Unlike cINN (which requires a Gaussian prior) or standard diffusion models, SBUnfold uses the reconstructed event itself as the starting point for the generation process. This allows it to learn local corrections rather than global distribution shifts.
• Robustness to Low Statistics: By training on simulation pairs and using the reconstructed event as a conditional prior, SBUnfold maintains performance even when the number of available real data events is extremely low (e.g., 1,000 events), a regime where OmniFold fails.
• High-Dimensional Unbinned Capability: The method operates on unbinned, high-dimensional data, avoiding the information loss inherent in histogram-based methods.

4. Results

The authors evaluated SBUnfold on a synthetic Z+jets dataset (1.6 million events) using six physics observables (Jet Mass, Jet Width, Constituent Multiplicity, $N$-subjettiness, $\ln \rho$, and $z_g$).

• Comparison with cINN (Generative):
  • SBUnfold achieved lower Earth Mover's Distance (EMD) and Triangular Discriminator scores across most distributions.
  • It showed superior performance in reproducing sharp features (e.g., the cutoff at $z_g=0$) because it leverages the reconstructed prior, whereas cINN must generate these features from a smooth Gaussian prior.
• Comparison with OmniFold (Discriminative):
  • In the standard regime (600k data points), SBUnfold performed comparably or better.
  • Crucial Test (Low Data): When the number of data events was reduced to 1,000, OmniFold's performance degraded significantly (EMD values disagreed by >2 standard deviations). SBUnfold and cINN remained stable because they rely on simulation for training. SBUnfold specifically showed the most robust performance in this low-statistics regime.
• Migration Matrix: Analysis of the migration matrix showed that SBUnfold learns to apply small, diagonal-dominant corrections, confirming it acts as a precise correction mechanism rather than a global reshaping.

5. Significance and Outlook

• Scientific Impact: SBUnfold offers a powerful tool for future high-luminosity collider experiments (like the HL-LHC) where specific rare processes may have very few observed events. It enables precise, high-dimensional differential cross-section measurements in data-scarce regions.
• Methodological Advancement: It bridges the gap between generative and discriminative unfolding, demonstrating that diffusion models can be adapted to learn conditional transport between arbitrary distributions without tractable priors.
• Future Work:
  • Extending the method to include the Maximization (M) step (iterative refinement) to further reduce systematic biases.
  • Adapting the architecture to handle permutation-invariant data (e.g., sets of particles) rather than fixed-order feature vectors.
  • Application to image-like datasets and full phase-space unfolding.

In summary, SBUnfold represents a state-of-the-art advancement in particle physics unfolding, providing a robust, high-fidelity solution that excels particularly in scenarios with limited experimental data.