Neural ODEs gradient calculation for multiple time steps

Question

I was reading the paper on Neural ODEs (here) and was wondering if anyone could offer some insight on calculation of the gradient of the loss function.

If we are only considering 2 time points, $t_0,t_1$, I understand how the adjoint method works. However, what confuses me is when the loss function involves multiple time points, say $t_0,t_1,t_2$.

The paper says (on p. 15) that the adjoint step for each intervals $[t_1,t_2]$ and $[t_0,t_1]$ can be performed and that the obtained gradients can be summed. I find this confusing as well as Figure 2 of the paper (page 2).

Using the paper's notation, I understand that $\mathbf{a}(t) = \frac{dL}{d \mathbf{z}(t)}$ and $\mathbf{a}_t(t) = \frac{dL}{dt(t)}$ need to be computed first on the interval $[t_1,t_2]$ then use these results, together with an adjustment of $\frac{dL}{d \mathbf{z}(t_1)}, \frac{dL}{dt(t_1)}$, to compute the quantities $\mathbf{a}(t),\mathbf{a}_t(t)$ on the interval $[t_0,t_1]$.

Following the code in this blog, algorithmically, we have:

on the time interval $[t_1,t_2]$: $$\mathbf{a}(t_1) = \mathbf{a}(t_2) - \int_{t_2}^{t_1} \text{some integrand} \,dt$$ and $$\mathbf{a}_t(t_1) = \mathbf{a}_t(t_2) - \int_{t_2}^{t_1} \text{some integrand} \,dt.$$

On the other hand, for the interval $[t_0,t_1]$, the initial conditions are updated so that

$$\mathbf{a}(t_0) = \mathbf{a}(t_1) + \frac{dL}{d \mathbf{z}(t_1)} - \int_{t_1}^{t_0} \text{some integrand} \,dt$$ and $$\mathbf{a}_t(t_0) = \mathbf{a}_t(t_1) -\frac{dL}{d t(t_1)} - \int_{t_1}^{t_0} \text{some integrand} \,dt.$$ Can anyone help me understand/show mathematically why the adjustments to the gradient computation have to be done like this?

Answer 1

From the paper:

Most ODE solvers have the option to output the state $\mathbf{z}(t)$ at multiple times. When the loss depends on these intermediate states, the reverse-mode derivative must be broken into a sequence of separate solves, one between each consecutive pair of output times (Figure 2). At each observation, the adjoint must be adjusted in the direction of the corresponding partial derivative ∂$L$/∂$\mathbf{z}(t_i)$.

The adjustments of $ \frac{dL}{d \mathbf{z}(t_1)}$ and $-\frac{dL}{d t(t_1)} $ are exactly what they're talking about here. These are the blue vertical dashed lines in Figure 2. I think you already understood all of this, and your question is why adjustments are needed for calculating the loss properly a.k.a why the second sentence quoted is true.

I think the adjustments are needed so that errors down the line don't affect loss earlier. Say we're approximating $z(t) = t^2$ at $t = 0, 1, 2$ and our current estimate of the dynamics outputs (1, 2, 5) compared with (0, 1, 4). The model is off by one for the last timestep, but the dynamics between timesteps is accurate, so you don't want it to adjust the loss for each point—just the last one. Subtracting the adjustments as described above would make it so that the loss is zero for $t = 0$ and $t=1$.

I haven't seen this paper before so my understanding may be flawed, so lmk if that's any help.