Building Normalizing Flows with Stochastic Interpolants

Introduction

Contemporary generative models have primarily been designed around the construction of a map between two probability distributions that transform samples from the first into samples from the second. While progress has been from various angles with tools such as implicit maps (Goodfellow et al., 2014; Brock et al., 2019), and autoregressive maps (Menick & Kalchbrenner, 2019; Razavi et al., 2019; Lee et al., 2022), we focus on the case where the map has a clear associated probability flow. Advances in this domain, namely from flow and diffusion models, have arisen through the introduction of algorithms or inductive biases that make learning this map, and the Jacobian of the associated change of variables, more tractable. The challenge is to choose what structure to impose on the transport to best reach a complex target distribution from a simple one used as base, while maintaining computational efficiency.

One convenient way to represent this time-continuous map is to define it as the flow associated with the ordinary differential equation (ODE)

where the dot denotes derivative with respect to $t$ and $v_{t}(x)$ is the velocity field governing the transport. This is equivalent to saying that the probability density function $\rho_{t}(x)$ defined as the pushforward of the base $\rho_{0}(x)$ by the map $X_{t}$ satisfies the continuity equation (see e.g. (Villani, 2009; Santambrogio, 2015) and Appendix A)

and the inference problem becomes to estimate a velocity field such that (3) holds.

Here we propose a solution to this problem based on introducing a time-differentiable interpolant

A useful instance of such an interpolant that we will employ is

though we stress the framework we propose applies to any $I_{t}(x_{0},x_{1})$ satisfying (4) under mild additional assumptions on $\rho_{0}$, $\rho_{1}$, and $I_{t}$ specified below. Given this interpolant, we then construct the stochastic process $x_{t}$ by sampling independently $x_{0}$ from $\rho_{0}$ and $x_{1}$ from $\rho_{1}$, and passing them through $I_{t}$:

We refer to the process $x_{t}$ as a stochastic interpolant.

Under this paradigm, we make the following key observations as our main contributions in this work:

The probability density $\rho_{t}(x)$ of $x_{t}$ connecting the two densities, henceforth referred to as the interpolant density, satisfies (3) with a velocity $v_{t}(x)$ which is the unique minimizer of a simple quadratic objective. This result is the content of Proposition 1 below, and it can be leveraged to estimate $v_{t}(x)$ in a parametric class (e.g. using deep neural networks) to construct a generative model through the solution of the probability flow equation (2), which we call InterFlow.

By specifying an interpolant density, the method therefore separates the tasks of minimizing the objective from discovering a path between the base and target densities. This is in contrast with conventional maximum likelihood (MLE) training of flows where one is forced to couple the choice of path in the space of measures to maximizing the objective.

We show that the Wasserstein-2 ($W_{2}$) distance between the target density $\rho_{1}$ and the density $\hat{\rho}_{1}$ obtained by transporting $\rho_{0}$ using an approximate velocity $\hat{v}_{t}$ in (2) is controlled by our objective function. We also show that the value of the objective on $\hat{v}_{t}$ during training can be used to check convergence of this learned velocity field towards the exact $v_{t}$.

We show that our approach can be generalized to shorten the path length of the interpolant density and optimize the transport by additionally maximizing our objective over the interpolant $I_{t}(x_{0},x_{1})$ and/or adjustable parameters in the base density $\rho_{0}$.

By choosing $\rho_{0}$ to be a Gaussian density and using (5) as interpolant, we show that the score of the interpolant density, $\nabla\log\rho_{t}$, can be explicitly related to the velocity field $v_{t}$. This allows us to draw connection between our approach and score-based diffusion models, providing theoretical groundwork for future exploration of this duality.

We demonstrate the feasibility of the method on toy and high dimensional tabular datasets, and show that the method matches or supersedes conventional ODE flows at lower cost, as it avoids the need to backpropagate through ODE solves. We demonstrate our approach on image generation for CIFAR-10 and ImageNet 32x32 and show that it scales well to larger sizes, e.g. on the 128$\times$128 Oxford flower dataset.

Early works on exploiting transport maps for generative modeling go back at least to Chen & Gopinath (2000), which focuses on normalizing a dataset to infer its likelihood. This idea was brought closer to contemporary use cases through the work of Tabak & Vanden-Eijnden (2010) and Tabak & Turner (2013), which devised to expressly map between densities using simple transport functions inferred through maximum likelihood estimation (MLE). These transformations were learned in sequence via a greedy procedure. We detail below how this paradigm has evolved in the case where the map is represented by a neural network and optimized accordingly.

Discrete and continuous time flows. The first success of normalizing flows with neural network parametrizations follow the work of Tabak & Turner (2013) with a finite set of steps along the map. By imposing structure on the transformation so that it remains an efficiently invertible diffeomorphism, the models of Rezende & Mohamed (2015); Dinh et al. (2017); Huang et al. (2018); Durkan et al. (2019) can be optimized through maximum likelihood estimation at the cost of limiting the expressive power of the representation, as the Jacobian of the map must be kept simple to calculate the likelihood. Extending this to the continuous case allowed the Jacobian to be unstructured yet still estimable through trace estimation techniques (Chen et al., 2018; Grathwohl et al., 2019; Hutchinson, 1989). Yet, learning this map through MLE requires costly backpropagation through numerical integration. Regulating the path can reduce the number of solver calls (Finlay et al., 2020; Onken et al., 2021), though this does not alleviate the main structural challenge of the optimization. Our work uses a continuous map $X_{t}$ as well but allows for direct estimation of the underlying velocity. While recent work has also considered simulation-free training by fitting a velocity field, these works present scalability issues (Rozen et al., 2021) and biased optimization (Ben-Hamu et al., 2022), and are limited to manifolds. Moreover, (Rozen et al., 2021) relies on interpolating directly the probability measures, which can lead to unstable velocities.

