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

Setup the ODE

• Set an appropriate grid
• Define placeholder y
• Define c
• Compute cy
• Compute dydt
• Define initial condition
• Define start time
• Define end time
• Define number of steps

Once the variables have been set up, the trajectory can be drawn like so:

• Draw the trajectory

• One can evaluate y at any point on the trajectory, and also move along it.
• One can highlight the entire trajectory using the highlight command.

---

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 a feel of it is by drawing multiple trajectories for different values of etc. Using native broadcasting semantics, it can be done easily by setting any of these variables to an array:

Array variable	Command
Initial conditions	y0 = \linspace 0 2 5
Parameter values	c = \linspace 0 1 5
Start times	t0 = \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:

4.1 Setup

• Define grid size
• Define c
• Define start time
• Define end time
• Define number of steps
• Define target value
• Define placeholder y
• Compute cy
• Compute dydt:
• Define adjoint placeholder a
• Compute adjoint dadt:
• Define gray color
• Define red color
• Define the target point and color it red
• Next to target , add text annotation

---

We maintain a separate initial condition for each iteration as y0it{k}.

Iteration 0

• Define initial condition y0it0
• Draw a point at the new starting point
• Solve the forward ODE
• Evaluate y at : we would like this point to be equal to the target point.
• Extract y-coordinate at

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 trajectory to gray
• Set previous backward trajectory to gray
• Solve the forward ODE
• Evaluate y at
• 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
• Highlight forward and backward trajectories

• Update the starting point after iteration 1

Iteration 2

• Set previous forward trajectory to gray
• Set previous backward trajectory 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
• Highlight forward and backward trajectories

• Update the starting point after iteration 2

Iteration 3

• Set previous forward trajectory to gray
• Set previous backward trajectory 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
• Highlight forward and backward trajectories

• Update the starting point after iteration 3

---

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 Setup

• Start time for forward trajectory
• End time for forward trajectory
• Start time for backward trajectory: we need a different variable than t-final due to the way ODE solver is implemented. In one's mental model, it can be assumed to be the same as t-final.
• Start time for second backward ODE: same as t-final but a separate variable.
• End time for second backward ODE: same as t-start but a separate variable.
• Number of steps
• Target value
• Initial condition for forward trajectory
• Initial condition for the second backward trajectory
• Forward trajectory placeholder
• State adjoint trajectory placeholder: for the first backward trajectory
• Parameter adjoint trajectory placeholder: for the second backward trajectory
• Value of parameter c before iteration 0
• Define colors
• Create target point and set it to red
• Next to target , add text annotation

Iteration 0

• Calculate c*y at iteration 0
• Define the ODE for the forward trajectory
• 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.

Setup

• Define colors
• Parameter c
• Forward initial condition : kept fixed — here we learn the start time.
• Forward end / adjoint start T
• Number of steps
• Target value
• Learning rate
• Forward trajectory placeholder
• Calculate c*y
• Define the forward ODE dy/dt = -c·y: autonomous, no explicit t.
• Adjoint trajectory placeholder
• Define the adjoint ODE da/dt = c·a: since .
• Compute : constant because the dynamics are autonomous.
• Create target point and set it to red
• Next to target , add text annotation

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

Iteration 0

• Start time
• 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 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).