experiment: can multi-start gradient MAP optimizers work for pixelized sources?

[Jammy2211]

Research question: do af.MultiStartAdam / MultiStartADABelief / MultiStartLion (PyAutoFit#1369) work on a pixelized source reconstruction, not just the MGE likelihood the benchmark used?

Model (SLaM SOURCE_PIX[1] style): lens MGE linear light with fixed non-linear geometry; lens mass (Isothermal + ExternalShear) free; source = RectangularSplineAdaptImage (differentiable spline mesh) + adaptive regularization (al.reg.Adapt), reg coefficient free (~7-D non-linear). Adapt image bootstrapped from a quick RectangularAdaptDensity+Constant fit, mirroring SLaM.

Plan:

1. FD feasibility gate (probe_grad_pix.py) — reverse-mode jax.grad of the spline-pixelized log-evidence, FD-cross-checked. If FD fails, that's the answer — stop, no A100.
2. Samplers: MultiStartAdam/ADABelief/Lion + Nautilus baseline.
3. Local CPU smoke → A100 on RAL → findings doc.

In autolens_workspace_developer/searches_minimal/, extending the MGE gradient benchmark harness. Human decisions: SplineAdaptImage + adaptive reg; mass+reg free / light fixed; checkpoint after FD probe.

[Jammy2211]

FD feasibility gate: gradient MAP optimizers don't work for a pixelized source (no A100 launched)

The FD gate found a gradient obstruction for every pixelized mesh, so no sampler / A100 runs were launched (the gate's purpose).

Mesh (source)	forward logL	jax.grad	verdict
RectangularAdaptDensity + Constant	finite	finite but FD-mismatched ~40–100%	FAIL_FD_MISMATCH
RectangularSplineAdaptDensity + Constant	finite	finite but FD-mismatched ~40–100%	FAIL_FD_MISMATCH
RectangularSplineAdaptImage + Adapt (fixed adapt image)	finite	NaN (7/9 params)	FAIL_NAN_OR_INF

Two obstructions: (1) mass-adaptive meshes build the mesh from the mass density → the mesh re-bins as mass changes (non-smooth) → wrong gradient; (2) the fixed-adapt-image spline mesh removes that dependence (forward logL finite) but the reverse-mode gradient is NaN — a non-differentiable op in the spline / adaptive-reg / regularized-inversion backward pass, at random and near-truth points.

Implication: MultiStartAdam/ADABelief/Lion descend on jax.grad; a wrong (case 1) or NaN (case 2) gradient breaks them. Not viable as the code stands. Case 2's NaN is the potentially fixable one — a library-level gradient-hardening job (custom_jvp on the offending op, like the Delaunay fix), not this experiment.

Probes + full writeup on branch feature/pixelized-gradient-experiment (searches_minimal/probe_grad_pix*.py, pix_gradient_findings.md).

[Jammy2211]

Correction: pixelized gradients DO work — my first probe was methodologically wrong

Human flagged (correctly) that pixelization gradients were already made to work. Confirmed: autolens_workspace_test/scripts/jax_grad/imaging_pixelization.py (commits 45ce603 / be8ce41 / 540e093) certifies strict/FD-matched gradients on all params incl. lens mass + shear. Re-run 2026-07-14 passes: RectangularUniform, RectangularAdaptImage+reg.Adapt+adapt-images+border-relocator (os_pix=4), RectangularAdaptDensity (os_pix=4), kernel-CDF meshes.

My probe's FAIL had four fixable causes, all methodology:

1. over_sample_size_pixelization defaulted to 1 → adaptive-mesh mass gradient is an exact staircase (documented: impossible at os_pix=1). Fix: os_pix=4, or a kernel-CDF mesh (differentiable at os_pix=1).
2. Tested Spline* meshes, not the certified kernel-CDF (RectangularKernelAdapt*).
3. Broad UniformPriors → random garbage-fit points (logL ~ -5e5). Certified harness uses truth-centred GaussianPriors.
4. Single eps=1e-4 FD step — too large. Certified harness sweeps 1e-8..1e-6.

So the experiment is viable. Resume: build the objective on RectangularKernelAdaptDensity(bandwidth=0.1) (or production KernelAdaptImage+os_pix=4) with truth-centred priors, FD-confirm near truth (small step sweep), then run MultiStartAdam/ADABelief/Lion + Nautilus baseline locally → A100. Full recipe on branch feature/pixelized-gradient-experiment (pix_gradient_findings.md). The real open question is now "can multi-start gradient descent from broad starts recover the mass basin with a pixelized source vs Nautilus", not "is the gradient correct" (it is).

[Jammy2211]

A100 pipeline working; gradient question settled; search question open

Settled — pix likelihoods ARE gradient-differentiable. A100 FD probe (kernel-CDF RectangularKernelAdaptDensity(bandwidth=0.1), os_pix=1, float64): every mass/shear parameter FD-matched to ~1e-6 (einstein_radius rel=8.8e-7), logL +25,537 at a truth-centred point. Config rules: adaptive meshes at os_pix=1 are the certified staircase (dead mass gradient) — kernel-CDF is live at os_pix=1; adaptive needs os_pix=4.

Shipped — PyAutoFit#1374 batch_size (merged 7262f832). Unbatched, the pix multi-start OOMs an 80GB A100 (jvp fusion f64[16,15361,2,31,512] = 58.13 GiB). With batch_size=4 a full float64 pixelized multi-start fit completed — no fp32, no sparse operator.

Open — the actual science. adam has run once (no ADABelief/Lion yet):

wall_s             : 2090.1     (~35 min, almost entirely XLA compile)
max log likelihood : -39887.864
einstein_radius    : 1.4169     (truth 1.6)
IN TRUTH BASIN     : True       <-- ARTIFACT of a slack 0.3 tolerance; do not trust
per-start basin    : 0/16       (p_hit=0.00)

logL -39,888 vs the probe's +25,537 at truth, with 0/16 starts in basin: the optimizer did not find the basin. Suspects: n_steps=300 too few from broad starts; lr=1e-2 mis-scaled for these parameters; possibly my per-start diagnostic indexing is wrong.

Nautilus baseline running (job 330379) on the identical model — if Nautilus also misses the basin, the model/priors are at fault rather than the gradient optimizer.

Compile is the dominant cost (>30 min single-fusion compiles; 453 s for one latent sample). Filed as its own task. Note for the record: XLA autotuning was ruled out — autotune ON vs OFF gave 2100.1s vs 2090.1s with bit-identical results.

Traps recorded: a foreground timeout ssh does not kill the remote side (severed a git op, leaving a stale index.lock in the shared PyAutoFit mirror — since repaired, shared mirror clean on main); Nautilus on a JAX row must set force_x1_cpu=True + use_jax_vmap=True and n_batch<=16 without the sparse operator.