Neural ODEs

In the previous paragraph, I essentially claimed that even if the loss is a function of some black box integrator, i.e. $\mathcal{L}\left(I\left[f_{\theta}\right]\right)$, we can train $f_{\theta}$ with backpropagation. How is that even possible? If we don't know how to make the gradients flow back from $I$, aren't we stuck? The adjoint sensitivity method, introduced by Pontryagin in the early 60s, allows us to compute the gradient by solving a second ODE, backward in time. Consider the following problem: given

$$\mathcal{L}\left(\text{ODESolve}(\mathbf{z}(t_0), f, t_0, t_1, \theta)\right)$$

, where $\text{ODESolve}$ is a callable forward integrator, $t_0$ is the initial time, $t_1$ is the final time, $\mathbf{z}(t)$ is the state, and $f$ is the current dynamics network parameterized by $\theta$, compute $\dfrac{\partial \mathcal{L}}{\partial \theta}$.

First, define the adjoint, $\mathbf{a}(t)$, as the gradient of the loss with respect to the state $\mathbf{z}(t)$:

$$\mathbf{a}(t) = \dfrac{\partial \mathcal{L}}{\partial \mathbf{z}(t)}$$

The key identity that allows us to find $\dfrac{\partial \mathcal{L}}{\partial \theta}$ is:

$$\dfrac{\partial \mathcal{L}}{\partial \theta} = - \displaystyle\int_{t_1}^{t_0} \mathbf{a}(t)^{T} \dfrac{\partial f(\mathbf{z}(t), t, \theta)}{\partial \theta} dt$$

This tells us that integrating $-\mathbf{a}(t)^T \dfrac{\partial f(\mathbf{z}(t), t, \theta)}{d\theta}$ backward in time allows us to retrieve the derivative of the loss function with respect to the parameters. Now as long as we know how to evaluate the trajectory of $\mathbf{a}(t)$ backward in time, we can get our gradients for optimization. Thus, it is of interest to know the dynamics of $\mathbf{a}(t)$. The following identity gives us just that:

$$\dot{\mathbf{a}}(t) = -\mathbf{a}(t)^{T}\dfrac{\partial f(\mathbf{z}(t), t, \theta)}{\partial \mathbf{z}}$$

One small point about the notation here, is that in the previous equation defining $\mathbf{a}(t)$, we had the notation $\partial \mathbf{z}(t)$, whereas here we have $\partial \mathbf{z}$. The way I interpret this is that, because $\mathcal{L}$ is a functional of the entire trajectory $\mathbf{z}([t_0, t_1])$, we need to specify which $\mathbf{z}$ along the trajectory we're nudging. On the other hand, $f$ is just a regular function with three inputs, meaning there is only one $\mathbf{z}$ to nudge to begin with.

Armed with the knowledge of how both $\mathbf{z}$ and $\mathbf{a}$ evolve backward in time, we run our numerical integrator on three sets of dynamics:

$$(f, -\mathbf{a}(t)^T \dfrac{\partial{f}}{\partial \mathbf{z}}, -\mathbf{a}(t)^T \dfrac{\partial{f}}{\partial \theta})$$

with the third coordinate being of interest for the actual loss gradient. Another interesting point to note here is that since the third coordinate depends on both $\mathbf{a}(t)$ and $\mathbf{z}(t)$, there is a choice to be made of whether to use the updated or previous values of those functions when discretizing the evolution.

I initially considered providing justifications for the above identities, but decided against it to prevent this post from getting excessively long. Interested readers can refer to Vaibhav Patel's blog post deriving these identities.

If the internal steps of the numerical integrator are known and written to be autodifferentiable, it is possible to compute the gradients by classic backpropagation. In fact, as we will see, though this alternative is not as mathematically elegant, it might actually be preferable in certain scenarios. In the original paper, the adjoint sensitivity method is proposed as a more memory-efficient approach with significantly lower numerical error. However, in practice, I found that the adjoint sensitivity method took a very long time, and that the standard backprop approach yielded results that were just as good, if not better, for the simple problems posed in my experiments.

