From 03f7fb5a043824cd1b3f7fb62a60a3f149aa551a Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Tue, 21 Jul 2026 19:03:16 +0100 Subject: [PATCH] feat(mle): auto-convergence early-stopping for multi-start gradient searches Add a MultiStartGradientConvergence settings object (mirroring AutoCorrelationsSettings) so MultiStartAdam/ADABelief/Lion/Prodigy stop early once the global-best figure-of-merit plateaus, defaulting on so users no longer hand-tune n_steps. n_steps remains a hard ceiling. The plateau check runs per-step (fom_history updates every step) while checkpointing stays at the iterations_per_full_update boundary, so early-stop works with the default single-chunk cadence. Skipped when resurrect=True (pixelized regime). Carries stop_reason in search_internal for resume and the phase-2 results contract. Phase 1 of 2 (results-contract + aggregator hardening follows). Refs #1406. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01LSEf7D5urqPgvffKBqzPsF --- autofit/__init__.py | 3 + .../mle/multi_start_gradient/convergence.py | 105 ++++++++++++++++++ .../search/mle/multi_start_gradient/search.py | 69 ++++++++++-- .../search/mle/test_multi_start_gradient.py | 90 +++++++++++++++ 4 files changed, 258 insertions(+), 9 deletions(-) create mode 100644 autofit/non_linear/search/mle/multi_start_gradient/convergence.py diff --git a/autofit/__init__.py b/autofit/__init__.py index a5b11d8a8..ba74c6265 100644 --- a/autofit/__init__.py +++ b/autofit/__init__.py @@ -92,6 +92,9 @@ from .non_linear.search.mle.multi_start_gradient.search import MultiStartADABelief from .non_linear.search.mle.multi_start_gradient.search import MultiStartLion from .non_linear.search.mle.multi_start_gradient.search import MultiStartProdigy +from .non_linear.search.mle.multi_start_gradient.convergence import ( + MultiStartGradientConvergence, +) from .non_linear.paths.abstract import AbstractPaths from .non_linear.paths import DirectoryPaths from .non_linear.paths import DatabasePaths diff --git a/autofit/non_linear/search/mle/multi_start_gradient/convergence.py b/autofit/non_linear/search/mle/multi_start_gradient/convergence.py new file mode 100644 index 000000000..e588f70fc --- /dev/null +++ b/autofit/non_linear/search/mle/multi_start_gradient/convergence.py @@ -0,0 +1,105 @@ +import numpy as np + +from autofit.non_linear.test_mode import is_test_mode + + +class MultiStartGradientConvergence: + def __init__( + self, + check_for_convergence: bool = True, + window: int = 50, + rtol: float = 1.0e-4, + atol: float = 1.0e-3, + min_steps: int = 100, + ): + """ + Settings for the auto-convergence (early-stopping) check of the multi-start + gradient searches (``MultiStartAdam`` / ``MultiStartADABelief`` / + ``MultiStartLion`` / ``MultiStartProdigy``). + + This mirrors the ``AutoCorrelationsSettings`` precedent that lets the + ensemble MCMC samplers (Emcee / Zeus) terminate before their full step + budget: the search still has a hard ``n_steps`` ceiling (it never runs + forever), but when ``check_for_convergence`` is ``True`` it stops early once + the global-best figure-of-merit has plateaued. + + The search minimises ``-2 * log_posterior`` (a chi-squared-like quantity), + so the tracked global best figure-of-merit ``best_fom`` is monotonically + non-increasing. Convergence is declared when the improvement over the + trailing ``window`` of the best-fom history is within tolerance:: + + improvement = fom_history[-window] - fom_history[-1] # >= 0 + converged = improvement <= atol + rtol * abs(fom_history[-window]) + + evaluated only once at least ``max(min_steps, window)`` steps have been + taken, so early transient plateaus (before the population has descended) + cannot trigger a false stop. + + Scope: this is the parametric-source (MGE / Sersic) regime, where a + global-best plateau genuinely means converged. It is deliberately **not** + applied when ``resurrect=True`` (the pixelized regime), whose best-fom + climbs in long plateaus punctuated by breakthrough jumps that plateau + detection would false-stop on; there the search leans on the ``n_steps`` + ceiling instead. + + Parameters + ---------- + check_for_convergence + Whether the global-best figure-of-merit is checked to terminate the + search early. If ``True`` the search may stop before ``n_steps`` has + been reached; if ``False`` all ``n_steps`` steps are taken. + window + The number of trailing steps of the best-fom history over which the + plateau (improvement within tolerance) is measured. A longer window + requires the best-fom to be flat for longer before stopping, trading + a few extra steps for robustness against noise. + rtol + The relative tolerance on the best-fom improvement over the window. + atol + The absolute tolerance on the best-fom improvement over the window. + min_steps + The minimum number of steps that must be taken before the convergence + check is allowed to terminate the search, so the population has time to + descend out of its broad starts before a plateau can stop it. + """ + self.check_for_convergence = check_for_convergence + self.window = window + self.rtol = rtol + self.atol = atol + self.min_steps = min_steps + + if is_test_mode(): + self.window = 1 + self.min_steps = 1 + + def check_if_converged(self, fom_history) -> bool: + """ + Whether the multi-start gradient search has converged: the global-best + figure-of-merit has plateaued over the trailing ``window``. + + Parameters + ---------- + fom_history + The history of the global best figure-of-merit, one entry per step + (monotonically non-increasing). Convergence is only assessed once at + least ``max(min_steps, window)`` entries are present. + """ + if not self.check_for_convergence: + return False + + fom_history = np.asarray(fom_history) + + if fom_history.size < max(self.min_steps, self.window): + return False + + latest = fom_history[-1] + past = fom_history[-self.window] + + # Both non-finite (no finite basin found yet) -> not converged; the plateau + # of the actual descent is what we want to detect, not a flat +inf run. + if not np.isfinite(latest) or not np.isfinite(past): + return False + + improvement = past - latest + + return bool(improvement <= self.atol + self.rtol * abs(past)) 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 05e23887c..f94c7785d 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -12,6 +12,9 @@ from autofit.non_linear.initializer import AbstractInitializer from autofit.non_linear.samples.sample import Sample from autofit.non_linear.samples.samples import Samples +from autofit.non_linear.search.mle.multi_start_gradient.convergence import ( + MultiStartGradientConvergence, +) class AbstractMultiStartGradient(AbstractMLE): @@ -38,6 +41,7 @@ def __init__( start_lower_limit: float = 0.15, start_upper_limit: float = 0.85, resurrect: bool = False, + convergence: Optional[MultiStartGradientConvergence] = None, initializer: Optional[AbstractInitializer] = None, iterations_per_full_update: int = None, iterations_per_quick_update: int = None, @@ -112,6 +116,16 @@ def __init__( searchable at all. (Even so, on such landscapes a nested sampler still wins decisively — resurrection makes gradient MAP *viable* there, not competitive.) + convergence + Auto-convergence (early-stopping) settings. When + ``check_for_convergence`` is ``True`` (the default) the search stops + early once the global-best figure-of-merit has plateaued, so users do + not have to hand-tune ``n_steps`` — which remains a hard ceiling / max + budget (the search never runs forever). ``None`` builds the default + ``MultiStartGradientConvergence()``. The check is deliberately skipped + when ``resurrect=True`` (the pixelized regime, whose best-fom climbs in + breakthrough jumps that a plateau check would false-stop on); there the + search leans on the ``n_steps`` ceiling. """ super().__init__( @@ -136,6 +150,9 @@ def __init__( self.start_lower_limit = start_lower_limit self.start_upper_limit = start_upper_limit self.resurrect = resurrect + self.convergence = ( + convergence if convergence is not None else MultiStartGradientConvergence() + ) self.logger.debug(f"Creating {self.optax_method} MultiStartGradient Search") @@ -223,6 +240,7 @@ def batched_value_and_grad(params): fom_history = list(search_internal["fom_history"]) total_steps = int(search_internal["total_steps"]) n_resurrections = int(search_internal.get("n_resurrections", 0)) + stop_reason = search_internal.get("stop_reason", None) self.logger.info( "Resuming MultiStartGradient search (previous samples found)." @@ -246,6 +264,7 @@ def batched_value_and_grad(params): fom_history = [] total_steps = 0 n_resurrections = 0 + stop_reason = None self.logger.info( f"Starting new {self.optax_method} MultiStartGradient search " @@ -256,13 +275,26 @@ def batched_value_and_grad(params): # ``resurrect`` is on); seeded independently of the broad-start draw. resurrect_rng = np.random.default_rng(1) - while total_steps < self.n_steps: + # ``n_steps`` is the hard ceiling / max budget; ``stop_reason`` becomes + # ``"converged"`` if the auto-convergence check stops the search early + # (this also short-circuits the loop on a resumed, already-converged run). + while total_steps < self.n_steps and stop_reason != "converged": steps_remaining = self.n_steps - total_steps iterations = min( self.iterations_per_full_update or self.n_steps, steps_remaining ) + # Convergence is assessed every step (``fom_history`` updates every + # step) rather than only at the ``iterations_per_full_update`` boundary, + # so early-stopping works even with the default single-chunk update + # cadence (``iterations_per_full_update=None``). Checkpointing and the + # expensive ``perform_update`` still happen only at the boundary (or on + # convergence), so the extra work per step is a cheap NumPy plateau + # check. The check is skipped when ``resurrect`` is on (pixelized + # regime), where a plateau does not mean converged. + converged = False + for _ in range(iterations): foms, grads = batched_value_and_grad(params) @@ -299,7 +331,18 @@ def batched_value_and_grad(params): updates, opt_state = step_update(grads, opt_state, params, foms) params = optax.apply_updates(params, updates) - total_steps += iterations + total_steps += 1 + + if not self.resurrect and self.convergence.check_if_converged( + fom_history + ): + converged = True + break + + if converged: + stop_reason = "converged" + elif total_steps >= self.n_steps: + stop_reason = "max_steps" search_internal = { "params": np.asarray(params), @@ -309,17 +352,27 @@ def batched_value_and_grad(params): "fom_history": np.asarray(fom_history), "total_steps": total_steps, "n_resurrections": n_resurrections, + "stop_reason": stop_reason, } self.paths.save_search_internal(obj=search_internal) + # A converged (or ceiling-reached) boundary is the final update, so it + # runs with ``during_analysis=False``; intermediate boundaries do not. + is_final = converged or total_steps >= self.n_steps self.perform_update( model=model, analysis=analysis, - during_analysis=total_steps < self.n_steps, + during_analysis=not is_final, fitness=fitness, search_internal=search_internal, ) + if converged: + self.logger.info( + f"{self.optax_method} MultiStartGradient converged early at " + f"step {total_steps} (best figure-of-merit plateaued)." + ) + self.logger.info(f"{self.optax_method} MultiStartGradient sampling complete.") return search_internal, fitness @@ -374,9 +427,9 @@ def _build_optimizer(self, optax, jax): @jax.jit def step_update(grads, opt_state, params, values): - return jax.vmap( - lambda g, s, p, v: optimizer.update(g, s, p, value=v) - )(grads, opt_state, params, values) + return jax.vmap(lambda g, s, p, v: optimizer.update(g, s, p, value=v))( + grads, opt_state, params, values + ) else: @@ -490,9 +543,7 @@ def samples_via_internal_from( # Fitness.call returns -2 * log_posterior, so log_posterior = -0.5 * fom. best_log_posterior = -0.5 * float(search_internal["best_fom"]) log_likelihood_list = [best_log_posterior - log_prior_list[0]] - log_likelihood_list += [ - np.nan for _ in range(len(parameter_lists) - 1) - ] + log_likelihood_list += [np.nan for _ in range(len(parameter_lists) - 1)] weight_list = [1.0] + [0.0] * (len(parameter_lists) - 1) 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 8ffe4e737..c105c4299 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 @@ -114,6 +114,96 @@ def test__batch_size_is_carried_to_every_rule(): assert cls(batch_size=8).batch_size == 8 +def test__convergence_default_is_on_and_carried(): + """Auto-convergence is a shared base-class setting, on by default (so users do + not hand-tune ``n_steps``), and a custom settings object is carried through.""" + for cls in ( + af.MultiStartAdam, + af.MultiStartADABelief, + af.MultiStartLion, + af.MultiStartProdigy, + ): + default = cls().convergence + assert isinstance(default, af.MultiStartGradientConvergence) + assert default.check_for_convergence is True + # documented defaults for the plateau check. + assert default.window == 50 + assert default.min_steps == 100 + assert default.rtol == pytest.approx(1.0e-4) + assert default.atol == pytest.approx(1.0e-3) + + custom = af.MultiStartGradientConvergence(check_for_convergence=False) + assert cls(convergence=custom).convergence is custom + + +def test__check_if_converged__plateau_stops_climbing_does_not(): + # rtol/atol both zero: converges only on an exactly-flat window. + convergence = af.MultiStartGradientConvergence( + window=3, min_steps=3, rtol=0.0, atol=0.5 + ) + + # Flat global-best over the window -> converged. + assert convergence.check_if_converged([10.0, 10.0, 10.0]) is True + + # Still descending over the window -> not converged. + assert convergence.check_if_converged([10.0, 5.0, 1.0]) is False + + +def test__check_if_converged__respects_min_steps_and_window_length(): + convergence = af.MultiStartGradientConvergence( + window=3, min_steps=5, rtol=0.0, atol=1.0 + ) + + # Fewer than max(min_steps, window) entries -> never terminates early. + assert convergence.check_if_converged([10.0, 10.0, 10.0]) is False + assert convergence.check_if_converged([10.0, 10.0, 10.0, 10.0]) is False + + # Once min_steps entries exist and the window is flat -> converged. + assert convergence.check_if_converged([10.0] * 5) is True + + +def test__check_if_converged__relative_tolerance(): + convergence = af.MultiStartGradientConvergence( + window=3, min_steps=3, rtol=1.0e-4, atol=0.0 + ) + + # Improvement 0.1 over a best-fom ~1000 is within rtol * 1000 = 0.1 -> converged. + assert convergence.check_if_converged([1000.1, 1000.05, 1000.0]) is True + + # A large improvement over the same window is not within tolerance. + assert convergence.check_if_converged([1002.0, 1001.0, 1000.0]) is False + + +def test__check_if_converged__disabled_and_non_finite(): + # check_for_convergence=False never terminates, even on a flat window. + disabled = af.MultiStartGradientConvergence( + check_for_convergence=False, window=3, min_steps=3 + ) + assert disabled.check_if_converged([10.0, 10.0, 10.0]) is False + + # A non-finite window (no finite basin found yet) is not a real plateau. + convergence = af.MultiStartGradientConvergence(window=3, min_steps=3, atol=1.0) + assert convergence.check_if_converged([np.inf, np.inf, np.inf]) is False + + +def test__dict_round_trip__convergence(): + # The convergence settings must survive serialisation so a resumed search + # keeps the same early-stopping behaviour. + search = af.MultiStartAdam( + n_starts=4, + convergence=af.MultiStartGradientConvergence( + check_for_convergence=False, window=7, min_steps=9 + ), + ) + restored = from_dict(to_dict(search)) + + assert isinstance(restored, af.MultiStartAdam) + assert isinstance(restored.convergence, af.MultiStartGradientConvergence) + assert restored.convergence.check_for_convergence is False + assert restored.convergence.window == 7 + assert restored.convergence.min_steps == 9 + + def test__dict_round_trip(): dictionary = to_dict(af.MultiStartLion(n_starts=7, n_steps=33)) restored = from_dict(dictionary)