Automatic differentiation is too slow

[hersle]

Automatic differentiation (with ForwardDiff.jl) does not perform as well as it should. There seems to be type instabilities when solving ODEProblems generated by ModelingToolkit related to the (combined?) use of dual numbers, splines, reinitialization, etc.

[hersle]

SymBoltz should use SciMLSensitivity.

See also Chris' comment in the announcement post discussion.

[hersle]

Experimenting without MTK:

using OrdinaryDiffEq, Zygote, SciMLSensitivity, ForwardDiff, FiniteDiff

function f(du, u, p, t)
    a = u[1]
    Ωr0, Ωm0, ΩΛ0 = p[1], p[2], p[3]
    a′ = √(Ωr0/a^4 + Ωm0/a^3 + ΩΛ0) * a^2
    du[1] = a′
end

prob = ODEProblem(f, [NaN], (0.0, 2.0), [NaN, NaN])

function a_final(x; sensealg = nothing)
    Ωr0, Ωm0 = x[1], x[2]
    ΩΛ0 = 1 - Ωr0 - Ωm0
    p = [Ωr0, Ωm0, ΩΛ0]
    u0 = [1e-5]
    newprob = remake(prob; u0, p)    
    sol = solve(newprob, Tsit5(); save_everystep = false, sensealg)
    return sol.u[end][1]
end

a_final_1(x) = a_final(x; sensealg = SciMLSensitivity.ForwardSensitivity())
a_final_2(x) = a_final(x; sensealg = SciMLSensitivity.ForwardDiffSensitivity())
a_final_3(x) = a_final(x; sensealg = SciMLSensitivity.GaussAdjoint())
a_final_4(x) = a_final(x; sensealg = SciMLSensitivity.InterpolatingAdjoint())
a_final_5(x) = a_final(x; sensealg = SciMLSensitivity.QuadratureAdjoint())
a_final_6(x) = a_final(x; sensealg = SciMLSensitivity.ReverseDiffAdjoint())

x0 = [1e-5, 0.3]
a_final(x0)
@time Zygote.gradient(a_final_1, x0)
@time Zygote.gradient(a_final_2, x0)
@time Zygote.gradient(a_final_3, x0)
@time Zygote.gradient(a_final_4, x0)
@time Zygote.gradient(a_final_5, x0)
@time Zygote.gradient(a_final_6, x0)
@time ForwardDiff.gradient(a_final, x0)
@time FiniteDiff.finite_difference_gradient(a_final, x0)

  0.001261 seconds (3.07 k allocations: 116.812 KiB)
  0.001561 seconds (3.45 k allocations: 154.250 KiB)
  0.001713 seconds (4.30 k allocations: 199.938 KiB)
  0.001538 seconds (3.80 k allocations: 158.016 KiB)
  0.002294 seconds (9.87 k allocations: 319.609 KiB)
  0.002102 seconds (14.49 k allocations: 588.094 KiB)
  0.000078 seconds (89 allocations: 4.750 KiB)
  0.000071 seconds (328 allocations: 16.000 KiB)

[hersle]

Experimenting with MTK after last updates:

using ModelingToolkit, SymbolicIndexingInterface, OrdinaryDiffEq, Zygote, SciMLSensitivity, ForwardDiff, FiniteDiff

@independent_variables t
D = Differential(t)
@variables a(t)
@parameters Ωr0 Ωm0 ΩΛ0
eqs = [D(a) ~ √(Ωr0/a^4 + Ωm0/a^3 + ΩΛ0) * a^2]
@mtkbuild M = System(eqs, t)

prob = ODEProblem(M, [M.a => 1e-5, M.Ωr0 => NaN, M.Ωm0 => NaN, M.ΩΛ0 => NaN], (0.0, 2.0))

setter = setsym_oop(prob, [M.Ωr0, M.Ωm0, M.ΩΛ0])
function a_final(x; sensealg = nothing)
    Ωr0, Ωm0 = x[1], x[2]
    ΩΛ0 = 1 - Ωr0 - Ωm0
    u0, p = setter(prob, [Ωr0, Ωm0, ΩΛ0])
    newprob = remake(prob; u0, p)    
    sol = solve(newprob, Tsit5(); save_everystep = false, sensealg)
    return sol.u[end][1]
end

a_final_1(x) = a_final(x; sensealg = SciMLSensitivity.ForwardSensitivity())
a_final_2(x) = a_final(x; sensealg = SciMLSensitivity.ForwardDiffSensitivity())
a_final_3(x) = a_final(x; sensealg = SciMLSensitivity.GaussAdjoint())
a_final_4(x) = a_final(x; sensealg = SciMLSensitivity.InterpolatingAdjoint())
a_final_5(x) = a_final(x; sensealg = SciMLSensitivity.QuadratureAdjoint())
a_final_6(x) = a_final(x; sensealg = SciMLSensitivity.ReverseDiffAdjoint())

x0 = [1e-5, 0.3]
a_final(x0)
#@time Zygote.gradient(a_final_1, x0) # fails
@time Zygote.gradient(a_final_2, x0)
@time Zygote.gradient(a_final_3, x0)
@time Zygote.gradient(a_final_4, x0)
@time Zygote.gradient(a_final_5, x0)
#@time Zygote.gradient(a_final_6, x0) # fails
@time ForwardDiff.gradient(a_final, x0)
@time FiniteDiff.finite_difference_gradient(a_final, x0)

  0.005766 seconds (18.64 k allocations: 1.417 MiB, 0.01% compilation time)
  0.020760 seconds (46.00 k allocations: 3.468 MiB, 0.01% compilation time)
  0.013520 seconds (46.25 k allocations: 3.445 MiB, 0.01% compilation time)
  0.014456 seconds (48.86 k allocations: 3.574 MiB, 0.01% compilation time)
  0.000102 seconds (120 allocations: 12.203 KiB)
  0.000103 seconds (458 allocations: 42.062 KiB)

[hersle]

This is now much better with the latest LinearSolve sparsity + ForwardDiff improvements:

using SymBoltz, ForwardDiff
M = ΛCDM()
pars = parameters_Planck18(M)
prob = CosmologyProblem(M, pars)

l = 25:25:2500
jl = SphericalBesselCache(l)
diffpars = [M.c.Ω₀, M.b.Ω₀]
probgen = parameter_updater(prob, diffpars)
DlTT(θ) = spectrum_cmb(:TT, probgen(θ), jl; normalization = :Dl)
θ = [pars[par] for par in diffpars]
@time DlTTs = DlTT(θ)
@time ∂DlTTs_∂θ_ad = ForwardDiff.jacobian(DlTT, θ)

  0.289164 seconds (647.73 k allocations: 1.263 GiB)
  1.360755 seconds (913.83 k allocations: 3.470 GiB, 9.83% gc time)

The ForwardDiff time is now roughly the same time it would take to run the equivalent finite differences.

Reverse-mode AD is still left for future work.