Precondition the fit parameters, and restart a stalled minimizer

[davidwalter2]

Problem

A Z dilepton in-situ muon efficiency fit with 2112 unconstrained Chebyshev
coefficients took 145 minutes, and most of that time bought nothing: 20 of
78 iterations returned a bit-identical loss while consuming 57 minutes, and the
worst single iteration spent 27 minutes to arrive back where it started.

The cause is conditioning. trust-krylov solves its subproblem with a Krylov
method whose inner iteration count grows like sqrt(kappa), and the
correlation condition number of that parameter block is 6.1e4. Nothing in
rabbit currently scales or preconditions parameters.

This PR adds two independent things: a preconditioner, and recovery from a
stalled minimizer.

1. --precondition (off by default)

Reparameterise theta = theta_ref + T y with T = L^-T from the Cholesky of a
reference Hessian restricted to a selected block, so T^T H0 T = I there.

The important design point: this is preconditioning obtained as a change of
variables, so the minimizer is left completely untouched. scipy's
trust-krylov (GLTR/_trlib) accepts no user preconditioner and does not need
to — a spherical trust region in y is an H0-aligned ellipsoid in theta,
which is exactly the desired effect.

It is confined to the three scipy callbacks in fit(); self.x holds physical
parameters everywhere else, so the postfit Hessian, covariance, impacts and
pulls are untouched and need no mapping back.

The transform is block diagonal, and by default the blocks are found from
the reference matrix rather than named by hand: threshold its correlations and
take connected components. On the in-situ fit that recovers 233-240 components
where the parameterisation has 240 (step, eta, charge) blocks, knowing nothing
about parameter names. Two reasons that matters more than tidiness:

• Cholesky is O(m^3), so 240 blocks of ~12 cost ~4e-5 of one block of 2112.
• A union of individually fine groups can be singular. On the combined W+Z
  fit the 1921 ABCD fake parameters were not factorisable and took the 2112
  in-situ ones down with them, leaving the fit unpreconditioned. With blocks, an
  unusable block is skipped and the rest still apply.

New options: --precondition, --preconditionParams (names, regexes or
systematic groups; defaults to the unconstrained parameters, where it pays off —
constrained nuisances are already normalised by their prior),
--preconditionBlocks {auto,expressions,none} with --preconditionBlockThreshold,
--preconditionFrom {hessian,gaussnewton}, --preconditionRidge.

auto is the default so that --precondition on its own does something sensible
without the user having to know which of their parameters are badly correlated;
expressions (one block per --preconditionParams entry) and none (one joint
block, the original behaviour) are there when that is known.

The ridge comes from the block's spectrum

A block that is not positive definite gets a ridge before the Cholesky. That
ridge used to escalate by powers of a hundred (1e-8 -> 1e-6 -> 1e-4 -> 1e-2),
which steps straight over what a mildly indefinite block needs: on the combined
W+Z fit two blocks holding 168 parameters were skipped as "not factorisable"
when their smallest eigenvalues were only 2-3% of max|diag| and nothing was
near zero — 1e-4 was too small and 1e-2 was the next rung up.

The ridge required to restore definiteness is set by the most negative
eigenvalue, so it is computed rather than guessed. The caller's value is still
tried first, since a positive definite block needs nothing more and the Cholesky
is cheap; only if that fails is the extra O(m^3) eigendecomposition worth it.

Measured on that fit: 69 of 71 blocks factorised before, 71 of 71 after, so
all 4033 selected parameters are now preconditioned. Splitting the failing
blocks further was the alternative and a worse one — they only became
factorisable at a correlation threshold of 0.5, which would have discarded every
correlation between 0.1 and 0.5, the bulk of what the transform is for.

The blocks are re-derived at every restart along with the factorisations, since
the correlation structure moves with the fit: over five builds of one 4D fit the
count went 233, 240, 230, 229, 213 with the median block pinned at 12 while the
largest grew 73 -> 204, i.e. clusters merge as the minimum is approached.

