Solving neural ODEs from first principles

Run the editors in order, top to bottom. First some helper functions are defined:

• solve(f, t_start, y_start, t_end) solves $\frac{dy}{dt} = f(t, y)$ with scipy's solve_ivp. When t_end comes before t_start, it runs backwards in time.
• at(sol, t) is the value of a solution at time $t$.
• draw(sol, color) draws a solution as a curve in the $(t, y)$ plane, growing from its start to its end.
• fresh() clears the canvas and zooms in on the part of the plane this page uses.

import numpy as np
from scipy.integrate import solve_ivp

blue, orange, green, gray = '#2563eb', '#ea580c', '#16a34a', '#a8a29e'

def solve(f, t_start, y_start, t_end):
    return solve_ivp(f, (t_start, t_end), [y_start], dense_output=True, rtol=1e-8, atol=1e-10)

def at(sol, t):
    return float(sol.sol(t)[0])

def grow(s, sol, steps):
    ts = sol.t[0] + np.minimum(np.linspace(0, 1, steps), s) * (sol.t[-1] - sol.t[0])
    return np.stack([ts, sol.sol(ts)[0]], axis=1)

def draw(sol, color, steps=100):
    s = pie.scalar(0)
    path = pie.lift(grow, pie.Array[pie.Point])(s, sol, steps, hidden=True)
    lines = pie.line(path[:-1], path[1:], stroke=color)
    pie.translate_scalar_to(s, 1, animate=True)
    return lines

def fresh():
    pie.clear()
    pie.scalar(120, out='unit')
    pie.scalar(-0.5, out='originX')
    pie.scalar(-1, out='originY')

Consider the ODE

Draw its trajectory:

fresh()
c, y0, t0, t1 = 0.5, 1.0, 0.0, 2.0

def f(t, y):
    return -c * y

traj = solve(f, t0, y0, t1)
draw(traj, blue)

One can evaluate $y$ at any point on the trajectory:

print(at(traj, 1.0))

One can also move along the trajectory:

def position(t):
    return (t, at(traj, t))

t = pie.scalar(t0, out='t')
pie.lift(position, pie.Point)(t, fill='crimson')
pie.translate_scalar_to(t, t1, animate=True)

1. Drawing many trajectories

One can ask how the trajectory varies with input variables like the initial condition $y_0$, parameter $c$ or the start time $t_0$. One way to get an intuition is by drawing multiple trajectories for different values of $y_0$ etc.

Initial conditions $y_0$ from 0 to 2:

fresh()
pie.text('Trajectories for multiple values of y0', 1, -0.4, 14)
for value in np.linspace(0, 2, 5):
    draw(solve(lambda t, y: -0.5 * y, 0, value, 2), blue)

Parameter values $c$ from 0 to 1:

fresh()
pie.text('Trajectories for multiple values of c', 1, -0.4, 14)
for value in np.linspace(0, 1, 5):
    draw(solve(lambda t, y: -value * y, 0, 1, 2), blue)

Start times $t_0$ from -2 to 0:

fresh()
pie.text('Trajectories for multiple values of t0', 0, -0.4, 14)
for value in np.linspace(-2, 0, 5):
    draw(solve(lambda t, y: -0.5 * y, value, 1, 2), blue)

2. Defining the problem

Let's say one wants to reach the point $(t,y) = (2,0.2)$, drawn in red, when the trajectory ends at $t_1$.

fresh()
draw(traj, blue)
target = 0.2
pie.point(t1, target, fill='red')
pie.text('target', t1 + 0.2, 0.2, 12)

Currently there is a gap between the actual value and the target:

y1 = at(traj, t1)
top = pie.point(t1, y1, fill=blue)
pie.line(top, pie.point(t1, target, hidden=True), stroke='#d97706')
pie.text('gap', t1 + 0.2, y1, 12)

Is there a way to use this gap as a loss and use gradient descent to update one of the input variables above to reduce the loss? Yes.

One can calculate the loss and its gradient at the end of a trajectory and backpropagate it through the trajectory to modify the inputs to the ODE.

