fix: LM damping parity + stall guards for potential-correction iterative engines - #676
Merged
Merged
Conversation
…ive engines The iterative engines' Marquardt mu*diag(H) damping diverged from the reference implementation's mu*I, changing recovered signals at fixed iteration budgets (PyAutoLens#672). damping= option added to both engines (imaging default restored to identity; interferometer keeps marquardt); rejected steps below tol now return as converged and consecutive rejections are capped, ending the end-of-budget rejection storm. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This was referenced Jul 31, 2026
Collaborator
Author
|
Workspace PR (phase 3 acceptance test): PyAutoLabs/autolens_workspace_test#243 |
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.
Summary
Closes the potential-correction JAX-vs-Python discrepancy investigation of #672 (upstream author report: "for the same set of regularization parameters ... the recovered mass-perturbation signal differs between my original pure Python code and the new JAX-ported version").
Side-by-side on the reference 200x200 demo dataset (identical inputs, cold start, the author's demo-2 hyper-params, n_iter=5) localized the divergence to the LM damping matrix: the reference implementation damps
H + mu*I(in both its numpy and JAX paths), while the port used the scale-invariant Marquardt formH + mu*diag(H)everywhere. The two are not interchangeable at a fixed iteration budget — identical initial costs (5.0276e6, cost functions agree exactly), but after 5 iterations the reference engine reaches cost 6.2694e3 while the Marquardt-damped port sits at 1.0367e4 with a noise dkappa (corr 0.024). Withmu*Ipatched in, the port matches the reference trajectory to 4 significant figures at every iteration (6.2696e3) — the implementations are equivalent; the damping form was the entire discrepancy. One-shot (single-step) inversion needed no fix: machine-precision parity (evidence diff 6e-7, dkappa corr 1.000000).The investigation also exposed an end-of-budget failure mode in both implementations: at a cost minimum no decreasing step exists, so the inner LM loop rejects trial steps (each a full Jacobian rebuild, ~minutes at 200x200) until mu exceeds 1e15 — 47 wasted minutes and, for the reference implementation on a 15GB machine, an OOM death mid-storm.
Changes:
dense_util.solve_lm_step_fromgainsdamping="identity"|"marquardt"(function default"marquardt", unchanged — the JAX-parity smoke scripts that call it directly are unaffected).IterFitDpsiSrcImaginggainsdamping(default"identity"— restores the reference implementation's behavior: cold-start converges in a few iterations at the author's published hyper-params) andmax_consecutive_rejections(default 10).IterFitDpsiSrcInterferometergains the same two options withdamping="marquardt"retained as its default (visibility-weighted curvatures ~1e11 are why the scale-invariant form exists; the research: realistic-uv iterative potential-correction recovery campaign #627-certified uv recipe is unchanged).tolnow returns as converged (growing mu only shrinks the step further — previously the tol check fired only on accepted steps), andmax_consecutive_rejectionsbounds the rejection storm.API Changes
Changed Behaviour
al.pc.IterFitDpsiSrcImagingnow defaults to identity LM damping (damping="identity"), matching the reference implementation of Cao et al. 2025. Cold-start recovery at published hyper-params works again without the feat: x0 warm start for the imaging iterative engine (parity) #630 warm start (the warm start remains supported and is still the certified recipe underdamping="marquardt"). Passdamping="marquardt"for the previous behavior.tolor aftermax_consecutive_rejectionsconsecutive rejections, instead of rejecting untilmu > 1e15. States returned at these exits are the same current-best states the mu-ceiling exit returned; runs that previously burned a rejection storm at convergence now return up to tens of minutes sooner.Added
dampingandmax_consecutive_rejectionsparameters onal.pc.IterFitDpsiSrcImagingandal.pc.IterFitDpsiSrcInterferometer;dampingparameter onal.pc.dense_util.solve_lm_step_from.Test Plan
test_autolens/full suite: 492 passed (includes new identity-damping solve tests, damping plumb-through tests, and a warm-start-at-optimum stall-guard test bounding Jacobian rebuilds).autolens_workspace_testimaging/subhalo_recovery.pyrun against this branch under the exact CI smoke env — passes with identical one-shot evidence.Generated by the PyAutoLabs agent workflow.
🤖 Generated with Claude Code