Continuous Normalizing Flows

Continuous normalizing flows operate in a very simple setup. You have \(x_0\sim\rho_0\), how can we transform this sample, via an invertible transformation, to a sample \(x_1\sim\rho_1\) where \(\rho_1\) is a distribution we want to sample from (which we have access to empirically).

First, let’s think of a deterministic function \(f_\theta(x)\) such that \(f_\theta(x_0)\sim\rho_1\) given that \(x_0\sim\rho_0\). This is how normalizing flows started where \(f_\theta\) is parametrized by a neural network. To train such a model, a crucial quantity to access is the likelihood of an empirical sample \(\hat{\rho}_1(x_1)\) if it were to be generated by the model. Notice here that we’re using the notation \(\hat{\rho}_1\) to designate the \(f_\theta(x_0)\sim\hat{\rho}_1, x_0\sim\rho_0\).

The computation of this quantity involves the Jacobian as you can easily check

\[\rho_0(x_0) = \hat{\rho}_1(x_1) \Big|\frac{\partial f_\theta(x_0)}{\partial x_0}\Big| \implies \hat{\rho}_1(x_1) = \rho_0(f_\theta^{-1}(x_1))\Big|\frac{\partial f_{\theta}}{\partial x}(f_\theta^{-1}(x_1))\Big|^{-1}\]

Now, since \(x_1\) is a high-dimensional vector (e.g, image), then the Jacobian is huge. And thus, people had to come up with clever architectures to make the computation of the Jacobian feasible. But this in turn restricted the class of functions \(f_\theta(x)\). The solution is to extend this framework to Continuous Normalizing Flows (CNF).

In CNF, instead of asking: “How can we transform \(x_0\rightarrow x_1\) in one step?”. We ask “How can I deterministically transform \(x_0\) continuously in time according to a dynamical system described by an ODE \(\dot{x}_t = f_\theta(x_t, t)\) where at \(t=1\) we have \(x_1\sim\rho_1\) ?”. The beauty behind this formulation is that we have decomposed the complexity of a single jump from \(x_0\rightarrow x_1\) into a time-varying velocity field that continuously, and gradually takes us there. More beautifully, the math works out so that we actually avoid computing the Jacobian. Let’s see how this is the case.

We now switch to using \(p_t(x)\) for the model’s time-evolving density under the flow, with \(p_0 = \rho_0\). At \(t=1\), \(p_1\) is the model’s push-forward density, which training aims to match to the data distribution \(\rho_1\).

First, let’s think about \(x_t = X(t, x_0)\) at time \(t\) and notice that, again, according to the change of measure formula we have

\[p_t(x_t) =p_0(x_0)\Big| \frac{\partial x_t}{\partial x_0}\Big|^{-1} \implies \mathrm{log}(p_t(x_t)) = \mathrm{log}(p_0(x_0)) - \mathrm{log}(|J(t)|)\]

where \(J(t)\) is the time-varying Jacobian, and we think of it, for now, as function of \(t\). Ok, now what? In such a setup (a deterministic ODE connecting two measures), there are two equivalent mathematical views: the Lagrangian view and the Eulerian view. For a rough mental model, take a look at Eulerian & Lagrangian Views.

In our case, the Eulerian view is represented by the probability flow PDE which is computationally prohibitive but mathematically very important. The Lagrangian view, on the other hand, is much more tractable computationally and is the backbone of likelihood computation for CNF training.

So let’s move forward with the Lagrangian view and compute the total derivative of density (or log-density) in time. We have

\[\frac{d\mathrm{log}(p_t(x_t))}{dt} = -\frac{d\mathrm{log}(|J(t)|)}{dt} = -\frac{1}{|J(t)|}\frac{d|J(t)|}{dt}\]

Now, let’s compute the total derivative of the Jacobian. To that end, note that for any matrix \(A\) we have

\[|A| = \sum_j A_{i, j}C_{i, j}\ \ \ \forall i \implies \frac{\partial|A|}{\partial A_{i, j}} = C_{i, j}\]

