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.

[Jammy2211]

Nautilus arbiter COMPLETE (job 330513, 2026-07-16, 1h32m — 1h10m compile + 22m sampling):

• Converged cleanly: 27,840 calls, N_eff 1233, logZ +17345.4, f_live 0.01.
• Max logL +17418.9 at einstein_radius = 1.3102 (truth 1.6). Inside the slack 0.3 basin tol by 0.01 — per this issue's own warning, do not over-read that flag.
• Key tension: the FD probe evaluated +25537 at a truth-centred point that is inside this model's space (mass and reg both free), so converged Nautilus's best sample is ~8k nats below a provably-existing better point. Reads like the pixelized-source degeneracy (source deforms to absorb a slightly-wrong mass) making the r_E≈1.31 solution the dominant posterior mode, with the truth peak narrow/low-mass at n_live=100.
• For the optimizer question: Adam's −39888 is 57k nats below converged Nautilus — the search question is decisively live, and the honest bar for any optimizer is now +17419, not the slack r_E tolerance.
• Compile cache is now warm on RAL for this graph; the 1h10m fusion is a one-time cost. Job 330513's log: pixgrad_logs/samp_pixgrad_nautilus_330513.log. The isolated pixgrad_pyautofit worktree is now safe to delete (#1374 merged, job finished).
• Follow-up runs on the lr hypothesis proceed under Experiment: learning-rate-free multi-start optimizers (MGE + pixelized) #101 (stacked branch).

[Jammy2211]

Nautilus baseline landed (RAL job 330513, re-run of 330379 which timed out; read while resuming the compile-time research #71 in autolens_profiling):

wall_s             : 5517.8   (of which ~4200 s = ONE XLA compile of input_reduce_fusion)
max log likelihood : 17418.9
einstein_radius    : 1.3102   (truth 1.6)
N_eff              : 1233     log Z: +17345.4

Reading per this task's resume note: Nautilus reaches logL +17419 — far above adam's −39888, but below the FD probe's +25537 at the truth-centred point, and r_E 1.31 vs truth 1.6. So Nautilus ALSO fails to reach the truth basin → the suspect shifts to model/priors, not the gradient optimizer. ("IN TRUTH BASIN: True" in the log is the known slack-tolerance artifact — ignore it.)

Compile-time context for the wall number is being handled in autolens_profiling#71 (the 1h10m single-fusion compile was ~76% of this job's wall time).

[Jammy2211]

Shipped via PR #102 (merged). Search-question diagnosis and continuation in #101 / PR #103. Closing.