Stochastic Interpolants

These notes are based on Albergo, Boffi, and Vanden-Eijnden, “Stochastic Interpolants: A Unifying Framework for Flows and Diffusions.”

The main idea behind Stochastic Interpolants is that you want to learn a transport, which could be deterministic or stochastic, between two probability densities \(\rho_0, \rho_1\). The Stochastic Interpolant is simply a continuous-time stochastic process that bridges the two and is of the form:

where \((x_0, x_1)\) is sampled from a coupling with density \(\rho(x_0, x_1)\) and \(z \sim \mathcal{N}(0, 1) \perp (x_0, x_1)\). Note that you can sample \(I_t\) very easily. The most important realization is that the time-dependent marginal density of \(I_t\), denoted by \(p_t(x)\), satisfies a transport equation (giving an ODE-based generative model) and also a Fokker-Planck equation for any choice of diffusion coefficient (giving a family of SDE-based generative models with the same marginals). So the same \(p_t\) can be generated by infinitely many processes, an ODE, or any of a family of SDEs, and all we have to do is empirically learn a few important properties of \(I_t\), namely its velocity field and score function, which together determine all of them.

Derivation of the Transport Equation

Let’s first take the simpler case of \(\gamma(t) = 0\), meaning no added noise, so \(I_t\) is a deterministic interpolation of the coupling. Note that this case is still “stochastic” because the randomness still comes from the coupling \((x_0, x_1)\sim\rho\), just not from \(z\). This sharpens the contrast with CNFs, where the randomness is only in \(x_0\) and then the flow is deterministic.

One thing we’d like to prove is that the time-dependent density \(p_t(x)\) satisfies a Transport Equation

for some velocity \(v_t(x)\). Notice here that for a given point \((x, t)\), there is no longer a unique velocity for the particles moving in time for various pairs of \((x_0, x_1)\). For example, for the linear interpolant \(I_t = (1-t)x_0 + tx_1\), the velocity along a sample path is \(x_1 - x_0\), which depends on the specific coupling draw, so two paths passing through the same \((x, t)\) can have different velocities. And so we can’t apply the nimble toggling of Eulerian and Lagrangian Formulations. Moreover, we don’t have a nice bijective transformation \(\mathbb{R}^d \rightarrow \mathbb{R}^d\) along a particle path like in the CNF setup so it wouldn’t be a good idea to talk about Jacobians. We’ll try a different approach instead.

Let’s start by forming an intuitive picture of how one would calculate \(p_t(x)\). Given a region \(\mathcal{R}\subset \mathbb{R}^d\), notice that

Now, to get the density \(p_t(x)\), we imagine shrinking the region \(\mathcal{R}\) until it collapses to \(\{x\}\). Of course, plugging \(\mathcal{R} = \{x\}\) would just give \(0\). But it does point us to the intuition that we should replace the indicator function \(1_{I_t(x_0, x_1)\in \mathcal{R}}\) with an object that picks out density values when integrated against. Only one object comes to mind: the Dirac Delta Function. Thus, we should expect the following result:

Indeed, this relationship turns out to be correct. Ok, this is good. But what we want, to prove the transport equation, is \(\partial_t p_t(x)\), which requires us to differentiate inside \(\delta(\cdot)\). But the Dirac Delta Function is not really a function, it’s a distribution. For our discussion, a distribution is something we integrate a function against, such as \(\rho(x_0, x_1)\), to spit out a real number such as \(p_t(x)\). In other words, it’s a functional. So we can think of the following equation like so:

where the functional \(F_{t, x}\) is parametrized by \((t,x)\) and what we want is

but it’s not clear how to find \(\partial_t F_{t, x}(\rho(x_0, x_1))\).

Fourier Transforms come to the rescue! It turns out that we can represent

Intuitively, for \(x=0\), we’re integrating a constant over an infinite range so it blows up. Otherwise, the integrand cycles on the complex plane and things cancel out. Of course, this is not a rigorous argument. Rigorously, this is the statement that \(\delta\) is the inverse Fourier transform of the constant function, in the distributional sense. In our case, I like to think of the integral formula as notational convenience to manipulate objects containing \(\delta\) in their formulations, like \(F_{t, x}\) above and compute their derivatives, which are defined and do make sense. Let’s verify this. We’ll try to compute \(\partial_t p_t(x)\) using this substitution.

This implies that

Now, notice that

This means we can write \(\partial_t p_t(x) = -\nabla \cdot j_t(x)\) where

The object \(j_t(x)\) can look intimidating. But the point is that we were able to express \(\partial_t p_t(x)\) as the divergence of another function of \((t, x)\). Finally, define

and you have your transport equation:

Now of course, we still don’t have a clear physical intuition on what \(v_t(x)\) is or, most importantly, how to empirically estimate it from data. That’s what we’ll do next.

