Generative models on phase space

This paper introduces generative models, including diffusion and flow matching variants, that are explicitly constructed to remain confined to the massless N-particle Lorentz-invariant phase space manifold at every sampling step, thereby ensuring exact energy-momentum conservation and providing a clear framework for studying particle correlations in high-energy physics.

The Gist

The Big Picture: Teaching AI to Play by Physics Rules

Imagine you are teaching a robot to paint pictures of a busy city street.

• The Problem: If you just show the robot thousands of photos, it might learn to paint a car floating in the sky or a building with no foundation. It learns the look of the data, but it doesn't understand the rules of reality (like gravity or laws of physics).
• The Goal: In particle physics, scientists use computers to simulate collisions of subatomic particles. These simulations are incredibly complex. They want to use AI (specifically "Generative Models") to speed this up. But if the AI makes a mistake where energy disappears or momentum isn't conserved, the simulation is useless.
• The Solution: This paper introduces a new way to train AI so that it is physically impossible for it to break the rules of physics. It doesn't just "guess" the right answer; it is forced to stay on the "highway" of valid physics at every single step of its thinking process.

---

The Core Concept: The "q-Space" Shortcut

To understand their trick, let's use an analogy of a molded clay sculpture.

1. The Problem with Standard AI (The "p-Space" Approach)

Imagine you have a lump of clay (the data) that must always be shaped like a perfect sphere (the laws of physics: energy and momentum conservation).

• Standard AI tries to learn this by starting with a ball of clay and adding random noise to it, then trying to smooth it back out.
• The Flaw: When the AI adds noise, the clay might get squished into a cube or a star shape. When it tries to smooth it back, it might accidentally leave a dent or a wobble. The final result looks like a sphere, but it's not a perfect sphere. In physics, even a tiny wobble means energy is lost, which breaks the simulation.

2. The Paper's Solution: The "q-Space" Map

The authors use a clever mathematical trick called RAMBO (Random Momenta Everywhere) to create a "magic map."

• The Analogy: Imagine you have a flat, infinite sheet of paper (this is q-space). On this paper, you can draw any shape you want without worrying about rules.
• The Transformation: There is a magical machine (the RAMBO algorithm) that takes any shape you draw on the paper and instantly squishes, stretches, and folds it into a perfect sphere (this is p-space, or real physical phase space).
• The Magic: If you draw a straight line on the paper, the machine turns it into a perfect curve on the sphere. If you draw a circle on the paper, it becomes a perfect sphere.
• The Strategy: Instead of teaching the AI to smooth out the clay directly (where it might make mistakes), the authors teach the AI to draw on the flat paper (q-space). Because the "magic machine" (RAMBO) is perfect, whatever the AI draws on the paper gets turned into a perfectly valid physics event when it comes out the other side.

In short: They moved the AI's "brain" to a place where it doesn't have to worry about the rules, because the rules are automatically applied by the machine that converts the AI's output into reality.

---

How the AI Learns: The "De-Noising" Process

The paper uses two types of AI: Diffusion Models and Flow Matching. Here is how they work in this new system:

The Diffusion Model (The "Reverse Snowstorm")

Imagine a room filled with snowflakes (pure randomness).

1. Forward Process: You take a beautiful, complex sculpture (a particle collision) and slowly throw snow at it until it's completely buried and unrecognizable.
2. The Goal: Teach an AI to look at a snow-covered blob and figure out how to remove the snow to reveal the sculpture underneath.
3. The Innovation: In this paper, the "snow" isn't just random noise. It's a specific type of noise that, when the AI removes it, leaves behind a perfectly uniform distribution of particles (like a perfectly smooth, featureless sphere).
4. The Result: As the AI "de-noises" the data, it builds up correlations between particles. Because the whole process happens in the "magic map" (q-space), the final sculpture is guaranteed to be a perfect sphere (conserving energy and momentum).

The Flow Matching Model (The "River Current")

Imagine a river flowing from a calm lake (simple noise) to a raging waterfall (complex data).

1. The Goal: Teach the AI to predict the direction of the current at every point so you can swim upstream from the waterfall back to the lake.
2. The Innovation: Just like with diffusion, they set up the river so that the "lake" is a perfect, uniform distribution of particles.
3. The Result: The AI learns the "current" (the rules of physics) and can swim upstream to generate new, valid particle collisions.

