Generative models on phase space

TL;DR

This paper introduces generative models confined to the Lorentz-invariant phase space manifold, improving physical constraint adherence and interpretability in high-energy physics data sampling.

Contribution

The authors develop models that are inherently confined to the phase space manifold, ensuring exact physical constraints during sampling.

Findings

Models successfully learn distributions with complex singularity structures.

The approach enhances interpretability and physical consistency of generated data.

Demonstrated on simulated jet data with promising results.

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.