On Kinetic Optimal Probability Paths for Generative Models

Neta Shaul, Ricky T. Q. Chen, Maximilian Nickel, Matthew Le, Yaron Lipman
Proceedings of the 40th International Conference on Machine Learning, PMLR 202:30883-30907, 2023.

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 $1$ in the general case of arbitrary normalized dataset consisting of $n$ samples in $d$ dimension as $n/\sqrt{d}\rightarrow 0$. A consequence of this result is that the Conditional Optimal Transport (Cond-OT) path becomes kinetic optimal as $n/\sqrt{d}\rightarrow 0$. We further support this theory with empirical experiments on ImageNet.

Cite this Paper

BibTeX
@InProceedings{pmlr-v202-shaul23a, title = {On Kinetic Optimal Probability Paths for Generative Models}, author = {Shaul, Neta and Chen, Ricky T. Q. and Nickel, Maximilian and Le, Matthew and Lipman, Yaron}, booktitle = {Proceedings of the 40th International Conference on Machine Learning}, pages = {30883--30907}, year = {2023}, editor = {Krause, Andreas and Brunskill, Emma and Cho, Kyunghyun and Engelhardt, Barbara and Sabato, Sivan and Scarlett, Jonathan}, volume = {202}, series = {Proceedings of Machine Learning Research}, month = {23--29 Jul}, publisher = {PMLR}, pdf = {https://proceedings.mlr.press/v202/shaul23a/shaul23a.pdf}, url = {https://proceedings.mlr.press/v202/shaul23a.html}, 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 $1$ in the general case of arbitrary normalized dataset consisting of $n$ samples in $d$ dimension as $n/\sqrt{d}\rightarrow 0$. A consequence of this result is that the Conditional Optimal Transport (Cond-OT) path becomes kinetic optimal as $n/\sqrt{d}\rightarrow 0$. We further support this theory with empirical experiments on ImageNet.} }
Endnote
%0 Conference Paper %T On Kinetic Optimal Probability Paths for Generative Models %A Neta Shaul %A Ricky T. Q. Chen %A Maximilian Nickel %A Matthew Le %A Yaron Lipman %B Proceedings of the 40th International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2023 %E Andreas Krause %E Emma Brunskill %E Kyunghyun Cho %E Barbara Engelhardt %E Sivan Sabato %E Jonathan Scarlett %F pmlr-v202-shaul23a %I PMLR %P 30883--30907 %U https://proceedings.mlr.press/v202/shaul23a.html %V 202 %X 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 $1$ in the general case of arbitrary normalized dataset consisting of $n$ samples in $d$ dimension as $n/\sqrt{d}\rightarrow 0$. A consequence of this result is that the Conditional Optimal Transport (Cond-OT) path becomes kinetic optimal as $n/\sqrt{d}\rightarrow 0$. We further support this theory with empirical experiments on ImageNet.
APA
Shaul, N., Chen, R.T.Q., Nickel, M., Le, M. & Lipman, Y.. (2023). On Kinetic Optimal Probability Paths for Generative Models. Proceedings of the 40th International Conference on Machine Learning, in Proceedings of Machine Learning Research 202:30883-30907 Available from https://proceedings.mlr.press/v202/shaul23a.html.

Related Material