Score-based flows. Adjacent research has made use of diffusion processes, commonly the Ornstein-Uhlenbeck (OU) process, to connect the target $\rho_{1}$ to the base $\rho_{0}$. In this case the transport is governed by a stochastic differential equation (SDE) that is evolved for infinite time, and the challenge of learning a generative model can be framed as fitting the reverse time evolution of the SDE from Gaussian noise back to $\rho_{1}$ (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021b). Doing so indirectly learns the velocity field by means of learning the score function $\nabla\log\rho_{t}(x)$, using the Fischer divergence instead of the MLE objective. While this approach has shown great promise to model high dimensional distributions (Rombach et al., 2022; Hoogeboom et al., 2022), particularly in the case of text-to-image generation (Ramesh et al., 2022; Saharia et al., 2022), there is an absence of theoretical motivation for the SDE-flow framework and the complexity it induces. Namely, the SDE must evolve for infinite time to connect the distributions, the parameterization of the time steps remains heuristic (Xiao et al., 2022), and the criticality of noise, as well as the score, is not absolutely apparent (Bansal et al., 2022; Lu et al., 2022). In particular, while the objective used in score-based diffusion models was shown to bound the Kullback-Leibler divergence (Song et al., 2021a), actual calculation of the likelihood requires one to work with the ODE probability flow associated with the SDE. This motivates further research into effective, ODE-driven, approaches to learning the map. Our approach can be viewed as an alternative to score-based diffusion models in which the ODE velocity is learned through the interpolant $x_{t}$ rather than an OU process, leading to greater simplicity and flexibility (as we can connect any two densities exactly over a finite time interval).

Bridge-based methods. Heng et al. (2021) propose to learn Schroedinger bridges, which are a entropic regularized version of the optimal transportation plan connecting two densities in finite time, using the framework of score-based diffusion. Similarly, Peluchetti (2022) investigates the use of bridge processes, i.e. SDE whose position is constrained both at the initial and final times, to perform exact density interpolation in finite time. A key difference between these approaches and ours is that they give diffusion-based models, whereas our method builds a probability flow ODE directly using a quadratic loss for its velocity, which is simpler and shown here to be scalable.

Interpolants. Co-incident works by Liu et al. (2022); Lipman et al. (2022) derive an analogous optimization to us, with a focus on straight interpolants, also contrasting it with score-based methods. Liu et al. (2022) describe an iterative way of rectifying the interpolant path, which can be shown to arrive at an optimal transport map when the procedure is repeated ad infinitum (Liu, 2022). We also propose a solution to the problem of optimal transport that involves optimizing our objective over the stochastic interpolant.

2 Notations and assumptions

A few additional technical assumptions on $\rho_{0}$, $\rho_{1}$ and $I_{t}$ are listed in Appendix B. Given any function $f_{t}(x_{0},x_{1})$ we denote

its expectation over $t$, $x_{0}$, and $x_{1}$ drawn independently from the uniform density on $$,$\rho_{0}$, and$\rho_{1}$, respectively. We use$\nabla$ to denote the gradient operator.

Stochastic Interpolants and Associated Flows

Our first main theoretical result can be phrased as follows:

The stochastic interpolant $x_{t}$ defined in (6) with $I_{t}(x_{0},x_{1})$ satisfying (4) has a probability density $\rho_{t}(x)$ that satisfies the continuity equation (3) with a velocity $v_{t}(x)$ which is the unique minimizer over $\hat{v}_{t}(x)$ of the objective

In addition the minimum value of this objective is given by

Proposition 1 is proven in Appendix B under Assumption B.1. As this proof shows, the first statement of the proposition remains true if the expectation over $t$ is performed using any probability density $\omega(t)>0$, which may prove useful in practice. We now describe some primary facts resulting from this proposition, itemized for clarity:

The objective $G(\hat{v})$ is given in terms of an expectation that is amenable to empirical estimation given samples $t$, $x_{0}$, and $x_{1}$ drawn from $\rho_{0}$, $\rho_{1}$ and $U()$. Below, we will exploit this property to propose a numerical scheme to perform the minimization of $G(\hat{v})$.

While the minimizer of the objective $G(\hat{v})$ is not available analytically in general, a notable exception is when $\rho_{0}$ and $\rho_{1}$ are Gaussian mixture densities and we use the trigonometric interpolant (5) or generalization thereof, as discussed in Appendix C.

The minimal value in (10) achieved by the objective implies that a necessary (albeit not sufficient) condition for $\hat{v}=v$ is

In our numerical experiments we will monitor this quantity. This minimal value also suggests to maximize $G(v)=\min_{\hat{v}}G(\hat{v})$ with respect to additional control parameters (e.g. the interpolant) to shorten the $W_{2}$ length of the path $\{\rho_{t}(x):t\in\}$. In Appendix D, we show that this procedure achieves optimal transport under minimal assumptions.

The last bound in (10), which is proven in in Lemma B.2, implies that the path length is always finite, even if it is not the shortest possible.

Let us now provide some intuitive derivation of the statements of Proposition 1:

Continuity equation. By definition of the stochastic interpolant $x_{t}$ we can express its density $\rho_{t}(x)$ using the Dirac delta distribution as

Since $I_{t=0}(x_{0},x_{1})=x_{0}$ and $I_{t=1}(x_{0},x_{1})=x_{1}$ by definition, we have $\rho_{t=0}=\rho_{0}$ and $\rho_{t=1}=\rho_{1}$, which means that $\rho_{t}(x)$ satisfies the boundary conditions at $t=0,1$ in (3). Differentiating (12) in time using the chain rule gives

Therefore if we introduce the velocity $v_{t}(x)$ via

we see that we can write (13) as the continuity equation in (3).

Variational formulation. Using the expressions in (12) and (14) for $\rho_{t}(x)$ and $j_{t}(x)$ shows that we can write the objective (9) as

Since $\rho_{t}(x)$ and $j_{t}(x)$ have the same support, the minimizer of this quadratic objective is unique for all $(t,x)$ where $\rho_{t}(x)>0$ and given by (15).

Minimum value of the objective. Let $v_{t}(x)$ be given by (15) and consider the alternative objective

where we expanded the square and used the identity $v_{t}(x)\rho_{t}(x)=j_{t}(x)$ to get the second equality. The objective $G(\hat{v})$ given in (16) can be written in term of $H(\hat{v})$ as

The equality in (10) follows by evaluating (18) at $\hat{v}_{t}(x)=v_{t}(x)$ using $H(v)=0$.

where we used the shorthand notation $I_{t}=I_{t}(x_{0},x_{1})$ and $\partial_{t}I_{t}=\partial_{t}I_{t}(x_{0},x_{1})$. Evaluating this inequality at $\hat{v}=v$ using $H(v)=0$ we deduce

