Aperitif

This article shows how to optimize the input variables of a differential equation using gradient descent. The crux is computing the gradients of the loss w.r.t these variables without backpropagating through the solver - this is the adjoint sensitivity method popularized by the Neural ODEs paper. It is also a poorly understood method in the ML community in my experience.

Three different components of a simple ODE are optimized. I work through it the way one would with pen and paper, prioritizing correctness over efficiency. Doing so illuminates some core ideas that efficient machinery tends to hide. The proofs lean on perturbation theory and are, IMHO, easier to follow.

Consider the ODE

initial condition
parameter
start time
end time

• Draw its trajectory
• One can evaluate at any point on the trajectory.
• One can also move along the trajectory.

---

1. Using broadcasting with ODEs

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

Array variable	Command
Initial conditions	y0 = gm.linspace(0, 2, 5)
Parameter values	c = gm.linspace(0, 1, 5)
Start times	t0 = gm.linspace(-2, 0, 5)

---

2. Defining the problem

Let's say one wants to reach the point when the trajectory ends at .

Currently there is a gap between the actual value and the target. 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 w.r.t the output , , is trivial. The tricky part is calculating the gradient w.r.t variables like:

Variable	Gradient
Initial value
Parameter
Start time

---

3. Gradient w.r.t the initial condition

The gradient of the loss w.r.t at any time is . It is trivial to calculate and our goal is to calculate . This problem has a closed form solution:

3.1 Proof

• For small , we have
• Note that in the equation above, is the parameter and is the identity matrix.

Applying chain rule, we have:

Taking the limit , we get the adjoint equation:

This is an ODE for that can be solved backwards in time from to obtain .

Once we have the gradient w.r.t , we can update it using gradient descent and repeat the process.

---

4. Optimizing

Let's start with the original equation: The state adjoint equation is given by:

Maintain a separate initial condition for each iteration : y0it{k}.

4.1 Iteration 0

• Solve the forward ODE
• Evaluate y at : we would like this point to be equal to the target point.

The loss function is defined to be so that . Now it is time to set up the adjoint equation and run it backwards in time (from to ):

• Compute initial adjoint aT0:
• Solve the backward adjoint ODE: solve the backward adjoint equation
• Evaluate adjoint at
• Extract gradient grad0: grad0 is the
• Gradient descent step: this is the gradient descent step with learning rate of .
• Update the new starting point: note that the new start point shifted a little bit.

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

• Run a forward trajectory to calculate
• Calculate
• Run a backward trajectory to calculate

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

Iteration 1

• Set previous forward and backward trajectories to gray.
• Solve the forward ODE
• Evaluate y at
• These steps have no visual change:
  • Extract y-coordinate at
  • Compute initial adjoint aT1:
• Solve the backward adjoint ODE
  • Evaluate adjoint at
  • Extract gradient grad1: grad1 is the
  • Gradient descent step
• Update the starting point after iteration 1
• Highlight forward and backward trajectories

Iteration 2

• Set previous forward and backward trajectories to gray.
• Solve the forward ODE
• Evaluate y at
  • Extract y-coordinate at
  • Compute initial adjoint aT2:
• Solve the backward adjoint ODE
  • Evaluate adjoint at
  • Extract gradient grad2: grad2 is the
  • Gradient descent step
• Update the starting point after iteration 2
• Highlight forward and backward trajectories

Iteration 3

• Set previous forward and backward trajectories to gray.
• Solve the forward ODE
• Evaluate y at
  • Extract y-coordinate at
  • Compute initial adjoint aT3:
• Solve the backward adjoint ODE
  • Evaluate adjoint at
  • Extract gradient grad3: grad3 is the
  • Gradient descent step
• Update the starting point after iteration 3
• Highlight forward and backward trajectories

---

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 , the adjoint equation is quite similar to that for y:

• : the parameter adjoint, i.e. , representing how sensitively the loss changes with respect to the parameter .
• : the state adjoint, i.e. , representing how sensitively the loss changes with respect to the state .

5.1 Proof

• Define an augmented variable that stacks the state with the parameter. Since the parameter is constant in time (), its dynamics are

• For small , exactly as before:

• Define the augmented adjoint , where and . Applying the chain rule identically to the earlier argument gives the augmented adjoint equation:

• The Jacobian of is block-structured because the parameter rows of are zero:

• Substituting and expanding the two blocks:

The first block recovers the state adjoint equation from before, while the second block gives the parameter adjoint equation:

• Since the loss has no explicit dependence on the parameter at the final time (it only enters through the trajectory), the boundary condition is . Integrating backwards from to accumulates the parameter gradient:

---

Coming back to the original ODE , the parameter adjoint equation is:

To solve this equation, for each time step , we will need the values and .

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

5.2 Iteration 0

• Run the forward ODE
• Get the end point of the trajectory
• y at time T at iteration 0
• Get the gradient of loss w.r.t y(T):
• Define the ODE for the first backward trajectory
• Run the first backward ODE
• Get the start point of the state adjoint trajectory.
• Sample the point a(t) at T: a-samp0
• Get the end point of the forward trajectory.
• Sample the point y(t) at T: y-samp0

It is important to note that a-samp0 and y-samp0 will be updated for different values of time t as the parameter adjoint in the steps below is run.

• Define the ODE for the second backward trajectory
• Run the second backward ODE
• Get the point at the end of the trajectory
• Get the y-coordinate, which is the gradient
• Perform gradient descent step and update the text

Here is a quick overview of what we just did:

• Run a forward trajectory to calculate .
• Calculate .
• Run a backward trajectory to calculate .
• Run another backward trajectory to calculate .

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

Iteration 1

