Generative models on phase space

Zachary Bogorad; Ibrahim Elsharkawy; Yonatan Kahn; Andrew J. Larkoski; Noam Levi

hep-phApr 2, 2026

Generative models on phase space

Authors: Zachary Bogorad, Ibrahim Elsharkawy, Yonatan Kahn, Andrew J. Larkoski, Noam Levi

Organizations: Theory Division, Fermi National Accelerator Laboratory, Batavia, IL, USA · Department of Physics, University of Toronto and Vector Institute, Toronto, ON, Canada and NERSC, Lawrence Berkeley National Laboratory, Berkeley, California, USA · Department of Physics, University of Toronto and Vector Institute, Toronto, ON, Canada · American Physical Society, Hauppauge, New York, USA · Tel Aviv University, Tel Aviv, Israel

Abstract

Deep generative models such as diffusion and flow matching are powerful machine learning tools capable of learning and sampling from high-dimensional distributions. They are particularly useful when the training data appears to be concentrated on a submanifold of the data embedding space. For high-energy physics data, consisting of collections of relativistic energy-momentum 4-vectors, this submanifold can enforce extremely strong physically-motivated priors, such as energy and momentum conservation. If these constraints are learned only approximately, rather than exactly, this can inhibit the interpretability and reliability of such generative models. To remedy this deficiency, we introduce generative models which are, by construction, confined at every step of their sampling trajectory to the manifold of massless N-particle Lorentz-invariant phase space in the center-of-momentum frame. In the case of diffusion models, the "pure noise" forward process endpoint corresponds to the uniform distribution on phase space, which provides a clear starting point from which to identify how correlations among the particles emerge during the reverse (de-noising) process. We demonstrate that our models are able to learn both few-particle and many-particle distributions with various singularity structures, paving the way for future interpretability studies using generative models trained on simulated jet data.

Figures & tables

Figure 1: ( Left ) Dalitz plot of 100,000 samples from the muon decay distribution ( 15 ). ( Center ) Dalitz plot of 100,000 samples from the reverse process of the diffusion model trained on the data in the left panel. ( Right ) Uniform phase space for comparison.

Figure 2: Distributions of the logarithm of the theoretical Dalitz plot PDF for the true muon decay distribution and for our generated distribution.

Figure 3: Energy distributions for the muon decay matrix element.

Figure 4: Distributions of the Rosenblatt transformation parameters $u_{1}=16E_{1}^{3}(1-E_{1})$ and $u_{2}=v^{2}(3-4v)/E_{1}^{2}/(3-4E_{1})$ with $v=E_{1}+E_{2}-1/2$ for the true muon decay distribution and for our generated distribution. In the true distribution, $u_{1}$ and $u_{2}$ should be independently uniformly distributed on $[0,1]$ .

Figure 5: Angular distributions for the muon decay matrix element.

Figure 6: Dalitz plots of 500,000 samples from the $e^{+}e^{-}\to q\bar{q}g$ distribution as generated by MadGraph , with varying cutoff $m^{2}_{\rm cut}$ on the two-particle invariant mass.

Figure 7: As in Fig. 1 , but for the $q\bar{q}g$ distribution with $m^{2}_{\rm cut}=0.01$ .

Figure 8: Single-particle energy distributions comparing diffusion model samples to the ground truth distribution of

e^{+}e^{-}\to q\bar{q}g

events for various values of

m^{2}_{\rm cut}

Figure 9: Single-particle angular distributions comparing diffusion model samples to the ground truth distribution of

e^{+}e^{-}\to q\bar{q}g

events for

m^{2}_{\rm cut}=10^{-3}

Figure 10: ( Left ) Distribution of

\log_{10}\tau

for 500,000

q\bar{q}g

events generated from MadGraph , compared to the analytic prediction in Eq. ( 19 ). ( Right ) Distribution of

\tau

for 100,000 samples from diffusion models trained on the

q\bar{q}g

events from the left panel. For comparison, we also show the distribution for uniform sampling of phase space, along with its analytic prediction in Eq. ( 20 ).

Figure 11: ( Left ) Energy distributions for the APS events sampled from Eq. ( 21 ), for various values of the cutoff

\xi_{m}

. ( Right ) Comparison of

q

-space and

p

-space distributions for

\xi_{m}=40

Figure 12: Distributions of

\tau={\rm min}\{p_{I}\cdot p_{J}\}

