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

Jul 17, 2026

Learning Standard Model structure from LHC data with Riemannian flow matching

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.

Midori Kato, Kevin A. Urquía-Calderón, Inar Timiryasov +1

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

Ask PDF

Apr 4, 2026

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

We address the problem of recovering a time-varying 4D distribution from a sparse sequence of 2D projections - analogous to novel-view synthesis from sparse cameras, but applied to the 4D transverse phase space density

ρ(x,p_x,y,p_y)

of charged particle beams. Direct single shot measurement of this high-dimensional distribution is physically impossible in real particle accelerator systems; only limited 1D or 2D projections are accessible. We propose PhaseFlow4D, a feedback-guided latent diffusion model that reconstructs and tracks the full 4D phase space from incomplete 2D observations alone, with built-in hard physics constraints. Our core technical contribution is a 4D VAE whose decoder generates the full 4D phase space tensor, from which 2D projections are analytically computed and compared against 2D beam measurements. This projection-consistency constraint guarantees physical correctness by construction - not as a soft penalty, but as an architectural prior. An adaptive feedback loop then continuously tunes the conditioning vector of the latent diffusion model to track time-varying distributions online without retraining. We validate on multi-particle simulations of heavy-ion beams at the Facility for Rare Isotope Beams (FRIB), where full physics simulations require

\sim

6 hours on a 100-core HPC system. PhaseFlow4D achieves accurate 4D reconstructions 11000

\times

faster while faithfully tracking distribution shifts under time-varying source conditions - demonstrating that principled generative reconstruction under incomplete observations transfers robustly beyond visual domains.

Alexander Scheinker, Alexander Plastun, Peter Ostroumov

Los Alamos National Laboratory Los Alamos 87545, NM, USA · Facility for Rare Isotopes Beams Michigan State University, East Lansing, MI 48824, USA

4D Reconstruction Physics-Informed Diffusion Models Latent Diffusion Models High-Energy Physics

Ask PDF

Sep 30, 2026

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

Physics-constrained generative models aim to generate physical fields that match a target distribution and satisfy prescribed constraints. However, enforcing these constraints often increases sampling costs through iterative corrections or training costs through residual optimization and trajectory unrolling. To address this issue, we introduce Preconditioned Manifold one-step Flow Matching (PMosFM), a preconditioned manifold matching framework for one-step physics-constrained generation. By encoding constraints in a manifold decoder, PMosFM learns transport in intrinsic coordinates without separate residual losses or terminal residual unrolling. A geometric preconditioner rescales coordinates using the decoder-induced metric, while a regularized covariance transform approximately whitens the interpolation-state inputs. A finite-interval objective couples velocity supervision with consistency between decoded endpoints in physical space. We show that exact parameterization removes residual-induced Gauss--Newton curvature, that geometric and covariance effects separate in a local conditioning bound, and that physical flow-map error bounds endpoint distributional error. Controlled ablations examine conditioning, and experiments evaluate optimizer-update time and memory footprint. At inference, PMosFM uses one neural transport evaluation followed by physical decoding. Experiments across benchmarks show lower training and sampling time than the multi-step baselines at comparable physical and distributional fidelity. Code and datasets will be released publicly.

Zhangyong Liang, Haibin Ling

Department of Artificial Intelligence Westlake University

Constrained Generative Modeling Physics-Informed Generative Modeling Flow Matching One-Step Generative Modeling Flow-Based Generative Modeling

Ask PDF