SciML / SciML/SciMLSensitivity.jl

Continuous adjoints add d(sum(u0))/dp when u0 depends on parameters through initialization

Open
#1,663 0 comments 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.