SciML / SciML/SciMLSensitivity.jl
Continuous adjoints add d(sum(u0))/dp when u0 depends on parameters through initialization
Nobody has claimed this yet.
- Dominant language
- Julia
- Stars
- 396
- Forks
- 90
- Avg merge
- 23h 1m
- Merged PRs (30d)
- 49
Description
For a ModelingToolkit problem whose initial values depend on parameters, InterpolatingAdjoint, GaussAdjoint and QuadratureAdjoint return the correct gradient plus the gradient of sum(u0) with respect to the parameters. ForwardDiffSensitivity and the default sensealg are correct; BacksolveAdjoint is wrong in a different way.
using ModelingToolkit, OrdinaryDiffEq, SciMLSensitivity, Zygote, ForwardDiff
using ModelingToolkit: t_nounits as t, D_nounits as D
using SymbolicIndexingInterface: setp_oop
@parameters a b
@variables y1(t) y2(t) y3(t)
eqs = [D(y1) ~ -(a + b) * y1 + y2, D(y2) ~ -(a + b) * y2 + y3, D(y3) ~ -(a + b) * y3]
sys = mtkcompile(System(eqs, t; name = :sys,
initial_conditions = [y1 => a * cos(0.1), y2 => a * cos(0.2), y3 => a * cos(0.3), a => 1.2, b => 2.1]))
prob = ODEProblem(sys, nothing, (0.0, 1.0))
set_p = setp_oop(prob, [a, b])
sa = InterpolatingAdjoint(autojacvec = ReverseDiffVJP(true))
loss(ps) = sum(Array(solve(remake(prob; p = set_p(prob, ps)), Tsit5(); saveat = 0.1, abstol = 1e-8, reltol = 1e-8, sensealg = sa))[:, 1]) # sum(u(0))
ForwardDiff.gradient(loss, [1.2, 2.1]) # [2.9304, 0.0], = sum(cos.([0.1, 0.2, 0.3]))
Zygote.gradient(loss, [1.2, 2.1])[1] # [5.8608, 0.0], twice that
A loss over the whole trajectory carries the same excess of 2.9304. The same happens with u0 given as a function of ps in remake; a plain ODEProblem without ModelingToolkit and u0 = f(p) is exact.
It comes from concrete_solve.jl: igs = _init_originator_gradient(originator, init_loss, tunables) with init_loss(t) = sum(nu0) is accumulated into dp unconditionally, so the contribution of u0 is seeded with a cotangent of ones instead of the adjoint state at t0. #1643 mentions this as known; I did not find an issue for it.
SciMLSensitivity 7.119.7, ModelingToolkit 11.42.1, SciMLBase 3.53.3, Zygote 0.7.13, Julia 1.12.5.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the reproducible ModelingToolkit example and inspect concrete_solve.jl, especially the _init_originator_gradient call and the adjoint entry points named in the issue. Compare the continuous adjoints with ForwardDiffSensitivity and the default sensealg; done means no extra sum(u0) contribution remains for parameter-dependent initialization, including function-based u0 and whole-trajectory losses.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 58/100