feat(search): add multi-start gradient MAP searches (Adam/ADABelief/Lion)#1370
Merged
Conversation
…ion) Promote the benchmark-winning multi-start first-order gradient MAP optimizer to first-class PyAutoFit searches. Adds AbstractMultiStartGradient(AbstractMLE) and the concrete MultiStartAdam / MultiStartADABelief / MultiStartLion rules, running N broad multi-starts vmapped over the af.Fitness seam with a fixed self-normalised optax update per start, returning best-basin MAP + per-start diagnostics through the standard samples/result contract. optax added to the jax extra. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This was referenced Jul 14, 2026
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.
Summary
Promotes the benchmark-winning multi-start gradient MAP optimizer (autolens_workspace_developer PR#96+#98, Phase-3 complete — nothing beat wide multi-start Adam on the kinked lens likelihood) to first-class PyAutoFit searches. Adds
AbstractMultiStartGradient(AbstractMLE)and three concrete rule-classes, mirroring the existingAbstractBFGS → BFGS/LBFGSidiom. Each runs N broad multi-starts in parallel (jax.vmap) on the unconstrained parameterization, with a fixed self-normalised optax update per start, returning the best-basin MAP point plus per-start basin diagnostics through the standard search/samples/result contract (save/resume included).Phase 1 of the promotion: core search + samples/result. Follow-ons (not in this PR): autofit_workspace_test JAX truth-basin validation (library-first gate), Phase 2 (config/packaged defaults), Phase 3 (autofit_workspace examples). Resolves the library portion of #1369.
API Changes
Additive only — three new searches on the public API, no removed/renamed/changed symbols:
af.MultiStartAdam(certified best),af.MultiStartADABelief(tied),af.MultiStartLion(sign-based, smaller default lr).optaxadded to the[jax]optional-dependency extra (lazy import; clear error if absent).See full details below.
Test Plan
test_autofit/full suite — 1479 passed, 1 skipped (incl. 4 newtest_multi_start_gradientplumbing tests;mle/optimizesubtree 7/7)MultiStartAdamrecovers a truth basin exactly on a convex toy JAX analysisto_dict/from_dict) round-tripsFull API Changes (for automation & release notes)
Added
af.MultiStartAdam(name=, path_prefix=, n_starts=48, n_steps=300, learning_rate=1e-2, start_lower_limit=0.15, start_upper_limit=0.85, ...)— multi-start Adam MAP searchaf.MultiStartADABelief(...)— same, ADABelief rule (default lr 1e-2)af.MultiStartLion(...)— same, Lion rule (sign-based, default lr 1e-3)autofit.non_linear.search.mle.multi_start_gradient.search.AbstractMultiStartGradient(AbstractMLE)— shared base owning the vmapped multi-start optax loopoptax— new[jax]extra dependencyMigration
Generated by the PyAutoLabs agent workflow.