# Nested gradient computations with differential equation solver

**URL:** <https://discourse.julialang.org/t/nested-gradient-computations-with-differential-equation-solver/79042>\
**Category:** New to Julia\
**Tags:** question\
**Created:** [April 5, 2022, 12:04pm UTC](https://discourse.julialang.org/t/nested-gradient-computations-with-differential-equation-solver/79042 "2022-04-05T12:04:22Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![quentinf00](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/quentinf00/32/35215_2.png) [@quentinf00](https://discourse.julialang.org/u/quentinf00)\
**Post date:** [April 5, 2022, 12:04pm UTC](https://discourse.julialang.org/t/nested-gradient-computations-with-differential-equation-solver/79042/1 "2022-04-05T12:04:22Z")

</div>

Hi, I am exploring the possibilities of combining Julia’s differential equations solvers and automatic differentiation capabilities, and I am stuck while trying to backpropagate twice through an ode solver.

Here is some code to reproduce the problem :

```julia
module Minimal

using Zygote
using DifferentialEquations
using DiffEqSensitivity

t = 1:0.1:10
gt = sin.(t)
x0 = gt + randn(size(gt)...) * 0.2
η0 = 0.01

function dsin(du, u, p, t)
	du[1] = cos(t)
end

function innergrad()
	g = gradient(x0) do x
		prob = ODEProblem(dsin, [x[1]], (1, 10))
		sol = DifferentialEquations.solve(
			prob, RK4(), saveat=collect(1:0.1:10),
		)

		solx = [xx[1] for xx in sol.u]

		sum((solx - x) .^ 2)
	end
	g[1]
end

function outergrad()
	gradient(η0) do η
		g = innergrad()
		sum( (gt - (x0 - η*g)).^2 )
	end
end
end

Minimal.innergrad() # OK
Minimal.outergrad() # KO

```

Here are my installed packages (`Pkg.status()`):

```julia
  [a077e3f3] DiffEqProblemLibrary v4.15.0
  [41bf760c] DiffEqSensitivity v6.71.0
  [0c46a032] DifferentialEquations v7.1.0
  [61744808] DynamicalSystems v2.1.8
  [587475ba] Flux v0.12.9
  [7073ff75] IJulia v1.23.2
  [a98d9a8b] Interpolations v0.13.5
  [2b0e0bc5] LanguageServer v4.2.0
  [30363a11] NetCDF v0.11.4
  [1dea7af3] OrdinaryDiffEq v6.7.1
  [91a5bcdd] Plots v1.27.3
  [438e738f] PyCall v1.93.1
  [e88e6eb3] Zygote v0.6.37

```

and julia version: `Version 1.7.2`  
Any help greatly appreciated 🙂

---

<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:** [April 5, 2022, 12:41pm UTC](https://discourse.julialang.org/t/nested-gradient-computations-with-differential-equation-solver/79042/2 "2022-04-05T12:41:42Z")

</div>

> [@quentinf00](#):
>
> and I am stuck while trying to backpropagate twice through an ode solver.

Why backprop twice? Second order of an ODE solver is essentially always faster to do via forward-over-reverse, and that exists. It’ll happen automatically when using Newton:

[https://diffeqflux.sciml.ai/dev/examples/second\_order\_adjoints/](https://diffeqflux.sciml.ai/dev/examples/second_order_adjoints/)

and uses the following:

[https://diffeq.sciml.ai/stable/analysis/sensitivity/#Second-Order-Sensitivity-Analysis-via-second\_order\_sensitivities](https://diffeq.sciml.ai/stable/analysis/sensitivity/#Second-Order-Sensitivity-Analysis-via-second_order_sensitivities)

And note, when you do forward-over-reverse, you can do Hessian-free operations, i.e. `Hv`, to do Newton-Krylov without computing Hessians.

---

<div class="post-metadata">

**Author:** ![goerch](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/goerch/32/29122_2.png) [@goerch](https://discourse.julialang.org/u/goerch)\
**Post date:** [April 5, 2022, 1:07pm UTC](https://discourse.julialang.org/t/nested-gradient-computations-with-differential-equation-solver/79042/3 "2022-04-05T13:07:00Z")

</div>

The following seems to work for me

```julia
module Minimal

using Zygote
using ForwardDiff
using DifferentialEquations
using DiffEqSensitivity

t = 1:0.1:10
gt = sin.(t)
x0 = gt + randn(size(gt)...) * 0.2
η0 = 0.01

function dsin(du, u, p, t)
	du[1] = cos(t)
end

function innergrad()
	g = gradient(x0) do x
		prob = ODEProblem(dsin, [x[1]], (1, 10))
		sol = DifferentialEquations.solve(
			prob, RK4(), saveat=collect(1:0.1:10),
		)

		solx = [xx[1] for xx in sol.u]

		sum((solx - x) .^ 2)
	end
	g[1]
end

function outergrad()
	ForwardDiff.derivative(η0) do η
		g = innergrad()
		sum( (gt - (x0 .- η*g)).^2 )
	end
end
end

Minimal.innergrad() # OK
Minimal.outergrad() # KO

```

---

<div class="post-metadata">

**Author:** ![quentinf00](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/quentinf00/32/35215_2.png) [@quentinf00](https://discourse.julialang.org/u/quentinf00)\
**Post date:** [April 5, 2022, 1:10pm UTC](https://discourse.julialang.org/t/nested-gradient-computations-with-differential-equation-solver/79042/4 "2022-04-05T13:10:27Z")

</div>

[EDIT]: code updated with working solution

@ChrisRackauckas , Thank you for your reply. Maybe just to give you some additional context: the idea behind my question is to optimize for the hyperparameters of a variational data assimilation scheme (4dvar), maybe my code was too minimal… (see below for the full code).

- The inner differentiation is part of a gradient descent to minimize a variational cost
- I want to then optimize for some hyperparameters of the gradient descent, i.e. backpropagating through the gradient descent

I’m not sure if your solution still applies, I need to read up to better understand the different term you introduced ( forward-over-reverse, Hessian free, Newton Krylov…) . I’ll try some stuff and come back here with a solution or some additional questions 🙂

```julia
module Tmp
using Zygote
using ForwardDiff
using ReverseDiff
using DifferentialEquations
using DiffEqSensitivity
using Plots

# 4d variational cost
function varcost(t, df, obs_idx, x, y)
        prob = ODEProblem(df, [x], (1, 10))
        sol = DifferentialEquations.solve(
                prob, RK4(), saveat=collect(t),
	)

        solx = [xx[1] for xx in sol.u]
        obs_cost = sum((solx[obs_idx] - y) .^ 2)

        obs_cost 
end

# Gradient descent solver
function varsolve(x0, df, t, xvarcost, niter, η)
	x = deepcopy(x0)
	for i in 1:niter
		g = gradient(xvarcost, x)[1]
		x = x - η * g
	end
        prob = ODEProblem(df, [x], (1, 10))
        sol = DifferentialEquations.solve( prob, RK4(), saveat=collect(t),)

        [xx[1] for xx in sol.u]
end

function outer_fit_rev(η0, ηvarsolve, loss, niter=1, lr=0.001)
	η = deepcopy(η0)
	for i in 1:niter
		gη = ReverseDiff.gradient(loss∘ηvarsolve∘first, [η], ReverseDiff.GradientConfig([η]))[1]
		η = η - lr * gη
	end
	η
end
function outer_fit_fwd(η0, ηvarsolve, loss, niter=1, lr=0.001)
	η = deepcopy(η0)
	for i in 1:niter
		gη = ForwardDiff.derivative(loss∘ηvarsolve, η)[1]
		η = η - lr * gη
	end
	η
end
end

using Plots
using Statistics

t = 1:0.1:10
gt = sin.(t)
sub_samp = 3
ss_idx = 1:sub_samp:length(gt)
ss_gt = gt[ss_idx]
obs = ss_gt + randn(size(ss_gt)...) * 0.2

x0 = ss_gt[1] - 0.3

function dsin(du, u, p, t)
	du[1] = cos(t)
end
xvarcost = (x) -> Tmp.varcost(t, dsin, ss_idx, x, obs)

niterinner = 1
η0 = 0.01

x_η0 = Tmp.varsolve(x0, dsin, t, xvarcost, niterinner, η0)

niterouter = 2
lr = 0.001

ηvarsolve = (η)-> Tmp.varsolve(x0, dsin, t, xvarcost, niterinner, η)
training_loss = (x) -> mean((x - gt).^2)

@time η = Tmp.outer_fit_rev(η0, ηvarsolve, training_loss, niterouter, lr)

x_η = Tmp.varsolve(x0, dsin, t, xvarcost, niterinner, η)

plot(t, gt, label="gt")
scatter!(t[ss_idx], obs, label="obs")
plot!(t, x_η0, label="x_η0")
plot!(t, x_η, label="x_η")

```

---

<div class="post-metadata">

**Author:** ![quentinf00](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/quentinf00/32/35215_2.png) [@quentinf00](https://discourse.julialang.org/u/quentinf00)\
**Post date:** [April 5, 2022, 1:56pm UTC](https://discourse.julialang.org/t/nested-gradient-computations-with-differential-equation-solver/79042/5 "2022-04-05T13:56:51Z")

</div>

> [@goerch](#):
>
> `using ForwardDiff`

Great !! indeed this solves my problem thanks a lot, but to be honest I don’t really understand why, could you maybe provide some insights or some documentation relating to the impact of this import ?

---

<div class="post-metadata">

**Author:** ![goerch](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/goerch/32/29122_2.png) [@goerch](https://discourse.julialang.org/u/goerch)\
**Post date:** [April 5, 2022, 2:36pm UTC](https://discourse.julialang.org/t/nested-gradient-computations-with-differential-equation-solver/79042/6 "2022-04-05T14:36:26Z")

</div>

> [@quentinf00](#):
>
> Great !! indeed this solves my problem thanks a lot, but to be honest I don’t really understand why, could you maybe provide some insights or some documentation relating to the impact of this import ?

Whenever I encounter one of these incomprehensible error messages with `Zygote` I try to understand the problem using alternate [AD packages](https://github.com/JuliaDiff).

Another test shows that

```julia
function outergrad()
	ReverseDiff.gradient([η0]) do η
		g = innergrad()
		sum( (gt - (x0 .- η.*g)).^2 )
	end
end

```

seems to work here too, but `Zygote` still has problems to manage this.

---

<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:** [April 5, 2022, 3:38pm UTC](https://discourse.julialang.org/t/nested-gradient-computations-with-differential-equation-solver/79042/7 "2022-04-05T15:38:32Z")

</div>

How many hyperparameters do you have? 1000? If you have like 20 hyperparameters, then using forward-mode will be faster on that. Reverse mode is for very specific things and really shouldn’t be thought of as a hammer.

---

<div class="post-metadata">

**Author:** ![quentinf00](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/quentinf00/32/35215_2.png) [@quentinf00](https://discourse.julialang.org/u/quentinf00)\
**Post date:** [April 5, 2022, 6:54pm UTC](https://discourse.julialang.org/t/nested-gradient-computations-with-differential-equation-solver/79042/8 "2022-04-05T18:54:41Z")

</div>

In the code above only 1 parameter but my end goal would be to train some models with a few hundred thousand parameters.  
I’ve played around with Flux.jl, and this is why went with Zygote.jl by default.  
For your information, I am just starting out in Julia and have been using pytorch so far therefore I haven’t had to think too much about the AD part of my work so all your comments are very much appreciated 😉

---

<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:** [April 5, 2022, 8:40pm UTC](https://discourse.julialang.org/t/nested-gradient-computations-with-differential-equation-solver/79042/9 "2022-04-05T20:40:15Z")

</div>

Not parameters, hyperparameters.

---

<div class="post-metadata">

**Author:** ![quentinf00](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/quentinf00/32/35215_2.png) [@quentinf00](https://discourse.julialang.org/u/quentinf00)\
**Post date:** [April 6, 2022, 8:10am UTC](https://discourse.julialang.org/t/nested-gradient-computations-with-differential-equation-solver/79042/10 "2022-04-06T08:10:05Z")

</div>

I apologize for the lack of clarity… Let me try again:

I have a parametrized function f that performs a gradient descent on a variational cost (that includes some differential eqs integration)

f\_{\theta}(y) = \hat{x} = \underset{x}{argmin} J(y,x)

I want to learn the parameters \Theta of this function.

\Theta = \underset{\theta}{argmin}L(f\_\theta)

In my example so far the only parameter of this function was the gradient step used to minimize J. I used the terminology hyperparameter (probably a little bit abusively) because I thought the problem was similar to optimizing for the learning rate in a classical ML problem.

In my target usecase this parametrized function contains a couple neural networks with a few hundred thousand parameters. They are not really hyperparameters but optimizing them do require nested gradient descent.

I hope this was clearer. And thanks again for taking the time