---

Why This Matters: The "Traffic Jam" of Particles

In high-energy physics (like at the Large Hadron Collider), particles often fly off in specific patterns. Sometimes, they fly in a "jet" where many particles are packed tightly together.

• The Singularity Problem: Sometimes, the math says a particle should have zero energy or be perfectly aligned with another. This creates a "singularity" (a mathematical infinity). Standard AI gets confused here and produces garbage data.
• The Paper's Success: The authors tested their method on these tricky, "singular" situations.
  • They found that even though the AI didn't perfectly learn the messy, zero-energy tail (which is physically impossible to measure anyway), it perfectly learned the important parts of the distribution.
  • It learned the "shape" of the traffic jam correctly, ensuring that the total energy and momentum of the jam were conserved exactly.

The Takeaway: "Physics for AI, AI for Physics"

The paper concludes with a beautiful two-way street:

1. AI for Physics: By forcing the AI to stay on the "physics highway" (using q-space), scientists can trust the simulations. They don't have to waste time checking if the AI broke the laws of thermodynamics.
2. Physics for AI: By using these rigid, perfect physics rules as a testbed, scientists can better understand how AI learns. They can see exactly what the AI is learning (the correlations between particles) versus what is just random noise.

Summary Metaphor:
Think of the AI as a student taking a test.

• Old Way: The student is given a blank sheet of paper and told "Draw a car." They might draw a car with wheels on the roof. It looks like a car, but it won't drive.
• New Way (This Paper): The student is given a car-shaped stencil (the q-space map). They can draw whatever they want inside the stencil. When they lift the stencil, the result is guaranteed to be a car with wheels on the bottom. The student is free to be creative with the details (the paint job, the speed), but the fundamental structure is always correct.

This allows scientists to generate millions of realistic particle collision events instantly, with the absolute guarantee that the laws of physics are never violated.

Technical Summary

1. Problem Statement

Deep generative models (such as Diffusion Models and Flow Matching) have shown success in high-energy physics (HEP) for simulating particle collisions. However, a critical limitation exists: standard models trained on 3-momentum vectors ($p$-space) do not exactly satisfy the physical constraints of energy and momentum conservation.

• The Manifold Constraint: HEP data (collections of $N$ relativistic 4-vectors) lives on a specific submanifold of the ambient space defined by Lorentz invariance and conservation laws (specifically, $\sum p_i^\mu = p_{beam}^\mu$).
• The Failure Mode: Current models learn these constraints only approximately. This leads to generated events that violate conservation laws, reducing their interpretability and reliability for scientific discovery.
• The Challenge: Enforcing these constraints exactly is difficult because the phase space manifold has a complex geometry (a $(3N-4)$-dimensional submanifold in $R^{3N}$), and standard diffusion/flow processes assume an ambient Euclidean space.

2. Methodology: $q$-Space Generative Modeling

The authors propose a framework called $q$-space generative modeling. Instead of performing the generative process directly on physical momenta ($p$-space), they map the data to an auxiliary, unconstrained space ($q$-space) where the generative process occurs, and then map the results back to physical space.

Core Mechanism: The RAMBO Algorithm

The method leverages the RAMBO (Random Momenta Better) algorithm, a 40-year-old technique for sampling massless $N$-particle phase space.

• The Mapping: RAMBO maps an unconstrained set of $N$ 3-vectors $Q = \{q_I\}$ (drawn from a specific distribution) to a physical phase space point $P = \{p_I\}$ that strictly satisfies energy-momentum conservation.
• The Transformation: The map involves a Lorentz boost to the center-of-momentum (CM) frame and a homogeneous rescaling to set the total energy to unity.
• The Advantage: In $q$-space, there are no constraints. The generative model learns a distribution in $R^{3N}$, and the RAMBO map acts as a nonlinear projection that enforces exact conservation laws upon mapping back to $p$-space.

Implementation Details

The authors adapt two types of generative models to this framework:

1. Diffusion in $q$-space:

   • Forward Process: Instead of adding Gaussian noise to $p$-vectors, the model adds noise to $q$-vectors. The "pure noise" endpoint is not a standard Gaussian but the uniform distribution on phase space (mapped from a specific distribution $pref(Q)$ in $q$-space).
   • Score Function: The model learns the score function $\nabla \log p_t(Q)$ in $q$-space.
   • Reverse Process: The model denoises $Q$ and maps the trajectory back to $P$ at every step, ensuring the generated sample is always on the physical manifold.
   • Prior: The prior is the uniform distribution on phase space, allowing for a clear interpretation of correlations learned during the reverse process.