Optimizing the transport. It is natural to ask whether our stochastic interpolant construction can be amended or generalized to derive optimal maps. Here, we state a positive answer to this question by showing that the maximizing the objective $G(\hat{v})$ in (9) with respect to the interpolant yields a solution to the optimal transport problem in the framework of Benamou & Brenier (2000). This is proven in Appendix D, where we also discuss how to shorten the path length in density space by optimizing adjustable parameters in the base density $\rho_{0}$. Experiments are given in Appendix H. Since our primary aim here is to construct a map $T=X_{t=1}$ that pushes forward $\rho_{0}$ onto $\rho_{1}$, but not necessarily to identify the optimal one, we leave the full investigation of the consequences of these results for future work, but state the proposition explicitly here. In their seminal paper, Benamou & Brenier (2000) showed that finding the optimal map requires solving the minimization problem

The minimizing coupling $(\rho_{t}^{*},\phi_{t}^{*})$ for gradient field $v_{t}^{*}(x)=\nabla\phi_{t}^{*}(x)$ is unique and satisfies:

In the interpolant flow picture, $\rho_{t}(x)$ is fixed by the choice of interpolant $I_{t}(x_{0},x_{1})$, and in general $\rho_{t}(x)\neq\rho_{t}^{*}(x)$. Because the value of the objective in (22) is equal to the minimum of $G(\hat{v})$ given in (10), a natural suggestion to optimize the transport is to maximize this minimum over the interpolant. Under some assumption on the Benamou-Brenier density $\rho^{*}_{t}(x)$ solution of (23), this procedure works. We show this through the use of interpolable densities as discussed in Mikulincer & Shenfeld (2022) and defined in D.1.

Assume that (i) the optimal density function $\rho_{t}^{*}(x)$ minimizing (22) is interpolable and (ii) (23) has a classical solution. Consider the max-min problem

where $G(\hat{v})$ is the objective in (9) and the maximum is taken over interpolants satisfying (4). Then a maximizer of (24) exists, and any maximizer $I_{t}^{*}(x_{0},x_{1})$ is such that the probability density function of $x^{*}_{t}=I_{t}^{*}(x_{0},x_{1})$, with $x_{0}\sim\rho_{0}$ and $x_{1}\sim\rho_{1}$ independent, is the optimal $\rho_{t}^{*}(x)$, the mimimizing velocity is $v_{t}^{*}(x)=\nabla\phi_{t}^{*}(x)$, and the pair $(\rho_{t}^{*}(x),\phi_{t}^{*}(x))$ satisfies (23).

The proof of Proposition 2 is given in Appendix D, along with further discussion. Proposition 2 relies on Lemma D.3 that reformulates (22) in a way which shows this problem is equivalent to the max-min problem in (24) for interpolable densities.

The following result shows that the objective in (17) controls the Wasserstein distance between the target density $\rho_{1}$ and the the density $\hat{\rho}_{1}$ obtained as the pushforward of the base density $\rho_{0}$ by the map $\hat{X}_{t=1}$ associated with the velocity $\hat{v}_{t}$:

Let $\rho_{t}(x)$ be the exact interpolant density defined in (12) and, given a velocity field $\hat{v}_{t}(x)$, let us define $\hat{\rho}_{t}(x)$ as the solution of the initial value problem

where $H(\hat{v})$ is the objective function defined in (17).

The proof of Proposition 3 is given in Appendix E: it leverages the following bound on the square of W-2 distance

where $X_{t}$ is the flow map solution of (2) with the exact $v_{t}(x)$ defined in (15) and $\hat{X}_{t}$ is the flow map obtained by solving (2) with $v_{t}(x)$ replaced by $\hat{v}_{t}(x)$.

2 Link with score-based generative models

The following result shows that if $\rho_{0}$ is a Gaussian density, the velocity $v_{t}(x)$ can be related to the score of the density $\rho_{t}(x)$:

Assume that the base density $\rho_{0}(x)$ is a standard Gaussian density $N(0,\text{Id})$ and suppose that the interpolant $I_{t}(x_{0},x_{1})$ is given by (5). Then the score $\nabla\log\rho_{t}(x)$ is related to velocity $v_{t}(x)$ as

The proof of this proposition is given in Appendix F. The first formula for $t\in[0,1)$ is based on a direct calculation using Gaussian integration by parts; the second formula at $t=1$ is obtained by taking the limit of the first using $v_{t=1}(x)=0$ from (B.20) and l’Hôpital’s rule. It shows that we can in principle resample $\rho_{t}$ at any $t\in$ using the stochastic differential equation in artificial time $\tau$ whose drift is the score $\hat{s}_{t}(x)$ obtained by evaluating (28) on the estimated $\hat{v}_{t}(x)$:

Similarly, the score $\hat{s}_{t}(x)$ could in principle be used in score-based diffusion models, as explained in Appendix F. We stress however that while our velocity is well-behaved for all times in $$, as shown in (B.20), the drift and diffusion coefficient in the associated SDE are singular at$t=0,1$.

Practical Implementation and Numerical Experiments

The objective detailed in Section 2 is amenable to efficient empirical estimation, see (I.1), which we utilize to experimentally validate the method. Moreover, it is appealing to consider a neural network parameterization of the velocity field. In this case, the parameters of the model $\hat{v}$ can be optimized through stochastic gradient descent (SGD) or its variants, like Adam (Kingma & Ba, 2015). Following the recent literature regarding density estimation, we benchmark the method on visualizable yet complicated densities that display multimodality, as well as higher dimensional tabular data initially provided in Papamakarios et al. (2017) and tested in other works such as Grathwohl et al. (2019). The 2D test cases demonstrate the ability to flow between empirical densities with no known analytic form. In all cases, numerical integration for sampling is done with the Dormand–Prince, explicit Runge-Kutta of order (4)5 (Dormand & Prince, 1980). In Sections 3.1-3.4 the choice of interpolant for experimentation was selected to be that of (5), as it is the one used to draw connections to the technique of score based diffusions in Proposition 4. In Section H we use the interpolant (B.2) and optimize $a_{t}$ and $b_{t}$ to investigate the possibility and impact of shortening the path length.

An intuitive first test to benchmark the validity of the method is sampling a target density whose analytic form is known or whose density can be visualized for comparison. To this end, we follow the practice of choosing a few complicated 2-dimensional toy datasets, namely those from (Grathwohl et al., 2019), which were selected to differentiate the flexibility of continuous flows from discrete time flows, which cannot fully separate the modes. We consider anisotropic curved densities, a mixture of 8 separated Gaussians, and a checkerboard density. The velocity field of the interpolant flow is parameterized by a simple feed forward neural network with ReLU Nair & Hinton (2010) activation functions. The network for each model has 3 layers, each of width equal to 256 hidden units. Optimizer is performed on $G(\hat{v})$ for 10k epochs. We plot a kernel density estimate over 80k samples from both the flow and true distribution in Figure 2. The interpolant flow captures all the modes of the target density without artificial stretching or smearing, evincing a smooth map.

