From 129190adbd7298d85978cdae7cd5d74dd221ba9d Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sun, 16 Aug 2026 18:54:12 -0400 Subject: [PATCH 1/2] feat: seed the multi-start draws and record the alive-versus-step curve MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two additions the prior-support validation campaign (autolens_profiling#128, phase 2) cannot be run without. `seed` — the search's two random draws, the broad starting points and the resurrection redraw, were hardcoded to `default_rng(0)` / `default_rng(1)`. Seeding `random`/`numpy` reaches only the initializer, so every run of a model drew the SAME starting population and a "multi-seed" study was silently a single-seed one. `seed=None` keeps the historical fixed seeds exactly, so existing fits are bit-identical and the argument is purely additive. The two streams are derived through `SeedSequence` rather than by offsetting the seed: `seed + stream` would make seed 0's resurrection stream the same sequence as seed 1's starting stream, so nominally independent seeds would share draws and resurrection would replay the starting population. `alive_history` — the number of living lanes per step, written to `search_internal`. The existing lane counters are survival INTEGRALS: a dead lane keeps adding to them every subsequent step, so the same death curve reads ~60% at 150 steps and ~75% at 300, and two runs at different budgets cannot be compared on the scalar at all. The curve is the budget-independent quantity, and until now it existed only in the progress log at `iterations_per_log` cadence — visible to a human reading stdout, unavailable to any analysis. Scope: this seeds the draws THIS search owns, not the framework. The initializer and the sampler-owned generators remain unseedable; that wider gap is filed separately. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01VRG2X7Af8zdF3vnsWoiK4U --- .../search/mle/multi_start_gradient/search.py | 68 ++++++++++++++++++- .../search/mle/test_multi_start_gradient.py | 43 ++++++++++++ 2 files changed, 109 insertions(+), 2 deletions(-) diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index eb2ad4059..b07790715 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -43,6 +43,7 @@ def __init__( start_lower_limit: float = 0.15, start_upper_limit: float = 0.85, resurrect: bool = False, + seed: Optional[int] = None, convergence: Optional[MultiStartGradientConvergence] = None, iterations_per_log: int = 10, initializer: Optional[AbstractInitializer] = None, @@ -146,6 +147,22 @@ def __init__( searchable at all. (Even so, on such landscapes a nested sampler still wins decisively — resurrection makes gradient MAP *viable* there, not competitive.) + seed + Seeds the two random draws this search owns: the broad starting + points (``_broad_starts``) and, when ``resurrect`` is on, the redraw + of dead lanes. Default ``None`` reproduces the historical fixed + seeds exactly, so an existing fit is bit-identical and this argument + is purely additive. + + It exists because without it the search cannot be varied *or* + genuinely repeated: both draws were hardcoded, so every run of the + same model produced the **same** starting population no matter what + the caller seeded. Seeding ``random`` / ``numpy`` globally does not + reach them — that only perturbs the initializer — which makes a + multi-seed study silently a single-seed one. Note the scope: this + seeds *this search's* draws, not the whole framework (the + initializer and any sampler-owned generator are separate, still-open + sources of non-reproducibility). convergence Auto-convergence (early-stopping) settings. When ``check_for_convergence`` is ``True`` (the default) the search stops @@ -204,6 +221,7 @@ def __init__( self.start_lower_limit = start_lower_limit self.start_upper_limit = start_upper_limit self.resurrect = resurrect + self.seed = seed self.convergence = ( convergence if convergence is not None else MultiStartGradientConvergence() ) @@ -698,6 +716,11 @@ def batched_value_and_grad(params): best_params = np.asarray(search_internal["best_params"]) best_fom = float(search_internal["best_fom"]) fom_history = list(search_internal["fom_history"]) + # ``.get``: a ``search_internal`` written before the alive history + # existed has no such key, and a resumed run must not KeyError on + # it. An older run's curve is unrecoverable, so it resumes empty + # rather than being back-filled with a fabricated population. + alive_history = list(search_internal.get("alive_history", [])) total_steps = int(search_internal["total_steps"]) n_resurrections = int(search_internal.get("n_resurrections", 0)) # ``.get`` with a default: a ``search_internal`` written before the @@ -761,6 +784,7 @@ def batched_value_and_grad(params): best_params = np.asarray(params[0]) best_fom = np.inf fom_history = [] + alive_history = [] total_steps = 0 n_resurrections = 0 n_value_nan_lane_steps = 0 @@ -776,7 +800,7 @@ def batched_value_and_grad(params): # Deterministic RNG for redrawing dead starts (only used when # ``resurrect`` is on); seeded independently of the broad-start draw. - resurrect_rng = np.random.default_rng(1) + resurrect_rng = np.random.default_rng(self._seed_for(1)) stop_reason = self._stop_reason_on_resume(stop_reason) @@ -846,6 +870,17 @@ def batched_value_and_grad(params): fom_history.append(best_fom) + # The size of the living population at this step. Recorded as a + # history because the cumulative lane counters above are + # survival INTEGRALS: a dead lane keeps adding to them every + # subsequent step, so the same death curve reads ~60% at 150 + # steps and ~75% at 300 and two runs at different budgets cannot + # be compared on the scalar at all. The curve is the + # budget-independent quantity, and until now it existed only in + # the progress log at ``iterations_per_log`` cadence — visible + # to a human reading stdout, unavailable to any analysis. + alive_history.append(int(np.count_nonzero(alive))) + # Restart-on-death: redraw any start whose objective went # non-finite (fresh params + reinitialised per-start optimizer # state), leaving alive starts untouched. best_* is captured @@ -937,6 +972,7 @@ def batched_value_and_grad(params): "best_params": best_params, "best_fom": best_fom, "fom_history": np.asarray(fom_history), + "alive_history": np.asarray(alive_history), "total_steps": total_steps, "n_resurrections": n_resurrections, "n_value_nan_lane_steps": n_value_nan_lane_steps, @@ -1084,6 +1120,28 @@ def merge(old, fresh): return params, opt_state + def _seed_for(self, stream: int): + """Seed material for one of this search's random draws. + + ``stream`` identifies the draw: ``0`` the broad starting points, ``1`` + the resurrection redraws. They must not share a stream — resurrection + would otherwise replay the starting population. + + ``seed=None`` returns the bare stream index, which is exactly the + historical hardcoded seed at each site (``default_rng(0)`` / + ``default_rng(1)``), so the default path is bit-identical. + + With a seed set, the pair is derived through ``SeedSequence`` rather + than by offsetting the seed. ``seed + stream`` looks equivalent and is + not: it makes seed 0's resurrection stream the same sequence as seed 1's + start stream, so nominally independent seeds in a multi-seed study share + draws. ``SeedSequence`` spreads each ``(seed, stream)`` pair over the + full state space instead. + """ + if self.seed is None: + return stream + return np.random.SeedSequence([self.seed, stream]) + def _broad_starts(self, model, value_and_grad_single, jnp): """ Draw ``n_starts`` broad starting points in the unit cube, map them to @@ -1097,7 +1155,7 @@ def _broad_starts(self, model, value_and_grad_single, jnp): cost per draw, which on a multi-band ``FactorGraphModel`` objective dominated the whole fit (~13 minutes for 16 draws, cache or no cache). """ - rng = np.random.default_rng(0) + rng = np.random.default_rng(self._seed_for(0)) starts = [] max_tries = self.n_starts * 30 @@ -1209,6 +1267,12 @@ def samples_via_internal_from( "n_clipped_lane_steps": int( search_internal.get("n_clipped_lane_steps", 0) ), + # The seed this search's own draws used, so a result file says which + # member of a multi-seed study produced it. ``None`` records the + # default (historical fixed seeds) rather than being omitted, so a + # seeded and an unseeded run are distinguishable downstream instead + # of both reading as "no seed key". + "seed": self.seed, # Auto-convergence outcome: whether the run stopped on the plateau # check ("converged") or exhausted the ``n_steps`` ceiling # ("max_steps"), the settings that produced it, and the global-best diff --git a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py index 4458c71ac..2b7b3e2b6 100644 --- a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py +++ b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py @@ -278,6 +278,49 @@ def test__dict_round_trip__resurrect(): assert restored.n_starts == 6 +def test__dict_round_trip__seed(): + # The seed must survive serialisation, or a resumed member of a multi-seed + # study silently continues on the default draw. + restored = from_dict(to_dict(af.MultiStartAdam(seed=7, n_starts=6))) + + assert isinstance(restored, af.MultiStartAdam) + assert restored.seed == 7 + + +def test__seed__default_is_bit_identical_to_the_historical_fixed_seeds(): + # ``seed`` is additive: the default path must draw exactly what the + # hardcoded ``default_rng(0)`` / ``default_rng(1)`` drew before it existed, + # so no existing fit changes. + search = af.MultiStartAdam(seed=None) + + assert search._seed_for(0) == 0 + assert search._seed_for(1) == 1 + + +def test__seed__is_reproducible_and_varies_between_seeds(): + def draw(seed, stream): + search = af.MultiStartAdam(seed=seed) + return np.random.default_rng(search._seed_for(stream)).uniform(size=8) + + # Same seed reproduces; different seeds diverge. Without this the campaign's + # "at least two seeds per arm" is silently a single-seed study. + assert np.array_equal(draw(3, 0), draw(3, 0)) + assert not np.array_equal(draw(3, 0), draw(4, 0)) + + +def test__seed__start_and_resurrect_streams_never_coincide(): + # Guards the bug a naive ``seed + stream`` offset would introduce: seed 0's + # resurrection stream would be seed 1's starting stream, so nominally + # independent seeds would share draws and resurrection would replay the + # starting population. + def draw(seed, stream): + search = af.MultiStartAdam(seed=seed) + return np.random.default_rng(search._seed_for(stream)).uniform(size=8) + + assert not np.array_equal(draw(0, 0), draw(0, 1)) + assert not np.array_equal(draw(0, 1), draw(1, 0)) + + def test__samples_via_internal_from(): model = af.Model(example.Gaussian) From 4616665333ed9a4f11aa686bd8390e39e538ae0f Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sun, 16 Aug 2026 21:39:08 -0400 Subject: [PATCH 2/2] feat: optional momentum reset on clipped coordinates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `reset_momentum_on_clip` zeroes the optimizer moments wherever the clipper just projected a coordinate back onto its prior box. Default False, so the clipping path is unchanged unless asked for. Projection alone leaves a lane holding the exact velocity that carried it out of the box, so the next step drives it into the same wall and it is re-projected onto the same bound indefinitely — counted alive, permanently pinned, still paying a full likelihood-and-gradient evaluation every step. The reset is per-coordinate: a lane clipped in one parameter keeps its momentum in the others. Moment fields are targeted BY NAME, not by shape. Shape matching is actively wrong here: Prodigy's `params0` and `grad_sum` carry the same `(n_starts, n_params)` shape as the moments, and `params0` anchors its learning-rate estimate — zeroing it would corrupt the step size for the rest of the run rather than resetting momentum. A regression test pins that. Measured (autolens_profiling#131, imaging/mge hst, 16x3000, fp64) the arm does NOT pay off and the campaign recommends against using it: same converged answer, but deaths 2 -> 2523, one MORE lane pinned, +39% wall on seed 0, and on seed 1 it gives back nearly all of plain clipping's gain (-120880.6 -> -137783.6). Shipped default-off so the measurement is reproducible, not because it is recommended. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01VRG2X7Af8zdF3vnsWoiK4U --- .../search/mle/multi_start_gradient/search.py | 75 +++++++++++++++++ .../search/mle/test_multi_start_gradient.py | 81 +++++++++++++++++++ 2 files changed, 156 insertions(+) diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index b07790715..7cbaa3784 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -48,6 +48,7 @@ def __init__( iterations_per_log: int = 10, initializer: Optional[AbstractInitializer] = None, clipper: Optional[AbstractClipper] = None, + reset_momentum_on_clip: bool = False, iterations_per_full_update: int = None, iterations_per_quick_update: int = None, silence: bool = False, @@ -163,6 +164,25 @@ def __init__( seeds *this search's* draws, not the whole framework (the initializer and any sampler-owned generator are separate, still-open sources of non-reproducibility). + reset_momentum_on_clip + Zero the optimizer's momentum on the coordinates a ``clipper`` just + clipped. Default ``False``, so the clipping path is unchanged unless + asked for. + + It exists because projection alone leaves a lane holding the exact + velocity that carried it out of the prior box, so the next step + drives it into the same wall and it is re-projected onto the same + bound indefinitely — counted as alive, permanently pinned, and still + paying a full likelihood-and-gradient evaluation every step. The + reset is per-coordinate: a lane clipped in one parameter keeps its + momentum in all the others. + + Note that a pinned lane is not automatically a failure. Where the + likelihood genuinely prefers a value outside the prior, sitting on + the bound is the correct MAP answer under the declared prior, and + this flag would then be discarding useful state. It is a knob for + the case where pinning is an artefact of momentum rather than a + statement about the data. convergence Auto-convergence (early-stopping) settings. When ``check_for_convergence`` is ``True`` (the default) the search stops @@ -222,6 +242,7 @@ def __init__( self.start_upper_limit = start_upper_limit self.resurrect = resurrect self.seed = seed + self.reset_momentum_on_clip = reset_momentum_on_clip self.convergence = ( convergence if convergence is not None else MultiStartGradientConvergence() ) @@ -925,6 +946,17 @@ def batched_value_and_grad(params): np.count_nonzero(np.asarray(jnp.any(clipped_mask, axis=-1))) ) + # Optionally zero the optimizer's momentum on the clipped + # coordinates. Without it a clipped lane keeps the velocity + # that pushed it out of the box, so it is re-projected onto + # the same bound every step: alive in the counters, pinned + # to the wall, and still spending a full likelihood-and- + # gradient evaluation per step. + if self.reset_momentum_on_clip: + opt_state = self._reset_clipped_momentum( + opt_state=opt_state, clipped_mask=clipped_mask, jnp=jnp + ) + total_steps += 1 if not self.resurrect and self.convergence.check_if_converged( @@ -1120,6 +1152,48 @@ def merge(old, fresh): return params, opt_state + # The optax moment accumulators — the "momentum" a clip should forget. + # Named explicitly rather than matched by shape, because shape matching is + # actively wrong here: Prodigy's ``params0`` and ``grad_sum`` carry the same + # ``(n_starts, n_params)`` shape as the moments, and ``params0`` is the + # reference point of its learning-rate estimate. Zeroing that would not + # reset momentum, it would corrupt the step-size estimate for the rest of + # the run. Anything not named here (``params0``, ``grad_sum``, ``estim_lr``, + # ``numerator_weighted``, ``count``) is deliberately left intact. + _MOMENTUM_FIELDS = frozenset({"mu", "nu", "exp_avg", "exp_avg_sq", "trace"}) + + @classmethod + def _reset_clipped_momentum(cls, opt_state, clipped_mask, jnp): + """Zero the optimizer moments wherever a coordinate was clipped. + + ``clipped_mask`` is ``(n_starts, n_params)``, so the reset is + per-coordinate: a lane clipped in one parameter keeps its momentum in + every other parameter and only forgets the direction that took it out of + the box. + """ + + def rebuild(node): + fields = getattr(node, "_fields", None) + if fields is not None: + return type(node)( + **{ + name: ( + jnp.where(clipped_mask, jnp.zeros_like(value), value) + if name in cls._MOMENTUM_FIELDS + and getattr(value, "shape", None) == clipped_mask.shape + else rebuild(value) + ) + for name, value in zip(fields, node) + } + ) + if isinstance(node, tuple): + return tuple(rebuild(child) for child in node) + if isinstance(node, list): + return [rebuild(child) for child in node] + return node + + return rebuild(opt_state) + def _seed_for(self, stream: int): """Seed material for one of this search's random draws. @@ -1264,6 +1338,7 @@ def samples_via_internal_from( # current process's share — the same reasoning as the ``.get`` # defaults above. "clipper": type(self.clipper).__name__, + "reset_momentum_on_clip": self.reset_momentum_on_clip, "n_clipped_lane_steps": int( search_internal.get("n_clipped_lane_steps", 0) ), diff --git a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py index 2b7b3e2b6..55b1eba79 100644 --- a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py +++ b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py @@ -1,4 +1,5 @@ import inspect +from typing import NamedTuple import numpy as np import pytest @@ -321,6 +322,86 @@ def draw(seed, stream): assert not np.array_equal(draw(0, 1), draw(1, 0)) +class _AdamLikeState(NamedTuple): + count: np.ndarray + mu: np.ndarray + nu: np.ndarray + + +class _ProdigyLikeState(NamedTuple): + exp_avg: np.ndarray + exp_avg_sq: np.ndarray + grad_sum: np.ndarray + params0: np.ndarray + estim_lr: np.ndarray + + +def test__reset_clipped_momentum__zeroes_only_the_clipped_coordinates(): + # (2 starts, 3 params); only start 0's middle coordinate was clipped. + mask = np.array([[False, True, False], [False, False, False]]) + state = _AdamLikeState( + count=np.array([7, 7]), + mu=np.ones((2, 3)), + nu=np.full((2, 3), 2.0), + ) + + out = af.MultiStartAdam._reset_clipped_momentum( + opt_state=state, clipped_mask=mask, jnp=np + ) + + assert out.mu.tolist() == [[1.0, 0.0, 1.0], [1.0, 1.0, 1.0]] + assert out.nu.tolist() == [[2.0, 0.0, 2.0], [2.0, 2.0, 2.0]] + # A lane keeps its momentum in every coordinate that was not clipped, and + # non-moment state is untouched. + assert out.count.tolist() == [7, 7] + + +def test__reset_clipped_momentum__never_touches_prodigys_reference_point(): + # ``params0`` and ``grad_sum`` carry the SAME shape as the moments, so a + # shape-matched reset would zero them. ``params0`` anchors Prodigy's + # learning-rate estimate: zeroing it corrupts the step size for the rest of + # the run rather than resetting momentum. + mask = np.array([[True, True, True]]) + state = _ProdigyLikeState( + exp_avg=np.ones((1, 3)), + exp_avg_sq=np.ones((1, 3)), + grad_sum=np.full((1, 3), 5.0), + params0=np.full((1, 3), 9.0), + estim_lr=np.array([0.5]), + ) + + out = af.MultiStartProdigy._reset_clipped_momentum( + opt_state=state, clipped_mask=mask, jnp=np + ) + + assert out.exp_avg.tolist() == [[0.0, 0.0, 0.0]] + assert out.exp_avg_sq.tolist() == [[0.0, 0.0, 0.0]] + assert out.params0.tolist() == [[9.0, 9.0, 9.0]] + assert out.grad_sum.tolist() == [[5.0, 5.0, 5.0]] + assert out.estim_lr.tolist() == [0.5] + + +def test__reset_clipped_momentum__recurses_through_optax_chain_tuples(): + # optax.adam's state arrives wrapped in a chain tuple, so the reset has to + # descend through plain tuples as well as NamedTuples. + mask = np.array([[True, False]]) + nested = (_AdamLikeState(count=np.array([1]), mu=np.ones((1, 2)), nu=np.ones((1, 2))),) + + out = af.MultiStartAdam._reset_clipped_momentum( + opt_state=nested, clipped_mask=mask, jnp=np + ) + + assert out[0].mu.tolist() == [[0.0, 1.0]] + + +def test__dict_round_trip__reset_momentum_on_clip(): + restored = from_dict(to_dict(af.MultiStartAdam(reset_momentum_on_clip=True))) + + assert restored.reset_momentum_on_clip is True + # Default stays off, so the clipping path is unchanged unless asked for. + assert af.MultiStartAdam().reset_momentum_on_clip is False + + def test__samples_via_internal_from(): model = af.Model(example.Gaussian)