2. Flow Matching in $q$-space:

   • The model learns a velocity field $v_\theta(Q, t)$ that transports a prior distribution to the data distribution in $q$-space.
   • While the authors found that using the uniform phase space prior was difficult to train, they successfully used a Gaussian prior in $q$-space.
   • Crucially, the mapping back to $p$-space ensures exact conservation.

Data Augmentation

To handle the non-uniqueness of the RAMBO map (multiple $q$-vectors can map to the same $p$-vector), the authors employ data augmentation. They randomly boost and rescale training data points into $q$-space using different parameters $(b, x)$. This "fills out" the $q$-space distribution, helping the model learn permutation-invariant features relevant to jets.

3. Key Contributions

• Exact Constraint Satisfaction: The primary contribution is a generative framework that guarantees exact energy-momentum conservation by construction, eliminating the need for post-hoc correction or soft penalties.
• Interpretability via Uniform Prior: By setting the forward process endpoint to the uniform distribution on phase space, the reverse process explicitly reveals how physical correlations (e.g., soft/collinear singularities) emerge from a structureless state.
• Permutation Equivariance: The framework naturally handles unordered sets of particles (essential for QCD jets) by using permutation-equivariant architectures (Point-Edge Transformers) in $q$-space.
• Bridging Physics and AI: The paper demonstrates how physical symmetries can be embedded in the generative process (the trajectory) rather than just the network architecture, offering a new paradigm for "Physics for AI."

4. Results

The authors validated their method on two scales:

A. Low-Dimensional (3-Particle) Examples

• Smooth Distribution (Muon Decay): The model successfully learned the smooth matrix element for $\mu^- \to e^- \nu_\mu \bar{\nu}_e$. It achieved excellent agreement with ground truth in energy distributions and Dalitz plots (Wasserstein-1 distance $\approx 0.0016$).
• Nearly-Singular Distribution ($e^+e^- \to q\bar{q}g$): The model learned a distribution with soft and collinear singularities.
  • While the model struggled to perfectly reproduce the unphysical low-energy tail (which depends on the infrared cutoff), it perfectly learned the physically relevant $\tau$ distribution (an IRC-safe observable) away from the singularity.
  • The model correctly identified the kinematic endpoint $\tau = 1/6$, determined purely by phase space geometry.

B. High-Dimensional Example (10-Particle)

• APS Matrix Element: The authors tested on a 10-particle distribution exhibiting antenna pole structures (simulating QCD parton showers).
• Performance:
  • Diffusion Models: Learned the $\tau$ distribution correctly over a dynamic range of 3 orders of magnitude, matching the analytic leading-log prediction.
  • Flow Matching: Outperformed diffusion models, learning the distribution over 9 orders of magnitude in $\tau$.
  • Conservation: Unlike $p$-space models which exhibited energy/momentum violations comparable to the particle energies, the $q$-space models satisfied conservation to machine precision.

Comparison with $p$-Space Models

Table I in the paper shows that $p$-space models (diffusion and flow matching) fail to conserve energy and momentum, with violations often exceeding 10% of the median particle energy. In contrast, $q$-space models have zero violation.

5. Significance and Outlook

• Reliability for HEP: This approach provides a robust tool for generating simulated jet data where physical laws are strictly obeyed, crucial for precision measurements and new physics searches.
• Diagnostic Tool for AI: By constraining the generative process to a known physical manifold, researchers can use the reverse trajectory to diagnose what latent structures (hierarchical, compositional) the neural network is actually learning, separating physical correlations from artifacts.
• Scalability: The method is a "drop-in" replacement for existing generative models operating on 3-momenta. It scales to high particle multiplicities ($N \sim 200$) without requiring complex architectural changes to enforce symmetries.
• Future Work: The authors suggest extending this to massive particles and applying the framework to real experimental jet data to study the hierarchical structure of QCD showers.

In summary, the paper introduces a mathematically rigorous and practically effective method to embed the fundamental symmetries of high-energy physics directly into the generative process, solving the long-standing issue of conservation law violations in AI-driven particle physics simulations.