Note that calculating the gradient of a scalar like loss $L$ w.r.t the output $y_1$, $\frac{\partial L}{\partial y_1}$, is trivial. The tricky part is calculating the gradient w.r.t variables like:

3. Gradient w.r.t the initial condition $y_0$

The gradient of the loss w.r.t $y$ at any time $t$ is $\frac{\partial L}{\partial y(t)} := a(t)$. It is trivial to calculate $a(t_1) = \frac{\partial L}{\partial y(t_1)}$ and our goal is to calculate $a(t_0) = \frac{\partial L}{\partial y(t_0)}$. This problem has a closed form solution:

4. Optimizing $y_0$

Let's start with the original equation: $$\frac{dy}{dt} = f(y) = -cy = -0.5y$$ The state adjoint equation is given by: $$\frac{da}{dt} = -a \frac{\partial f}{\partial y} = 0.5a$$

In code, f is the right side of the ODE and adjoint is the right side of the state adjoint equation:

fresh()
c, t0, t1, target = 0.5, 0.0, 2.0, 0.2
pie.point(t1, target, fill='red')
pie.text('target', t1 + 0.2, 0.2, 12)

def f(t, y):
    return -c * y

def adjoint(t, a):
    return c * a

4.1 Iteration 0

• Start from the current initial condition $y_0 = 1$ and solve the forward ODE:

y0 = 1.0
y0_node = pie.scalar(y0, out='y0')
pie.point(t0, y0_node, fill=blue)
pie.text('current y0', t0 - 0.5, y0_node, 12)
fwd = solve(f, t0, y0, t1)
fwd_line = draw(fwd, blue)

• Evaluate y at $t_1$: we would like this point to be equal to the red target point.

y1 = at(fwd, t1)
end = pie.point(t1, y1, fill=blue)
print(y1)

The loss function is defined to be $\frac{1}{2} (y_1 - \text{target})^2$ so that $\frac{dL}{dy_1} = y_1 - \text{target}$. Now it is time to set up the adjoint equation and run it backwards in time (from $t_1$ to $t_0$):

• Compute the initial adjoint $a_1 = \frac{dL}{dy_1} = y_1 - \text{target}$:

a1 = y1 - target
print(a1)

• Solve the backward adjoint equation, from $t_1$ to $t_0$:

adj = solve(adjoint, t1, a1, t0)
adj_line = draw(adj, orange)

• Evaluate the adjoint at $t_0$. This is $\frac{dL}{dy_0}$:

grad = at(adj, t0)
print(grad)

• Take a gradient descent step with a learning rate of $1$. Note that the start point shifts a little bit:

y0 = y0 - grad
pie.translate_scalar_to(y0_node, y0, animate=True)
print(y0)

This completes one iteration. Here is an overview of what we did in the first iteration:

• Run a forward trajectory (blue) to calculate $y_1$
• Calculate $a(t_1) = \frac{\partial L}{\partial y(t_1)}$
• Run a backward trajectory (orange) to calculate $\frac{\partial L}{\partial y(t_0)}$

Now let's run it for a few more iterations.

def step_y0():
    global y0, fwd_line, adj_line, end
    pie.set_stroke(fwd_line, gray)
    pie.set_stroke(adj_line, gray)
    pie.set_fill(end, gray)

    fwd = solve(f, t0, y0, t1)
    fwd_line = draw(fwd, blue)
    y1 = at(fwd, t1)
    end = pie.point(t1, y1, fill=blue)

    adj = solve(adjoint, t1, y1 - target, t0)
    adj_line = draw(adj, orange)
    grad = at(adj, t0)

    y0 = y0 - grad
    pie.translate_scalar_to(y0_node, y0, animate=True)
    print(f'y1 = {y1:.4f}, dL/dy0 = {grad:.4f}, new y0 = {y0:.4f}')

step_y0()

step_y0()

step_y0()

One can see the initial condition converging towards the optimal value and the end point of the trajectory also converging towards the target point.

