research: BlackJAX adaptive-tempered SMC with gradient inner kernel (MGE)

[Jammy2211]

Overview

Stage (a) of the JAX-native posterior sampler wave (2026-07-16 deep-research shortlist). Upgrade the existing gradient-free blackjax_smc.py (adaptive-tempered SMC with an RWM inner kernel, sampling in unit-cube space via a non-differentiable pure_callback bridge) to a gradient inner kernel (MALA first, HMC behind a flag) that samples in physical-parameter space, where jax.grad of the MGE imaging likelihood is already certified OK_HMC_VIABLE (probe_grad.py). SMC is stage (a) because it yields log-evidence for free (sum of tempering log_likelihood_increment) while using gradients — the one gradient sampler that also produces logZ. Scope: MGE parametric likelihood only; pixelized work is a deferred follow-up. First stage of a 5-stage wave — issued alone; do not bulk-issue b–e.

Plan

• Add a gradient inner-kernel SMC variant sampling in physical space (differentiable; no per-step host roundtrip), replacing the cube uniform log-prior with a JAX-native physical-space log-prior.
• Start with MALA (cheapest gradient kernel), then HMC behind a flag, both via blackjax.adaptive_tempered_smc with inner_kernel_tuning for per-temperature step-size adaptation.
• Preserve SMC log-evidence recovery and report it against the Nautilus / nss_jit / RWM-SMC baselines.
• Deliverable: runnable script + smc_gradient_findings.md. Workspace-only, no library edits.

Detailed implementation plan

Affected Repositories

• autolens_workspace_developer (primary)

Branch Survey

Repository	Current Branch	Dirty?
./autolens_workspace_developer	main	dirty (pre-existing dev-scratch under searches_minimal/; target files blackjax_smc.py/_setup.py/probe_grad.py are clean) — work runs in an isolated worktree off origin/main

Suggested branch: feature/blackjax-smc-gradient-kernel
Concurrent worktree alongside the parked pix-gradient-slogdet-revalidation claim (different files; git supports N worktrees per repo).

Implementation Steps

1. New searches_minimal/blackjax_smc_grad.py (sibling to blackjax_smc.py, reusing _setup + _metrics.MLTracker):
   • log_likelihood(params) in physical space — model.instance_from_vector(vector=params, xp=jnp) → analysis.log_likelihood_function(instance=…) (the probe_grad.py form; fully differentiable, no pure_callback).
   • log_prior(params) JAX-native: iterate model.priors_ordered_by_id, Gaussian→-0.5*((x-mu)/sigma)**2, Uniform→where(in-bounds,-log(width),-inf), LogUniform→log-space form; NaN masking via jnp.where at degenerate points.
   • Inner kernel: MALA (blackjax.mala) first, then HMC behind a flag, fed to blackjax.adaptive_tempered_smc via mcmc_step_fn/mcmc_init_fn, wrapped with blackjax.smc.inner_kernel_tuning for per-temperature step-size (HMC: + mass matrix / trajectory length).
   • Initial particles drawn from the prior in physical space (host-side vector_from_unit_vector at init only; hot path stays JAX-native).
2. Time the one-shot value_and_grad compile separately (mirrors existing scripts; nss_grad baseline compile ≈ 1291 s is the figure to report/beat).
3. Run on the standard MGE problem; write output/blackjax_smc_grad_summary.txt, regenerate the output/comparison.txt row, author smc_gradient_findings.md (physical-vs-cube design note, MALA-vs-HMC, logZ vs Nautilus, gradient-kernel viability verdict). RAL A100 fp32 GPU pass if warranted; laptop CPU fp64 for wiring smoke.

Constraints (from prior campaigns)

