Adapting Noise to Data: Generative Flows from 1D Processes

AI summary

Overview

Research area: Generative modeling — specifically flow matching (FM) and diffusion models, combined with optimal transport and one-dimensional probability/quantile theory (stat.ML).

Technical level: Advanced. The paper builds on Wasserstein spaces of probability measures, continuity equations, mean-reverting stochastic processes, and optimal transport couplings; it assumes comfort with measure-theoretic probability and differential equations.

Scope: The paper proposes learning the latent noise distribution of a flow-based generative model directly from data by parameterizing per-coordinate one-dimensional quantile functions and optimizing them with the Wasserstein distance, then demonstrates the approach on synthetic 2D distributions, MNIST, CIFAR-10, and (per the abstract) heavy-tailed weather data.

What This Paper Is About

Flow-based generative models normally start from a standard Gaussian latent, which is a poor fit when the data is heavy-tailed, compactly supported, or otherwise far from Gaussian. Existing fixes (for example Student-$t$ noise with a tunable degrees-of-freedom parameter, or $\alpha$-stable noise) still require a human to pick and hand-tune the noise family. The authors instead learn the noise itself: they restrict each coordinate's noise to an arbitrary 1D distribution, represent it by a monotone quantile function, and fit it to the data by minimizing the Wasserstein distance between noise and data, all inside the standard flow matching training loop.

Key Contributions

1. A general construction of multidimensional flows from 1D flows. The authors decompose a $d$-dimensional noising process into independent one-dimensional components (the structure used by Gaussian noise and Brownian motion) and show how to lift any 1D process plus an associated mean-reverting process into an accessible conditional flow for flow matching in arbitrary dimension.
2. New 1D noising processes beyond Brownian motion. They examine the physics-inspired Kac process (linked to the damped wave equation) and an MMD gradient flow (linked to a uniform process), which yield compactly supported noise and better regularity of the FM velocity field. Neither has a trivial multidimensional analogue: the Kac PDE is no longer mass-conserving in dimension $d \geq 3$, and the existence of the MMD gradient flow in multiple dimensions is unclear because the MMD is not convex.
3. Learning the one-dimensional noise distribution from data. The central contribution: parameterizing each coordinate's latent through a quantile function (implemented with a rational quadratic spline) and learning it by minimizing $W_2^2$ between the data and the generated noise distribution, jointly with the flow matching velocity field.
4. Extensions to few-step and alternative constructions. The paper defines quantile interpolants that satisfy an identity and composition property, fitting the framework of consistency-style models such as inductive moment matching (IMM), and describes an alternative multidimensional construction based on radially symmetric processes (Appendix D).

Main Findings

• Learned noise helps on heavy-tailed targets. In the Neal's funnel comparison, four latents were tested: uniform on $[-1,1]$, standard Gaussian, Student-$t$ with parameters $(20,4)$ (inspired by the choice in Pandey et al., 2024), and the learned distribution. The two heavy-tailed noises performed significantly better than the compact and Gaussian latents. In a separate funnel experiment using a capacity-constrained network (three layers, width 64, no positional embeddings), the compact latent performed worst, followed by the Gaussian, while the learned latent adapted to the target's heavy tails.
• Transport paths get shorter. On the checkerboard distribution, the method learns a latent approximating a uniform distribution over the target's support; combined with optimal-transport coupling, the resulting transport paths are substantially shorter than starting from a Gaussian, and the vector field converges faster during training.
• Capacity-constrained image models improve most. On MNIST, using channel widths of 8, 16 and 32 (with $\lambda=1$, $\beta=0.1$), the learned latent achieved significantly lower FIDs at every capacity level, which the authors attribute to the noise distribution no longer carrying redundant information the network must model.
• The learned latent captures per-pixel marginal structure. On MNIST, pixel marginals are strongly structured (center pixels frequently active, border pixels nearly always zero); the learned quantiles concentrate around the essentially-always-black values and spread around zero where the data is more variable.
• Computational overhead is small. With $K=32$ spline bins and $d=3072$ (CIFAR-10), the total parameter count is $\mathcal{O}(Kd)$, roughly 300k parameters. On CIFAR-10 with an NVIDIA RTX 5090, the paper measures approximately 2.7% overhead during joint training and 0.5% after freezing; the quantile is frozen after 55k iterations.
• CIFAR-10 training efficiency. The paper reports that its method reached the best validation FID after 320k steps, while the baseline took 340k; FID was evaluated over 5 seeds with the mean reported, and visualized samples were generated with 100 Euler steps. A table covers regularization weights $\beta = 0.2, 0.3, 0.5, 0.8, 1.0$ plus a baseline, reporting pairs of values per row (for example 7.81 and 4.75 at $\beta=0.2$, and 8.42 and 4.63 for the baseline) alongside a $-\log\det$ diagnostic; the column labels in the provided text reference 20 and 100 Euler steps, and the exact metric-to-column assignment is not unambiguous from the extracted text.
• The independence assumption is a real limitation. For a 2D Gaussian mixture with nine unevenly weighted modes, the learned latent cannot reproduce the target's joint distribution and is explicitly not the product of the correct marginals; instead it approximates a distribution whose components can no longer independently improve the transport cost. The independence constraint restricts the latent to per-coordinate adaptation, and cross-dimensional dependencies are handled by the velocity field through the target displacement $(y-x)$.

Methodology in Plain English

The starting point is the observation that standard flow matching and diffusion add independent noise to each coordinate. The authors take that per-coordinate independence seriously: they let the latent be the product of $d$ independent 1D distributions, so the whole multidimensional noising process is built from 1D pieces. Because every 1D distribution can be written as a quantile function (an increasing function on the unit interval), each piece can be parameterized directly and simply by a single monotone rational quadratic spline per coordinate, with learned per-coordinate scale and bias.

To make these 1D pieces fit a diffusion-style "data plus noise" schedule, the authors use a mean-reverting process $X_t = f(t)X_0 + N_{g(t)}$ with schedules $f(0)=1, f(1)=0$ and $g(0)=0, g(1)=1$. Setting $f(t)=1-t$ and $g(t)=t^2$ recovers standard flow matching, and the choices $f(t)=\exp(-h(t)/2)$, $g(t)=1-\exp(-h(t))$ with $h(t)=\int_0^t \beta_{\min}+s(\beta_{\max}-\beta_{\min}),ds$ and $\beta_{\min}=0.1$, $\beta_{\max}=20$ recover the processes used in score-based generative models.

The noise is then learned. The training objective combines three terms: the conditional flow matching loss, a Wasserstein-distance loss $W_2^2(\mu_0, \nu_\phi)$ between the data distribution $\mu_0$ and the quantile-induced latent $\nu_\phi$, and an entropy regularizer equal to the expected log-determinant of the quantile's Jacobian (the differential entropy of the latent), written as $\mathcal{L}(\theta,\phi) = \mathcal{L}_{\mathsf{CFM}}(\theta,\phi) +

Authors’ abstract

The default Gaussian latent in flow-based generative models poses challenges when learning certain distributions such as heavy-tailed ones. We introduce a general framework for learning data-adaptive parametric prior distributions (latent noise) using one-dimensional quantile functions, optimized via the Wasserstein distance between noise and data. The quantile-based prior parameterization naturally adapts to both heavy-tailed and compactly supported distributions and shortens transport paths. Numerical results on heavy-tailed weather and image datasets confirm the method's flexibility and effectiveness achieved with negligible computational overhead.