for ground-truth APS events (solid) and samples from trained diffusion models (dashed). For comparison, uniform events on 10-particle phase space and the analytic distributions Eqs. ( 23 )–( 24 ) are also shown. A zoom-in of the large-

\tau

region is shown at right.

Figure 13: Comparison of flow matching models in

p

-space and

q

-space with a

q

-space diffusion model, for the

N=10

APS distribution with

\xi_{m}=30

	$\|\sum_{I}E_{I}-E_{\rm CM}\|/\tilde{E}$	$\|\sum_{I}p_{I,x}\|/\tilde{E}$	$\|\sum_{I}p_{I,y}\|/\tilde{E}$	$\|\sum_{I}p_{I,z}\|/\tilde{E}$
$p$ -space diffusion	$0.2585$	$0.1264$	$0.1220$	$0.0820$
$p$ -space flow matching	$0.0094$	$0.0023$	$0.0024$	$0.0023$

Table 1: Mean conservation violations for generated events (

N=10

particles,

\xi_{m}=30

), expressed as fractions of the median single-particle energy

\tilde{E}

in each event. Generative models in

q

-space satisfy conservation to machine precision (as does the ground-truth data), while generative models in

p

-space exhibit energy and momentum conservation violations that are comparable to the median particle energy.

Figure 14: Per-event energy and momentum violation normalized by the median single-particle energy

\tilde{E}

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 15: Energy (left) and angular (center and right) distributions of particle 3 for varying data augmentation strategies.

Figure 16: Training data in

q

-space for the different data augmentation strategies in Fig. 15 .

Figure 17: Distributions of the logarithm of the theoretical Dalitz plot PDF for the true muon decay distribution and for generated distributions with

N_{\rm mult}=

1 or 2. The model with

N_{\rm mult}=2

is a worse match to the true distribution even for the high-probability events.

Figure 18: Rosenblatt transformations as in Fig. 4 , but for diffusion models trained with $N_{\rm mult}=$ 1 or 2. The non-uniformity of both $u_{1}$ and $u_{2}$ is evident for $N_{\rm mult}=2$ .

Observable	$N_{\rm mult}=1$	$N_{\rm mult}=2$
$E_{1}$	$\mathbf{0.73\times 10^{-3}}$	$2.4\times 10^{-3}$
$E_{2}$	$\mathbf{0.35\times 10^{-3}}$	$1.2\times 10^{-3}$
$E_{3}$	$\mathbf{1.6\times 10^{-3}}$	$5.7\times 10^{-3}$
$\ln\left(p_{\rm muon}^{(2)}\right)$	$\mathbf{17\times 10^{-3}}$	$78\times 10^{-3}$
$\cos\theta_{3}$	$60\times 10^{-3}$	$\mathbf{3.8\times 10^{-3}}$
$\cos\theta_{\rm EP}$	$\mathbf{2.2\times 10^{-3}}$	$3.8\times 10^{-3}$

Appendix

Table 2: The Wasserstein-1 distances between various projections of the event distributions generated using

N_{\rm mult}=1

2

and using the theoretical event distribution; all variables are defined in the main text. Note that Wasserstein distances carry units and, correspondingly, scale with the variation of the underlying variable, so comparisons between distinct rows must be made with caution; the two columns, however, can be safely compared. The lower value in each row is bolded.

	FID inv	FID AE	AUC inv	AUC AE	KS $(E_{3})$	KS $(\cos\theta_{3})$
default	$\mathbf{4.8\times 10^{-5}}$	$0.052\pm 0.002$	$\mathbf{0.500\pm 0.008}$	$0.527\pm 0.011$	$\mathbf{0.0270}$	$0.0667$
case 1: none	$1.8\times 10^{-4}$	$0.076\pm 0.003$	$0.503\pm 0.004$	$0.545\pm 0.007$	$0.0353$	$0.0551$
case 2: $N_{\rm mult}=2$	$6.9\times 10^{-5}$	$0.029\pm 0.001$	$0.486\pm 0.005$	$0.511\pm 0.007$	$0.0355$	$0.0214$
case 3: $N_{\rm mult}=10$	$9.2\times 10^{-5}$	$0.033\pm 0.001$	$0.501\pm 0.008$	$0.509\pm 0.012$	$0.0480$	$0.0175$
case 4: cont.	$1.8\times 10^{-4}$	$\mathbf{0.019\pm 0.002}$	$0.513\pm 0.009$	$\mathbf{0.494\pm 0.010}$	$0.0480$	$\mathbf{0.0143}$

