# Slow backward gradients in ODEs

**URL:** <https://discourse.julialang.org/t/slow-backward-gradients-in-odes/78333>\
**Category:** Modelling & Simulations\
**Tags:** question, diffeq\
**Created:** [March 23, 2022, 2:10pm UTC](https://discourse.julialang.org/t/slow-backward-gradients-in-odes/78333 "2022-03-23T14:10:32Z")\
**Posts on this page:** 8\
**Page:** 1

<div class="post-metadata">

**Author:** ![deismic](https://avatars.discourse-cdn.com/v4/letter/d/cdc98d/32.png) [@deismic](https://discourse.julialang.org/u/deismic)\
**Post date:** [March 23, 2022, 2:10pm UTC](https://discourse.julialang.org/t/slow-backward-gradients-in-odes/78333/1 "2022-03-23T14:10:33Z")

</div>

Hi all,

I am new to Julia and I am trying to solve an ODE and obtain its gradients wrt parameters `p`. The ODE is the Hodgkin-Huxley model, a standard model in neuroscience (full ODE at the end of this message).

I can solve the ODE as follows:

```julia
p = [70.0,10.0,0.05];
u0 = [-59;0.2;0.9;0.1];
tspan = [0.0,600.0];
prob = ODEProblem(my_ode,u0,tspan,p)
t_solve = @elapsed sol = solve(prob,Euler(),dt=0.01);
println("Time to solve: ", t_solve)

```

And I can get its gradients with:

```julia
function G(p)
  tmp_prob = remake(prob,u0=convert.(eltype(p),u0),p=p)
  sol = sum(solve(tmp_prob,Euler(),dt=0.01,sensealg=SensitivityADPassThrough()))
end
t_grad = @elapsed grad = ReverseDiff.gradient(G,p)
println("Time to obtain grad: ", t_grad)

```

On my machine, it takes about `0.05` seconds to solve the ODE, but it takes about `3.5` seconds to obtain the gradient.

My question is:

- Am I doing anything obviously wrong?
- Why is backpropagation 70 times more expensive than the solve? E.g. in neural networks, the gradient usually costs about 1-4 times as much as the forward pass.

Note: I know that I should not be using `Euler` and that `ForwardDiff` is faster in my current example (the model I am eventually interested in has \>50 parameters). I am just curious if there is some way of improving the efficiency of backward mode gradients. I have tried using `Zygote`, but the problem persists.

Thanks a lot for your help!

Code for ODE:

```julia
using ReverseDiff
using DifferentialEquations: DifferentialEquations
using DifferentialEquations

function my_ode(du, u, p, t)
    E1 = u[1]
    m1 = u[2]
    h1 = u[3]
    n1 = u[4]
    
    V_T = -63.0
    g_Na1, g_Kd1, g_leak1 = p
    
    alpha_act1 = -0.32 * (E1 - V_T - 13) / (exp(-(E1 - V_T - 13) / 4) - 1)
    beta_act1 = 0.28 * (E1 - V_T - 40) / (exp((E1 - V_T - 40) / 5) - 1)
    du[2] = (alpha_act1 * (1.0 - m1)) - (beta_act1 * m1)

    alpha_inact1 = 0.128 * exp(-(E1 - V_T - 17) / 18)
    beta_inact1 = 4 / (1 + exp(-(E1 - V_T - 40) / 5))
    du[3] = (alpha_inact1 * (1.0 - h1)) - (beta_inact1 * h1)

    alpha_kal1 = -0.032 * (E1 - V_T - 15) / (exp(-(E1 - V_T - 15) / 5) - 1)
    beta_kal1 = 0.5 * exp(-(E1 - V_T - 10) / 40)
    du[4] = (alpha_kal1 * (1.0 - n1)) - (beta_kal1 * n1)
    
    area = 20000.0
    I_Na1 = -(E1 - 50) * g_Na1 * area * (m1^3) * h1
    I_K1 = -(E1 - (-90)) * g_Kd1 * area * n1^4
    I_leak1 = -(E1 - (-65)) * g_leak1 * area
    
    I_ext = 1.0 * area
    du[1] = (I_leak1 + I_K1 + I_Na1 + I_ext) / (1 * area)
end

```

---

<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:** [March 23, 2022, 2:29pm UTC](https://discourse.julialang.org/t/slow-backward-gradients-in-odes/78333/2 "2022-03-23T14:29:35Z")

</div>

Interesting. Benchmarking in global scope and with `@elapsed` is always error prone. If I have done everything correctly

```julia
function test1()
    p = [70.0,10.0,0.05];
    u0 = [-59;0.2;0.9;0.1];
    tspan = [0.0,600.0];
    prob = ODEProblem(my_ode,u0,tspan,p)
    sol = solve(prob,Euler(),dt=0.01);
end

@btime test1()

function G(p)
    u0 = [-59;0.2;0.9;0.1];
    tspan = [0.0,600.0];
    prob = ODEProblem(my_ode,u0,tspan,p)
    tmp_prob = remake(prob,u0=convert.(eltype(p),u0),p=p)
    sol = sum(solve(tmp_prob,Euler(),dt=0.01,sensealg=SensitivityADPassThrough()))
end

function test2()
    p = [70.0,10.0,0.05];
    grad = ReverseDiff.gradient(G,p)
end

@btime test2()

```

yields

```julia
  16.679 ms (240307 allocations: 24.58 MiB)
  2.251 s (24660957 allocations: 977.22 MiB)

```

for me. You could also `const`-ify your global variables. `@elapsed` is potentially measuring compilation and runtime.

Edit: corrected mistakes.

---

<div class="post-metadata">

**Author:** ![deismic](https://avatars.discourse-cdn.com/v4/letter/d/cdc98d/32.png) [@deismic](https://discourse.julialang.org/u/deismic)\
**Post date:** [March 23, 2022, 3:52pm UTC](https://discourse.julialang.org/t/slow-backward-gradients-in-odes/78333/3 "2022-03-23T15:52:30Z")

</div>

Thanks a lot! The difference indeed seems to stem from working in global scope (substituting `btime` with `elapsed` in your example also works fine). Thank you!

---

<div class="post-metadata">

**Author:** ![deismic](https://avatars.discourse-cdn.com/v4/letter/d/cdc98d/32.png) [@deismic](https://discourse.julialang.org/u/deismic)\
**Post date:** [March 23, 2022, 4:01pm UTC](https://discourse.julialang.org/t/slow-backward-gradients-in-odes/78333/4 "2022-03-23T16:01:04Z")

</div>

Unfortunately, I have to reopen this. There is a bug in the answer: In the function `G(p)`, the input `p` is unused and overwritten. After fixing this bug, the runtime for the gradient is as high as posted originally

---

<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:** [March 23, 2022, 4:14pm UTC](https://discourse.julialang.org/t/slow-backward-gradients-in-odes/78333/5 "2022-03-23T16:14:29Z")

</div>

Two more data points: as you mentioned `ForwardDiff` is good here:

```julia
function test3()
    p = [70.0,10.0,0.05];
    grad = ForwardDiff.gradient(G,p)
end

@btime test3()

```

with

```julia
 24.930 ms (240307 allocations: 48.38 MiB)

```

`Zygote` fails with a compile error `promotion of types Static.False and Int64 failed to change any arguments`.

---

<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:** [March 23, 2022, 9:19pm UTC](https://discourse.julialang.org/t/slow-backward-gradients-in-odes/78333/6 "2022-03-23T21:19:53Z")

</div>

One more update (got Zygote working):

```julia
function test1()
    p = [70.0,10.0,0.05];
    u0 = [-59;0.2;0.9;0.1];
    tspan = [0.0,600.0];
    prob = ODEProblem(my_ode,u0,tspan,p)
    sol = solve(prob,Euler(),dt=0.01,save_at=600.0);
end

@show test1()

G1(p) = begin
    u0 = convert.(eltype(p), [-59;0.2;0.9;0.1]);
    tspan = [0.0,600.0];
    prob = ODEProblem(my_ode,u0,tspan,p)
    sol = sum(solve(prob,Euler(),dt=0.01,save_at=600.0))
end

G2(p) = begin
    u0 = convert.(eltype(p), [-59;0.2;0.9;0.1]);
    tspan = [0.0,600.0];
    prob = ODEProblem(my_ode,u0,tspan,p)
    sol = sum(solve(prob,Euler(),dt=0.01,save_at=600.0,sensealg=SensitivityADPassThrough()))
end

function test2()
    p = [70.0,10.0,0.05];
    grad = ReverseDiff.gradient(G2,p)
end

@show test2()

function test3()
    p = [70.0,10.0,0.05];
    grad = ForwardDiff.gradient(G1,p)
end

@show test3()

function test4()
    p = [70.0,10.0,0.05];
    grad = Zygote.gradient(G1,p)
end

@show test4()

@btime test1()
@btime test2()
@btime test3()
@btime test4()

```

yields

```julia
  17.118 ms (240445 allocations: 24.59 MiB)
  2.133 s (24661093 allocations: 977.24 MiB) # ReverseDiff
  25.753 ms (240487 allocations: 48.40 MiB) # ForwardDiff
  168.403 ms (1859063 allocations: 183.53 MiB) # Zygote

```

so Zygote does slightly better here. I also tried to optimize the `ReverseDiff` case using `ReverseDiff.@forward` without much success. Haven’t checked yet if `Zygote` has a similar mechanism. One more observation: possible selections of `sensealg` seem to be dependent from AD.

---

<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:** [March 24, 2022, 8:57am UTC](https://discourse.julialang.org/t/slow-backward-gradients-in-odes/78333/7 "2022-03-24T08:57:21Z")

</div>

First of all, this model is tiny. If you go by the results of [[1812.01892] A Comparison of Automatic Differentiation and Continuous Sensitivity Analysis for Derivatives of Differential Equation Solutions](https://arxiv.org/abs/1812.01892), you need about 100 ODEs to overcome the cost of reverse-mode AD in the best case, in the worst case even more than 1000. Reverse-mode just generally has a higher overhead but better scaling. So in this case, it shouldn’t be surprising that ForwardDiff is best.

For the other result, Zygote plugs into the DiffEqSensitivity setup [Local Sensitivity Analysis (Automatic Differentiation) · DifferentialEquations.jl](https://diffeq.sciml.ai/stable/analysis/sensitivity/) and so it ends up “not being too bad” because it’s using a hand-written adjoint rule, where all of the internal work is probably delegated to a vjp defined by Enzyme to handle the mutations well.

I’m not sure if ReverseDiff.jl ends up hitting the same rules? I think we defined a connection here [https://github.com/SciML/DiffEqBase.jl/blob/master/src/reversediff.jl#L47-L50](https://github.com/SciML/DiffEqBase.jl/blob/master/src/reversediff.jl#L47-L50) but I’d need to double check if that actually catches it. The problem with ReverseDiff.jl though is that if you use it in `Array{TrackedReal}` accidentally (instead of `TrackedArray`), you can miss dispatches like this and get some huge performance drops. I think the only real application for ReverseDiff.jl at the user level is for compiled tape mode on cases where `Array{TrackedReal}` is used to avoid errors (i.e. to handle mutation).

---

<div class="post-metadata">

**Author:** ![deismic](https://avatars.discourse-cdn.com/v4/letter/d/cdc98d/32.png) [@deismic](https://discourse.julialang.org/u/deismic)\
**Post date:** [March 31, 2022, 10:59am UTC](https://discourse.julialang.org/t/slow-backward-gradients-in-odes/78333/8 "2022-03-31T10:59:16Z")

</div>

Fantastic, thanks a lot for the detailed reply!