Skip to content

Ridge Levenberg–Marquardt

RidgeLevenbergMarquardt solves underdetermined interpolation problems by minimizing a ridge objective whose small-ridge limit is the minimum-seminorm interpolant. Where LevenbergMarquardt encodes the selection in a parameter-space metric (the algorithm's damping geometry), this solver encodes it in the objective — every inner problem is a well-posed nonlinear least squares that any solver can attack with textbook guarantees.

Problem statement

The flattened parameter vector splits into two blocks, \(x = [x_m;\, x_f] \in \mathbb R^p\): the metric block \(x_m \in \mathbb R^{n_m}\), covered by a user positive-definite metric \(W \succ 0\) (a Metric), and the free block \(x_f \in \mathbb R^{n_f}\), unpenalized (\(n_f = p - n_m \ge 0\) is inferred from the iterate at init). Given an underdetermined residual \(r(x) \in \mathbb R^m\) (\(m < p\)), the target is the minimum-seminorm interpolant

\[ x^\dagger \;=\; \operatorname*{argmin}_x \; \|x_m\|_W^2 \quad\text{s.t.}\quad r(x) = 0 , \]

with \(\|v\|_W = \sqrt{v^\top W v}\), identified under the condition that no null direction of \(J\) lives entirely in the free block (\(\ker J \cap (\{0\} \times \mathbb R^{n_f}) = \{0\}\) at the solution; Eldén 1982 — the full-space penalty \(\mathrm{blockdiag}(W, 0)\) is PSD by design). For kernel coefficient problems \(\|x_m\|_W^2\) is the squared RKHS norm, so \(x^\dagger\) is the minimum-RKHS-norm interpolant.

The ridge family and why the objective carries the selection

The solver minimizes, for a strictly positive weight \(\lambda\) (the ridge),

\[ F_\lambda(x) \;=\; \|r(x)\|^2 + \lambda\,\|x_m\|_W^2 . \]

Under the identification condition (plus standard local regularity of an isolated constrained minimizer) each fixed-\(\lambda\) problem has an isolated minimizer \(x_\lambda\), and

\[ x_\lambda \;=\; x^\dagger + O(\lambda) \quad\text{as } \lambda \downarrow 0 \]

— classical nonlinear Tikhonov regularization (Engl–Kunisch–Neubauer 1989; Engl–Hanke–Neubauer 1996, Ch. 10). The selection lives in the objective on purpose: plain Gauss–Newton on an underdetermined system converges to some interpolant, and characterizing which one requires either an explicit null-space correction (Campbell–Kunkel–Bobinyec 2012, relaxed in Pes–Rodriguez 2022) or an affine residual (Izmailov–Solodov 2026). With the ridge objective no such characterization is needed — the minimizer itself carries the selection, at any solver accuracy.

The whitened change of variables

The metric enters through an invertible factor \(F\) with \(W = F^\top F\); \(F\) should be upper triangular, and the canonical example is the upper Cholesky factor, F = jnp.linalg.cholesky(K, upper=True). The solver extends it by the identity on the free block and runs entirely in the whitened variable

\[ y = \bar F x, \qquad \bar F = \operatorname{blockdiag}(F,\, I_{n_f}), \qquad \tilde J = J \bar F^{-1}, \qquad E = \operatorname{blockdiag}(I_{n_m},\, 0) , \]

where the penalty is simply \(\lambda \|y_m\|^2\). This is an exact, constant linear change of variables — the chain rule is applied by hand once, never by differentiating through a solve — and neither \(W\) nor \(\bar F\) is ever materialized: the algorithm only applies the metric's factor ops (\(F v\), \(F^{-1} v\), \(F^{-\top} v\)).

Everywhere in the implementation the equivalent augmented-residual view is used: standard Euclidean LM on

\[ R(y) = \begin{bmatrix} r(x) \\ \sqrt{\lambda}\, y_m \end{bmatrix}, \qquad A = \begin{bmatrix} \tilde J \\ \sqrt{\lambda}\, [\,I_{n_m} \mid 0\,] \end{bmatrix} \]

— the rows of \(F\) are the appended penalty equations, and in \(y\) they are the constant stack \([\,I \mid 0\,]\). Each LM step solves

\[ \left(\tilde J^\top \tilde J + \lambda E + \mu I\right) \delta_y = -\left(\tilde J^\top r + \lambda\,[\,y_m;\, 0\,]\right), \qquad \delta_x = \bar F^{-1} \delta_y , \]

(equivalently, the least-squares problem with \(\sqrt{\mu}\,I\) damping rows appended to \(A\)); the iterate is stored in \(x\), and trial penalties use the linearity \(\bar F(x + \delta_x) = y + \delta_y\), so no second factor application is ever needed.

Three properties follow from the penalty rows being constant in \(y\):

  • the gradient identity \(A^\top R = \tilde J^\top r + \lambda [y_m; 0]\) is exact;
  • the penalty rows have zero second directional derivative, so geodesic acceleration (Transtrum–Sethna 2012) uses the true residual only — \(R_{vv} = [\,f_{vv};\,0\,]\) — and the standard formulas apply unchanged;
  • accept/reject compares the plain scalar objective \(F_\lambda\) (Marquardt 1963; Moré 1978; Nocedal–Wright 2006, Ch. 10), with the trust region posed in the whitened geometry: \(\mu\|\delta_y\|^2 = \mu\,(\|\delta_{x_m}\|_W^2 + \|\delta_{x_f}\|^2)\).

The conditioning payoff is structural: the whitened normal matrix \(\tilde J^\top \tilde J + \lambda E\) has a clean spectral floor at \(\lambda\) on the metric block regardless of \(W\)'s conditioning — no interaction between the kernel's conditioning and the Gram product — which is what keeps the default Cholesky() path accurate at deep ridge. (Measured on the production DAE kernel drivers at \(\lambda = 3{\times}10^{-12}\): 25–35% fewer LM steps than the unwhitened x-space formulation at ~8× lower per-step cost than its QR fallback — 1.8–3× faster wall-clock with identical solution quality; that experiment is why the whitened path is the only one.)

The two scalars are fully decoupled:

  • \(\mu\) (damping) is the trust-region parameter — it moves every step by accept/reject exactly as in stock LM;
  • \(\lambda\) (ridge) is the selection weight — carried as traced state, changed only by init() or a solve callback, annealed monotonically toward a positive floor. \(\lambda\) never enters the metric's factorization, so continuation composes unchanged.

\(\lambda = 0\) is out of contract everywhere (constructor validation, callback contract, continuation floors), so \(\tilde J^\top \tilde J + \lambda E\) is positive definite under the identification condition at every reachable state.

Ridge continuation

The homotopy \(\lambda \downarrow \lambda_{\min}\) is a callback, not solver machinery — AnnealRidge is the shipped schedule, and its docstring documents the mechanics for composing it into a callback of your own:

from nlls_gram import AnnealRidge

anneal = AnnealRidge(ridge_floor=1e-8, decrease=0.1)
result = solver.solve(x0, callback=anneal, user_state=anneal.init_state(),
                      atol=1e-8, max_steps=500)

The callback multiplies lm_state.ridge by decrease whenever the inner fixed-\(\lambda\) problem is approximately stationary (info.grad_norm below grad_rtol relative to its reference at the current level), never below ridge_floor. Two pacing notes from production use: the per-level references compound (each level's reference is the gradient it entered with), so a schedule that freezes below the noise floor wants a wider decrease (e.g. 0.01 — larger jumps keep the references generous) or the opt-in stall_rtol stagnation advance; and with the conjunctive stopping rule, choose atol between the ridge-floor residual and the last intermediate level's residual (they differ by roughly 1 / decrease) — the intermediate levels are stationary too, and atol is what rules them out, so the solve can only stop at the floor. Solving each level out and passing to the limit is the nonlinear Tikhonov path above; annealing per stationarity event is the iteratively regularized Gauss–Newton method (Bakushinskii 1992; Blaschke–Neubauer–Scherzer 1997; Kaltenbacher–Neubauer–Scherzer 2008), whose theory wants exactly this monotone, boundedly geometric schedule. A callback ridge change is treated as a problem change: that step's convergence test is suppressed (its diagnostics were computed at the old \(\lambda\)) and the ridge-keyed factorization caches invalidate.

Stopping: the two-phase picture

A ridge solve has two phases. Phase 1 drives the residual to its floor — fast, a handful of Gauss–Newton-quality steps. Phase 2 slides the iterate along the interpolation set, resolving the null-space (selection) component while \(\|r\|\) stays essentially constant. A pure-residual test is blind to phase 2. The toy problem \(r(x) = x_1 - 1\) with the identity metric from \(x_0 = (0, 3)\) makes it concrete: the residual floors at \((1, 3)\), and everything that happens afterwards — the slide to the ridge minimizer \((1/(1{+}\lambda),\, 0)\) — is invisible to \(\|r\|\). Stopping on the residual alone returns an answer whose second coordinate is wrong by 3, at machine-perfect residual.

The stopping geometry is the whitened one: info.step_norm (bounded by xtol) is \(\|\delta_y\|\) — the W-norm of the step on the metric block — and info.grad_norm (bounded by gtol) is the whitened stationarity

\[ \left\| \tilde J^\top r + \lambda\,[\,y_m;\, 0\,] \right\| \]

— the dual \(W^{-1}\)-norm of the half-gradient on the metric block. Objective values (loss, resid_loss, penalty_value) are invariant under the change of variables.

Phase 2 is what gtol exists to detect, and it should be calibrated, not guessed. At a ridge minimizer the gradient vanishes as an exact cancellation of its two terms,

\[ \tilde J^\top r \;=\; -\lambda\,[\,y_m;\, 0\,] , \]

each of magnitude \(\approx \lambda\,\|y_m\| = \lambda \sqrt{q(x)}\) with \(q(x) = \|x_m\|_W^2\) the penalty value. The calibrated bound asks that the computed sum be below \(c\) times that common magnitude — i.e. that the two terms cancel to relative accuracy \(c\):

\[ \|\tilde J^\top r + \lambda [y_m; 0]\| \;<\; c \cdot \lambda\,\sqrt{q}, \qquad c \sim 10^{-3} . \]

It is the standard relative-residual criterion of numerical linear algebra (\(\|Ax - b\| \le \mathrm{tol}\,\|b\|\)) applied to the stationarity equation, and the link to the selection error is first-order: along null-space directions the whitened Hessian is \(\lambda E\), so a gradient of norm \(g\) leaves a relative null-space displacement of exactly \(g / (\lambda\sqrt q)\). A gtol at \(c\,\lambda\sqrt q\) resolves the selection to \(\sim c\) relative accuracy; a gtol orders of magnitude above it leaves it visibly unresolved (measured on the asset pricing driver: a loose gtol left the free scalar \(p_0\) off by 310% at a machine-perfect residual; the calibrated one, 0.1%). The solver reports \(\sqrt q\) directly — info.penalty_grad_norm equals sqrt(info.penalty_value) — and the solution's seminorm is usually known to an order of magnitude before any pilot run, so the recipe is simply

gtol = 1e-3 * ridge * jnp.sqrt(q)   # q = the solution's squared (RKHS) seminorm
result = solver.solve(x0, max_steps=400, gtol=gtol, atol=2e-8)

The same formula is a scaled dual-feasibility criterion in the constrained-optimization sense — the ridge problem is the quadratic-penalty form of \(\min \|x_m\|_W^2\) s.t. \(r = 0\) with multiplier estimate \(\nu = r/\lambda\), and codes like IPOPT scale their KKT-residual tests by multiplier magnitudes in exactly this way. Note the consequences: gtol must be re-calibrated when the residual scaling, the metric scaling, or the ridge level changes (all three move \(\lambda\sqrt q\)), and during continuation a single absolute gtol can only be right at one level — pair it with atol placement (above) so the solve stops only at the floor.

Floors and precision. The achievable grad_norm is bounded by \(\varepsilon_{\text{mach}}\) and the conditioning of the augmented stack (which grows like \(1/\sqrt{\lambda}\)); gtol must sit above that noise floor or the solve dies loudly at MAX_STEPS. Two regimes follow. On well-conditioned problems at \(\lambda \gtrsim \sqrt{\varepsilon}\), the calibrated \(10^{-3}\lambda\sqrt q\) sits comfortably above the floor and the recipe applies as stated (in float32 the same structure holds with \(\varepsilon \sim 10^{-7}\); use the ridge=None default \(\lambda = \sqrt{\varepsilon_{f32}}\) and a looser \(c\)). Deep continuation floors on ill-scaled Jacobians (\(\lambda\) far below \(\sqrt{\varepsilon}\)) can push the noise floor above the calibrated value — there gtol becomes a measured constant instead: run once, watch where info.grad_norm flattens, and set gtol just above it. And one floor-choice rule: pushing the ridge floor lower makes the answer worse past a point — the total error \(O(\lambda) + O(\varepsilon_{\text{mach}}/\lambda)\) is minimized near \(\lambda^* \sim \sqrt{\varepsilon_{\text{mach}}}\), about 1e-8 in float64; only go below that with a measured gtol.

The Metric interface

A Metric supplies the factor through four ops, each taking a metric-block vector — or a matrix whose leading axis is size (columns batched) — and a SolverContext carrying everything the solver knows at the call site (the flat iterate x, the live LMState, args, p; the shipped metrics ignore it, a custom metric may key off it):

  • factor_apply(v, ctx) — \(F v\)
  • factor_solve(v, ctx) — \(F^{-1} v\)
  • factor_solve_transpose(v, ctx) — \(F^{-\top} v\) (the workhorse: \(\tilde J^\top\) assembly, gradient pullbacks, the AD rule)
  • norm(v, ctx) — \(\|v\|_W\) (vectors only; must match \(\|F v\|_2\) to floating-point accuracy — the solver compares objective values built from both forms)

Two contracts matter. The factor must be exact: the identity penalty block in \(y\) is hardcoded, so an approximate factor silently changes the objective (unlike a CG preconditioner, which may be sloppy). And metrics hash by identity: construct one at setup scope and reuse it — a rebuilt equal-config metric keys a fresh compiled solve loop.

Shipped metrics:

  • IdentityMetric(size) — \(F = I\): plain ridge \(\|r\|^2 + \lambda\|x_m\|^2\), every op a passthrough, with no special-casing anywhere in the solver.
  • RepeatedFactorMetric(F, repeats=J) — repeats copies of one upper-triangular block factor on the metric block, \(W = \operatorname{blockdiag}(F^\top F, \ldots, F^\top F)\). The constructor takes the factor, not the Gram matrix (callers typically already hold it): F = jnp.linalg.cholesky(K, upper=True). All repeated blocks and all batched columns share a single triangular product/solve.

The metric instance rides in lm_state.metric, so a callback may replace it mid-solve by constructing a new instance of the same type (see Callbacks); the solver treats the change as a problem change — the metric defines this objective — suppressing that step's convergence test and invalidating the whitened caches. The gradient and the implicit-AD rule treat \(F\) as constant between callback actions: the state-dependence of a callback-refreshed factor is not differentiated.

Kernel instantiation

For kernel collocation with \(J\) repeated coefficient blocks (one per equation stack) and \(s\) unpenalized structural scalars in the free block,

\[ W = \operatorname{blockdiag}(K, \ldots, K), \qquad q(x) = \sum_j \alpha_j^\top K \alpha_j , \]

built by RepeatedFactorMetric(jnp.linalg.cholesky(K, upper=True), repeats=J) — one Cholesky, batched triangular ops over the repeated blocks, no epsilon shift on the scalars (they carry no penalty at all, unlike the metric formulation's \(\varepsilon I\)). Each inner LM step is then a kernel ridge regression of the relinearized equations, exactly the Gauss–Newton scheme of Chen–Hosseini–Owhadi–Stuart (2021) with the ridge kept explicit. An \(O(N)\) state-space (Matérn) factor pairing with the matrix-free CG path is planned for a later release; at the problem sizes this package targets the dense factor serves all three linear solvers.

Choose the free block deliberately. The identification condition — whatever the residual does not pin, the metric must — is easy to violate by moving too much into the free block, and it can hold marginally while the selection still needs a block's seminorm in the objective. The experiments that motivated this design produced a clear negative result: on the DAE drivers, removing the penalty from coefficient blocks whose levels are pinned by initial conditions left the residual converged but let the unpenalized paths drift enormously between collocation points (seminorms exploding by 2–80×, ground-truth checks failed). Keep every kernel coefficient block under the metric; reserve the free block for genuinely structural scalars, and verify the choice empirically — the free coordinates must be reproducible across perturbed starts (multi-start agreement is the practical test).

Linear solvers

linear_solver takes a typed config — Cholesky(), QR(), or CG(preconditioner, tol=..., atol=..., maxiter=...) — so each method's knobs live on its own config and cannot be passed with another (the configs hash by value: equal configs share one compiled solve loop). All three solve \((\tilde J^\top \tilde J + \lambda E + \mu I)\,\delta_y = -g\) and share one factorization (or inner-solve setup) between the velocity and geodesic-acceleration solves.

linear_solver Method Cost per update When
Cholesky() (default) dense normal equations; \(\tilde J^\top = \bar F^{-\top} J^\top\) by one batched factor solve, then \(G = \tilde J^\top \tilde J + \lambda E\) (\(\lambda\) added on the metric-block diagonal only) cached across rejected steps (a reject re-factors in \(p^3/3\) without the GEMM or the \(\tilde J\) materialization) \(mp^2\) GEMM + \(n_m\)-block triangular solve (both skipped on reject) + \(p^3/3\) default; the \(\lambda\) spectral floor keeps it accurate at deep ridge
QR() QR of the augmented stack \([\tilde J;\ \sqrt{\lambda}\,[I \mid 0]]\) cached per \((x, \lambda)\); each damping update re-factors \([R;\sqrt{\mu}I]\) and solves by corrected semi-normal equations with one refinement pass (Björck 1987; Björck 1996, §6.6.5; the damping-row structure is Moré 1978's) \((m{+}n_m)p^2\) QR (skipped on reject) + \(2p^3/3\)-ish refactor extreme tiny-\(\lambda\)/tiny-\(\mu\) regimes, where even the whitened normal equations square a large \(\|\tilde J\|^2/\mu\)
CG(preconditioner, ...) matrix-free preconditioned CG on the damped whitened normal operator itself — the same SPD system Cholesky() factors, with products through jvp/vjp and the metric's factor ops; the required preconditioner (a typed Preconditioner; IdentityPreconditioner() opts out) sits in CG's M slot iterations × (one \(\tilde J\) and one \(\tilde J^\top\) product + one preconditioner apply) Jacobians too large to materialize, given a structured preconditioner; the \(\lambda\) floor on the metric block helps the spectrum

Everything runs at the residual dtype — there is no promotion knob. The selection resolution is bounded by the problem dtype either way (the stationarity test reads the gradient, which lives at the residual dtype; at tiny \(\lambda\) the float32 gradient noise \(\sim 10^{-7}/\lambda\) bounds the selection regardless of any wider factorization), and QR() is the in-dtype fix for extreme conditioning. A float32 program that genuinely needs float64 selection should run the solve in float64.

Preconditioners

The CG config requires a typed Preconditioner in both roles: a subclass implementing apply(v, damping, ctx), an SPD approximation of \((\tilde J^\top \tilde J + \lambda E + \mu I)^{-1}\) applied with the live damping (zero in the AD role), receiving the same SolverContext as the metric ops. IdentityPreconditioner() is the explicit opt-out; a custom one is a small dataclass:

@dataclass(frozen=True, eq=False)
class JacobiPreconditioner(Preconditioner):
    diagonal: jax.Array

    def apply(self, v, damping, ctx):
        return v / (self.diagonal + damping)

A preconditioner changes the CG iteration path, never the subproblem — approximations are safe (unlike the metric factor, which must be exact).

For repeated interacting blocks (multiple "agents" coupled through shared equations), BlockEigenPreconditioner(families, permutation) is the shipped workhorse: a block-diagonal approximation over a chosen grouping of the whitened coordinates, eigendecomposed in the constructor and applied with the damping-analytic shift \(\Lambda + \texttt{ridge\_weight}\cdot\lambda + \mu\) (metric-block families carry the live ridge read from ctx.lm_state.ridge, free-block families are damping-only). The eigenbasis rides in the carried instance (lm_state.preconditioner).

Adaptive refreshes through the solve callback. A callback rebuilds the instance from the live iterate by calling the constructor; staleness only moves the CG iteration path, so a refresh is never a problem change. The natural policy pairs refreshes with AnnealRidge — rebuild exactly when the anneal advances a level, gated with lax.cond so the eigendecomposition is paid only when it fires (a jnp.where merge would pay it every step):

anneal = AnnealRidge(ridge_floor=1e-11)

def callback(ctx):
    action = anneal(ctx)
    advanced = action.lm_state.ridge < ctx.lm_state.ridge
    precond = jax.lax.cond(
        advanced,
        lambda: BlockEigenPreconditioner(build_families(ctx.x), PERMUTATION),
        lambda: ctx.lm_state.preconditioner,
    )
    return dataclasses.replace(
        action,
        lm_state=dataclasses.replace(action.lm_state, preconditioner=precond),
    )

solver = RidgeLevenbergMarquardt(
    residual_fn, metric=metric,
    linear_solver=CG(BlockEigenPreconditioner(build_families(x0), PERMUTATION),
                     tol=1e-10, maxiter=2500),
)
result = solver.solve(x0, callback=callback, user_state=anneal.init_state(),
                      max_steps=120, gtol=1e-11)

The AD-role CG applies the carried instance at zero damping — a callback-refreshed eigenbasis is exactly the near-solution build the tangent solve wants. ad_solver=None inherits it at the AD-default tolerance and unbounded iteration budget; ad_solver=CG(None, tol=..., maxiter=...) inherits it while pinning the knobs; an explicit ad_solver=CG(instance, ...) uses that instance as-is (hooks marked requires_positive_damping fall back to unpreconditioned). One measured calibration note: with a family layout whose blocks carry the same ridge floor as the operator, the CG iteration count SATURATES as the ridge anneals down (it does not grow like \(1/\sqrt{\lambda}\)) — budget maxiter for the saturated count, since a truncated inner solve stalls the endgame. Two scope cautions: save_steps=True records a full copy of args per step, and a parallel multi_start vmaps the callback so a lax.cond rebuild lowers to a select that evaluates BOTH branches in every lane — keep both out of preconditioned-CG production runs.

Implicit differentiation

solve(...).x has a custom implicit rule with respect to p: Gauss–Newton differentiation of the ridge stationarity, posed on the whitened variable —

\[ \left(\tilde J^\top \tilde J + \lambda E\right) \dot y = -\tilde J^\top \frac{\partial r}{\partial p}\,\dot p , \qquad \dot x = \bar F^{-1} \dot y , \]

with \(\lambda\) frozen (stop-gradient) at the returned state's ridge — the continuation schedule's and the multi-start selection's dependence on p are deliberately ignored — and no damping in the AD matrix. ad_solver=Cholesky() assembles and factors; CG(preconditioner, tol=..., atol=..., maxiter=...) is matrix-free CG on the same operator (positive definite does not mean well conditioned — unpreconditioned CG degrades as \(\lambda\) shrinks). ad_solver=None (the default) matches the forward path: Cholesky() for the dense forwards, and a CG forward keeps both its family and its preconditioner — the typed apply is damping-analytic, so the forward hook serves the undamped system at zero damping exactly (requires_positive_damping hooks fall back to unpreconditioned); the AD tolerance and iteration budget stay at the AD defaults unless an explicit ad_solver=CG(...) pins them.

The contract, stated plainly: exact differentiation carries two extra terms (\(\sum_i r_i \nabla^2 r_i\) in the matrix, \((\partial J^\top/\partial p)r\) on the right), both dropped. Both are exactly zero for an affine residual and first order in \(\|r\|\) otherwise — so the rule is exact in the interpolating limit in absolute terms. On a genuinely underdetermined curved system, however, the null-space block of the exact tangent carries the constraint-curvature term \(\sum_i \nu_i \nabla^2 r_i\) with multiplier \(\nu = r/\lambda\), which does not vanish relative to the \(\lambda\)-scaled penalty as \(\lambda \downarrow 0\): tangent components read off null-space directions retain a Gauss–Newton bias proportional to the residual curvature. This is the same approximation the metric solver's frozen-projector AD rules make. Failed statuses return exact zero tangents and evaluate the masked tangent program at stop-gradient copies of the original inputs and the initial ridge.

API reference

nlls_gram.Metric

Positive-definite metric W, via factor callbacks.

x = [x_m; x_f] splits into the metric block x_m -- the leading size coordinates, covered by W -- and a free block x_f. The metric is supplied through callbacks for an invertible factor F with W = F'F, the whitened variable being F x_m and ||v||_W = ||F v||_2. F should be upper triangular; the canonical example is F = jnp.linalg.cholesky(K, upper=True). The solver extends it to F_bar = blockdiag(F, sqrt(free_scale) I) over the whole vector and never materializes either.

Subclasses implement the ops on metric-block vectors (or matrices whose LEADING axis is size; columns are batched):

  • factor_apply(v, ctx): F v
  • factor_solve(v, ctx): F^{-1} v
  • factor_solve_transpose(v, ctx): F^{-T} v
  • norm(v, ctx): ||v||_W, defaulted to ||factor_apply(v)||_2. Provided for callers; the solvers measure in the whitened variable and never call it.

Every op receives a :class:~nlls_gram.SolverContext carrying the solver's live state, so an exotic metric can key off the iterate through ctx.x and ctx.lm_state.

Metric instances are JAX PYTREES: array fields are traced leaves, the type plus its static fields are structure. Every concrete class must be registered with :func:~nlls_gram.register_pytree_dataclass -- the solvers reject unregistered instances. The instance rides inside the solver state (lm_state.metric), so a solve callback replaces the metric by constructing a new instance of the same type (same static fields, same leaf shapes and dtypes) -- pure traced ops, no recompilation. Equal-config instances with fresh arrays share one compiled solve loop. dataclasses.replace works too: metric constructors only validate and derive shapes, so re-running them under trace is cheap.

free_scale weights the free block in the whitened variable: 1.0 (the default) leaves it Euclidean. The ridge solver never penalizes the free block whatever the scale -- free_scale only changes its trust-region geometry -- while for the metric solver it IS that block's damping weight. It is a traced leaf, canonicalized by each constructor to the factor's float dtype, so changing it never recompiles.

Contracts: the factor must be EXACT. The solver hardcodes the identity penalty block in the whitened variable, so an approximate factor silently changes the objective (unlike a CG preconditioner, which may be sloppy). The ridge weight never enters the factorization, so ridge continuation composes unchanged. How a subclass fulfills the ops -- prefactorized storage, factorize-in-__init__, fully matrix-free -- is its constructor's business, and the constructor must be traceable when the metric is rebuilt inside a jitted callback.

Source code in src/nlls_gram/metrics.py 
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
class Metric:
    """Positive-definite metric ``W``, via factor callbacks.

    ``x = [x_m; x_f]`` splits into the metric block ``x_m`` -- the leading
    ``size`` coordinates, covered by ``W`` -- and a free block ``x_f``. The
    metric is supplied through callbacks for an invertible factor ``F`` with
    ``W = F'F``, the whitened variable being ``F x_m`` and
    ``||v||_W = ||F v||_2``. ``F`` should be upper triangular; the canonical
    example is ``F = jnp.linalg.cholesky(K, upper=True)``. The solver extends
    it to ``F_bar = blockdiag(F, sqrt(free_scale) I)`` over the whole vector
    and never materializes either.

    Subclasses implement the ops on metric-block vectors (or matrices whose
    LEADING axis is ``size``; columns are batched):

    - ``factor_apply(v, ctx)``: ``F v``
    - ``factor_solve(v, ctx)``: ``F^{-1} v``
    - ``factor_solve_transpose(v, ctx)``: ``F^{-T} v``
    - ``norm(v, ctx)``: ``||v||_W``, defaulted to ``||factor_apply(v)||_2``.
      Provided for callers; the solvers measure in the whitened variable and
      never call it.

    Every op receives a :class:`~nlls_gram.SolverContext` carrying the
    solver's live state, so an exotic metric can key off the iterate through
    ``ctx.x`` and ``ctx.lm_state``.

    Metric instances are JAX PYTREES: array fields are traced leaves, the
    type plus its static fields are structure. Every concrete class must be
    registered with
    :func:`~nlls_gram.register_pytree_dataclass` -- the solvers reject
    unregistered instances. The instance rides inside the solver state
    (``lm_state.metric``), so a ``solve`` callback replaces the metric by
    constructing a new instance of the same type (same static fields, same
    leaf shapes and dtypes) -- pure traced ops, no recompilation. Equal-config
    instances with fresh arrays share one compiled solve loop.
    ``dataclasses.replace`` works too: metric constructors only validate and
    derive shapes, so re-running them under trace is cheap.

    ``free_scale`` weights the free block in the whitened variable: ``1.0``
    (the default) leaves it Euclidean. The ridge solver never penalizes the
    free block whatever the scale -- ``free_scale`` only changes its
    trust-region geometry -- while for the metric solver it IS that block's
    damping weight. It is a traced leaf, canonicalized by each constructor to
    the factor's float dtype, so changing it never recompiles.

    Contracts: the factor must be EXACT. The solver hardcodes the identity
    penalty block in the whitened variable, so an approximate factor silently
    changes the objective (unlike a CG preconditioner, which may be sloppy).
    The ridge weight never enters the factorization, so ridge continuation
    composes unchanged. How a subclass fulfills the ops -- prefactorized
    storage, factorize-in-``__init__``, fully matrix-free -- is its
    constructor's business, and the constructor must be traceable when the
    metric is rebuilt inside a jitted callback.
    """

    size: int
    free_scale: float = 1.0

    def factor_apply(self, v, ctx):
        """``F v`` for a metric-block vector or leading-axis-batched matrix."""
        raise NotImplementedError

    def factor_solve(self, v, ctx):
        """``F^{-1} v`` for a metric-block vector or batched matrix."""
        raise NotImplementedError

    def factor_solve_transpose(self, v, ctx):
        """``F^{-T} v`` for a metric-block vector or batched matrix."""
        raise NotImplementedError

    def norm(self, v, ctx):
        """``||v||_W = ||F v||_2`` for a metric-block vector."""
        return jnp.linalg.norm(self.factor_apply(v, ctx))

factor_apply(v, ctx)

F v for a metric-block vector or leading-axis-batched matrix.

Source code in src/nlls_gram/metrics.py 
90
91
92
def factor_apply(self, v, ctx):
    """``F v`` for a metric-block vector or leading-axis-batched matrix."""
    raise NotImplementedError

factor_solve(v, ctx)

F^{-1} v for a metric-block vector or batched matrix.

Source code in src/nlls_gram/metrics.py 
94
95
96
def factor_solve(self, v, ctx):
    """``F^{-1} v`` for a metric-block vector or batched matrix."""
    raise NotImplementedError

factor_solve_transpose(v, ctx)

F^{-T} v for a metric-block vector or batched matrix.

Source code in src/nlls_gram/metrics.py 
 98
 99
100
def factor_solve_transpose(self, v, ctx):
    """``F^{-T} v`` for a metric-block vector or batched matrix."""
    raise NotImplementedError

norm(v, ctx)

||v||_W = ||F v||_2 for a metric-block vector.

Source code in src/nlls_gram/metrics.py 
102
103
104
def norm(self, v, ctx):
    """``||v||_W = ||F v||_2`` for a metric-block vector."""
    return jnp.linalg.norm(self.factor_apply(v, ctx))

nlls_gram.SolverContext dataclass

What the solver knows at a metric, preconditioner, or linear-solver call site -- the inner algebra's context, as opposed to :class:LMContext, which a per-step user callback receives.

Fields are None where the call site has nothing to offer:

  • x: the current FLATTENED iterate (the whole parameter vector, not just the metric block).
  • lm_state: the live :class:LMState (damping, ridge, caches, and the carried metric/preconditioner instances). In the implicit-AD rule this is the returned state under stop_gradient -- inert conditioning data, like the ridge.
  • args / p: the residual's auxiliary data and differentiation parameters as passed to solve/update.
Source code in src/nlls_gram/lm_types.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class SolverContext:
    """What the solver knows at a metric, preconditioner, or linear-solver
    call site -- the inner algebra's context, as opposed to
    :class:`LMContext`, which a per-step user callback receives.

    Fields are ``None`` where the call site has nothing to offer:

    - ``x``: the current FLATTENED iterate (the whole parameter vector, not
      just the metric block).
    - ``lm_state``: the live :class:`LMState` (damping, ridge, caches, and
      the carried metric/preconditioner instances). In the implicit-AD rule
      this is the returned state under ``stop_gradient`` -- inert
      conditioning data, like the ridge.
    - ``args`` / ``p``: the residual's auxiliary data and differentiation
      parameters as passed to ``solve``/``update``.
    """

    x: Any = None
    lm_state: Any = None
    args: Any = None
    p: Any = None

nlls_gram.IdentityMetric dataclass

Bases: Metric

The identity metric W = I on size coordinates -- plain ridge for the ridge solver, Euclidean damping for the metric solver.

F = I: every factor op is the identity and norm is the Euclidean norm, with no special-casing anywhere downstream.

Source code in src/nlls_gram/metrics.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
@dataclass(frozen=True, eq=False)
class IdentityMetric(Metric):
    """The identity metric ``W = I`` on ``size`` coordinates -- plain ridge for
    the ridge solver, Euclidean damping for the metric solver.

    ``F = I``: every factor op is the identity and ``norm`` is the Euclidean
    norm, with no special-casing anywhere downstream.
    """

    size: int
    free_scale: float = 1.0

    def __post_init__(self):
        if self.size < 0:
            raise ValueError("size must be nonnegative")
        object.__setattr__(
            self,
            "free_scale",
            _canonical_free_scale(self.free_scale, jnp.result_type(float)),
        )

    def factor_apply(self, v, ctx):
        _check_leading_size(v, self.size)
        return v

    def factor_solve(self, v, ctx):
        _check_leading_size(v, self.size)
        return v

    def factor_solve_transpose(self, v, ctx):
        _check_leading_size(v, self.size)
        return v

    def norm(self, v, ctx):
        _check_leading_size(v, self.size)
        return jnp.linalg.norm(v)

nlls_gram.RepeatedFactorMetric dataclass

Bases: Metric

repeats copies of one block factor: W = blockdiag(F'F, ...).

F is an upper-triangular invertible block factor -- e.g. jnp.linalg.cholesky(K, upper=True) for a positive-definite Gram matrix K, giving the repeated kernel seminorm sum_j alpha_j' K alpha_j over repeats coefficient blocks. The constructor takes the FACTOR, not K (callers typically already hold it); a positive-SEMIdefinite K needs a shift first: jnp.linalg.cholesky(K + epsilon * I, upper=True). Triangularity and positive definiteness are assumed, not validated.

All repeated blocks (and all batched columns) share a single triangular product or solve: the ops reshape the metric block into the columns of one (block, repeats * cols) matrix, so no repeated factor or full block diagonal is ever formed. size = repeats * F.shape[0].

Source code in src/nlls_gram/metrics.py 
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
@dataclass(frozen=True, eq=False)
class RepeatedFactorMetric(Metric):
    """``repeats`` copies of one block factor: ``W = blockdiag(F'F, ...)``.

    ``F`` is an upper-triangular invertible block factor -- e.g.
    ``jnp.linalg.cholesky(K, upper=True)`` for a positive-definite Gram matrix
    ``K``, giving the repeated kernel seminorm ``sum_j alpha_j' K alpha_j``
    over ``repeats`` coefficient blocks. The constructor takes the FACTOR, not
    ``K`` (callers typically already hold it); a positive-SEMIdefinite ``K``
    needs a shift first: ``jnp.linalg.cholesky(K + epsilon * I, upper=True)``.
    Triangularity and positive definiteness are assumed, not validated.

    All repeated blocks (and all batched columns) share a single triangular
    product or solve: the ops reshape the metric block into the columns of one
    ``(block, repeats * cols)`` matrix, so no repeated factor or full block
    diagonal is ever formed. ``size = repeats * F.shape[0]``.
    """

    F: jax.Array
    repeats: int = field(default=1, kw_only=True)
    free_scale: float = field(default=1.0, kw_only=True)
    size: int = field(init=False)

    def __post_init__(self):
        F = jnp.asarray(self.F)
        if F.ndim != 2 or F.shape[0] != F.shape[1] or F.shape[0] == 0:
            raise ValueError("F must be a nonempty square matrix")
        dtype = jnp.result_type(F, 1.0)
        if not jnp.issubdtype(dtype, jnp.floating):
            raise TypeError("F must have a real floating-point dtype")
        if self.repeats < 1:
            raise ValueError("repeats must be a positive integer")
        object.__setattr__(self, "F", F.astype(dtype))
        object.__setattr__(self, "size", self.repeats * F.shape[0])
        object.__setattr__(
            self, "free_scale", _canonical_free_scale(self.free_scale, dtype)
        )

    def _map_blocks(self, block_op, v):
        _check_leading_size(v, self.size)
        block_size = self.F.shape[0]
        trailing_shape = v.shape[1:]
        packed = jnp.moveaxis(
            v.reshape((self.repeats, block_size) + trailing_shape), 0, 1
        ).reshape(block_size, -1)
        return jnp.moveaxis(
            block_op(packed).reshape((block_size, self.repeats) + trailing_shape),
            0,
            1,
        ).reshape((self.size,) + trailing_shape)

    def factor_apply(self, v, ctx):
        return self._map_blocks(lambda m: mm(self.F, m), v)

    def factor_solve(self, v, ctx):
        return self._map_blocks(
            lambda m: jsp_linalg.solve_triangular(self.F, m, lower=False), v
        )

    def factor_solve_transpose(self, v, ctx):
        return self._map_blocks(
            lambda m: jsp_linalg.solve_triangular(self.F.T, m, lower=True), v
        )

nlls_gram.Preconditioner

SPD preconditioner for the solvers' CG paths.

apply(v, damping, ctx) returns an SPD approximation of the damped operator's inverse -- (J~'J~ + ridge E + damping I)^{-1} in parameter space under :class:~nlls_gram.CG, (J~J~' + damping I)^{-1} in residual space under :class:~nlls_gram.GramCG. In the forward role it sits in CG's M slot with the live damping; in the ad_solver role the implicit system is undamped and damping is zero. ctx is the same :class:~nlls_gram.SolverContext the metric factor ops receive.

Preconditioner instances are JAX PYTREES: array fields are traced leaves, the type plus its static fields are structure. Every concrete class must be registered with :func:~nlls_gram.register_pytree_dataclass -- the solvers reject unregistered instances. The instance rides inside the solver state (lm_state.preconditioner), so a solve callback refreshes it by calling the CONSTRUCTOR again with fresh arrays -- same type, same leaf shapes and dtypes, pure traced ops, no recompilation. Never dataclasses.replace a preconditioner: replace re-runs __init__, re-paying any eigendecomposition or sketch, and construction-time-only inputs are not stored to re-supply. A stale instance only changes the CG iteration path, never the converged step, so refreshing rarely (or never) is always safe.

Implement a custom one as a small registered frozen dataclass::

@dataclass(frozen=True, eq=False)
class JacobiPreconditioner(Preconditioner):
    diagonal: jax.Array

    def apply(self, v, damping, ctx):
        return v / (self.diagonal + damping)

register_pytree_dataclass(JacobiPreconditioner, data_fields=("diagonal",))

Subclasses whose apply divides by the live damping must set requires_positive_damping = True; the constructor rejects them for the AD role, where damping is zero.

Source code in src/nlls_gram/preconditioners.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
class Preconditioner:
    """SPD preconditioner for the solvers' CG paths.

    ``apply(v, damping, ctx)`` returns an SPD approximation of the damped
    operator's inverse -- ``(J~'J~ + ridge E + damping I)^{-1}`` in parameter
    space under :class:`~nlls_gram.CG`, ``(J~J~' + damping I)^{-1}`` in
    residual space under :class:`~nlls_gram.GramCG`. In the forward role it
    sits in CG's ``M`` slot with the live damping; in the ``ad_solver`` role
    the implicit system is undamped and ``damping`` is zero. ``ctx`` is the
    same :class:`~nlls_gram.SolverContext` the metric factor ops receive.

    Preconditioner instances are JAX PYTREES: array fields are traced
    leaves, the type plus its static fields are structure. Every concrete
    class must be registered with
    :func:`~nlls_gram.register_pytree_dataclass` -- the solvers reject
    unregistered instances. The instance rides inside the solver state
    (``lm_state.preconditioner``), so a ``solve`` callback refreshes it by
    calling the CONSTRUCTOR again with fresh arrays -- same type, same leaf
    shapes and dtypes, pure traced ops, no recompilation. Never
    ``dataclasses.replace`` a preconditioner: ``replace`` re-runs
    ``__init__``, re-paying any eigendecomposition or sketch, and
    construction-time-only inputs are not stored to re-supply. A stale
    instance only changes the CG iteration path, never the converged step,
    so refreshing rarely (or never) is always safe.

    Implement a custom one as a small registered frozen dataclass::

        @dataclass(frozen=True, eq=False)
        class JacobiPreconditioner(Preconditioner):
            diagonal: jax.Array

            def apply(self, v, damping, ctx):
                return v / (self.diagonal + damping)

        register_pytree_dataclass(JacobiPreconditioner, data_fields=("diagonal",))

    Subclasses whose ``apply`` divides by the live damping must set
    ``requires_positive_damping = True``; the constructor rejects them for
    the AD role, where damping is zero.
    """

    requires_positive_damping = False

    def apply(self, v, damping, ctx):
        raise NotImplementedError

nlls_gram.IdentityPreconditioner dataclass

Bases: Preconditioner

The identity map as an explicit "no preconditioner" choice for CG.

Nobody should run Krylov methods without thinking about preconditioning, so opting out is an explicit, greppable decision rather than a silent default. Stateless: every instance compares equal, so equal CG configs share one compiled solve loop.

Source code in src/nlls_gram/preconditioners.py 
70
71
72
73
74
75
76
77
78
79
80
81
@dataclass(frozen=True)
class IdentityPreconditioner(Preconditioner):
    """The identity map as an explicit "no preconditioner" choice for ``CG``.

    Nobody should run Krylov methods without thinking about preconditioning,
    so opting out is an explicit, greppable decision rather than a silent
    default. Stateless: every instance compares equal, so equal ``CG``
    configs share one compiled solve loop.
    """

    def apply(self, v, damping, ctx):
        return v

nlls_gram.BlockEigenPreconditioner dataclass

Bases: Preconditioner

Block-diagonal eigenbasis preconditioner over grouped whitened coordinates.

The workhorse for structured whitened normal operators J~'J~ + ridge E + damping I built from repeated interacting blocks (multiple "agents" coupled through shared equations): approximate the operator by a block-diagonal matrix over a chosen grouping of the whitened coordinates, eigendecompose each block, and apply the exact inverse of the shifted approximation

v  ->  V ((V' v) / (Lambda + ridge_weight * ridge + damping)) V'

per block -- analytic in both the live damping (traced; it changes per LM step) and the live ridge (read from ctx.lm_state.ridge, so ridge continuation composes with no refresh).

families is a sequence of (blocks, ridge_weight) pairs: blocks has shape (groups, size, size) -- the stacked diagonal blocks of J~'J~ restricted to that family's coordinate groups, in permuted order. Families in the metric block set ridge_weight = 1 (their diagonal carries the ridge spectral floor); free-block families set 0 (damping-only, so the zero-damping AD role applies their plain inverse -- positive definite whenever the free block is identified). Blocks are symmetrized and eigendecomposed HERE, once; positive semidefiniteness is assumed, not validated (entries may be traced). permutation is the 1-D integer array reordering the flattened whitened parameter vector into family-major order (v_permuted = v[permutation]); families are consumed in sequence and must cover it exactly. jnp.arange(n) serves when the natural layout already is family-major.

The constructor is fully traceable, so a solve callback refreshes the preconditioner from the live iterate by constructing a new instance -- gate the rebuild with jax.lax.cond (e.g. on a ridge-continuation level advance) so the eigendecomposition is paid only when it fires; a jnp.where merge would pay it every step.

Source code in src/nlls_gram/preconditioners.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
@dataclass(frozen=True, eq=False)
class BlockEigenPreconditioner(Preconditioner):
    """Block-diagonal eigenbasis preconditioner over grouped whitened
    coordinates.

    The workhorse for structured whitened normal operators
    ``J~'J~ + ridge E + damping I`` built from repeated interacting blocks
    (multiple "agents" coupled through shared equations): approximate the
    operator by a block-diagonal matrix over a chosen grouping of the
    whitened coordinates, eigendecompose each block, and apply the exact
    inverse of the shifted approximation

        v  ->  V ((V' v) / (Lambda + ridge_weight * ridge + damping)) V'

    per block -- analytic in both the live ``damping`` (traced; it changes per
    LM step) and the live ``ridge`` (read from ``ctx.lm_state.ridge``, so
    ridge continuation composes with no refresh).

    ``families`` is a sequence of ``(blocks, ridge_weight)`` pairs:
    ``blocks`` has shape ``(groups, size, size)`` -- the stacked diagonal
    blocks of ``J~'J~`` restricted to that family's coordinate groups, in
    permuted order. Families in the metric block set ``ridge_weight = 1``
    (their diagonal carries the ``ridge`` spectral floor); free-block
    families set ``0`` (damping-only, so the zero-damping AD role applies
    their plain inverse -- positive definite whenever the free block is
    identified). Blocks are symmetrized and eigendecomposed HERE, once;
    positive semidefiniteness is assumed, not validated (entries may be
    traced). ``permutation`` is the 1-D integer array reordering the
    flattened whitened parameter vector into family-major order
    (``v_permuted = v[permutation]``); families are consumed in sequence and
    must cover it exactly. ``jnp.arange(n)`` serves when the natural layout
    already is family-major.

    The constructor is fully traceable, so a ``solve`` callback refreshes
    the preconditioner from the live iterate by constructing a new instance
    -- gate the rebuild with ``jax.lax.cond`` (e.g. on a ridge-continuation
    level advance) so the eigendecomposition is paid only when it fires; a
    ``jnp.where`` merge would pay it every step.
    """

    families: InitVar[Any]
    permutation: jax.Array
    eigenvectors: tuple = field(init=False)
    eigenvalues: tuple = field(init=False)
    ridge_weights: tuple = field(init=False)
    inverse_permutation: jax.Array = field(init=False)

    def __post_init__(self, families):
        permutation = jnp.asarray(self.permutation)
        if permutation.ndim != 1 or not jnp.issubdtype(permutation.dtype, jnp.integer):
            raise ValueError("permutation must be a 1-D integer array")
        eigenvectors, eigenvalues, ridge_weights = [], [], []
        covered = 0
        for blocks, ridge_weight in families:
            blocks = jnp.asarray(blocks)
            if blocks.ndim != 3 or blocks.shape[1] != blocks.shape[2]:
                raise ValueError(
                    "each family's blocks must have shape (groups, size, size); "
                    f"got {blocks.shape}"
                )
            symmetrized = 0.5 * (blocks + jnp.swapaxes(blocks, 1, 2))
            values, vectors = jnp.linalg.eigh(symmetrized)
            # eigh of a numerically PSD block can return tiny negative
            # eigenvalues; clamped at zero the apply shift stays positive for
            # any positive ridge/damping (and the zero-damping AD role stays
            # SPD whenever the family itself is).
            eigenvalues.append(jnp.maximum(values, 0.0))
            eigenvectors.append(vectors)
            ridge_weights.append(jnp.asarray(ridge_weight, dtype=blocks.dtype))
            covered += blocks.shape[0] * blocks.shape[1]
        if covered != permutation.shape[0]:
            raise ValueError(
                f"families cover {covered} coordinates but the permutation "
                f"has {permutation.shape[0]}"
            )
        object.__setattr__(self, "permutation", permutation)
        object.__setattr__(self, "eigenvectors", tuple(eigenvectors))
        object.__setattr__(self, "eigenvalues", tuple(eigenvalues))
        object.__setattr__(self, "ridge_weights", tuple(ridge_weights))
        object.__setattr__(self, "inverse_permutation", jnp.argsort(permutation))

    def apply(self, v, damping, ctx):
        # LevenbergMarquardt carries no ridge, so its metric-block families
        # shift by the damping alone.
        carried = ctx.lm_state.ridge
        ridge = (
            jnp.zeros((), v.dtype)
            if carried is None
            else jnp.asarray(carried, dtype=v.dtype)
        )
        permuted = v[self.permutation]
        pieces = []
        offset = 0
        for V, values, ridge_weight in zip(
            self.eigenvectors, self.eigenvalues, self.ridge_weights, strict=True
        ):
            groups, size = V.shape[0], V.shape[1]
            segment = permuted[offset : offset + groups * size]
            offset += groups * size
            shift = ridge_weight * ridge + damping
            coefficients = jnp.einsum(
                "gab,ga->gb", V, segment.reshape(groups, size), precision=HIGHEST
            )
            pieces.append(
                jnp.einsum(
                    "gab,gb->ga", V, coefficients / (values + shift), precision=HIGHEST
                ).reshape(-1)
            )
        return jnp.concatenate(pieces)[self.inverse_permutation].astype(v.dtype)

References

  • Bakushinskii, A. B. (1992). "The problem of the convergence of the iteratively regularized Gauss–Newton method." Comput. Math. Math. Phys. 32(9), 1353–1359.
  • Björck, Å. (1987). "Stability analysis of the method of seminormal equations for linear least squares problems." Linear Algebra Appl. 88–89, 31–48.
  • Björck, Å. (1996). Numerical Methods for Least Squares Problems. SIAM.
  • Blaschke (Kaltenbacher), B., A. Neubauer, and O. Scherzer (1997). "On convergence rates for the iteratively regularized Gauss–Newton method." IMA J. Numer. Anal. 17(3), 421–436.
  • Campbell, S. L., P. Kunkel, and K. Bobinyec (2012). "A minimal norm corrected underdetermined Gauß–Newton procedure." Appl. Numer. Math. 62(5), 592–605.
  • Chen, Y., B. Hosseini, H. Owhadi, and A. M. Stuart (2021). "Solving and learning nonlinear PDEs with Gaussian processes." J. Comput. Phys. 447, 110668.
  • Eldén, L. (1982). "A weighted pseudoinverse, generalized singular values, and constrained least squares problems." BIT 22, 487–502.
  • Engl, H. W., M. Hanke, and A. Neubauer (1996). Regularization of Inverse Problems. Kluwer.
  • Engl, H. W., K. Kunisch, and A. Neubauer (1989). "Convergence rates for Tikhonov regularisation of non-linear ill-posed problems." Inverse Problems 5(4), 523–540.
  • Izmailov, A. F., and M. V. Solodov (2026). "Local convergence of the Gauss–Newton methods for constrained nonlinear equations." Comput. Optim. Appl. (doi:10.1007/s10589-026-00801-4).
  • Kaltenbacher, B., A. Neubauer, and O. Scherzer (2008). Iterative Regularization Methods for Nonlinear Ill-Posed Problems. de Gruyter.
  • Marquardt, D. W. (1963). "An algorithm for least-squares estimation of nonlinear parameters." J. SIAM 11(2), 431–441.
  • Moré, J. J. (1978). "The Levenberg–Marquardt algorithm: implementation and theory." In Numerical Analysis (Dundee 1977), Lecture Notes in Math. 630, Springer, 105–116.
  • Nocedal, J., and S. J. Wright (2006). Numerical Optimization, 2nd ed. Springer, Ch. 10.
  • Pes, F., and G. Rodriguez (2022). "A doubly relaxed minimal-norm Gauss–Newton method for underdetermined nonlinear least-squares problems." Appl. Numer. Math. 171, 233–248.
  • Transtrum, M. K., and J. P. Sethna (2012). "Improvements to the Levenberg–Marquardt algorithm for nonlinear least-squares minimization." arXiv:1201.5885.