The punch line of this write up is to gather intuition on why we had to use Fourier Transforms to derive the Transport Equation. The writeup above is not strictly rigorous, and the Stochastic Interpolants paper starts from Fourier Transforms (or characteristic function), and moves in reverse. Though when I first looked at the proof, I was baffled at the thought of moving into the complex domain and using Fourier Transforms so suddenly. And this is my way of building intuition of how I would have come up with their approach if I were in their shoes.

Visualizations

Here are some examples of learned transport from a 2D Gaussian to various target distributions:

Interpreting and Estimating The Velocity Field

Ok, now that we have this object \(v_t(x)\), let’s gather some intuition on what it represents. First, recall the expression

which can be intuitively interpreted as aggregating the density of \(\rho(x_0, x_1)\) through the filter \(x = I_t(x_0, x_1)\) which gives us \(p_t(x)\). The expression of \(j_t(x)\)

is doing something very similar. Instead of simply summing up \(\rho(x_0, x_1)\) through the filter, it’s computing a weighted sum \(\partial_t I_t(x_0, x_1)\) through the filter with \(\rho(x_0, x_1)\) being the weights. This should immediately hint to us that \(j_t(x)\) resembles an expectation of \(\partial_tI_t(x_0, x_1)\) over some distribution. Now, note that as it stands, the weights \(\rho(x_0, x_1)\) filtered through \(\delta(x-I_t(x_0, x_1))\) don’t add up to \(1\). So to have a legitimate conditional expectation, we should re-normalize by dividing by the sum of the weights over the filter which is nothing but \(p_t(x)\)! This gives us exactly \(v_t(x)= j_t(x)/p_t(x)\)! So \(v_t(x)\) is indeed an expectation, but what’s the distribution? Well, it is nothing but \(\rho(x_0, x_1)\) restricted to \(I_t(x_0, x_1) = x\) and re-normalized so it’s the conditional expectation:

This is a beautiful result, not only because it’s very intuitive. But also, it has a close resemblance to Continuous Normalizing Flows. In the case of CNF, the velocity field of a particle was a unique deterministic function of \((x, t)\). In the case of Stochastic Interpolants, the velocity field is not unique given \((x, t)\) and there’s still inherent uncertainty to \(\partial_t I_t\). However, the problem naturally solves itself by taking an expectation!

Finally, our last task is to come up with the algorithm to estimate \(v_t(x)\) empirically. Because \(v_t(x)\) is a conditional expectation, and conditional expectations are characterized as \(L^2\) minimizers, the loss function is indeed a simple \(L^2\) expression. Regressing \(\partial_t I_t\) against a function of \(I_t\) under squared loss recovers \(v_t\) exactly. To see this clearly, let’s zoom out to the general setup of regression, where we have two variables \((X, Y)\sim \rho(x, y)\) and we’d like to minimize

over all functions \(f\) of \(X\). Notice that

Now, given that we can choose \(f(X)\) whatever we want given a fixed value of \(X\), optimizing \(J(f)\) reduces to pointwise optimization of \(f(X)\) for every \(X\). And we can think of the inner expectation as

for some constant \(a\) that we can choose to be whatever we want. By differentiating this with respect to \(a\) and setting to \(0\), we see that

All in all this tells us that the optimal function is \(f^{*} = \mathbb{E}[Y \mid X]\).

Going to \(v_t(x)\), by drawing this analogy with regression, we can set

And our objective becomes

where in the last step we realize that the joint distribution of \((I_t, \partial_t I_t)\) is generated by sampling \((x_0, x_1)\). And we’re done! This last equation can be computed empirically. Of course, during training we sample time \(t\) from a uniform distribution so the overall objective is then

To recap, what we’ve done essentially is to

1. Define a Stochastic Interpolant of our choosing that we can sample efficiently \(I_t = I_t(x_0, x_1)\)
2. Prove that time-dependent probability distribution \(p_t\) satisfies the Transport Equation
3. Finding an expression for the velocity field \(v_t(x)\) as a conditional expectation
4. Coming up with a simple quadratic objective \(G(\hat{v})\) where \(v = \mathrm{argmin}\, G(\hat{v})\)
5. And now we have a generative model \(x_0 \sim \rho_0, \dot{x}_t = \hat{v}(x_t, t)\), where we integrate this ODE from \(t=0\) to \(t=1\) to generate a sample from (approximately) \(\rho_1\)

The end product, \(\hat{v}(x, t)\) is the same as CNFs. But whereas in CNFs we didn’t have any restrictions on the velocity field and required two ODE integrations every training iteration. Here

1. We pre-specified the interpolation path before we even start training, and
2. We don’t require ODE integrations to learn \(\hat{v}\), just a simple quadratic loss