5. Optimizing parameter c

For a parameter p and an ODE $\frac{dy}{dt} = f(t,y,p)$, the adjoint equation is quite similar to that for y:

• $a_p$: the parameter adjoint, i.e. $\frac{dL}{dp}$, representing how sensitively the loss changes with respect to the parameter $p$.
• $a$: the state adjoint, i.e. $\frac{dL}{dy}$, representing how sensitively the loss changes with respect to the state $y$.

Coming back to the original ODE $\frac{dy}{dt} = f(t,y,c) = -cy$, the parameter adjoint equation is:

To solve this equation, for each time step $t$, we will need the values $a(t)$ and $y(t)$.

• Sampling $y(t)$: these values can be sampled from the forward trajectory. This will be our first (and only) forward trajectory per iteration.
• Sampling $a(t)$: we still need to run the backward adjoint equation to get these values. This will be our first backward trajectory per iteration.
• Sampling $a_p(t)$: now that we have the trajectories to sample a(t) and y(t), we can run the backward trajectory to get $a_p(t_0)$. This will be our second backward trajectory per iteration.

This time c is the variable being learned. Its current value is shown on the canvas:

fresh()
y0, t0, t1, target = 1.0, 0.0, 2.0, 0.2
pie.point(t1, target, fill='red')
pie.text('target', t1 + 0.2, 0.2, 12)

c = 0.5
c_node = pie.scalar(c, out='c')

def c_label(c):
    return f'c = {c:.4f}'

pie.text(pie.lift(c_label, str)(c_node), 1, 1, 14)

def f(t, y):
    return -c * y

def adjoint(t, a):
    return c * a

5.2 Iteration 0

• Run the forward ODE:

fwd = solve(f, t0, y0, t1)
fwd_line = draw(fwd, blue)

• Get the end point of the trajectory, and the gradient of the loss w.r.t y(T): $a_T = \frac{\partial L}{\partial y_T} = y_T - \text{target}$

y1 = at(fwd, t1)
end = pie.point(t1, y1, fill=blue)
a1 = y1 - target
print(a1)

• Run the first backward ODE, the state adjoint $\frac{da}{dt} = ca$:

adj = solve(adjoint, t1, a1, t0)
adj_line = draw(adj, orange)

• Run the second backward ODE, the parameter adjoint $\frac{da_p}{dt} = ay$, starting from $a_p(t_1) = 0$. At every time $t$ the solver asks for, the right side reads $a(t)$ from the first backward trajectory and $y(t)$ from the forward trajectory:

def accumulator(t, a_p):
    return [at(adj, t) * at(fwd, t)]

acc = solve(accumulator, t1, 0.0, t0)
acc_line = draw(acc, green)

• Get the value at the end of the trajectory, which is the gradient $\frac{dL}{dc}$:

grad = at(acc, t0)
print(grad)

• Perform the gradient descent step and update the text:

c = c - grad
pie.translate_scalar_to(c_node, c, animate=True)

Here is a quick overview of what we just did:

• Run a forward trajectory (blue) to calculate $y(t_1)$.
• Calculate $a(t_1) = \frac{\partial L}{\partial y(t_1)}$.
• Run a backward trajectory (orange) to calculate $a(t)$.
• Run another backward trajectory (green) to calculate $a_p(t_0)$.

Now we just repeat the same procedure for a few more iterations.

def step_c():
    global c, fwd, adj, fwd_line, adj_line, acc_line, end
    pie.set_stroke(fwd_line, gray)
    pie.set_stroke(adj_line, gray)
    pie.set_stroke(acc_line, gray)
    pie.set_fill(end, gray)

    fwd = solve(f, t0, y0, t1)
    fwd_line = draw(fwd, blue)
    y1 = at(fwd, t1)
    end = pie.point(t1, y1, fill=blue)

    adj = solve(adjoint, t1, y1 - target, t0)
    adj_line = draw(adj, orange)

    acc = solve(accumulator, t1, 0.0, t0)
    acc_line = draw(acc, green)
    grad = at(acc, t0)

    c = c - grad
    pie.translate_scalar_to(c_node, c, animate=True)
    print(f'y1 = {y1:.4f}, dL/dc = {grad:.4f}, new c = {c:.4f}')

