On Kinetic Optimal Probability Paths for Generative Models
Neta Shaul, Ricky T. Q. Chen, Maximilian Nickel, Matthew Le, Yaron Lipman
Abstract
Recent successful generative models are trained by fitting a neural network to an a-priori defined tractable probability density path taking noise to training examples. In this paper we investigate the space of Gaussian probability paths, which includes diffusion paths as an instance, and look for an optimal member in some useful sense. In particular, minimizing the Kinetic Energy (KE) of a path is known to make particles' trajectories simple, hence easier to sample, and empirically improve performance in terms of likelihood of unseen data and sample generation quality. We investigate Kinetic Optimal (KO) Gaussian paths and offer the following observations: (i) We show the KE takes a simplified form on the space of Gaussian paths, where the data is incorporated only through a single, one dimensional scalar function, called the data separation function. (ii) We characterize the KO solutions with a one dimensional ODE. (iii) We approximate data-dependent KO paths by approximating the data separation function and minimizing the KE. (iv) We prove that the data separation function converges to in the general case of arbitrary normalized dataset consisting of samples in dimension as . A consequence of this result is that the Conditional Optimal Transport (Cond-OT) path becomes kinetic optimal as . We further support this theory with empirical experiments on ImageNet.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers20
- FlowMM: Generating Materials with Riemannian Flow MatchingBenjamin Kurt Miller, Ricky T. Q. Chen, Anuroop Sriram, Brandon M. WoodICML 2024 · 101 citations
- Metric Flow Matching for Smooth Interpolations on the Data ManifoldKacper Kapusniak, Peter Potaptchik, Teodora Reu, Leo Zhang et al.NeurIPS 2024 · 89 citations
- D-Flow: Differentiating through Flows for Controlled GenerationHeli Ben-Hamu, Omri Puny, Itai Gat, Brian Karrer et al.ICML 2024 · 82 citations
- Fisher Flow Matching for Generative Modeling over Discrete DataOscar Davis, Samuel Kessler, Mircea Petrache, Ismail Ilkan Ceylan et al.NeurIPS 2024 · 79 citations
- Neural Flow Diffusion Models: Learnable Forward Process for Improved Diffusion ModellingGrigory Bartosh, Dmitry P. Vetrov, Christian Andersson NaessethNeurIPS 2024 · 49 citations
Builds on12
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Diffusion Models Beat GANs on Image SynthesisPrafulla Dhariwal, Alexander Quinn NicholNeurIPS 2021 · 13,211 citations
- High-Resolution Image Synthesis with Latent Diffusion ModelsRobin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser et al.CVPR 2022 · 13,123 citations
- Improved Denoising Diffusion Probabilistic ModelsAlexander Quinn Nichol, Prafulla DhariwalICML 2021 · 5,234 citations
- Elucidating the Design Space of Diffusion-Based Generative ModelsTero Karras, Miika Aittala, Timo Aila, Samuli LaineNeurIPS 2022 · 3,959 citations
Related papers
- Flow Matching for Generative ModelingYaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel et al.ICLR 2023 · 87 citations
- A Kinetic-Energy Perspective of Flow MatchingZiyun Li, Huancheng Hu, Soon Hoe Lim, Xuyu Li et al.ICML 2026
- Kernel Neural Optimal TransportAlexander Korotin, Daniil Selikhanovych, Evgeny BurnaevICLR 2023 · 10 citations
- Learning diffusion at lightspeedAntonio Terpin, Nicolas Lanzetti, Martín Gadea, Florian DörflerNeurIPS 2024 · 26 citations
- Understanding DDPM Latent Codes Through Optimal TransportValentin Khrulkov, Gleb V. Ryzhakov, Andrei Chertkov, Ivan V. OseledetsICLR 2023 · 12 citations
