Add Gauss-Newton gradient preconditioner (pyzag.preconditioning) - #31
Open
hugary1995 wants to merge 4 commits into
Open
Add Gauss-Newton gradient preconditioner (pyzag.preconditioning)#31hugary1995 wants to merge 4 commits into
hugary1995 wants to merge 4 commits into
Conversation
|
A data-driven, range-free alternative to RangeRescale for least-squares model
calibration: precondition the gradient with the Gauss-Newton curvature
H = J^T W J (or its diagonal), delta_theta ~ (H + lam*diag(H))^{-1} g.
GaussNewtonPreconditioner wraps any torch optimizer (Adam, SGD, ...) and only
reshapes the gradient, so the base optimizer is swappable. H is estimated from a
subsample of residual rows (one reverse-mode sweep each) and refreshed on a free
gain-ratio trigger -- or never (rho=None) for a fixed preconditioner, the
data-driven analogue of a static per-parameter rescaling. Non-finite steps are
skipped with a warning; an on_refresh callback / refresh_steps expose the
recompute schedule.
Includes unit tests (linear least-squares): exact-Newton in full mode, estimator
accuracy, gain-ratio and fixed-once refresh, and optimizer swappability.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
hugary1995
force-pushed
the
feat/gauss-newton-preconditioning
branch
from
July 31, 2026 12:57
62136f8 to
4fb02b3
Compare
…upport
Supersedes the initial preconditioner in this branch. The redesign is driven by
one finding: the original wrapped Adam, and **Adam is exactly invariant to
gradient preconditioning**. It divides each coordinate by its own running
gradient RMS, so multiplying the gradient by any fixed diagonal leaves the step
bit-identical -- measured on a quadratic with curvature spread 1e6. Wrapping Adam
was a silent no-op, so the constructor now rejects the Adam family outright.
pyzag/curvature.py (new)
CurvatureEstimator: the shared H = J^T W J extraction, reverse-mode only (one
sweep per sampled row), with optional cotangent grouping so a plate-structured
model can get exact per-member rows at nbatch-fold fewer sweeps. Defaults to
every row -- subsampling is an opt-in trade, and on the NEML2 calibration
nsub=8 produces an unusable estimate on 5 draws in 12.
pyzag/preconditioning.py
GaussNewtonCurvature adds caching, staleness and Levenberg-Marquardt damping;
GaussNewtonPreconditioner drives the step. It keeps nsub=8, which is safe here
precisely because damping bounds the step and the next refresh replaces a bad
estimate -- a frozen reparametrization has neither defence, hence the
asymmetric defaults.
gauss_newton_rescalers() is the second lever: measure H once and hand
CurvatureRescale scalers to any optimizer, Adam included, because a
reparametrization moves the optimizer's state into the scaled coordinates
along with the metric. It rejects an ill-conditioned estimate rather than
returning a scale that is wrong by orders of magnitude.
Two failure-handling fixes, both found on a real calibration and not by the
unit tests:
- A non-finite objective now counts as a rejection. Every `nan > tol` compare
is False, so a NaN read as an *accepted* step, poisoned the anchor, and froze
the damping; a run then burned its whole budget at a dead point while
reporting a stale loss as if it had converged.
- The last evaluable point is tracked apart from the quadratic-model anchor. A
rollback invalidates the model but not the point it just restored to, and
clearing both together left the following step with no fallback, so two
consecutive failures raised while a good point sat in the parameters.
pyzag/reparametrization.py
CurvatureRescale, plus a fix for root-module parameters whose names picked up a
leading dot and failed to resolve.
pyzag/stochastic.py
PyroGaussNewtonOptim + PreconditionedSVI. Pyro builds one optimizer per
parameter and SVI.step backwards before the optimizer runs, so neither the
cross-parameter curvature nor the pre-update loss is reachable through the
stock interfaces. gaussian_map_residual builds the GN residual from a trace
including the prior terms -- the hyper sites enter only through the prior, so a
likelihood-only residual would silently zero them.
Tests: 59 across the two preconditioning suites, pinning the Adam invariance
itself (bypassing the guard, so the premise stays measured), the NaN and
consecutive-failure recovery paths, and the estimator defaults. Full suite green:
160 passed, 745 subtests.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Three errors, none of them behavioural:
- C0412 ungrouped pyro imports, from how the merge resolved that block.
- E1129 on `with pyro.poutine.trace(...)`. A false positive -- `trace` is a
factory pylint cannot resolve the return type of. Switched to the concrete
`trace_messenger.TraceMessenger`, which constructs an identical object and
follows the precedent already set for the scale/mask handlers a few hundred
lines above ("Nested context managers required so pylint can resolve scale
and mask").
- E0606 `eps` possibly-unbound. Pre-existing rather than new -- it reproduces on
the previous commit, and only started failing CI because the workflow installs
pylint fresh and a newer release detects it. Also a false positive, since
`sample_noise_outside` is fixed in __init__ and the two assignments are
exclusive, but the exclusivity was invisible to a reader as well, so it is now
written as a sentinel rather than suppressed. `self.eps` is a PyroSample, so
the read location decides where the site is sampled; both reads stay exactly
where they were.
pylint 10.00/10, black clean, copyright clean. Suite still 160 passed, 745
subtests.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Findings from re-reading the branch after several redesigns: - `_run_scaled` in test_preconditioning.py was dead *and* broken: no callers, and it passed `apply=`, a constructor argument that went away with the step-scaling mode. It would have raised TypeError if anyone had called it. Removed. (`_badly_scaled`, which it used, is still used by two live tests.) - The comment standing in for a removed non-finite test referred to "the pre-fix contract" -- a test that only ever existed inside this PR's own history, so the reference means nothing once merged. Trimmed to the part that is still useful: where that coverage now lives. - `test_it_actually_optimizes` unpacked an optimizer it never used. - `reset_anchor` now says it deliberately leaves `_good` alone. That the quadratic-model anchor and the last evaluable point have different lifetimes is the subtlest thing in the module, and this is the method that embodies it. Also checked and clean: no TODO/debug leftovers, no surviving references to the removed step-scaling or Adam-wrapping designs, every public function reachable, all doc-referenced symbols present, toctree resolves, and documented defaults (estimator exact, preconditioner nsub=8) match the code. pylint 10.00/10, black clean, 59 preconditioning tests pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
hugary1995
force-pushed
the
feat/gauss-newton-preconditioning
branch
from
August 7, 2026 12:12
c08f835 to
75c6108
Compare
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
1 participant
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.
What
Adds
pyzag.preconditioning.GaussNewtonPreconditioner— a data-driven, range-free alternative toRangeRescalefor least-squares model calibration.Minimizing
L(θ) = ½ r(θ)ᵀ W r(θ), the raw gradient is badly scaled when parameters have heterogeneous magnitudes/sensitivities. Instead of a hand-picked per-parameter range, precondition the gradient with the Gauss–Newton curvatureH = JᵀWJ(or its diagonal):Δθ ∝ (H + λ·diag(H))⁻¹ g.Design
Jᵀvproducts (one reverse-mode sweep each, e.g. throughsolve_adjoint);His estimated from a subsample of residual rows.rho) refreshesHonly when the cached curvature goes stale;rho=Nonecomputes it once (a fixed preconditioner — the data-driven analogue of a static rescaling).min_refresh_intervalbounds cost;on_refresh/refresh_stepsexpose the schedule.Tests
test/test_preconditioning.py(linear least-squares, no recursive solve needed): full-mode == exact Newton in one step,H == AᵀA, estimator accuracy, gain-ratio avoids re-refresh on an exact quadratic, fixed-once (rho=None) computes once, non-finite warns + skips, optimizer swappability.Notes
The module is torch-only (no pyzag internals), so it applies unchanged across pyzag versions. Motivated by a NEML2 Kocks–Mecking viscoplastic calibration via pyzag's adjoint, where the diagonal GN preconditioner computed once matches a well-tuned
RangeRescale+Adam and beats a mistuned one — with no parameter ranges.🤖 Generated with Claude Code