2 Dataset Interpolation

As described in Section 2, the velocity field associated to the flow can be inferred from arbitrary densities $\rho_{0},\rho_{1}$ – this deviates from the score-based diffusion perspective, in which one distribution must be taken to be Gaussian for the training paradigm to be tractable.

In Figure 2, we illustrate this capacity by learning the velocity field connecting the anisotropic swirls distribution to that of the checkerboard. The interpolant formulation allows us to draw samples from $\rho_{t}$ at any time $t\in$, which we exploit to check that the velocity field is empirically correct at all times on the interval, rather than just at the end points. This aspect of interpolants is also noted in (Choi et al., 2022), but for the purpose of density ratio estimation. The above observation highlights an intrinsic difference of the proposed method compared to MLE training of flows, where the map that is the minimizer of $G(\hat{v})$ is not empirically known. We stress that query access to $\rho_{0}$ or $\rho_{1}$ is not needed to use our interpolation procedure since it only uses samples from these densities.

3 Tabular data for higher dimensional testing

A set of tabular datasets introduced by (Papamakarios et al., 2017) has served as a consistent test bed for demonstrating flow-based sampling and its associated density estimation capabilities. We continue that practice here to provide a benchmark of the method on models which provide an exact likelihood, separating and comparing to exemplary discrete and continuous flows: MADE (Germain et al., 2015), Real NVP (Dinh et al., 2017), Convex Potential Flows (CPF) (Huang et al., 2021), Neural Spline Flows (NSP) Durkan et al. (2019), Free-form continuous flows (FFJORD) (Grathwohl et al., 2019), and OT-Flow (Finlay et al., 2020). Our primary point of comparison is to other continuous time models, so we sequester them in benchmarking.

We train the interpolant flow model on each target dataset listed in Table 2, choosing the reference distribution of the interpolant $\rho_{0}$ to be a Gaussian density with mean zero and variance $I_{d}$, where $d$ is the data dimension. The architectures and hyperparameters are given in Appendix I. We highlight some of the main characteristics of the models here. In each case, sampling of the time $t$ was reweighted according to a beta distribution, with parameters $\alpha,\beta$ provided in the same appendix.

Results from the tabular experiments are displayed in Table 2, in which the negative log-likelihood averaged over a test set of held out data is computed. We note that the interpolant flow achieves better or equivalent held out likelihoods on all ODE based models, except BSDS300, in which the FFJORD outperforms the interpolant by $\sim 0.6\%$. We note upwards of $30\%$ improvements compared to baselines. Note that these likelihoods are achieved without direct optimization of it.

4 Unconditional Image Generation

To compare with recent advances in continuous time generative models such as DDPM (Ho et al., 2020), Score SDE(Song et al., 2021b), and ScoreFlow (Song et al., 2021a), we provide a demonstration of the interpolant flow method on learning to unconditionally generate images trained from the CIFAR-10 (Krizhevsky et al., 2009) and ImageNet 32$\times$32 datasets (Deng et al., 2009; Van Den Oord et al., 2016), which follows suit with ScoreFlow and Variational Diffusion Models (VDM) (Kingma et al., 2021). We train an interpolant flow built from the U-Net architecture from DDPM (Ho et al., 2020) on a single NVIDIA A100 GPU, which was previously impossible under maximum likelihood training of continuous time flows. Experimental details can be found in Appendix I. Note that we used a beta distribution reweighting of the time sampling as in the tabular experiments. Table 2 provides a comparison of the negative log likelihoods (NLL), measured in bits per dim (BPD) and Frechet Inception Distance (FID), of our method compared to past flows and state of the art diffusions. We focus our comparison against models which emit a likelihood, as this is necessary to compare NLL.

We compare to other flows FFJORD and Glow (Grathwohl et al., 2019; Kingma & Dhariwal, 2018), as well as to recent advances in score based diffusion in DDPM, DDPM++, VDM, Score SDE, ScoreFlow, and Soft Truncation (Ho et al., 2020; Nichol & Dhariwal, 2021; Kingma et al., 2021; Song et al., 2021b; a; Kim et al., 2022). We present results without data augmentation. Our models emit likelihoods, measured in bits per dim, that are competitive with diffusions on both datasets with a NLL of $2.99$ and $3.45$. Measures of FID are proximal to those from diffusions, though slightly behind the best results. We note however that this is a first example of this type of model, and has not been optimized with the training tricks that appear in many of the recent works on diffusions, like exponential moving averages, truncations, learning rate warm-ups, and the like. To demonstrate efficiency on larger domains, we train on the Oxford flowers dataset (Nilsback & Zisserman, 2006), which are images of resolution 128$\times$128. We show example generated images in Figure 3.

Discussion, Challenges, and Future work

We introduced a continuous time flow method that can be efficiently trained. The approach has a number of intriguing and appealing characteristics. The training circumvents any backpropagation through ODE solves, and emits a stable and interpretable quadratic objective function. This objective has an easily accessible diagnostic which can verify whether a proposed minimizer of the loss is a valid minimizer, and controls the Wasserstein-2 distance between the model and the target.

One salient feature of the proposed method is that choosing an interpolant $I_{t}(x_{0},x_{1})$ decouples the optimization problem from that of also choosing a transport path. This separation is also exploited by score-based diffusion models, but our approach offers better explicit control on both. In particular we can interpolate between any two densities in finite time and directly obtain the probability flow needed to calculate the likelihood. Moreover, we showed in Section 2 and Appendices D and G that the interpolant can be optimized to achieve optimal transport, a feature which can reduce the cost of solving the ODE to draw samples. In future work, we will investigate more thoroughly realizations of this procedure by learning the interpolant $I_{t}$ in a wider class of functions, in addition to minimizing $G(\hat{v})$.

The intrinsic connection to score-based diffusion presented in Proposition 4 may be fruitful ground for understanding the benefits and tradeoffs of SDE vs ODE approaches to generative modeling. Exploring this relation is already underway (Lu et al., 2022; Boffi & Vanden-Eijnden, 2022), and can hopefully provide theoretical insight into designing more effective models.