• First we set the colors of previous trajectories to gray, and then hide the points
• Calculate c*y at iteration 1
• Define the ODE for the forward trajectory
• Run the forward ODE
• Get the end point of the trajectory
• y at time T at iteration 1
• Get the gradient of loss w.r.t y(T)
• Define the ODE for the first backward trajectory
• Run the first backward ODE
• Get the start point of the state adjoint trajectory.
• Sample the point a(t) at T
• Get the end point of the forward trajectory.
• Sample the point y(t) at T
• Define the ODE for the second backward trajectory
• Run the second backward ODE
• Get the point at the end of the trajectory
• Get the y-coordinate, which is the gradient
• Highlight the forward trajectory, the first backward trajectory and the second backward trajectory

• Perform gradient descent step and update the text

Iteration 2

• Set the colors of previous trajectories to gray, and then hide the points
• Calculate c*y at iteration 2
• Define the ODE for the forward trajectory
• Run the forward ODE
• Get the end point of the trajectory
• y at time T at iteration 2
• Get the gradient of loss w.r.t y(T)
• Define the ODE for the first backward trajectory
• Run the first backward ODE
• Get the start point of the state adjoint trajectory.
• Sample the point a(t) at T
• Get the end point of the forward trajectory.
• Sample the point y(t) at T
• Define the ODE for the second backward trajectory
• Run the second backward ODE
• Get the point at the end of the trajectory
• Get the y-coordinate, which is the gradient
• Highlight the forward trajectory, the first backward trajectory and the second backward trajectory

• Perform gradient descent step and update the text

Iteration 3

• Set the colors of previous trajectories to gray, and then hide the points
• Calculate c*y at iteration 3
• Define the ODE for the forward trajectory
• Run the forward ODE
• Get the end point of the trajectory
• y at time T at iteration 3
• Get the gradient of loss w.r.t y(T)
• Define the ODE for the first backward trajectory
• Run the first backward ODE
• Get the start point of the state adjoint trajectory.
• Sample the point a(t) at T
• Get the end point of the forward trajectory.
• Sample the point y(t) at T
• Define the ODE for the second backward trajectory
• Run the second backward ODE
• Get the point at the end of the trajectory
• Get the y-coordinate, which is the gradient
• Highlight the forward trajectory, the first backward trajectory and the second backward trajectory

• Perform gradient descent step and update the text

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

An implementation detail

You might be thinking that running two backward ODEs sounds inefficient. In practice, one can create an "augmented" ODE where both and are calculated on the go. More on this below.

---

6. Optimizing start time

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

• : the state adjoint at the start time, i.e. , which we already obtain by integrating the adjoint equation backwards to .
• : the value of the dynamics at the start point, i.e. the velocity at .

6.1 Proof

Here we use a clever trick - perturb initial time to which leads to different values at time . From that time onwards, it can be seen as a perturbation in the "initial condition" at and the state adjoint equation can be applied to it.

• Hold the initial value fixed and shift the start time by a small . We compare two trajectories at the same later time :

• Advancing the start time therefore perturbs the state at by

• From onward both trajectories obey the same dynamics, so the effect on the loss is captured by the state adjoint :

• Dividing by and taking the limit gives the start-time gradient:

• For the original ODE , this specializes to

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

For each iteration , we define a unique start time t-init{k}.

6.1 Iteration 0

• Run the forward ODE
• Get the end point of the trajectory
• Get the gradient of loss w.r.t y(T):
• Run the backward adjoint ODE
• Get the point at
• Compute the gradient
• Highlight forward and backward trajectories.
• Perform gradient descent step
• Update start point and text

Iteration 1

• First we set the color of previous trajectories to gray, and then hide the points
• Run the forward ODE
• Get the end point of the trajectory
• y at time T
• Get the gradient of loss w.r.t y(T)
• Run the backward adjoint ODE
• Get the point at
• Sample
• Compute the gradient
• Scale by the learning rate
• Highlight forward and backward trajectories.

• Perform gradient descent step
• Update and start point and text

Iteration 2

• First we set the color of previous trajectories to gray, and then hide the points
• Run the forward ODE
• Get the end point of the trajectory
• y at time T
• Get the gradient of loss w.r.t y(T)
• Run the backward adjoint ODE
• Get the point at
• Sample
• Compute the gradient
• Scale by the learning rate
• Highlight forward and backward trajectories.

• Perform gradient descent step
• Update and start point and text

Iteration 3

• First we set the color of previous trajectories to gray, and then hide the points
• Run the forward ODE
• Get the end point of the trajectory
• y at time T
• Get the gradient of loss w.r.t y(T)
• Run the backward adjoint ODE
• Get the point at
• Sample
• Compute the gradient
• Scale by the learning rate
• Highlight forward and backward trajectories.

• Perform gradient descent step
• Update and start point and text

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 , etc. by hand, which is not always possible once f gets complicated. Two changes make this practical at scale.

7.1 Run all equations together (one augmented ODE)

Instead of one solve for the state adjoint and a second solve for the parameter accumulator , stack everything into a single state and solve it once, backwards in time:

• Why: one solver pass instead of two or three means fewer function evaluations and a single, consistent time grid for every component.
• Reference: this is exactly the backward pass of the Neural ODE paper, which stacks and integrates them together. The component here is their parameter-gradient term.

7.2 Use a VJP instead of hand-derived gradients

We wrote out and by hand. This does not scale once f has many parameters (e.g. a neural net). The adjoint only ever needs the products and i.e. a vector–Jacobian product (VJP) which autodiff computes directly without forming the full Jacobian.

• Why: a VJP costs about the same as one evaluation of f, regardless of the number of parameters, and removes error-prone manual derivatives.
• Reference: reverse-mode autodiff / VJPs are standard in JAX (jax.vjp) and PyTorch (torch.autograd.grad).