# Correct way of computing adjoints/gradients with dense solution of ODE

**URL:** https://discourse.julialang.org/t/correct-way-of-computing-adjoints-gradients-with-dense-solution-of-ode/121964
**Category:** Numerics
**Tags:** differentiation, sciml, ad, differentialequation
**Created:** [October 29, 2024, 10:39pm UTC](https://discourse.julialang.org/t/correct-way-of-computing-adjoints-gradients-with-dense-solution-of-ode/121964 "2024-10-29T22:39:26Z")
**Posts on this page:** 6
**Page:** 1

<div class="post-metadata">

### Author: ![facusapienza](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/facusapienza/32/24317_2.png) [@facusapienza](https://discourse.julialang.org/u/facusapienza)
#### Post date: [October 29, 2024, 10:39pm UTC](https://discourse.julialang.org/t/correct-way-of-computing-adjoints-gradients-with-dense-solution-of-ode/121964/1 "2024-10-29T22:39:26Z")

</div>

Hi all,

I am trying to use the interpolated solution of my ODE for adjoint sensitivity analysis. Since I have a very small system of ODEs (`n=3`), I am interested in evaluating the performance of continuous adjoint methods when using computing a dense solution during my forward pass. However, when evaluating `dense=true` (together with `save_everystep=true` and `saveat=nothing`) I cannot get the gradients computed in the reverse pass.

Here is a MWI:

```julia
using SciMLSensitivity, OrdinaryDiffEqCore, OrdinaryDiffEqTsit5, Zygote
using BenchmarkTools
using Random, Distributions
using Optimization, OptimizationOptimisers, OptimizationOptimJL

function fiip(du, u, p, t)
    du[1] = dx = p[1] * u[1] - p[2] * u[1] * u[2]
    du[2] = dy = -p[3] * u[2] + p[4] * u[1] * u[2]
end
p = [1.5, 1.0, 3.0, 1.0];
u0 = [1.0; 1.0];

abstol = reltol = 1e-8
tspan = (0.0, 10.0)

prob = ODEProblem(fiip, u0, tspan, p)

N = 50
times = sort(rand(sampler(Uniform(tspan[1], tspan[2])), N))

function loss(u0, p)
    sol = solve(prob, Tsit5(), u0 = u0, p = p, saveat = times, 
                abstol=abstol, reltol=reltol, 
                sensealg=QuadratureAdjoint(autojacvec=ReverseDiffVJP(true)))
    return sum(sum(sol.u))
end

@benchmark du0, dp = Zygote.gradient(loss, u0, p)

```

This works correctly, but it is saving the solution at `saveat = times`. On the other hand, when evaluating a dense solution like here:

```julia
function loss_dense(u0, p)
    sol_dense = solve(prob, Tsit5(), u0 = u0, p = p, 
                    dense=true, save_everystep=true,
                    abstol=abstol, reltol=reltol, 
                    sensealg=InterpolatingAdjoint(autojacvec=ReverseDiffVJP(true)))
    return sum(sol_dense(times))
end

@benchmark du0, dp = Zygote.gradient(loss_dense, u0, p)

```

I get the message

```julia
LoadError: Standard interpolation is disabled due to sensitivity analysis being
used for the gradients. Only linear and constant interpolations are
compatible with non-AD sensitivity analysis calculations. Either
utilize tooling like saveat to avoid post-solution interpolation, use
the keyword argument dense=false for linear or constant interpolations,
or use the keyword argument sensealg=SensitivityADPassThrough() to revert
to AD-based derivatives.

```

Now, the message is quite clear: it seems that the continuous adjoint here does not like the interpolation (same error when using direct continuous adjoint method for optimization with `AutoZygote()`). However, is this the expected behaviour? I am doing something fundamentally wrong or maybe the continuous adjoint method for higher order interpolations hasn’t been implemented?

My motivation here is to have a dense solution in the forward pass so I can compute my adjoints as fast as possible without worrying about memory use.

Thank you!!!

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [October 30, 2024, 6:58am UTC](https://discourse.julialang.org/t/correct-way-of-computing-adjoints-gradients-with-dense-solution-of-ode/121964/2 "2024-10-30T06:58:42Z")

</div>

> [@facusapienza](#):
>
> Now, the message is quite clear: it seems that the continuous adjoint here does not like the interpolation (same error when using direct continuous adjoint method for optimization with `AutoZygote()`). However, is this the expected behaviour? I am doing something fundamentally wrong or maybe the continuous adjoint method for higher order interpolations hasn’t been implemented?

It’s expected and it’s fundamental to the method. The continuous interpolation of most methods requires an alternative calculation with respect to the `k` of each step, since the dense interpolations are built by solving a (usually one order lower) bi\*k[i] summation. One, in order for the standard interpolation to be differentiable, you’d need to differentiate w.r.t. the `k`s of the problem, which is equivalent to the discrete adjoint formulation since that is the core internal step information. As such, it does not make sense to really add that as a default calculation in a continuous adjoint because that would add all of the calculations of the discrete adjoint to the continuous adjoint guaranteeing it would be the slowest of all. Thus if you need to use the standard dense output of the ODE solver, it really only make sense to use discrete adjoints.

However, you could build an alternative interpolation on the output. There’s pros and cons of this: it’s not really guaranteed to have the same accuracy for example. But if you take `sol.u` and give it to DataInterpolations.jl you could get an interpolation like a CubicSpline based only on the `sol.u`. It thus does not use internal derivative estimates, and therefore crucially may not be as robust to stiff behavior, but for some applications this could be sufficient. Since it’s only based on the `sol.u` it would only need to differentiate w.r.t. the saved values which thus propagates back into the adjoint as the delta in the `g` delta functions and is thus fine.

---

<div class="post-metadata">

### Author: ![facusapienza](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/facusapienza/32/24317_2.png) [@facusapienza](https://discourse.julialang.org/u/facusapienza)
#### Post date: [October 31, 2024, 12:01am UTC](https://discourse.julialang.org/t/correct-way-of-computing-adjoints-gradients-with-dense-solution-of-ode/121964/3 "2024-10-31T00:01:05Z")

</div>

Thank you for your response @ChrisRackauckas !

Ok, that makes sense, yes. I can see the problem. So, in general, the only alternative if I want to use a dense solution of the forward pass is using discrete adjoints? I guess I was trying to see how to setup the following option in the continuous adjoint method in cases where I don’t care about memory but I still want to use a continuous method instead of discrete.

 ![image](https://global.discourse-cdn.com/julialang/original/3X/d/0/d0f183d7832de8fbe7d3286d513394e73829a239.png)

If I understand you correctly, this right now just makes sense using the discrete adjoint method, or for the continuos adjoint method this would have to be done with a different interpolation technique that just depends of `sol.u` rather that the internal interpolation using during solving. Am I right?

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [October 31, 2024, 8:18am UTC](https://discourse.julialang.org/t/correct-way-of-computing-adjoints-gradients-with-dense-solution-of-ode/121964/4 "2024-10-31T08:18:33Z")

</div>

Yes, the interpolation is important for the method of the continuous adjoint because it allows for the reverse to not require using the same time steps as going forwards (by reinterpolating `u`) and allows for the integral to be handled more optimally by choosing a minimal set of points (by reinterpolating the `lambda`). However, taking the derivative with respect to the interpolation is a much more complicated process that the continuous adjoint method cannot do very easily, and in fact if you want that done you have to either add all of the differentiation parts of the discrete adjoint (in which case, you might as well do a discrete adjoint at that point) or you use a k-independent interpolation like a spline (which then has some accuracy trade-offs)

---

<div class="post-metadata">

### Author: ![facusapienza](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/facusapienza/32/24317_2.png) [@facusapienza](https://discourse.julialang.org/u/facusapienza)
#### Post date: [October 31, 2024, 6:14pm UTC](https://discourse.julialang.org/t/correct-way-of-computing-adjoints-gradients-with-dense-solution-of-ode/121964/5 "2024-10-31T18:14:15Z")

</div>

Makes sense.

So, just to be sure, in my original code,

```julia
function loss(u0, p)
    sol = solve(prob, Tsit5(), u0 = u0, p = p, saveat = times, 
                abstol=abstol, reltol=reltol, 
                sensealg=QuadratureAdjoint(autojacvec=ReverseDiffVJP(true)))
    return sum(sum(sol.u))
end

```

this is already using the linear interpolation (no the interpolation of the solver) for constructing the continuous adjoint or is not using any form of interpolation at all?

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [October 31, 2024, 6:17pm UTC](https://discourse.julialang.org/t/correct-way-of-computing-adjoints-gradients-with-dense-solution-of-ode/121964/6 "2024-10-31T18:17:18Z")

</div>

> [@facusapienza](#):
>
> this is already using the linear interpolation (no the interpolation of the solver) for constructing the continuous adjoint or is not using any form of interpolation at all?

Sorry, `saveat` is fine, it’s interpolating but as part of the adjoint process. It’s just post-solution interpolation that would require differentiating the `k`s.