step_c()

step_c()

step_c()

As you can see, the forward trajectory continues to converge towards the target point.

6. Optimizing start time $t_0$

Unlike $y_0$, $c$ or $p$, the start time $t_0$ is not a quantity carried inside the trajectory; it only sets where integration begins. Its gradient turns out to be especially simple:

• $a(t_0)$: the state adjoint at the start time, i.e. $\frac{\partial L}{\partial y(t_0)}$, which we already obtain by integrating the adjoint equation backwards to $t_0$.
• $f(t_0, y(t_0), p)$: the value of the dynamics at the start point, i.e. the velocity $\frac{dy}{dt}$ at $t_0$.

• For the original ODE $f(t,y,c) = -cy$, this specializes to$$\frac{\partial L}{\partial t_0} = -a(t_0)\,\big(-c\,y(t_0)\big) = c\,y(t_0)\,a(t_0)$$

Only one backward ODE is required: once the backward state adjoint pass reaches $t_0$, the gradient is just a product of the adjoint, the parameter and the state at the start point.

This time $t_0$ is the variable being learned, with a learning rate of $30$. Its current value is shown above the start point:

fresh()
c, y0, t1, target, lr = 0.5, 1.0, 2.0, 0.2, 30
pie.point(t1, target, fill='red')
pie.text('target', t1 + 0.2, 0.2, 12)

def f(t, y):
    return -c * y

def adjoint(t, a):
    return c * a

t0 = 0.0
t0_node = pie.scalar(t0, out='t0')
pie.point(t0_node, y0, fill=blue)

def t0_label(t0):
    return f't0 = {t0:.4f}'

pie.text(pie.lift(t0_label, str)(t0_node), t0_node, 1.25, 12)

6.2 Iteration 0

• Run the forward ODE:

fwd = solve(f, t0, y0, t1)
fwd_line = draw(fwd, blue)

• Get the end point of the trajectory, and the gradient of the loss w.r.t y(T): $a(T) = y(T) - \text{target}$

y1 = at(fwd, t1)
end = pie.point(t1, y1, fill=blue)
a1 = y1 - target
print(a1)

• Run the backward adjoint ODE $T \to t_0$:

adj = solve(adjoint, t1, a1, t0)
adj_line = draw(adj, orange)

• Compute the gradient $dL/dt_0 = a(t_0)\cdot(c\cdot y_0)$:

grad = at(adj, t0) * c * y0
print(grad)

• Perform the gradient descent step, and update the start point and text:

t0 = t0 - lr * grad
pie.translate_scalar_to(t0_node, t0, animate=True)
print(t0)

def step_t0():
    global t0, fwd_line, adj_line, end
    pie.set_stroke(fwd_line, gray)
    pie.set_stroke(adj_line, gray)
    pie.set_fill(end, gray)

    fwd = solve(f, t0, y0, t1)
    fwd_line = draw(fwd, blue)
    y1 = at(fwd, t1)
    end = pie.point(t1, y1, fill=blue)

    adj = solve(adjoint, t1, y1 - target, t0)
    adj_line = draw(adj, orange)
    grad = at(adj, t0) * c * y0

    t0 = t0 - lr * grad
    pie.translate_scalar_to(t0_node, t0, animate=True)
    print(f'y1 = {y1:.4f}, dL/dt0 = {grad:.4f}, new t0 = {t0:.4f}')

step_t0()

step_t0()

step_t0()

In just a few iterations, the end point of the trajectory is pretty close to the target already.

7. Optimizing the engineering part

So far we ran each equation as its own ODE solve, which is clear to follow but wasteful. We also cheated by working out the gradients $\frac{\partial f}{\partial y}$, $\frac{\partial f}{\partial c}$ etc. by hand, which is not always possible once f gets complicated. Two changes make this practical at scale.