where \(C_{i, j} = \mathrm{Adj}(J)_{j, i}\) are the matrix cofactors since \(C_{i, j}\) don’t depend on \(A_{i, j}\). Applying this to \(\lvert J(t)\rvert\), we conclude

\[\frac{d|J(t)|}{dt} = \sum_{i, j} \frac{\partial |J|}{\partial J_{i, j}}\frac{\partial J_{i, j}(t)}{\partial t} = \sum_{i, j} \mathrm{Adj}(J)_{j, i}\frac{\partial J_{i, j}(t)}{\partial t} = \mathrm{Tr}\Big[\mathrm{Adj}(J) \frac{\partial J(t)}{\partial t}\Big] = |J(t)|\mathrm{Tr}\Big[J^{-1} \frac{\partial J(t)}{\partial t}\Big]\]

This implies that

\[\frac{d\mathrm{log}(|J(t)|)}{dt} = \mathrm{Tr}\Big[J^{-1}\frac{\partial J(t)}{\partial t}\Big]\]

Now, notice that

\[\frac{\partial J(t)}{\partial t} = \frac{\partial}{\partial t}\Big(\frac{\partial x_t}{\partial x_0}\Big) = \frac{\partial }{\partial x_0}\Big(\frac{\partial x_t}{\partial t}\Big) = \frac{\partial }{\partial x_0} \dot{x}_t = \frac{\partial f_\theta(x_t, t)}{\partial x_0} = \frac{\partial f_\theta(x_t, t)}{\partial x_t}\frac{\partial x_t}{\partial x_0} = \nabla f_\theta(x_t, t) J(t)\]

Putting it all together we conclude that

\[\frac{d\mathrm{log}(|J(t)|)}{dt} = \mathrm{Tr}\Big[J^{-1} \nabla f_{\theta}(x_t, t) J\Big] = \mathrm{Tr}\Big[\nabla f_\theta(x_t, t)\Big] = \nabla \cdot f_\theta(x_t, t)\]

And we finally get the Lagrangian equation

\[\frac{d \mathrm{log}(p_t(x_t))}{dt} = - \nabla \cdot f_\theta(x_t, t)\]

which we can integrate to get the likelihood. Denoting \(X_{s,t}(x)\) as the flow map that takes a point \(x\) at time \(s\) to its position at time \(t\) under the ODE, we obtain the canonical equation

\[\mathrm{log}(p_1(x_1)) = \mathrm{log}(p_0(X_{1, 0}(x_1))) - \int_0^1 \nabla \cdot f_{\theta}(X_{1, \tau}(x_1), \tau)d\tau\]

which lets us compute the log-likelihood used in training of CNFs. Wonderful! One final detail, let’s connect it to the Eulerian view. We’ll start from the Transport equation and then work backward.

\[\partial_t p_t(x) + \nabla \cdot (p_t(x) f_\theta(t, x)) = 0 \iff \partial_t p_t(x) + \nabla p_t(x)\cdot f_\theta(x, t) + p_t(x) \nabla \cdot f_\theta(x, t) = 0\]

Now, we plug our results above and we conclude that this is true if and only if

\[\partial_t p_t(x) + \nabla p_t(x)\cdot f_\theta(x, t) -\frac{dp_t(x_t)}{dt} = 0\]

Note how we switched the notation from \(x\) to \(x_t\), because we now think of \(x_t = X(t, x_0)\) as a specific particle originating from propagating the dynamical system in time from a specific initial condition \(x_0\)!

Now,

\[\frac{dp_t(x_t)}{dt} = \partial_t p_t(x_t) + \frac{\partial p_t(x_t)}{\partial x_t} \dot{x}_t = \partial_t p_t(x_t) + \nabla p_t(x_t)\cdot f_\theta(x_t, t)\]

Thus, all terms cancel out and the Transport equation holds as expected!

The Formal Equivalence of Transport: Eulerian PDE to Lagrangian ODE

The argument is simple. We have an Eulerian equation. We represent it as a particle path. Since the equation holds along particle paths, and we can represent every \((x,t)\) as lying on some particle path, the equation holds for all \((x,t)\). So the Eulerian equation holds.