T is applied as a dense matvec against a cached L^-1 rather than a
triangular solve: same flops, but the solve is inherently sequential while the
matvec is a parallel GEMV, measured 15x faster at m=2112 (0.19 ms vs 2.9 ms per
HVP, 0.09% of a 215 ms HVP). The triangular solve is kept as a fallback.

2. --maxRestarts (on by default)

Independent of preconditioning, and useful for any trust-region fit.

scipy's _minimize_trust_region shrinks the trust radius 4x on every rejected
step with no lower bound, and holds it in a local variable. Once it has
collapsed the minimizer takes infinitesimal steps far from any minimum: the loss
freezes while the gradient is still large. A fresh minimize() call resets the
radius, which is why restarting from a stalled point resumes the descent — the
effect of chaining --externalPostfit by hand.

fit() now does this itself: on an early-stopping stall, restart from the
stalled point, keep restarting for as long as the loss keeps coming down, and
stop only once a restart no longer reduces it. The preconditioner is rebuilt at
each restart point, since the transform whitens the Hessian where it was built
and by the time the fit has stalled elsewhere that Hessian has changed
(measured: 3x worse conditioning at the restart point than at the start).

Measured: a fit frozen at loss 91596 reached 11252 through four automatic
restarts, matching the minimum the 145-minute run found.

Results

The cleanest measurement is the 4-dimensional fit (74864 bins, the distribution
the analysis actually intends to fit). Same input, same
-t 0 --noHessian --noEDM:

configuration	wall-clock	2*deltaNLL	final loss	EDM
no preconditioning	338 min	73118.71	36559.35666339957	5.3e-15
preconditioned, one block	19.1 min	73118.71	36559.35666339955	3.7e-16
preconditioned, auto blocks	50 min	73117.52	36558.7587	-

The same minimum, 17.7x faster for the single-block transform: it and the
unpreconditioned fit agree to 13 significant figures on the loss and exactly on
2*deltaNLL, with the preconditioned one converged an order of magnitude more
tightly (25 iterations instead of 93; the unpreconditioned run's last iterations
cost 3294 s each).

The third row is the current default (--precondition with auto blocks) and it
is not the same comparison: it walks a different path and ends in a different,
slightly deeper minimum (2deltaNLL lower by 1.2), taking 2.6x longer to get
there. So no speed ratio should be quoted for it. That auto-blocking finds a
deeper minimum than the joint transform is consistent with the 2D picture below,
where this likelihood is shown to have several minima a few units apart in
2deltaNLL and a reparameterisation changes which one is reached. Use
--preconditionBlocks none to reproduce the 19.1 min row.

The 2-dimensional projection of the same analysis, for reference:

configuration	wall-clock	2*deltaNLL	EDM
no preconditioning	145.5 min	22499.28	7.0e-11
--precondition	14.3 min	22504.78	7.3e-16
--precondition, no restart	4.6 min	22509.84	-

There the runs land on different minima — the likelihood has several separated
by ~3 in 2*deltaNLL, and a reparameterisation changes which one the optimiser
walks to, though it provably cannot move any of them. All are genuine minima:
note the preconditioned fits reach EDM ~1e-16 where the unpreconditioned one
stops at 7e-11, i.e. the speedup does not come at the cost of convergence. The
4D comparison above has no such ambiguity.

The same input also has a pre-existing offline basis-orthogonalisation applied
in the histmaker; that took 179 minutes to reach 22513.30, so preconditioning in
rabbit is ~12x faster than the analysis's previous best in 2D (~18x in 4D) and
removes the need to regenerate a basis file whenever the parameterisation
changes.

Things reviewers may want to push on

