Inflationary Flows


SDE-ODE Duality of diffusion-based models. The forward (noising) SDE defining the DBM (left) gives rise to a sequence of marginal probability densities whose temporal evolution is described by a Fokker-Planck equation (FPE, middle). But this correspondence is not unique: the probability flow ODE (pfODE, right) gives rise to the same FPE. That is, while both the SDE and the pfODE transform the data distribution in the same way, the former is noisy and mixing while the latter is deterministic and neighborhood-preserving. Both models require knowledge of the score function \( \nabla_\mathbf{x} \log p_t(\mathbf{x}) \), which can learned by training either.

In this project we exploited a previously established connection between the stochastic and probability flow ordinary differential equations (pfODEs) underlying Diffusion-Based Models (DBMs) to derive a new class of models, inflationary flows, that uniquely and deterministically map high-dimensional data to a lower-dimensional Gaussian distribution via ODE integration. This map is both invertible and neighborhood-preserving, with controllable numerical error, with the result that uncertainties in the data are correctly propagated to the latent space. We demonstrate how such maps can be learned via standard DBM training using a novel noise schedule and are effective at both preserving and reducing intrinsic data dimensionality. The result is a class of highly expressive generative models, uniquely defined on a low-dimensional latent space, that afford principled Bayesian inference.

Check out additional details by visiting our inflationary flows website! A full preprint for this project is available here, and code for reproducing the experiment results and training new models can be found here.

Finally, this work was also featured on the CoSyNe 2024 Workshop entitled “I Can’t Believe It’s Not Better”, check out the video for this talk here.