• Reverse-mode gradients only; NaN-gradient degenerate points masked at the log-prior.
• lax.map batch_size tiling (PyAutoFit#1374 lever) if particle vmap fan-out pressures memory.
• No silent None/NaN guards — mask at the prior, let genuinely bad likelihoods crash loudly.

Key Files

• searches_minimal/blackjax_smc.py — RWM baseline to upgrade from.
• searches_minimal/probe_grad.py — the gradient-certified physical-space likelihood form.
• searches_minimal/_setup.py — build_dataset / build_model / build_analysis(use_jax=True).
• searches_minimal/_metrics.py — MLTracker convergence metric.
• searches_minimal/output/comparison.txt — benchmark baseline (Nautilus / nss_jit / nss_grad rows).

Original Prompt

Click to expand starting prompt (full 5-stage wave)

Research: JAX-native posterior sampler wave — implement and benchmark the ranked sampler shortlist from the 2026-07-16 deep-research session on the standard searches_minimal problem in @autolens_workspace_developer, via the sampler_pipeline ingest-prototype-profile-promote flow. Ranked stages: (a) blackjax adaptive tempered SMC upgraded to a gradient inner kernel (MALA/HMC + inner_kernel_tuning; extends the existing blackjax_smc.py RWM smoke script; gives logZ from tempering increments), (b) ChEES-HMC many-chain adaptation, (c) MCLMC + adjusted MCLMC warm-started from the multi-start Adam basin, paired with harmonic (learned harmonic mean) for evidence, (d) flowMC (normalizing-flow global jumps; multimodality insurance), (e) jaxns cameo (independent GPU-native nested-sampling cross-check vs Nautilus). Constraints from prior campaigns: reverse-mode-only gradients, NaN-gradient degenerate points need masking, vmap fan-out needs lax.map batch_size tiling (Fit#1374 lever), RAL A100 float32. Deliverable per stage: findings doc + benchmark vs Nautilus baseline. Issue stages one at a time as predecessors ship — do not bulk-issue.

This issue = stage (a) only. Scope narrowed by the user to the MGE parametric test case (gradients confirmed working); pixelized deferred to a later stage.

[Jammy2211]

Progress — script complete, wiring proven locally, RAL CPU convergence run submitted

searches_minimal/blackjax_smc_grad.py written — physical-space adaptive-tempered SMC with a gradient inner kernel (MALA default, HMC --kernel hmc), JAX-native log-prior over the 15 params, --tune for inner_kernel_tuning. Debugging surfaced three real issues, all fixed and understood:

1. Non-finite gradient at the "prior mean" — the vector of prior .mean attributes is a degenerate MGE config; warm-up/finite-check moved to the cube-0.5 median (the probe_grad.py-certified point).
2. Non-finite gradient at the median too — the light/mass profile-centre singularity (1/r at centre (0,0), which the symmetric median sits exactly on; probe_grad.py only ever ran on GPU). Fixed with a custom_jvp that masks non-finite gradient entries to zero — the campaign's "NaN-gradient degenerate points need masking" constraint. Random particles almost never hit it; the mask keeps a kernel stable if one wanders near.
3. MALA "Pred must be scalar" — blackjax MALA takes a scalar step (no metric arg); my per-dimension array step broke the MH acceptance. Fixed by whitening the sampling space (z = params / prior_scale) so a single scalar step is well-conditioned; HMC gets an identity mass matrix.

Local 4-particle run reaches λ=1.0 with no NaN poisoning — wiring is proven. Sampling quality is meaningless at that scale (the laptop's ~15GB caps the vmapped MGE-inversion gradient at ~4 particles; 16 OOMs).

Convergence run: GPUs are queued full, so submitted a CPU job on RAL (ral partition, 921GB nodes → no memory ceiling) — job 330955, two arms (MALA fixed-step; MALA + inner_kernel_tuning) at 128 particles to confirm the sampler anneals into the high-likelihood region and yields a sensible log Z vs the Nautilus baseline. Representative A100 timing follows once GPUs free up.

blackjax 1.5 (matching local) installed isolated on RAL (its venv shipped an ancient 0.1.0b1.dev); the script runs there as an untracked file, no branch switch, so the concurrent pix-gradient-slogdet-revalidation checkout is undisturbed.

[Jammy2211]

Convergence run complete (RAL CPU, float64) — wiring proven, naive gradient SMC does NOT converge

Result: the machinery runs end-to-end in float64, but naive MALA adaptive-tempered SMC fails to converge on the MGE likelihood. This is the informative finding of stage (a) (full write-up committed as searches_minimal/smc_gradient_findings.md).

Why: the MGE likelihood is ~1000× sharper than the prior, so the prior→posterior tempering path needs the inner-kernel step to shrink ~1000× along the way. No static or naively-adaptive step tracks this — step-size study (RAL CPU, float64, 64 particles):

Step regime	Behaviour	Verdict
Fixed 0.02	acc 0.44→0.000 by step ~15; λ stalls 0.001; max-logL flat	collapse
Fixed 0.001 (tiny)	acc ~0.8 early; max-logL improves −165k→−121k; then collapses ~step 17, crashes	gradient IS useful, but no completion
Spread-adaptive (--tune, 2.38/√d·std)	step too big (~0.6); acc frozen 0.000; tempering force-jumps λ→1 → false "converged", garbage logZ	collapse (false positive)

The tiny-step arm improving −165k→−121k shows the gradient genuinely pulls toward better fits — the blocker is purely step-size scheduling across the extreme concentration, not the gradient itself.

Also fixed along the way (beyond the 3 earlier): (4) the inner_kernel_tuning callback signature (rng_key, state, info) + step recomputed from state.particles; and the float64 env trap — JAX defaults to float32 and JAX_ENABLE_X64 is not inherited by sbatch (one job wasted in silent float32 before I caught it).

Next levers (tomorrow): (1) warm-start particles from the multi-start Adam basin — shorten the tempering path so the step-scale gap is never crossed (most promising; the prompt's stage-c idea); (2) acceptance-feedback dual-averaging step control; (3) HMC + posterior-matched mass matrix; (4) A100 representative timing once GPUs free.

Committed to feature/blackjax-smc-gradient-kernel (local, cfcf893) — not shipped pending a converging configuration.

[Jammy2211]

Warm-start reframe: the solution IS found — one units bug left

Per user steer: these samplers are meant to be warm-started from a JAX optimizer's MLE, so cold prior-start benchmarking was never a representative test. Built a shared warm-start harness and re-ran (RAL CPU job 331035, float64, 64 particles).

How wrong cold-starting was:

	max log L
prior median	−159,736
every cold arm (all step regimes)	−158k … −180k
multi-start Prodigy optimum	+31,788

The cold runs never came within ~190,000 log-units of the solution; Prodigy closes the gap in ~200 steps. Yesterday's "acceptance collapse" was the sampler failing to cross a gap the optimizer crosses trivially.

Warm-started results:

• Prodigy max log L 31787.93; MLE einstein_radius = 1.5997 vs truth 1.6 — essentially exact.
• Laplace reference scale held (no fallback): posterior widths σ ~ 2e-4…2e-3.
• Tempering advances smoothly λ 0.006→1.0 in 15–16 steps (cold stalled at λ~1e-3); log Z ≈ 31701, a sane evidence number.

New infrastructure (_warm_start.py, shared across the whole wave): multi-start Prodigy → one cached artifact (MLE + reference scale + logL) so every sampler in stages (a)–(e) starts from the identical point. Reference scale via Laplace (inverse-Hessian diagonal), falling back to multi-start spread, then prior fraction.

Evidence preserved: rather than dropping particles at the MLE (which destroys SMC's logZ — valid only when tempering starts from a normalised distribution), we temper geometrically from a normalised Gaussian reference g at the MLE: logprior_fn := log g, loglikelihood_fn := log prior + log L − log g. Since ∫g = 1, logZ remains the true evidence, comparable to Nautilus. This required adding the prior's normalisation constant (cancels in MH, not in evidence).

Outstanding — acceptance still 0.000 (particles never move; the "posterior" is the reference reweighted). Cause found and fixed but not yet re-run: MALA proposes x + ε∇ + √(2ε)ξ, so ε is a squared length — proposal length is √(2ε), not ε. The auto-step used ε ≈ σ, giving proposal length 0.076 against σ=0.005 (~16× too wide). Correct: ℓ = 2.38·d^(−1/6)·σ, ε = ℓ²/2 (HMC's step is a length). Same trap fixed in the --tune rule.

Committed local 5be51d7 / 2b455e4 on feature/blackjax-smc-gradient-kernel (not pushed). Next run should show non-zero acceptance and a genuinely mixed posterior.

[Jammy2211]

RESOLVED — the gradient SMC now samples

Auto step size from POSTERIOR width (sigma=0.5, ref sigma=1): 0.2871
  step 1: λ=0.108  max log L=31770.41  acc rate=0.799
  step 2: λ=0.206  max log L=31780.07  acc rate=0.781
  step 3: λ=0.329                      acc rate=0.754
  step 4: λ=0.523                      acc rate=0.523
  step 5: λ=0.664  max log L=31781.57  acc rate=0.413
  step 6: λ=0.849                      acc rate=0.262
  step 7: λ=1.000                      acc rate=0.166
Best fit: einstein_radius = 1.5998   (truth 1.6)

Acceptance is healthy (was pinned at 0.000), particles genuinely move, max log L climbs toward the Prodigy MLE (31787.93), and einstein_radius recovers to 1.5998 against a truth of 1.6.

Three compounding causes, each measured rather than guessed:

1. Prior-whitening does not whiten the posterior. Measured 269× anisotropy in prior-whitened coordinates — einstein_radius has prior scale 8.0 but posterior std 2e-4 (σ_z 5.3e-5) vs 1.4e-2 for the loosest parameter. No scalar step can serve that: one tuned to the mean is ~88× too large for the tightest. → whiten by the warm-start (posterior) scale.
2. Diagonal whitening isn't enough — the posterior is correlated. Laplace covariance condition number 568, |r|=0.95 between two parameters. Diagonal whitening leaves a tilted ridge a spherical proposal walks off. → whiten with the Cholesky factor of the full covariance, adding log|det L| to the evidence.
3. The step targeted the reference width, not the posterior width. The reference is deliberately inflated, so the posterior is σ≈1/inflate — scaling to σ=1 overshoots by inflate². Measured: ε=1.148 → acc 0.00, ε=0.1 → acc 0.94. → auto step targets the posterior width (ε=0.287 at inflate=2).

Also fixed a bug that was confounding my own experiments: --step-size had a numeric default, so passing that same value explicitly was indistinguishable from not passing it and silently fell through to auto-scaling — a step-size bracket returned byte-identical arms. Default is now None.

Remaining refinement: acceptance declines 0.80→0.17 as λ→1, exactly what a fixed step does as the target sharpens — that's the job of --tune. RAL job 331058 now running all three arms at 128 particles (MALA auto-step, MALA+tune, HMC) with the warm start refreshed to carry the full covariance.

Committed a467fad + 6867762 on feature/blackjax-smc-gradient-kernel.

[Jammy2211]

Closing as superseded by the 2026-08-17 inference-methods programme (PyAutoLabs/autolens_profiling#134), which now owns sampler direction.

Stage (a) got its answer. Warm-started gradient SMC samples: acceptance 0.80→0.17 across tempering, einstein_radius 1.5998 vs truth 1.6, max logL ~31781 vs the Prodigy MLE 31787.93. The gradient path is certified OK_HMC_VIABLE (probe_grad.py). The write-up is searches_minimal/smc_gradient_findings.md on branch feature/blackjax-smc-gradient-kernel (tip 6867762), which stays pushed — nothing is lost by closing this.

Not carried forward: the open threads at park time — RAL job 331058 (three warm arms @128p: MALA auto-step, MALA --tune, HMC, testing whether --tune holds acceptance ~0.57 as λ→1), the logZ-vs-Nautilus comparison.txt row, and A100 rep-timing. Stages (b) ChEES-HMC, (c) MCLMC+harmonic, (d) flowMC and (e) jaxns were never issued and are not carried forward as scoped in the original wave.

Full context is preserved in PyAutoMind at complete/archive/shelved/jax_native_posterior_sampler_wave.md (PyAutoLabs/PyAutoMind#204), including the warm_start.json caveat (regenerable cache — re-run with --refresh if it predates the cov field, else it silently falls back to diagonal).

Closing as not planned rather than completed: the findings branch was never merged and the benchmark comparison was left open, so this should not read as delivered.

---

Generated by Claude Code