• gaussnewton is deliberately not the default. The Fisher matrix is PSD
  by construction and so cannot represent negative curvature. The exact Hessian
  of this fit has 33 negative eigenvalues at the starting point; whitening with
  Gauss-Newton left all 33 negative and the fit froze after 13 iterations, and
  raising the ridge did not help. The ridge on the exact Hessian turns out to be
  load-bearing for exactly that reason. Both docstrings record this.
  It does not generalise either way: with restarts it was the fastest option
  in 2D (5.7 min, EDM 6.5e-17) and clearly the worst in 4D (52 min and a higher
  NLL than hessian's 19.1 min), where the starting kappa is 4e8 against 1.3e5
  for the exact Hessian. Situational, hence not the default.
• The preconditioned runs still stop slightly above the unpreconditioned
  reference (22504.78 vs 22499.28), and the ordering across runs tracks how
  many restarts were taken. A pure reparameterisation cannot move the minimum,
  so this is about where each run stops, not where the minimum is. Note the
  obvious check — rerun with --earlyStopping -1 and let scipy converge — does
  not work: see below.
• Convergence thresholds are basis-dependent: gtol applies to ||grad||,
  which is not invariant under the transform, whereas EDM is.

scipy does not terminate on its own

Worth stating plainly, because it motivates the defaults. _minimize_trust_region
has no termination test for a collapsed trust radius, and maxiter defaults to
200*nparams. Measured on this 3603-parameter fit with --earlyStopping -1:
the loss froze after 63 seconds at 11427.90, and the following 720 iterations
changed nothing, running at ~10 iterations/second — which projects to twenty
hours before maxiter would end it. It was killed at 3 minutes.

So a stalled fit does not eventually converge, and it does not error out either:
it spins. That is why --earlyStopping is now on by default (20) — with restarts
it means "restart here", not "give up here" — and it is why the honest reading of
the table above is that every one of these numbers is a stopping point rather
than a converged minimum.

Backwards compatibility

• Preconditioning is off by default; existing fits are unchanged.
• Two default changes, both behaviour-affecting:
  --earlyStopping moves from -1 (off) to 20, and --maxRestarts defaults
  to -1 (restart while the loss keeps improving). Together a stall now means
  "restart and continue" where it previously meant "keep spinning" (see above).
  --earlyStopping -1 --maxRestarts 0 restores the old behaviour exactly.
  Happy to make either opt-in instead if the team prefers.
• Preconditioner logging is one summary line per build rather than one line per
  block, with the per-block detail moved to debug: with auto-blocking a run has
  hundreds of blocks, not the one it had when that line was written.
• A minimizer exception that is not a stall is now logged at warning level.
  Previously the bare except Exception swallowed everything, so a broken
  callback looked exactly like a converged fit.

Testing

44 tests across tests/test_preconditioner.py and tests/test_restart.py;
65 pass across the suite. The load-bearing ones:

• Invariance: a preconditioned fit reaches identical parameters,
  uncertainties and NLL to an unpreconditioned one, for both trust-krylov
  (hessp path) and trust-exact (dense hess path), on a deliberately
  ill-conditioned model (correlation condition number 1.3e3 -> 1.0).
• Algebra: whitening, chain rule vs finite differences, hessp against the
  dense transform, cached-inverse fast path against the triangular fallback.
• Degenerate blocks, singular blocks and an unformable reference matrix all fall
  back to running unpreconditioned rather than failing the fit.
• The rebuild-before-restart test forces the stall rather than relying on the
  model tripping a threshold, since TF's multithreaded CPU reductions are not
  bitwise reproducible and an earlier version of it was observed to both pass
  and fail on identical code.
• The spectrum-derived ridge is tested on a block deliberately built with its
  smallest eigenvalues at -3% of max|diag|, i.e. between the old ladder's
  rungs, which the previous schedule skipped; and on a rank-deficient block with
  an exact linear dependence.
• Auto-blocking recovers a hidden block structure exactly, merges two groups when
  a strong correlation is injected between them, isolates a parameter with no
  curvature, and percolates into one block below the threshold — the documented
  failure mode is tested rather than assumed away.

Not included

Absolute-eigenvalue whitening (V |Lambda|^-1/2) for indefinite Hessians, which
would subsume the ridge and handle the negative-curvature directions explicitly;
and passing the same factor as M= to the scipy.sparse.linalg.cg used in the
Hessian-free covariance path.

🤖 Generated with Claude Code