Appendix

Table 3: Compact comparison of the five muon-decay models shown in Fig. 15 . All quantities are evaluated on the same 2000-event held-out sample. “FID inv ” denotes a Fréchet distance computed on Lorentz-invariant features, while “FID AE ” uses the latent space of a shared autoencoder trained only on truth events; for the latter we quote the mean and standard deviation over two independent shared-AE trainings. “AUC inv ” and “AUC AE ” are linear classifier two-sample test AUCs in the same two feature spaces, with mean and standard deviation taken over repeated random train/test splits and classifier seeds; values closer to

0.5

indicate better agreement with the target distribution. The final two columns give KS statistics for the energy and angular marginals most visibly affected by the augmentation choice. Lower is better for all metrics shown, except that the AUC values are optimal at

0.5

; the best entry in each column is bolded. For case 1, one non-finite sample out of 2000 generated events is removed before computing these aggregate metrics.

Figure 19: The

q

-space distributions at the start ( left ) and end ( right ) of the forward process, for different diffusion schedules and data augmentation strategies.

Figure 20: Reverse trajectory for the

N_{\rm mult}=1

data augmentation strategy. The left panel shows the

q

-space distributions halfway through the reverse process

(t/T=0.5)

, while the center panel shows the endpoint of the reverse process. The right panel shows the

\tau

distribution analogous to Fig. 10 .

Figure 21: Snapshots from the reverse process in

q

-space ( left ) and

p

-space ( right ), for the case of

q\bar{q}g

data with no augmentation in

q

-space.

Figure 22: Forward process start and end distributions in

q

-space for the

N=10

APS distributions.

Explore similar work

Learning Standard Model structure from LHC data with Riemannian flow matching
Jul 17, 2026Midori Kato, Kevin A. Urquía-Calderón, Inar Timiryasov +1Constrained Generative Modeling Flow Matching
Niels Bohr Institute, University of Copenhagen, Jagtvej 155A, DK-2200, Copenhagen, Denmark
PhaseFlow4D: Physically Constrained 4D Beam Reconstruction via Feedback-Guided Latent Diffusion
Apr 4, 2026Alexander Scheinker, Alexander Plastun, Peter Ostroumov4D Reconstruction Physics-Informed Diffusion Models
Los Alamos National Laboratory Los Alamos 87545, NM, USA · Facility for Rare Isotopes Beams Michigan State University, East Lansing, MI 48824, USA
PMosFM: Preconditioned Manifold Matching for One-Step Physics-Constrained Generation
Sep 30, 2026Zhangyong Liang, Haibin LingConstrained Generative Modeling Physics-Informed Generative Modeling
Department of Artificial Intelligence Westlake University

Jul 17, 2026J/K move · Enter open · S save

Learning Standard Model structure from LHC data with Riemannian flow matching

Midori Kato, Kevin A. Urquía-Calderón, Inar Timiryasov, Oleg Ruchayskiy

Niels Bohr Institute, University of Copenhagen, Jagtvej 155A, DK-2200, Copenhagen, Denmark

Constrained Generative Modeling Flow Matching Conditional Flow Matching Riemannian Flow Matching High-Energy Physics

Open paper AskPDF arXiv

In this work we demonstrate that a single transformer-based generative model can capture Standard Model structure spanning five decades of invariant mass, from the sub-GeV regime to the TeV continuum, a range that no single Monte Carlo sample covers. To achieve this we design ShellFlow, a Riemannian conditional flow matching model that, given the recorded event composition, generates each particle on its on-shell manifold. Its only physics priors are the on-shell condition and the invariant-mass formula. The model is trained on

\sim 10^{9}

real

pp

collision events from the ATLAS Open Data 13~TeV release and told nothing else. From a single training run, the model learns to reproduce all of the following: intra-particle kinematics, the dilepton resonances (

J/ψ

Υ

Z

) at their PDG positions, the leptonic Weinberg angle, the

W

and top-quark masses, and inter-particle correlations that enter no training objective. A substantial fraction of the Standard Model is thus learnable directly from recorded collision data.

Abstract

Explore similar work

Learning Standard Model structure from LHC data with Riemannian flow matching

PhaseFlow4D: Physically Constrained 4D Beam Reconstruction via Feedback-Guided Latent Diffusion

PMosFM: Preconditioned Manifold Matching for One-Step Physics-Constrained Generation