We thank Gérard Ben Arous, Nick Boffi, Kyle Cranmer, Michael Lindsey, Jonathan Niles-Weed, Esteban Tabak for helpful discussions about transport. MSA is supported by the National Science Foundation under the award PHY-2141336. MSA is grateful for the hospitality of the Center for Computational Quantum Physics at the Flatiron Institute. The Flatiron Institute is a division of the Simons Foundation. EVE is supported by the National Science Foundation under awards DMR-1420073, DMS-2012510, and DMS-2134216, by the Simons Collaboration on Wave Turbulence, Grant No. 617006, and by a Vannevar Bush Faculty Fellowship.

References

Appendix A Background on transport maps and the continuity equation

The following result is standard and can be found e.g. in (Villani, 2009; Santambrogio, 2015)

Let $\rho_{t}(x)$ satisfy the continuity equation

Assume that $v_{t}(x)$ is $C^{1}$ in both $t$ and $x$ for $t\geq 0$ and globally Lipschitz in $x$. Then, given any $t,t^{\prime}\geq 0$, the solution of (A.1) satisfies

where $X_{s,t}$ is the probability flow solution to

In words, Lemma A.1 states that an evaluation of the PDF $\rho_{t}$ at a given point $x$ may be obtained by evolving the probability flow equation (2) backwards to some earlier time $t^{\prime}$ to find the point $x^{\prime}$ that evolves to $x$ at time $t$, assuming that $\rho_{t^{\prime}}(x^{\prime})$ is available. In particular, for $t^{\prime}=0$, we obtain

The assumed $C^{1}$ and globally Lipschitz conditions on $v_{t}$ guarantee global existence (on $t\geq 0$) and uniqueness of the solution to (2). Differentiating $\rho_{t}(X_{t^{\prime},t}(x))$ with respect to $t$ and using (2) and (A.1) we deduce

Integrating this equation in $t$ from $t=t^{\prime}$ to $t=t$ gives

Evaluating this expression at $x=X_{t,t^{\prime}}(x)$ and using the group properties (i) $X_{t^{\prime},t}(X_{t,t^{\prime}}(x))=x$ and (ii) $X_{t^{\prime},s}(X_{t,t^{\prime}}(x))=X_{t,s}(x)$ gives (A.2). Equation (A.4) can be derived by using (A.2) to express $\rho_{t}(x)$ in the integral at the left hand-side, changing integration variable $x\to X_{t^{\prime},t}(x)$ and noting that the factor $\exp\left(-\int_{t^{\prime}}^{t}\nabla\cdot v_{s}(X_{t,s}(x))ds\right)$ is precisely the Jacobian of this change of variable. The result is the integral at the right hand-side of (A.4). ∎

Appendix B Proof of Proposition 1

We will work under the following assumption:

The densities $\rho_{0}(x)$ and $\rho_{1}(x)$ are continuously differentiable in $x$; $I_{t}(x_{0},x_{1})$ is continuously differentiable in $(t,x_{0},x_{1})$ and satisfies (4) and (7); and for all $t\in$ we have

This structural assumption guarantees that the stochastic interpolant $x_{t}$ has a probability density and a probability current. As shown in Appendix C it is satisfied e.g. if $\rho_{0}$ and $\rho_{1}$ are Gaussian mixture densities and

where $a_{t}$ and $b_{t}$ are $C^{1}$ function of $t\in$ satisfying

The interpolant (4) is in this class for the choice

Our proof of Proposition 1 will rely on the following result that quantifies the probability density and the probability current of the stochastic interpolant $x_{t}$ defined in (6):

If Assumption B.1 holds, then the stochastic interpolant $x_{t}$ defined in (6) has a probability density function $\rho_{t}(x)$ given by

withe the probability current $j_{t}(x)$ given by

Note that (B.8) and (B.9) can be formally rewritten as (12) and (14) using the Dirac delta distribution.

By definition of $x_{t}$ in (6), the characteristic function of this random variable is

Under Assumption B.1, the Fourier inversion theorem implies that $x_{t}$ has a density $\rho_{t}(x)$ given by (B.5). Taking the time derivative of this density gives

Lemma B.1 shows that $\rho_{t}(x)$ satisfies the continuity equation (3) with the velocity field $v_{t}(x)$ defined in (15). It also show that the objective function $G(\hat{v})$ in (9) is well-defined, which implies the first part Proposition 1.

Then, using the pointwise identity $v_{t}(x)\rho_{t}(x)=j_{t}(x)$ as well as (B.8) and (B.9), we can write

where we use the shorthand $I_{t}=I_{t}(x_{0},x_{1})$ and $\partial_{t}I_{t}=\partial_{t}I_{t}(x_{0},x_{1})$. Therefore

Lemma B.2 implies that the objective $H(\hat{v})$ in (17) is well-defined, and the second part of statement of Proposition 1 follows from the argument given after (17). For the third part of the proposition we can then proceed as explained in main text, starting from the Poisson equation (B.21).

The interpolant density $\rho_{t}(x)$ and the current $j_{t}(x)$ are given explicitly in Appendix C in the case where $\rho_{0}$ and $\rho_{1}$ are both Gaussian mixture densities and we use the linear interpolant (B.2).

Notice that we can evaluate $v_{t=0}(x)$ and $v_{t=1}(x)$ more explicitly. For example, with the linear interpolant (B.2) we have

For the trigonometric interpolant that uses (B.4) these reduce to

Finally, we note that the result of Proposition 1 remains valid if we work with velocities that are gradient fields. We state this result as:

Interestingly, with $v_{t}(x)$ defined in (15) and $\phi_{t}(x)$ defined as the solution to (B.21), we have

If we assume that $\hat{v}_{t}(x)=\nabla\hat{\phi}_{t}(x)$, we can write the objective in (16) as

where we integrated by parts in $x$ to get the second equality, we used the continuity equation (13) to get the second, and integrated by parts in $t$ to get the third using $\rho_{t=0}=\rho_{0}$ and $\rho_{t=1}=\rho_{1}$. Since the objective in the last expression is an expectation with respect to $\rho_{t}$, we can evaluate it as

Appendix C The case of Gaussian mixture densities

Here we consider the case where $\rho_{0}$ and $\rho_{1}$ are both Gaussian mixture densities. We denote

Note that if all the covariance matrices are the same, $C_{i}^{0}=C_{j}^{1}=C$, with the trigonometric interpolant in (5) we have $C_{t}^{ij}=C$, which justifies this choice of interpolant.

