# Optimizating computational time of gradient on big linear UDEs
**URL:** https://discourse.julialang.org/t/optimizating-computational-time-of-gradient-on-big-linear-udes/108645
**Category:** Modelling & Simulations
**Created:** [January 10, 2024, 9:41pm UTC](https://discourse.julialang.org/t/optimizating-computational-time-of-gradient-on-big-linear-udes/108645 "2024-01-10T21:41:02Z")
**Posts on this page:** 12
**Page:** 1
### Author:  [@jarroyoe](https://discourse.julialang.org/u/jarroyoe)
#### Post date: [January 10, 2024, 9:41pm UTC](https://discourse.julialang.org/t/optimizating-computational-time-of-gradient-on-big-linear-udes/108645/1 "2024-01-10T21:41:02Z")
I have an high-dimensional UDE model that includes a linear model. A MWE of the linear component of the model is:
```julia
using ForwardDiff, DifferentialEquations, SciMLSensitivity
trainingData = rand(100,4)
p0 = rand(100,100)
function nn!(du,u,p,t)
diffs = [u[j]-u[i] for i in 1:100, j in 1:100]
du = sum(p.*diffs,dims=2)
end
function predict(p)
prob = ODEProblem(nn!,trainingData[:,1],(1.,4.),p)
Array(solve(prob,saveat=1.))
end
function loss(p)
pred = predict(p)
sum(abs2,pred .- trainingData)
end
@time ForwardDiff.gradient(loss,p0);
```
```julia
27.527524 seconds (1.45 M allocations: 122.377 GiB, 33.60% gc time, 0.32% compilation time: 100% of which was recompilation)
```
In the actual code, it takes forever to calculate a single gradient (after several minutes it hasn’t been calculated).
Could somebody advise me on ways to optimize this computational time to be more reasonable?
---
### Author:  [@stevengj](https://discourse.julialang.org/u/stevengj)
#### Post date: [January 10, 2024, 10:59pm UTC](https://discourse.julialang.org/t/optimizating-computational-time-of-gradient-on-big-linear-udes/108645/2 "2024-01-10T22:59:16Z")
> [@jarroyoe](#):
>
> `p0 = rand(100,100)`
You have 10^4 parameters — computing the gradient by forward-mode AD effectively involves solving the ODE 10^4 times. Using a reverse-mode (“adjoint”) algorithm, in contrast, will effectively involve solving the ODE _one_ additional time to get the gradient.
You should probably use an adjoint/reverse method in this regime, not ForwardDiff.
A classic reference on adjoint-method (reverse-mode/backpropagation) differentiation of ODEs (and generalizations thereof) is [Cao et al (2003)](https://epubs.siam.org/doi/10.1137/S1064827501380630) ([pdf](http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.65.455&rep=rep1&type=pdf)). See also the [SciMLSensitivity.jl package’s documention on reverse-mode AD](https://docs.sciml.ai/SciMLSensitivity/stable/getting_started/#Reverse-Mode-Automatic-Differentiation) for adjoint-method sensitivity analysis with DifferentialEquations.jl, along with Chris Rackauckas’s [notes from 18.337](https://rawcdn.githack.com/mitmath/18337/7b0e890e1211bfa253782f7862389aeaa321e8d7/lecture11/adjoints.html). There is a nice YouTube [lecture on adjoint sensitivity of ODEs](https://www.youtube.com/watch?v=k6s2G5MZv-I), again using a similar notation.
---
### Author:  [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [January 10, 2024, 11:18pm UTC](https://discourse.julialang.org/t/optimizating-computational-time-of-gradient-on-big-linear-udes/108645/3 "2024-01-10T23:18:09Z")
And also, since it’s a linear ODE, you can specialize this to simply use `u(t,p) = u0*exp(A(p)*t)` which then has a very simple derivative. SciMLSensitivity currently won’t specialize on linear ODEs but it should in the future, for now it’s quite straight forward to derive. As Steven says though, you will want to do this using the adjoint as forward mode will have O(np) scaling vs O(n+p) scaling of using the adjoint.
---
### Author:  [@stevengj](https://discourse.julialang.org/u/stevengj)
#### Post date: [January 10, 2024, 11:36pm UTC](https://discourse.julialang.org/t/optimizating-computational-time-of-gradient-on-big-linear-udes/108645/4 "2024-01-10T23:36:53Z")
> [@ChrisRackauckas](#):
>
> you can specialize this to simply use `u(t,p) = u0*exp(A(p)*t)` which then has a very simple derivative.
Beware that differentiating a matrix exponential (with respect to parameters of the matrix) is not as simple as many people expect (but ChainRules.jl and hence Zygote.jl can do it). See also [Differentiating random walk probability w.t.r. rate of jump - #14 by stevengj](https://discourse.julialang.org/t/differentiating-random-walk-probability-w-t-r-rate-of-jump/107116/14)
---
### Author:  [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [January 11, 2024, 1:21am UTC](https://discourse.julialang.org/t/optimizating-computational-time-of-gradient-on-big-linear-udes/108645/5 "2024-01-11T01:21:11Z")
Yeah I meant the code is relatively straightforward. If you put that into Zygote it should just work. Though it won’t work with ForwardDiff.jl for the reason that you mention, that the squaring and scaling algorithm used in Base is not differentiable and ForwardDiff is missing a specialized rule on it.
BTW this reminds me, Normally for this kind of function instead of recommending `exp` I’d normally recommend ExponentialUtilities.jl with `expv`, but there’s a missing derivative there for various reasons. Do you happen to know of a better trick than the one mentioned here [Add ChainRules rules · Issue #40 · SciML/ExponentialUtilities.jl · GitHub](https://github.com/SciML/ExponentialUtilities.jl/issues/40#issuecomment-749423055) ?
---
### Author:  [@jarroyoe](https://discourse.julialang.org/u/jarroyoe)
#### Post date: [January 11, 2024, 3:17am UTC](https://discourse.julialang.org/t/optimizating-computational-time-of-gradient-on-big-linear-udes/108645/6 "2024-01-11T03:17:41Z")
Thanks to both of you for your help! Two comments regarding the discussion:
- When trying to use Zygote, I get the following error on the MWE:
```julia
f = LLVM.Function("julia __mapreducedim__ 5966")
(gty, inst, v) = (LLVM.IntegerType[LLVM.IntegerType(i64)], LLVM.PHIInst(%129 = phi double addrspace(13)* [poison, %L89.L110.loopexit_crit_edge.us125.unr-lcssa.us.1.L110.us129.us.1_crit_edge], [%113, %L89.L110.loopexit_crit_edge.us125.unr-lcssa.us.1.thread], [%119, %L93.us121.epil.us.1]), LLVM.PoisonValue(0x000000006a63fb20))
f = LLVM.Function("julia __mapreducedim__ 6969")
(gty, inst, v) = (LLVM.IntegerType[LLVM.IntegerType(i64)], LLVM.PHIInst(%129 = phi double addrspace(13)* [poison, %L89.L110.loopexit_crit_edge.us125.unr-lcssa.us.1.L110.us129.us.1_crit_edge], [%113, %L89.L110.loopexit_crit_edge.us125.unr-lcssa.us.1.thread], [%119, %L93.us121.epil.us.1]), LLVM.PoisonValue(0x00000000086ed390))
┌ Warning: EnzymeVJP tried and failed in the automated AD choice algorithm with the following error. (To turn off this printing, add `verbose = false` to the `solve` call)
└ @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/Rm4xX/src/concrete_solve.jl:23
AssertionError: false
ERROR: UndefRefError: access to undefined reference
```
follow by a long stacktrace.
- Unfortunately my model is not just linear, but the linear part is the one that takes most of the space in memory as it has thousands of parameters.
A small neural network makes predictions over subsets of the data. Maybe a more reasonable MWE of what I am trying to do would be:
```julia
using Zygote, DifferentialEquations, SciMLSensitivity, Lux, Random, ComponentArrays
rng = Random.default_rng()
trainingData = rand(100,4)
p0 = rand(100,100)
chain = Lux.Chain(Lux.Dense(4,5),Lux.Dense(5,4))
ltup = Lux.setup(rng, chain)
ps = ltup[1]
st = ltup[2]
p = ComponentVector(model_params = ps, connectivityMatrix = p0)
function nn!(du,u,p,t)
nns = reduce(vcat,[first(chain(u[((i-1)*4+1):((i-1)*4+4)],p.model_params,st)) for i in 1:25])
diffs = [u[j]-u[i] for i in 1:100, j in 1:100]
du = nns .+ sum(p.connectivityMatrix*diffs,dims=2)
end
function predict(p)
prob = ODEProblem(nn!,trainingData[:,1],(1.,4.),p)
Array(solve(prob,saveat=1.))
end
function loss(p)
pred = predict(p)
sum(abs2,pred .- trainingData)
end
@time Zygote.gradient(loss,p);
```
---
### Author:  [@wsmoses](https://discourse.julialang.org/u/wsmoses)
#### Post date: [January 11, 2024, 7:32pm UTC](https://discourse.julialang.org/t/optimizating-computational-time-of-gradient-on-big-linear-udes/108645/7 "2024-01-11T19:32:35Z")
That error message is coming from Enzyme, and seemingly an older version (that code has been substantially improved since).
What is your package status [and thus version of Enzyme]
---
### Author:  [@stevengj](https://discourse.julialang.org/u/stevengj)
#### Post date: [January 11, 2024, 7:35pm UTC](https://discourse.julialang.org/t/optimizating-computational-time-of-gradient-on-big-linear-udes/108645/8 "2024-01-11T19:35:07Z")
> [@jarroyoe](#):
>
> When trying to use Zygote, I get the following error on the MWE:
Have you tried rewriting it in terms of a matrix exponential, as Chris recommended?
---
### Author:  [@jarroyoe](https://discourse.julialang.org/u/jarroyoe)
#### Post date: [January 11, 2024, 9:16pm UTC](https://discourse.julialang.org/t/optimizating-computational-time-of-gradient-on-big-linear-udes/108645/9 "2024-01-11T21:16:04Z")
Thanks William! I knew I’ve seen that error before. My status was
```julia
Status `~/.julia/environments/v1.9/Project.toml`
⌃ [7da242da] Enzyme v0.11.7
⌃ [e88e6eb3] Zygote v0.6.65
```
I upgraded and now it is
```julia
Status `/central/home/jarroyoe/.julia/environments/v1.8/Project.toml`
[7da242da] Enzyme v0.11.12
[e88e6eb3] Zygote v0.6.68
```
and I am getting a different error:
```julia
ERROR: BoundsError: attempt to access 4-element StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64} at index [0]
```
---
### Author:  [@jarroyoe](https://discourse.julialang.org/u/jarroyoe)
#### Post date: [January 11, 2024, 9:46pm UTC](https://discourse.julialang.org/t/optimizating-computational-time-of-gradient-on-big-linear-udes/108645/10 "2024-01-11T21:46:02Z")
I just did and there is a significant performance boost when doing `loss(p0)` and `ForwardDiff.gradient(loss,p0)`. I am getting errors on the adjoint though.
```julia
using Zygote, DifferentialEquations, SciMLSensitivity
trainingData = rand(100,4)
p0 = rand(100,100)
function predict(p)
A = [i==j ? -sum(p[i,:]) : p[i,j] for i in 1:100, j in 1:100]
reduce(hcat,[trainingData[:,1]'*exp.(A*i) for i in 0:3]')
end
function loss(p)
pred = predict(p)
sum(abs2,pred .- trainingData)
end
@time Zygote.gradient(loss,p0);
```
yields:
```julia
ERROR: MethodError: no method matching adjoint(::Nothing)
Closest candidates are:
adjoint(::Union{LinearAlgebra.QR, LinearAlgebra.QRCompactWY, LinearAlgebra.QRPivoted}) at /central/software/julia/1.8.5/share/julia/stdlib/v1.8/LinearAlgebra/src/qr.jl:517
adjoint(::Union{LinearAlgebra.Cholesky, LinearAlgebra.CholeskyPivoted}) at /central/software/julia/1.8.5/share/julia/stdlib/v1.8/LinearAlgebra/src/cholesky.jl:558
adjoint(::LinearAlgebra.Hessenberg) at /central/software/julia/1.8.5/share/julia/stdlib/v1.8/LinearAlgebra/src/hessenberg.jl:424
```
---
### Author:  [@stevengj](https://discourse.julialang.org/u/stevengj)
#### Post date: [January 11, 2024, 10:00pm UTC](https://discourse.julialang.org/t/optimizating-computational-time-of-gradient-on-big-linear-udes/108645/11 "2024-01-11T22:00:39Z")
> [@jarroyoe](#):
>
> `exp.(A*i)`
I thought you wanted a matrix exponential? This is an elementwise exponential, which is very much not the same thing as `exp(A*i)`.
(Also, `exp(A*i)` for `i = 0:3` is the same as `exp(A)^i` for `i = 0:3`, but computing `exp(A)` once and re-using it will be much faster. I’m also a little confused by why you are multiplying the initial condition as a row vector on the left, though of course it is possible to do this if you have transposed your system matrix. Make sure that your new calculation matches your ODE solution … it looks quite different from what you wrote before at first glance!)
> [@jarroyoe](#):
>
> `no method matching adjoint(::Nothing)`
It appears you are trying to calculate `nothing'` (`== adjoint(nothing)`) somewhere. The stacktrace will tell you where this is being called.
---
### Author:  [@jarroyoe](https://discourse.julialang.org/u/jarroyoe)
#### Post date: [January 11, 2024, 10:28pm UTC](https://discourse.julialang.org/t/optimizating-computational-time-of-gradient-on-big-linear-udes/108645/12 "2024-01-11T22:28:25Z")
Thanks for pointing that out, this was me messing up the math! This code works flawlessly for the linear case.
```julia
using Zygote, DifferentialEquations, SciMLSensitivity
trainingData = rand(100,4)
p0 = rand(100,100)
function predict(p)
A = [i==j ? -sum(p[i,:]) : p[i,j] for i in 1:100, j in 1:100]
reduce(hcat,[exp(A*i)*trainingData[:,1] for i in 0:3])
end
function loss(p)
pred = predict(p)
sum(abs2,pred .- trainingData)
end
@time Zygote.gradient(loss,p0);
```