Computation of initial adjoint for NODE

Question

I'm reading the paper Neural Ordinary Differential Equations and I have a simple question about adjoint method. When we train NODE, it uses a blackbox ODESolver to compute gradients through model parameters, hidden states, and time. It uses another quantity $\mathbf{a}(t) = \partial L / \partial \mathbf{z}(t)$ called adjoint, which also satisfies another ODE. As I understand, the authors build a single ODE that computes all the gradients $\partial L / \partial \mathbf{z}(t_{0})$ and $\partial L / \partial \theta$ by solving that single ODE. However, I can't understand how do we know the value $\partial L / \partial \mathbf{z}(t_1)$ which corresponds to the initial condition for the ODE corresponds to the adjoint. I'm using this tutorial as a reference, and it defines custom forward and backward methods for solving ODE. However, for the backward computation (especially ODEAdjoint class in the tutorial) we need to pass $\partial L / \partial \mathbf{z}$ for backpropagation, and this enables us to compute $\partial L / \partial \mathbf{z}(t_i)$ from $\partial L / \partial \mathbf{z}(t_{i+1})$, but we still need to know the adjoint value $\partial L / \partial \mathbf{z}(t_N)$. I do not understand well about how pytorch's autograd package works, and this seems to be a barrier to understand this. Could anyone explain how it operates, and where $\partial L / \partial \mathbf{z}(t_1)$ (or $\partial L / \partial \mathbf{z}(t_N)$ if this is more comfortable) comes from? Thanks in advance.

---

Here's my guess for the initial adjoint from simple example. Let $d\mathbf{z}/dt = Az$ be a 2-dim linear ODE with given $A \in \mathbb{R}^{2\times 2}$. If we use Euler's method as a ODE solver, then the estimate for $z(t_1)$ is explicitly given as $$\hat{\mathbf{z}}(t_1) = \mathrm{ODESolve}(\mathbf{z}(t_0), f, t_0, t_1, \theta))= \left(I + \frac{t_1 - t_0}{N}A\right)^{N} \mathbf{z}(t_0) $$ where $N$ is the number of steps for Euler's method (so that $h = (t_1 - t_0) /N$ is the step size). If we use MSE loss for training, then the loss will be $$ L(\mathbf{z}(t_1)) = \Bigl|\Bigl| \mathbf{z}_1 - \left(I + \frac{t_1 - t_0}{N}A\right)^N\mathbf{z}(t_0)\Bigr|\Bigr|_2^2 $$ where $\mathbf{z}_1$ is the true value at time $t_1$, which is $\mathbf{z}_1 = e^{A(t_1 - t_0)}\mathbf{z}(t_0)$. Since adjoint $\mathbf{a}(t) = \partial L / \partial \mathbf{z}(t)$ satisfies $$\frac{d\mathbf{a}(t)}{dt} = -\mathbf{a}(t)^{T} \frac{\partial f(\mathbf{z}(t), t, \theta)}{\partial \mathbf{z}} = \mathbf{0},$$ $\mathbf{a}(t)$ is constant and we get $\mathbf{a}(t_0) = \mathbf{a}(t_1)$. So we do not need to use augmented ODE for computing $\mathbf{a}(t)$. However, I still don't know what $\mathbf{a}(t_1) = \partial L / \partial \mathbf{z}(t_1)$ should be. If my understanding is correct, since $L = ||\mathbf{z}_1 - \mathbf{z}(t_1)||^{2}_{2}$, it seems that the answer might be $$ \frac{\partial L}{\partial \mathbf{z}(t_1)} = 2(\mathbf{z}(t_1) - \mathbf{z}_1). $$ However, this doesn't seem to be true: if it is, and if we have multiple datapoints at $t_1, t_2, \dots, t_N$, then the loss is $$ L = \frac{1}{N} \sum_{i=1}^{N}||\mathbf{z}_i -\mathbf{z}(t_i)||_{2}^{2} $$ and we may have $$ \frac{\partial L}{\partial \mathbf{z}(t_i)} = \frac{2}{N} (\mathbf{z}(t_i) - \mathbf{z}_i), $$ which means that we don't need to solve ODE associated to $\mathbf{a}(t)$.

Answer 1

First, a forward pass is done to obtain predictions of $z$, at every $t$. Then the adjoint state is run backward in time for every $t$. Which gives the learning impulse. So an initial run is done to obtain values of the dynamical system at all time points and the last of these is the initial point for the backward pass.

This picture depicts it.

To be more specific with regards to the question in your initial paragraph the input of $t_0$ and $t_1$ into $ODE\_Solve$ are time points at which the network is evaluated. In this specific case, the first point is $t_0$ which should always be the case and $t_1$ is the last point. But you can also insert more time points for example $N$ time points than you input the sequence $t_0, t_1, t_2, ..., t_N$ into $ODE\_Solve$. In that case, $t_N$ is the last time point.

This picture also shows a lot. But in this picture, a RNN is used to generate $z_{t_0}$. The RNN encoder is given the time points in reversed order: $(t_N, t_{N-1}, t_{N-2}, ..., t_0)$.

This is probably not enough to satisfy your question, hopefully, I can get more specific with respect to the math later.