Generative models that generalise rather than memorise
Deep generative models often face a fundamental tradeoff: high sample quality can come at the cost of memorisation, where the model reproduces its training data rather than generalising across the underlying data geometry.
Geometry-aware noise for flow matching
Carré du champ flow matching (CDC-FM) is a generalisation of flow matching (FM) that improves this quality-generalisation tradeoff by regularising the probability path with geometry-aware noise. Standard FM uses homogeneous, isotropic noise. CDC-FM replaces it with spatially varying, anisotropic Gaussian noise whose covariance captures the local geometry of the latent data manifold. We prove that this geometric noise can be optimally estimated from the data, and that the estimate scales to large datasets.
Results
We evaluated CDC-FM on diverse datasets (synthetic manifolds, point clouds, single-cell genomics, animal motion capture and images) and with various neural network architectures (MLPs, CNNs and transformers). CDC-FM consistently offers a better quality-generalisation tradeoff. The improvements over standard FM are largest in data-scarce regimes and on highly non-uniformly sampled datasets, which are common in AI for science applications.
Our work provides a mathematical framework for studying the interplay between data geometry, generalisation and memorisation in generative models, together with a robust and scalable algorithm that can be readily integrated into existing flow matching pipelines. The code is on GitHub, and the paper is also on arXiv.
See also Riemannian metric matching and diffusion processes on implicit manifolds, which build on the same carré du champ operator.