That’s the whole thing. For a more detailed treatment of the conditions under which this toggle works, see Eulerian & Lagrangian Views.

The Adjoint Method

So far, we have formulated the setup of CNFs. Essentially, up to time-reversal which won’t be important for our discussions, we have a particle \(x(t)\) that evolves according to an ODE

\[\dot{x}_t = f_\theta(x(t), t),\ x(0) = x_0\]

Now, during training, we have a loss function, call it \(L\) that’s evaluated at \(x_1 = x(1)\). And we’d like to compute the gradient \(\frac{d L}{d \theta}\). Let’s try to derive this from first principles.

As a first try, let’s try to naively compute this quantity. But, it’s important to think how to represent \(L\) as a function. We’re asking: “How does changing the parameters \(\theta\) change the final loss value \(L\)?” Well, the mechanism through which it happens, is that by changing \(\theta\), we change the function \(f_\theta\) which changes the velocity field of the ODE which changes the entire trajectory leading to \(x_1\). Given this, it’s not unreasonable to write

\[x_1 = x_0 + \int_0^1 f_\theta(x(t, \theta), t)dt\]

And notice here that we wrote \(x(t, \theta)\) because we must not forget that the position of the particle at time \(t\) doesn’t only depend on time but also on the ODE which is controlled by \(\theta\). Then we’d have

\[\frac{d L}{d\theta} = \frac{dL}{dx_1} \cdot \frac{d}{d\theta}\int_0^1f_\theta(x(t, \theta), t)dt = \frac{dL}{dx_1} \ \cdot \int_0^1 \frac{d}{d\theta} f_\theta(x(t, \theta), t)dt\]

Ok, now \(\frac{dL}{dx_1}\) can be readily computable so let’s focus on the second term

\[\frac{d}{d\theta} f_\theta (x(t, \theta), t) = \partial_\theta f_\theta + \partial_x f_\theta \frac{\partial x(t, \theta)}{\partial \theta} = J_\theta(t) + J_x(t) \frac{\partial x(t)}{\partial \theta }\]

where we denote \(J_x, J_\theta\) as the Jacobians of the neural network with respect to the input and parameters, respectively. Now, do note that the Jacobians are, in principle, computable through automatic differentiation (though let’s worry about their cost later!). But, we’d be curious how to compute the object \(S(t) = \partial x(t)/\partial \theta\) . Notice that when we’re expressing objects as a function of time, we’re doing some mental gymnastics of switching our view of tracing the trajectory of a single particle. This is not trivial and takes some time to get used to. It gave me a great deal of confusion at first.

Now, having the view of \(S(t)\), it’s natural to derive an ODE that describes it since then we can compute it through numerical integration. Let’s do that.

\[\dot{S}(t) = \frac{d}{dt} \frac{\partial x(t, \theta)}{\partial \theta} = \partial_{\theta, t}x(t, \theta) = \frac{d}{d \theta} f_\theta = J_\theta(t) + J_x(t) S(t)\]

So we’ve found the ODE. In fact, this result should not be surprising. From the above, we can see immediately that

\[\frac{d L}{d\theta} = \frac{d L }{dx_1}\int_0^1 \dot{S}(t)dt\]

But also, seen a different way

\[\frac{d L}{d\theta} = \frac{d L }{dx_1}\frac{dx_1}{d\theta} = \frac{d L }{dx_1} S(1)\]

So it should come as no surprise that

\[S(1) = \int_0^1 \dot{S}(t)dt, \ S(0) = 0\]

Wonderful. Now, in principle, this gives us an algorithm. But is it easily computable? The answer is no. And that’s because integration involves computation of huge Jacobian matrices \(J_\theta, J_x\). In fact, \(S(t)\) itself is huge and impractical to materialize. Let’s take a different approach.

Parameterizing the Question in Time

Recall that we were trying to answer the question: “How does changing the parameters \(\theta\), by changing the full ODE trajectory from time \(0\rightarrow 1\) change the loss function at the end?”. A common technique in setups of dynamical systems is to parametrize the question in time, with our original question being only a special case of it. Given this, let’s ask the question: “How does changing the parameters \(\theta\), by changing the ODE trajectory from \(t\rightarrow 1\) change the loss function at the end?” This is now a function of \(t\). Let’s call it \(g(t)\). And our original question boils down to simply finding \(g(0)\)!

