Training flow matching in q-space with a uniform phase-space prior

Develop and demonstrate a flow matching generative model in q-space that uses the uniform distribution on N-particle Lorentz-invariant phase space as the prior and can be successfully trained to generate samples, thereby achieving the same physically motivated prior used by the diffusion construction while maintaining exact energy–momentum conservation along the sampling trajectory.

Background

The paper introduces generative modeling directly on the manifold of massless N-particle Lorentz-invariant phase space by operating in an auxiliary q-space, enabling exact energy–momentum conservation throughout sampling. For diffusion, the authors set the pure-noise endpoint to be the uniform distribution on phase space (via the RAMBO map), providing a physically interpretable starting point and ensuring learned deviations correspond to true correlations.

They note that, in principle, the same idea should extend to flow matching by taking the uniform phase-space distribution as the prior in q-space. However, they were not able to train such flow matching models successfully and therefore resorted to a fitted Gaussian prior, which lacks a direct physical interpretation in p-space. Establishing a viable training procedure for flow matching with a uniform phase-space prior would align it with the diffusion setup and preserve interpretability of the prior.