The interpolant density $\rho_{t}(x)$ obtained by connecting the probability densities in (C.2) using the linear interpolant (B.2) is given by

and it satisfies the continuity equation $\partial_{t}\rho_{t}(x)+\nabla\cdot j_{t}(x)=0$ with the current

This velocity field is growing at most linearly in $x$, and when the mode of the Gaussian are well separated, in each mode it approximately reduces to

Using the Fourier representation in (C.1) and proceeding as in the proof of Lemma B.1, we deduce that $\rho_{t}(x)$ is given by

Performing the integral over $k$ gives (C.4). Taking the time derivative of this density gives

Appendix D Optimizing transport through stochastic interpolants

Using the velocity $v_{t}(x)$ in (15) that minimizes the objective (9) in the ODE (2) gives an exact transport map $T=X_{t=1}$ from $\rho_{0}$ to $\rho_{1}$. However, this map is not optimal in general, in the sense that it does not minimize

over all $\hat{T}$ such that $\hat{T}\sharp\rho_{0}=\rho_{1}$. It is easy to understand why: In their seminal paper Benamou & Brenier (2000) showed that finding the optimal map requires solving the minimization problem

As also shown in (Benamou & Brenier, 2000), the velocity minimizing (D.2) is a gradient field, $v^{*}_{t}(x)=\nabla\phi^{*}_{t}(x)$, and the minimizing couple $(\rho^{*}_{t},\phi^{*}_{t})$ is unique and satisfies

Minimizing (9) over gradient fields reduces the value of the objective in (D.2), but it does not yield an optimal map either—indeed the gradient velocity $\nabla\phi_{t}(x)$ with the potential $\phi_{t}(x)$ solution to (B.21) only minimizes the objective in (D.2) over all $\hat{v}_{t}(x)$ with $\hat{\rho}_{t}(x)=\rho_{t}(x)$ fixed, as explained after (B.23).

It is natural to ask whether our stochastic interpolant construction can be amended or generalized to derive optimal maps. This question is discussed next, from two complementary perspectives: via optimization of the interpolant $I_{t}(x_{0},x_{1})$, and/or via optimization of the base density $\rho_{0}$, assuming that we have some leeway to choose this density.

Since (10) indicates that the minimum of $G(\hat{v})$ is lower bounded by minus the value of the objective in (D.2), one way to optimize the transport is to maximize this minimum over the interpolant. Under some assumption on the Benamou-Brenier density $\rho^{*}_{t}(x)$ solution of (D.3), this procedure works, as we show now. Let us begin with a definition:

Interpolable densities form a wide class, as discussed e.g. in Mikulincer & Shenfeld (2022). These densities also are the ones that can be learned by score-based diffusion modeling discussed in Section 2.2. They are useful for our purpose because of the following result showing that any interpolable density can be represented as an interpolant density:

Let $\rho_{t}(x)$ be an interpolable density in the sense of Definition D.1. Then

satisfies (4) and is such that the stochastic interpolant defined in (6) satisfies $x_{t}\sim\rho_{t}$.

We stress that the interpolant in (D.4) is in general not the only one giving the interpolable density $\rho_{t}(x)$, and the actual value of the map $T_{t}$ plays no role in the results below.

so that the boundary condition in (4) are satisfied. Then observe that, by definition of $T_{t}$,

This implies that, if $x_{0}\sim\rho_{0}$, $x_{1}\sim\rho_{1}$, and they are independent,

is a Gaussian random variable with mean zero and covariance

i.e. it is a sample from $N(0,\text{Id})$. By definition of $T_{t}$, this implies that

is a sample from $\rho_{t}$ if $x_{0}\sim\rho_{0}$, $x_{1}\sim\rho_{1}$, and they are independent. ∎

Assume that (i) the optimal density function $\rho_{t}^{*}(x)$ minimizing (D.2) is interpolable and (ii) (D.3) has classical solution. Consider the max-min problem

where $G(\hat{v})$ is the objective in (9) and the maximum is taken over interpolants satisfying (4). Then a maximizer of (D.10) exists, and any maximizer $I_{t}^{*}(x_{0},x_{1})$ is such that the probability density function of $x^{*}_{t}=I_{t}^{*}(x_{0},x_{1})$, with $x_{0}\sim\rho_{0}$ and $x_{1}\sim\rho_{1}$ independent, is the optimal $\rho_{t}^{*}(x)$, the mimimizing velocity is $v_{t}^{*}(x)=\nabla\phi_{t}^{*}(x)$, and the pair $(\rho_{t}^{*}(x),\phi_{t}^{*}(x))$ satisfies (D.3).

The proof of Proposition D.2 relies on the following simple reformulation of (D.2):

The Benamou-Brenier minimization problem in (D.2) is equivalent to the min-max problem

In particular, under the conditions on $\rho_{0}$ and $\rho_{1}$ such that (D.2) has a minimizer, the optimizer $(\rho^{*}_{t},v^{*}_{t},j^{*}_{t})$ is unique and satisfies $v^{*}_{t}(x)=\nabla\phi^{*}_{t}(x)$, $j^{*}_{t}(x)=\nabla\phi^{*}_{t}(x)\rho^{*}_{t}(x)$ with $(\rho^{*}_{t},\phi^{*}_{t})$ solution to (D.3).

Since (D.11) is convex in $\hat{v}$ and concave in $(\hat{\rho},\hat{\jmath})$, the min-max and the max-min are equivalent by von Neumann’s minimax theorem. Considering the problem where we minimize over $\hat{v}_{t}(x)$ first, the minimizer must satisfy:

Since $\hat{\rho}_{t}(x)\geq 0$ and $\hat{\jmath}_{t}(x)$ have the same support by the constraint in (D.11), the solution to this equation is unique on this support. Using (D.12) in (D.11), we can therefore rewrite the max-min problem as

This problem is equivalent to the Benamou-Brenier minimization problem in (D.2).

To write the Euler-Lagrange equations for the optimizers of the min-max problem (D.11), let us introduce the extended objective

These equations imply that $v^{*}_{t}(x)=\nabla\phi^{*}_{t}(x)$, $j^{*}_{t}(x)=v^{*}_{t}(x)\rho^{*}_{t}(x)=\nabla\phi^{*}_{t}(x)\rho^{*}_{t}(x)$, with $(\rho^{*}_{t},\phi^{*}_{t})$ solution to (D.3), as stated in the lemma. Since the optimization problem in (D.11) is convex in $\hat{v}$ and concave in $(\hat{\rho},\hat{\jmath})$, its optimizer is unique and solves these equations. ∎

