feat: per-parameter step scaling for the multi-start gradient searches - #1485
Merged
Conversation
`ScalerPriorWidth` gives each parameter its own step scale, derived from its
prior, applied as a change of variables so the rule steps in `phi = theta / s`
while the objective is still evaluated at the physical `theta = s * phi`.
Default `ScalerNone`, skipped rather than applied as a multiply by ones, so the
default path and its compiled step are both unchanged.
WHY THIS FAMILY. Every other search in the library already adapts per-parameter
scale: the nested samplers propose in the unit cube, Emcee's stretch move is
affine-invariant, Zeus self-tunes per direction, BlackJAXNUTS adapts a MASS
MATRIX via window_adaptation, and L-BFGS accumulates curvature. MultiStart
{Adam,Lion,ADABelief,Prodigy} is the sole exception, and not by oversight: the
Adam rules normalise per-coordinate by GRADIENT magnitude rather than PARAMETER
scale, so with m/sqrt(v) ~ 1 every coordinate moves ~lr in PHYSICAL units, and
Prodigy adds only a single GLOBAL `d`. On the reference imaging/mge cell prior
widths span 40x, so a step sensible for einstein_radius (width 8.0) is a
wall-crossing for bulge.centre (width 0.2) — measured at 31.5% of lane-steps
clipped (autolens_profiling#131). This is therefore the diagonal preconditioner
the family lacks, the analogue of the NUTS mass matrix, not a novel hack.
WHAT IT COSTS. NUTS LEARNS its preconditioner from the geometry during warmup;
this one is static and prior-derived — cheap, no warmup, no new hazard, but a
prior width is a PROXY for the posterior scale and is poor where the likelihood
is far tighter than the prior. Documented in the module rather than hidden.
THE RULE is the physical extent of a unit-cube-sized step: Uniform -> width;
LogUniform -> sqrt(lo*hi)*log(hi/lo) at the median, NOT the physical width,
which for LogUniform(1e-4, 1e4) is wrong by a factor of 543; Gaussian and
TruncatedGaussian -> sigma, NOT the truncation width, which for ell_comps
(sigma 0.3 in a [-1,1] box) overstates the scale six-fold; LogGaussian ->
exp(mean)*sigma. Anything degenerate falls back to 1.0 and logs — the vector is
a DIVISOR, so a zero would send the whole population to infinity. Normalised to
a geometric mean of exactly 1, so only the RATIOS change and an A/B is not
confounded with an effective learning-rate change.
APPLIED AS A CHANGE OF VARIABLES, not `params += s * updates`. Prodigy estimates
`d` from the distance actually travelled (params0, grad_sum), so rescaling its
update externally would leave its own estimate inconsistent with the trajectory
it believes it took. Composed as `fitness.call(phi * scale)`, JAX's chain rule
yields the gradient w.r.t. phi and there is no manual gradient rescaling
anywhere. Optimising the density OF phi instead would fold a Jacobian in and
move the MAP — silently, with every counter looking healthy — so
`test_scaler.py` pins the invariance directly. The map is linear, so an optimum
that genuinely sits on a boundary stays finite; that is what killed the logit
reparameterisation this replaces.
`search_internal` stays entirely in PHYSICAL units (`params` multiplied on write,
divided on resume). The scaler does not enter the search identifier, so a scaled
and an unscaled arm share an output directory; a file whose units depended on an
invisible knob would resume as a silently wrong population rather than as an
error. `samples_via_internal_from` reads that array directly and has no scaler
to consult.
COMPLEMENTARY TO THE CLIPPER, not a replacement. Scaling makes reaching a wall
rarer, never impossible, and where the likelihood genuinely prefers a value
outside the prior the clipped lane on the bound is the correct MAP under the
declared prior. `Clipper.project` takes an optional `scale` and clips against
`bounds / s` in the caller's own coordinates, rather than round-tripping the
whole batch through physical space every step.
VISIBILITY. The derived vector is appended to the `model.info` FILE — not to the
`Model.info` property, which has no access to the search and so cannot render a
search-dependent block. `search.summary` gains a `Scaler = ...` line under the
same discipline as the clipper's: `ScalerNone` emits nothing, so an existing
run's summary is unchanged byte for byte. There is no per-step count to report;
scaling's effect is read in the clipped lane-step rate.
NOT offered on LBFGS, which takes a `clipper` but does not own its step loop — a
knob that accepts a value and ignores it is worse than no knob. A test pins that.
Verified end-to-end on the real JAX path (the library suite is NumPy-only by
house rule, so it does not execute `_fit`): MultiStartProdigy on a 40x-spread
model reaches the same optimum to ~1e-9 with scaling on and off, Prodigy's own
`d` estimates differ (so the scaler is not an accidental no-op), the stored
per-start parameters come back physical, and the clip count falls 2 -> 0.
Note when reading logs: under a scaler the reported `d` is in SCALED units,
because that is the space the rule steps in. Do not compare it across a scaled
and an unscaled arm; compare the clip rate.
Issue #1483.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…per-default policy Two clarifications from review, both about honesty rather than behaviour. SCALE DEFINITION. Mixing Uniform -> width with Gaussian -> sigma means one scaled unit does NOT carry a single consistent probabilistic interpretation across coordinates. That is fine for the ENGINEERING goal here (stop a step sized for a wide coordinate crossing a narrow one's wall, for which order-of-magnitude commensurability is enough) but it is not a claim to prior-standardised coordinates, and the module now says so. The unified definitions to reach for if that is ever wanted are named: the prior's actual standard deviation, a robust central quantile width, or the local inverse-CDF derivative dtheta/du at a reference point — the last being the continuous generalisation of what the uniform and log-uniform rules already do. None is implemented, deliberately: if the simple width/sigma version does not move the clip rate, a more principled one will not either, and the diagnosis is wrong. CLIPPER DEFAULT. Reframed from hygiene to constrained-optimizer semantics. A search advertising a posterior with hard prior support is solving a CONSTRAINED problem, so a state outside that support is infeasible, not merely poor. Projection does not alter the constrained optimum, makes the invariant explicit, prevents silent invalid trajectories, stops prior exits masking later pathologies, and is what makes the diagnostics interpretable. Hard-support enforcement is therefore the INTENDED default for the gradient/MLE searches; the reason it is not yet is empirical breadth, not mathematics — everything measured is MultiStartProdigy on one lens cell, and Adam/Lion/ADABelief are unmeasured while being MORE exposed (Adam steps ~lr in physical units in every coordinate, Lion exactly it). When it flips it must not be sold as a long-budget accuracy gain: at 16x3000 the answer does not move at all. The scaler and the clipper stay separate features regardless. Even at a clip rate driven from 31% to 0.1%, projection is retained as the last-line invariant. Issue #1483. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Jammy2211
changed the base branch from
feature/clipper-validation-campaign
to
main
August 17, 2026 17:17
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.
Closes #1483. Stacked on
feature/clipper-validation-campaign(#1482) — merge that first.Adds
ScalerPriorWidth: a per-parameter step scale derived from each prior,applied as a linear change of variables
phi = theta / swith the objective stillevaluated at the physical
theta = s * phi. Default off and bit-identical whenoff, exactly as
seed,alive_historyandreset_momentum_on_clipshipped.Read the measurement before the code: it is a NEGATIVE result
This was built to fix the prior-exit failure mode. It does not. Three of the
four pre-registered falsification conditions fired (autolens_profiling#133,
results/notes/clipper_campaign/RESULTS.md):prior_boxprior_box_scaledprior_boxprior_box_scaledagainst 15097, a 0.3% change across 48,000 lane-steps. Seed 1 went the wrong way.
exceed the Nautilus reference.
0 -> 5884, with lanes inside their boxes — the likelihood going non-finite.
It ships default-off so that row stays reproducible, which is the same reason
reset_momentum_on_clipshipped after measuring badly. The artefact worth keepingis the measurement, not the feature. Nothing in this PR changes any existing fit.
Why it fails, and what the real failure turned out to be
The scale is normalised to a geometric mean of 1 so only the ratios change — but
the aggregate propensity to reach a wall is set by the global step magnitude
relative to the box widths, which is exactly what that normalisation holds fixed.
Scaling redistributes wall contact rather than reducing it (
pinned_coords19 -> 6 on seed 0: same six pinned lanes, a third as many coordinates each).
Inspecting the winning parameters showed the actual failure is not wall contact at
all. Seed 1 converged to a degenerate "no lens" solution —
einstein_radiusdriven to 0.0000, the
U(0, 8)lower bound, with the source displaced ~1.6" toabsorb the image. Once
theta_E = 0there is no deflection and nothing pulls themass back. Scaling made this worse because
einstein_radiusdrew the largestscale in the vector by far (23.87 against a 0.298 floor), so it took the
biggest steps straight into that attractor.
What the code does
autofit/non_linear/scaler.py—AbstractScaler/ScalerNone(default) /ScalerPriorWidth, modelled onClipper. Rule: Uniform -> width; LogUniform ->sqrt(lo*hi)*log(hi/lo)at the median (the physical width is wrong by 543x forLogUniform(1e-4, 1e4)); Gaussian and TruncatedGaussian -> sigma, not thetruncation width (
ell_compsis sigma 0.3 in a [-1,1] box); LogGaussian ->exp(mu)*sigma. Degenerate priors fall back to 1.0 and log — the vector is adivisor, so a zero would send the whole population to infinity.
dfrom thedistance actually travelled, so rescaling its update externally would desync its
own estimate from the trajectory it believes it took. Composed as
fitness.call(phi * scale), so JAX's chain rule yields the phi-gradient andthere is no manual gradient rescaling anywhere.
search_internalstays in PHYSICAL units. The scaler does not enter thesearch identifier, so scaled and unscaled arms share an output directory; a file
whose units depended on an invisible knob would resume as a silently wrong
population rather than as an error.
model.infogains the derived scale block;search.summarygainsScaler = .... Both emit nothing when scaling is off, so existing artefacts areunchanged byte for byte.
LBFGS, which takes aclipperbut does not own its steploop — an inert knob is worse than no knob. A test pins that.
Testing
test_autofit/non_linear/test_scaler.py, pure NumPy per the house rule. Pins theMAP invariance directly (a constant diagonal Jacobian cannot move the optimum —
optimising the density of phi instead would move it silently), the scale-rule
table including the two "not the width" cases, geometric-mean normalisation, the
degenerate-prior fallback, and clipper composition.
Full suite: 1852 passed, 2 skipped.
The suite is NumPy-only by house rule, so it never executes
_fit— a green suiteproves nothing about this feature. Verified separately end-to-end on the real JAX
path: same optimum to ~1e-9 with scaling on and off, Prodigy's
destimatesgenuinely differ (so it is not an accidental no-op), stored per-start parameters
come back physical, and the clip count fell 2 -> 0 on a toy model.
Known limitation, stated in the module
Mixing Uniform -> width with Gaussian -> sigma means one scaled unit carries no
single probabilistic interpretation. That is acceptable for an engineering goal
(commensurable steps) but is not prior standardisation. The unified alternatives —
prior standard deviation, robust central quantile width, local inverse-CDF
derivative — are named and deliberately not implemented: measure the simple
version first, and it did not work.
🤖 Generated with Claude Code