Add opt-in PreserveDiffModel to keep the diff model across Parameter-only re-solves#371
Open
yeabbratz wants to merge 3 commits into
Open
Add opt-in PreserveDiffModel to keep the diff model across Parameter-only re-solves#371yeabbratz wants to merge 3 commits into
yeabbratz wants to merge 3 commits into
Conversation
…ly re-solves MOI.optimize! unconditionally discards Optimizer.diff, so the first differentiation after every re-solve re-instantiates the diff model and re-copies the whole problem, even when the only change since the last instantiation was MOI.Parameter values. Add an opt-in PreserveDiffModel optimizer attribute (default false, stock behavior unchanged): when enabled and only parameter values changed - every structural edit already resets diff to nothing, so a surviving diff is structurally in sync by construction - optimize! keeps the diff model and refreshes its point-dependent state instead: input sensitivities are cleared and the new solution is re-copied after each solve, and parameter values (their ConstraintSet and their slot in the primal point) are re-written at the next differentiation call. Results are identical to the full-rebuild path; a solve that ends in a non-differentiable status falls back to discarding diff. Only differentiation models that natively support MOI.Parameter can be preserved (a ParametricOptInterface layer inside diff would not propagate value updates); models opt in through _diff_model_supports_preserve, currently NonLinearProgram.Model.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #371 +/- ##
==========================================
+ Coverage 91.88% 91.92% +0.04%
==========================================
Files 17 17
Lines 2722 2763 +41
==========================================
+ Hits 2501 2540 +39
- Misses 221 223 +2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Author
|
Updated for review: applied Note: the reformatted files are all under |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
MOI.optimize!(::DiffOpt.Optimizer)unconditionally setsmodel.diff = nothing, so the firstdifferentiation after every re-solve re-instantiates the differentiation model and re-runs a
full
MOI.copy_toof the problem. That is the right default — a re-solve invalidates thepoint-dependent state of
diff— but when the only change since the last instantiation isMOI.Parametervalues (the standard parametric training / sensitivity loop:set_parameter_value→optimize!→ differentiate, repeatedly), the structure insidediffis still valid, and the re-instantiate +
copy_tois pure per-solve overhead.This PR adds an opt-in optimizer attribute, default
false, stock behavior unchanged:When enabled,
optimize!keepsdiffand refreshes exactly the point-dependent state insteadof discarding everything:
diffare cleared and the new primal/dual solution is re-copiedafter each solve (
_copy_dual— the same call the rebuild path runs aftercopy_to, fed bythe same source);
diffat the next differentiation call(
_refresh_parameters): both theirMOI.ConstraintSetand their slot in the primal point(
_copy_dualcopiesMOI.VariablePrimalof every variable, and for a parameter that is itsvalue).
Why a surviving
diffis structurally in sync by construction: every structural edit onDiffOpt.Optimizer(add/delete variable or constraint, function/set modification, objectivechange,
copy_to, …) already resetsdiff = nothing— this PR only narrows theMOI.ConstraintSetsetter forMOI.Parameterto not discard it when the attribute is enabled.So "only parameter values changed" needs no separate dirty tracking: it is the only mutation
that leaves
diffalive. A solve that ends in a non-differentiable termination status alsofalls back to discarding
diff.Scope: only differentiation models that natively support
MOI.Parametercan be preserved —with a ParametricOptInterface layer inside
diff, a value update would not propagate into thealready-copied problem data. Models opt in through the internal
_diff_model_supports_preservetrait; this PR enables it for
NonLinearProgram.Modelonly (where a parameter update is anexact value assignment). For every other configuration the attribute has no effect.
Stale sensitivities cannot be served: the re-solve path re-copies
x/y/s, clears the inputcaches, and the output caches (
forw_grad_cache/back_grad_cache) are overwritten by everydifferentiation call. A parameter-only re-solve with
diffpreserved yields results identicalto the full rebuild — the tests assert
==, not≈.Tests
test_preserve_diff_model_parameter_resolves: 4 rounds ofset_parameter_value→optimize!→ reverse on two identical models (one preserved, one stock): gradients==atevery round, and the preserved model provably reuses the same
diffobject; forward modechecked too.
test_preserve_diff_model_parameter_change_without_resolve: changing parameters anddifferentiating without re-solving matches stock (which rebuilds with the new parameter
values at the old solution).
test_preserve_diff_model_structural_change_falls_back: a structural edit discards thepreserved
diff; the rebuilt model matches a fresh one constructed in the final state.parameters.jlset.Benchmark
Median of 7 (after warmup) of the first
reverse_differentiate!after a Parameter-onlyre-solve (the solve itself untimed) — the diff-model re-instantiate +
copy_tois exactly theresidual this PR removes. Parameterized NLP with P parameters and P variables (script below),
Ipopt, M4 Pro, 1 thread, Julia 1.12.6, both columns on this branch:
PreserveDiffModelWith the attribute enabled, the first backward after a re-solve lands on the level of a
repeated backward on the same solve (0.31 / 1.30 / 45.5 ms in the same run) — i.e. the copy
residual is fully removed; what remains at large P is the dense sensitivity contraction (#369's
target) and the O(P) parameter refresh.
bench_persistent_evaluator.jl
Context
Same residual-removal series as #369 (single adjoint solve for the nonlinear reverse) and the
evaluator-cache-reuse PR (the two compose: preserving
diffskips the copy, reusing the cacheskips the tape rebuild). In a downstream differentiable-ACOPF training loop the combination was
validated bitwise against the rebuild path across a case14→case2000 ladder and yields ×12.6
cumulative per-backward together with #369 at 2000-bus scale.
🤖 Generated with Claude Code