On Kinetic Optimal Probability Paths for Generative Models
Neta Shaul Ricky T. Q. Chen Maximilian Nickel Matt Le Yaron Lipman
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 \emph{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 \emph{kinetic optimal} as
$n/\sqrt{d}\rightarrow 0$. We further support this theory with empirical
experiments on ImageNet.