Now, what is \(g(t)\) exactly? Given the intuition we have, the proper way to view \(L\) is as \(L = L(X(x, t, \theta))\), where \(X(t, x, \theta)\), answers the question: “If we start at position \(x\) at time \(t\) given ODE governed by \(\theta\) , where do we end up?”. And notice that

\[g(t) = (\partial_\theta L(X(x, t, \theta)))(t)\]

where, again, we first thought of the 3D function treating \(x, t, \theta\) as three independent variables then we switched our view to strictly being a function of time tracing the trajectory of a single particle. Finding a Lagrangian ODE for \(g(t)\) is promising because, unlike \(\partial x(t)/\partial \theta\), \(g(t)\) itself is not huge to materialize, so hopefully its ODE can be more efficiently computable. Let’s do it.

First, notice that

\[g(t) = \frac{d L}{d X} \partial_\theta X \implies \dot{g}(t) = \frac{d L}{d X} \frac{d}{dt}(\partial_\theta X(t, x(t), \theta)) = \frac{dL}{dX}(\partial_{\theta, t}X + \partial_{\theta, x}X \cdot f_\theta)\]

On the other hand, we’ll take advantage of another identity noticing that \(\dot{X} = 0\), which is true because as we propagate time forward and move along the particle path the final destination won’t change. So we have

\[0 = \dot{X} = \frac{d}{dt}X(t, x(t), \theta) = \partial_t X + \partial_x X f_\theta\]

We want to bridge this identity with what we have on \(g(t)\). So again, let’s switch our view to a 3D function of \((t, x, \theta)\) and take \(\partial/\partial \theta\) then multiply by \(\frac{d L}{dX}\), we then have

\[\partial_{\theta,t} X + \partial_{\theta,x}X \cdot f_\theta + \partial_x X J_\theta = 0 \implies \frac{d L }{d X}(\partial_{\theta, t} X + \partial_{\theta, x} X \cdot f_\theta) + \frac{d L }{dX}\partial_x X J_\theta = 0\]

Interpreting this equation as a function of \(t\) on the particle trajectory and matching with what we have so far we conclude that

\[\dot{g}(t) + \Big(\frac{d L}{d X}\partial_x X\Big)(t) J_\theta(t) = 0\]

We’re not there yet. But this looks promising because we don’t have to materialize any big matrices. Even though \(J_\theta\) is huge, here it’s multiplied by the vector:

\[a(t) = \Big(\frac{d L}{dX}\partial_x X\Big)(t) = \partial_x L(X(x, t, \theta))(t)\]

which we’ll call the adjoint vector! \(a(t)\) itself is not huge to materialize so it’s natural (again) to try to find an ODE describing it. Let’s do that!

\[\dot{a}(t) = \frac{d}{dt}\Big(\frac{dL}{dX}\partial_x X\Big) = \frac{dL}{dX}\Big(\partial_{x, t}X + \partial_{x, x}X f_\theta\Big)\]

This time, let’s differentiate the identity \(\dot{X}\) with respect to \(\partial_x\) and then multiply by \(\frac{dL}{dX}\) we find

\[0 = \partial_{t, x}X + \partial_{x, x}X f_\theta + \partial_x X J_x\implies 0 = \Big(\frac{dL}{dX}(\partial_{t, x}X + \partial_{x, x}X f_\theta)\Big)(t) + \Big(\frac{dL}{dX}\partial_xX\Big)(t) J_x(t) = 0\]

which is nothing but

\[\dot{a}(t) + a(t)J_x(t) =0\implies\dot{a}(t) = -a(t)J_x(t)\]

And we’re done! We now derived the two crucial ODEs

\[\boxed{\dot{g}(t) = -a(t)J_\theta(t), \ \dot{a}(t) = -a(t)J_x(t)}\]

Both \(a(t), g(t)\) are small to materialize and the integration of both ODEs involves Jacobian-vector products which can be computed efficiently in linear time.