We can reformulate the max-min problem (D.10) as

where the maximization is taken over probability density functions $\hat{\rho}_{t}(x)$ and probability currents $\hat{\jmath}_{t}(x)$ as in (B.5) and (B.7) with $I_{t}$ replaced by $\hat{I}_{t}$, i.e. formally given in terms of the Dirac delta distribution as

the max-min problem (D.16) is similar to the one considered in Lemma (D.3), except that the maximization is taken over probability density functions and associated currents that can be written as in (D.17). By Proposition D.1, this class is large enough to represent $\rho_{t}^{*}(x)$ since we have assumed that $\rho_{t}^{*}(x)$ is an interpolable density, and the statement of the proposition follows. ∎

Since our primary aim here is to construct a map $T=X_{t=1}$ that pushes forward $\rho_{0}$ onto $\rho_{1}$, but not necessarily to identify the optimal one, we can perform the maximization over $\hat{I}_{t}(x_{0},x_{1})$ in a restricted class (though of course the corresponding map is no longer optimal in that case). We investigate this option in numerical examples in Appendix H, using interpolants of the type (B.2) and maximizing over the functions $a_{t}$ and $b_{t}$, subject to $a_{0}=b_{1}=1$, $a_{1}=b_{0}=0$. In Section G we also discussion generalizations of the interpolant that can render the optimization of the transport easier to perform. We leave the full investigation of the consequences of Proposition D.2 for future work.

and a calculation similar to the one presented in Appendix C shows that the associated velocity field $v_{t}(x)$ is

This is the velocity giving the optimal transport map $X_{t}(x)=x+(m_{1}-m_{0})t$.

In (Liu et al., 2022; Liu, 2022), where an approach similar to ours using the linear interpolant $x_{t}=x_{0}(1-t)+x_{1}t$ is introduced, an alternative procedure is proposed to optimize the transport. Specifically, it is suggested to rectify the map $T=X_{t=1}$ learned by repeating the procedure using the new interpolant $x^{\prime}_{t}=x_{0}(1-t)+T(x_{0})t$ with $x_{0}\sim\rho_{0}$. As shown in (Liu, 2022) iterating on this rectification step yields successive maps that are getting closer to optimality. The main drawback of this approach is that it requires each of these maps (including the first one) to be learned exactly, i.e. we must have $T\sharp\rho_{0}=\rho_{1}$ to use the interpolant $x^{\prime}_{t}$ above. If the maps are not exact, which is unavoidable in practice, the procedure introduces a bias whose amplitude will grow with the iterations, leading to instability (as noted in Liu et al. (2022); Liu (2022)).

D.2 Optimizing the base density

Optimizing $\rho_{0}$ only makes practical sense if we do so in a restricted class, like that of Gaussian densities that we just discussed. Still, we may wonder whether optimizing $\rho_{0}$ over all densities would automatically give $\rho_{0}=\rho_{1}$ and $v_{t}=0$. If the interpolant is fixed, the answer is no, in general. Indeed, even if we set $\rho_{0}=\rho_{1}$, the interpolant density will still evolve, i.e. $\rho_{t}\not=\rho_{0}$ except at $t=0,1$, in general. This indicates that optimizing the interpolant in concert with $\rho_{0}$ is necessary if we want to optimize the transport.

Appendix E Proof of Proposition 3

To use the bound in (27), let us consider the evolution of

Using $\dot{X}_{t}(x)=v_{t}(X_{t}(x))$ and $\dot{\hat{X}}_{t}(x)=\hat{v}_{t}(\hat{X}_{t}(x))$, we deduce

Therefore, by Gronwall’s inequality and since $Q_{0}=0$ we deduce

Since $W_{2}^{2}(\rho_{1},\hat{\rho}_{1})\leq Q_{1}$ by (27), we are done. $\square$

Note that the proposition suggests to regularize $G(\hat{v})$ using e.g.

with some small $\lambda>0$. In the numerical results presented in the paper no such regularization was included.

Appendix F Proof of Proposition 4 and link with score-based diffusion models

Assume that the interpolant is of the type (B.2) so that $\partial_{t}I_{t}(x_{0},x_{1})=\dot{a}_{t}x_{0}+\dot{b}_{t}x_{1}$. For $t\in(0,1)$ let us write expression (14) for the probability current as

If $\rho_{0}(x_{0})=(2\pi)^{-d/2}e^{-\frac{1}{2}|x_{0}|^{2}}$, we have the identity $x_{0}\rho_{0}(x_{0})=-\nabla_{x_{0}}\rho_{0}(x_{0})$. Inserting this equality in the last integral in (F.1) and integrating by part using

Solving this expression in $\nabla\log\rho_{t}(x)$ and specializing it to the trigonometric interpolant with $a_{t}$, $b_{t}$ given in (B.4) gives the first equation in (28). The second one can be obtained by taking the limit of this first equation using $v_{t=1}(x)=0$ from (B.20) and l’Hôpital’s rule. $\square$

Note that (F.3) shows that, when the interpolant is of the type (B.2) and $\rho_{0}(x_{0})=(2\pi)^{-d/2}e^{-\frac{1}{2}|x_{0}|^{2}}$, the continuity equation (3) can also be written as the diffusion equation

