A Likely Geometry of Generative Models

Modern generative models based on diffusion or flow matching transport noise toward the data manifold, without representing the data manifold explicitly. We propose Distance Marching that explores an alternative design by explicitly modeling a distance field of the data manifold. A distance field offers a geometric view of generation as a descent process toward the data manifold using only the current position. The negative gradient of the distance field provides the direction, and the distance value can optionally determine the step size. However, sparse training data in high-dimensional space make it difficult to identify the closest points on the data manifold needed to train the distance field. To address this challenge, we design new loss functions that emphasize denoising targets from closer data points. Our analysis shows that, without time input, emphasizing closer targets reduces denoising ambiguity and the bias toward the data mean. This distance-field perspective naturally motivates two principled deterministic samplers, showing fast convergence in inference. As a result, our model is the only one without time input that remains competitive with time-conditioned models across datasets and architectures. Moreover, our distance prediction is also helpful for early stopping during sampling and for out-of-distribution detection. We hope distance-field modeling can open a promising direction for exploring generative modeling from a geometric perspective.

Flow matching (FM) has recently emerged as a promising framework for generative modeling due to its conceptual simplicity and strong empirical performance. In FM, samples are transported along a vector field parameterized by a neural network, inducing a probability path that evolves from a simple noise distribution to the target data distribution, governed by an ordinary differential equation (ODE). However, existing FM approaches predominantly rely on probability paths derived from optimal transport (OT) between Gaussian distributions, which may be suboptimal for capturing complex data with inhomogeneous structures such as heavy tail or sharp contrast. In this work, we generalize FM to the broader class of location-scale families for handling data inhomogeneity and introduce a novel class of probability paths defined as geodesics on the manifold of probability distributions. We name this approach probabilistic geodesic flow matching to distinguish it from prior geodesic (Riemannian) FM methods defined in input space. We argue that Euclidean OT-based paths are not necessarily optimal in probability space and may limit modeling flexibility. Through synthetic benchmarks and scientific datasets at different scales, we demonstrate that the proposed method more effectively captures complex distributions, leading to improved or comparable performance compared with SOTA geometry-motivated generative models.

Many scientific datasets, such as molecular conformational ensembles or single-cell tissue measurements, are naturally modeled as meta-distributions: distributions over probability measures on non-Euclidean domains. Existing generative methods largely assume Euclidean geometry and fail to capture this structure. We introduce Riemannian Wasserstein Entropic Flow Matching (RWEFM), a generative framework on the Wasserstein space $\mathcal{P}_2(\mathcal{M})$ of a Riemannian manifold $(\mathcal{M},g)$. RWEFM is trained by regressing a neural vector field onto Riemannian optimal transport velocities, using McCann displacement interpolations as conditional paths. We confirm theoretically that this construction leads to a valid flow matching approach on $\mathcal{P}_2(\mathcal{M})$ and introduce the Riemannian Entropic Map, a GPU-efficient approximation of the optimal transport map on manifolds. Our experiments show that by respecting the intrinsic geometry of the data, RWEFM can generate whole single-cell samples in hyperspherical latent spaces and protein conformational ensembles on the torus. As RWEFM requires only a geodesic distance and a projection operator, it is not restricted to manifolds with closed-form geometry, which we demonstrate by generating distributions on a general triangulated mesh.