One of the most natural applications of neural ODEs is to learn the propagation dynamics of physical systems by learning from trajectory data. Given a state vector $\mathbf{x}$ that evolves according to some hidden dynamics $f(\mathbf{x}, t)$, we can train a neural ODE to automatically interpolate meaningfully between observed points by training on observed data $(\mathbf{x}_{t_{k}})_{k = 0}^{k = N}$. The resulting neural ODE can then be used as a surrogate to simulate hypothetical trajectories. One important caveat is that the original method of adjoint sensitivity only outlines how we can compute the derivative of a loss that depends on the final state in the trajectory. In order to enforce trajectory consistency throughout the desired time-interval, we need to supervise with intermediate time-steps. The tricks necessary to incorporate intermediate points are beyond the scope of this article. In the following experiment, I trained a neural ODE to learn the dynamics of a Van der Pol oscillator, which is a dynamical system that exhibits very interesting trajectories and evolves according to the following ODE:

$$\ddot{x} ​=\mu (1 - x^2) \dot{x} - x$$

With $\mu$ set to $1$, I sampled the phase space coordinates $(x, \dot{x})$ over a time interval $[0, 20]$ with a step size of $0.1$ and a batch size of $16$, and trained a neural network for $5000$ iterations to learn propagation dynamics that result in trajectories that match these samples.

Another cool application mentioned in the original paper is generative modeling with normalizing flows. The paper shows that using a continuous time representation of normalizing flows with neural ODEs allows evaluating flow models in $O(M)$ time, where $M$ is the number of hidden units. This is contrasted with standard normalizing flows, which need $O(M^3)$ operations for the same task. This means that neural ODEs are capable of efficiently evaluating "wide" flow layers where the dynamics might be defined as sum of several functions. The paper coins the term "Continuous Normalizing Flows" (CNF) for these continuous-time flow models. In the following experiment, I trained a CNF to learn how to morph a standard multivariate normal distribution in 2D into a spiral distribution defined by $(x, y)$, where

$$x \sim r\cos{\theta} + \mathcal{N}(0, \sigma^2)$$ $$y \sim r\sin{\theta} + \mathcal{N}(0, \sigma^2)$$ $$\theta \sim 4\pi \cdot \mathcal{U}_{[0, 1)}$$ $$r = a + b\theta$$

The model generates a target distribution by first taking a sample $z_0 \sim \mathcal{N}(0, I_2)$ and then morphing it into a sample $x$ obtained by forward integrating $z_0$ with our dynamics neural network $f_{\theta}$. Some clever math from the original paper allows us to compute the log-probability of $x$ by finding the accumulated change in log-probability density by integrating negative the divergence of $f_{\theta}$ back in time. We train the network by sampling examples of $x$ from the target distribution defined above, then reverse integrating $(f_{\theta}, -\nabla \cdot f_{\theta})$ to find the corresponding basepoint $z$ and the accumulated change in log-probability density, and finally computing the parameter gradients. Interestingly enough, here the "forward pass" itself is a backward-in-time integration.

I trained the model to integrate with $200$ steps over the time interval $[0, 1]$, with $5000$ training loop iterations and a batch (resampled every time with the spiral generator) of size 1024.

The following visualizations demonstrate the learned transformation. In the first visual, we see the standard normal distribution in 2D morph into the learned distribution. In the second visual, we see the true spiral distribution superimposed onto the learned distribution.

As can be seen above, the CNF was able to roughly capture the shape of the target distribution. Generating a new sample is as simple as sampling from the standard normal distribution and forward propagating to get the corresponding sample in the learned distribution. Since it is normalizing flow, we also have the property that any sample from the learned probability distribution can be mapped back to a unique sample in the original standard normal distribution, which is done here with reverse integration.