Since we assume that $\dot{a}_{t}\leq 0$ and $\dot{b}_{t}\geq 0$ (see (B.3), the diffusion coefficient in this equation is negative

This means that (F.5) is well-posed backward in time, i.e. it corresponds to backward diffusion from $\rho_{t=1}=\rho_{1}$ to $\rho_{t=0}=\rho_{0}=(2\pi)^{-d/2}e^{-\frac{1}{2}|x_{0}|^{2}}$. Therefore, reversing this backward diffusion, similar to what is done in score-based diffusion models, gives an SDE that transforms samples from $\rho_{0}$ to $\rho_{1}$. Interestingly, these forward and backward diffusion processes arise on the finite time interval $t\in$; notice however that both the drift and the diffusion coefficient are singular at $t=1$. This is unlike the velocity $v_{t}(x)$ which is finite at $t=0,1$ and is given by (B.19).

Appendix G Generalized interpolants

Our construction can be easily generalized in various ways, e.g. by making the interpolant depend on additional latent variables to be averaged upon. This enlarges the class of interpolant density we can construct, which may prove useful to get simpler (or more optimal) velocity fields in the continuity equation (3). Let us consider one specific generalization of this type:

Suppose that we decompose both $\rho_{0}$ and $\rho_{1}$ as

This corresponds to splitting the samples from $\rho_{0}$ and $\rho_{1}$ into $K$ (soft) clusters, and only interpolating between samples in cluster $k$ in $\rho_{0}$ and samples in cluster $k$ in $\rho_{1}$. This clustering can either be done beforehand, based on some prior information we may have about $\rho_{0}$ and $\rho_{1}$, or be learned (more on this point below).

It is easy to see that the PDF $\rho_{t}$ of $x_{t}$ is formally given by

and that this density satisfies the continuity equation (B.11) for the current

Therefore this equation can be written as (3) with the velocity $v_{t}(x)$ which is the unique minimizer of a generalization of the objective (9). We state this as:

The stochastic interpolant $x_{t}$ defined in (G.2) with $I^{k}_{t}(x_{0},x_{1})$ satisfying (4) for each $k=1,\ldots,K$ has a probability density $\rho_{t}(x)$ that satisfies the continuity equation (3) with a velocity $v_{t}(x)$ which is the unique minimizer over $\hat{v}_{t}(x)$ of the objective

where we used the shorthand notations $I_{t}^{k}=I_{t}^{k}(x_{0},x_{1})$ and $\partial_{t}I_{t}^{k}=\partial_{t}I_{t}^{k}(x_{0},x_{1})$. In addition the minimum value of this objective is given by

As discussed in Section 2, the minimizer of the objective in equation (10) can be maximized with respect to the interpolant $I_{t}$ as a route toward optimal transport. We motivate this by choosing a parametric class for the interpolant $I_{t}$ and demonstrating that the max-min optimization in equation (D.10) can give rise to velocity fields which are easier to learn, resulting in better likelihood estimation. The 2D checkerboard example is an appealing test case because the transport is nontrivial. In this case, we train the same flow as in Section 3.1, with and without optimizing the interpolant. We choose a simple parametric class for the interpolant given by a Fourier series expansion

where parameters $\{\alpha,\beta\}_{m=1}^{M}$ define learnable $\hat{a}_{t},\hat{b}_{t}$ via

This is but one set of possible parameterizations. Another, for example, could be a rational quadratic spline, so that the endpoints for $\hat{a}_{t}$, $\hat{b}_{t}$ can be properly constrained as they are above in the Fourier expansion. In Figure H.1, we plot the log likelihood, the learned interpolants $\hat{a}_{t}$, $\hat{b}_{t}$ compared to their initializations, as well as the path length as it evolves over training epochs. With the learned interpolants, the path length is reduced, and the resultant velocity field under the same number of training epochs endows a model with better likelihood estimation, as shown in the left plot of the figure. For the path optimality experiments, $M=7$ Fourier coefficients were used to parameterize the interpolant. This suggests that minimizing the transport cost can create models which are, practically, easier to learn.

In addition to parameterizing the interpolant, one can also parameterize the base density $\rho_{0}$ in some simple parametric class. We show that including this in the min-max optimization given in equation (D.10) can further reduce the transport cost and improve the final log likelihood. The results are given in Figure H.2. We train an interpolant just as described above, but also allow the base density $\rho_{0}$ to be parameterized as a Gaussian with mean $\hat{\mu}$ and covariance $\hat{\Sigma}$. The inclusion of the learnable base density results in a significantly reduced path length, thereby bringing the model closer to optimal transport.

Appendix I Implementation Details for Numerical Experiments

Let $\{x_{0}^{i}\}_{i=1}^{N}$ be $N$ samples from the base density $\rho_{0}$, $\{x_{1}^{j}\}_{i=1}^{n}$ $n$ samples from the target density $\rho_{1}$, and $\{t_{k}\}_{k=1}^{K}$ $K$ samples from the uniform density on $$. Then an empirical estimate of the objective function in (9) is given by

The architectural information and hyperparameters of the models for the resulting likelihoods in Table 2 is presented in Table 3. ReLU (Nair & Hinton, 2010) activations were used throughout, barring the BSDS300 dataset, where ELU (Clevert et al., 2016) was used. Table formatting based on (Durkan et al., 2019).

In addition, reweighting of the uniform sampling of time values in the empirical calculation of (I.1) was done using a Beta distribution under the heuristic that the flow should be well trained near the target. This is in line with the statements under Proposition 1 that any weight $w(t)$ maintains the same minimizer.

The details for the image datasets are provided in Table 4. We built our models based off of the U-Net implementation provided by lucidrains public diffusion code, which we are grateful for https://github.com/lucidrains/denoising-diffusion-pytorch. We use the sinusoidal time embedding, but otherwise use the default implementation other than changing the U-Net dimension multipliers, which are provided in the table. Like in the tabular case, we reweight the time sampling to be from a beta distribution. All models were implemented on a single A100 GPU.

Below, we show that the results achieved in 3 are driven by a model that can train significantly more efficiently than the maximum likelihood approach to ODEs. Following that, we provide an illustration of the convergence requirements on the objective defined in (11).

We briefly describe the experimental setup for testing the computational efficiency of our model as compared to the FFJORD maximum likelihood learning method. We use the same network architectures for the interpolant flow and the FFJORD implementation, taking the architectures used in that paper. For the Gaussian case, this is a 3 layer neural network with hidden widths of 64 units; for the 43-dimensional MiniBooNE target, this is a 3 layer neural network with hidden widths of 860 units.

Figure I.1 shows a comparison of both the cost per training epoch and the convergence of the log likelihood across epochs. We take the architecture of the vector field as defined in the FFJORD paper for the 2-dimensional Gaussian and MiniBooNE, and use it to define the vector field for the interpolant flow. The left side of Figure I.1 shows that the cost per iteration is constant for the interpolant flow, while it grows for MLE based approaches as the ODE gets more complicated to solve. The speedup grows with dimension, 400x on MiniBooNE.

The right side of Figure I.1 shows that, under equivalent conditions, the interpolant flow can converge faster in number of training steps, in addition to being cheaper per step. The extent of this benefit is dependent on both hyperparameters and the dataset, so a general statement about convergence speeds is difficult to make. For the sake of this comparison, we averaged over 5 trials for each model and dataset, with variance shown shaded around the curves.