research: find what kills MGE multi-start lanes (62% value-NaN)

[Jammy2211]

Overview

The trapped-lane counter (PyAutoFit#1475 004f798, PyAutoGalaxy#572 695b27c) ran on the real imaging/mge cell and cleared the ell_comps plateau as a suspect — and surfaced a much larger effect nobody has characterised. On a production dataset/model/analysis via build_for_cell (16 starts x 150 steps, cloud CPU, 352s):

counter	lane-steps	share
n_value_nan_lane_steps	1498	62.42%
n_grad_nan_lane_steps	9	0.38%
n_constrained_lane_steps	0	0.00%
n_resurrections	0	—

The population fell from alive 16/16 to alive 2/16. Roughly seven of every eight starts die, and the best fit is being found by two survivors.

This contradicts a documented assumption: AbstractMultiStartGradient's resurrect docstring justifies the resurrect=False default by asserting the parametric MGE-class cell "has only the measure-zero singularity, so the apply_if_finite guard suffices". The measurement says otherwise, and the dominant mode is value-NaN, which apply_if_finite does not rescue — it only zeroes the step, so the lane stays dead and resurrect=False never redraws it. n_resurrections: 0 confirms that mechanism.

This task finds the cause. It does not change any search behaviour — see Boundary.

Plan

• Find where the non-finite likelihood comes from — which parameters/regions and which operation. This is the open question and it does not need a GPU.
• Measure whether resurrect=True recovers the budget on this cell, and what it costs.
• Confirm the effect at production budget, on GPU, across seeds — graded on the alive-versus-step curve, not on reproducing the scalar "62%".
• Record the finding in this repo's hazard index.

Detailed implementation plan

Affected repositories

• autolens_profiling (primary) — investigation, measurement, hazard-index entry
• PyAutoFit / PyAutoGalaxy — read only here; any fix is a separate task

Branch survey

Repository	Current branch	Dirty?
autolens_profiling	main (a34d6191)	clean
PyAutoMind	claude/mge-lane-death-ih8q92	clean

Suggested branch: research/mge-lane-death

Implementation steps, ordered by information per unit time

This ordering is deliberate and differs from the order the questions were first written down. The cause is the unanswered question and it is the cheap one; the GPU reproduction confirms an effect already measured once. Do not queue for a GPU before step 1 has been attempted.

1. Where the NaNs come from. Candidates: the mask edge, a Gaussian sigma collapsing, the linear inversion / NNLS solve, the sqrt at r=0. The existing ~6-minute 16x150 CPU run already produces all 1498 deaths — instrument it to dump each lane's parameter vector at the step its value first goes non-finite, and the distribution of those vectors is the answer. Minutes-long iteration, not a queued job.
2. Whether resurrect=True recovers the budget on this cell, and what it costs. The docstring says resurrection is for pixelized sources; this evidence suggests the parametric cell may need it too. Measurement of the candidate remedy, not adoption of it.
3. Whether it holds at production budget, on GPU, across seeds. One seed at reduced budget on CPU so far, and the float64 GPU path is different numerics. Log alive N/16 per step so curves can be overlaid across budgets.

Grade step 3 on the curve, not the scalar

62% is a survival integral, not a hazard rate: a value-NaN lane is frozen by apply_if_finite and never redrawn, so it keeps counting a value-NaN lane-step on every subsequent step. Hold the death curve fixed and run 300 steps instead of 150 and the same physics reports ~75%. A production run coming back above 62% is therefore not automatically a worse result. Compare alive-versus-step curves; use the scalar only within a fixed budget.

"When they happen" is already partly answered

Inverting the 1498 over 14 dead lanes with no resurrections gives sum(k_i) = 14*150 - 1498 = 602, i.e. a mean death step of ~43 of 150. Deaths concentrate in the first third — trajectories walking into a wall mid-descent, not bad initial draws that _broad_starts failed to filter. Step 1's per-lane death steps should confirm or overturn this; a disagreement is itself a finding, since it would imply lanes recovering and re-dying, which resurrect=False says they cannot.

Key files

• scripts/misc/searches/_setup.py — build_for_cell, the production dataset/model/analysis. Note it writes into dataset/ (rewrites the HST FITS, adds positions.json, emits results/simulators/*); it is not read-only.
• results/searches/multi_start_nan_accounting/ — where the existing NaN-accounting artefact lives; the new results belong alongside it.
• The hazard index in this repo — where the cause gets recorded.

Environment

• Python 3.12+ (autonerves); pip install jaxnnls is required for the JAX NNLS solver path and is not pulled in by default.
• Install autolens with --no-deps if using editable local autofit/autogalaxy, or the released wheels clobber them.
• Cell scripts honour SEARCHES_N_STARTS / SEARCHES_N_STEPS / SEARCHES_BATCH_SIZE / SEARCHES_DISABLE_VIZ.
• imaging/mge runs on CPU in ~6 min at 16x150. The pixelized mesh cells do not — that is JIT compile, not memory (batch_size=1 clears the OOM), and they need the GPU. Shrinking the source mesh is the wrong cost lever; the image-plane grid at mask_radius 3.5 dominates.
• On A100 set jax_enable_x64 explicitly — it is not inherited under sbatch, and float32 would understate the quantity under test.

A trap worth knowing before re-running

validate_zero.py (in the prompt) is the positive control that makes the 0 meaningful — without it, a zero cannot be distinguished from "nothing was watching". Two notes:

• It must probe with jnp arrays. Concrete Python floats trip validate_ell_comps and raise before the constraint is ever reached.
• Its four checks call discovery/evaluation/the predicate directly, so none of them exercises the search loop that accumulates the counter. What closes that gap is in the result JSON: the counters come back as integer 0/1498, not null, and the driver reads them via si.get(k) — which would have given null had _fit never written the key. A null in the counters block on a re-run means broken plumbing, not a clean cell. The two look identical in a summary table.

Original prompt

The full prompt — including the verbatim result JSON and both self-contained driver scripts (rerun_cell.py, validate_zero.py) — lives in PyAutoMind and travels with the task:

draft/research/autolens_profiling/mge_lane_death.md (permalink, 0b8895ae)

The scripts are deliberately not committed to this repo: they are throwaway measurement drivers, not part of the script tiers, and they hardcode a workspace root and a scratchpad output path. Copy them out to run.

Background on the counter itself: complete/2026/08/frozen-lane-counter.md — the full run, the detector-design finding, and the positive control.

[Jammy2211]

Cause found: the deaths are in the PRIOR term, not the likelihood

Reproduced and root-caused on cloud CPU. The likelihood never went non-finite in any run. Not once, in ~7200 lane-steps across three arms.

The mechanism

_fit builds its objective as Fitness(..., fom_is_log_likelihood=False, convert_to_chi_squared=True), i.e.

fom = -2 * (log_likelihood + sum(log_prior_list))

A UniformPrior returns log_prior = -inf outside its box. MultiStartGradient steps in physical parameter space through optax with no projection back onto that box. So when a lane crosses a hard prior edge the objective becomes ±inf, np.isfinite(foms) reads it as dead, and with resurrect=False it is never redrawn — it stays dead for every remaining step. That accumulation is the 62%.

Evidence

1. Reproduction (16x150, cloud CPU, 250s) — 1446 value-NaN lane-steps (60.25%), 18 gradient-NaN, 0 constrained, 0 resurrections, alive 16/16 -> 2/16, 14 lanes dead at steps [34,36,36,37,38,39,39,40,40,41,41,52,56,125]. Against the filed run's 1498/9/0/0 and the same 2/16.

The survival-integral identity holds exactly: sum(150 - k_i) = 14*150 - 654 = 1446, which is n_value_nan_lane_steps to the unit. That counter is precisely the area under the death curve — measured, not inferred.

2. Per-lane autopsy at the death vectors. Instrumented _nan_lane_counts (the params that produced foms is live in the caller's frame) to record each lane's death step and parameter vector. Then walked the whole pipeline at those vectors — positions penalty, deflections, convergence, model data, residuals, chi-squared, NNLS reconstruction, log-determinants, figure_of_merit:

verdict at death vector	lanes
finite likelihood, sum(log_prior) = -inf, 1-2 params outside a UniformPrior	11 / 14
parameters already NaN (the gradient-NaN path)	2 / 14
unexplained — every term re-evaluates finite eagerly	1 / 14

Every stage of the likelihood was finite at every death vector, figure_of_merit around -1.6e5. Only the prior term was -inf.

3. First causal test — refuted a narrower hypothesis. 10 of the 11 box exits were the shear, UniformPrior(-0.3, 0.3), so I widened only those two priors to (-1, 1). Deaths did not collapse — they moved later and got marginally worse. Widening one box just moves the wall.

4. Decisive test. Neutered the prior term itself (log_prior_list_from_vector -> zeros), leaving likelihood, model, data and stepping untouched:

arm	value-NaN	grad-NaN	constrained	dead	alive at end
baseline (reproduction)	1446 (60.25%)	18	0	14/16	2
shear box widened to ±1	1422 (59.25%)	36	13	15/16	1
prior term neutered	215 (8.96%)	27	667 (27.79%)	3/16	13

Value-NaN drops 85%; survivors go 2 -> 13. All three residual deaths have NaN parameters, i.e. the gradient-NaN path — not a likelihood death.

What this changes

• The resurrect docstring's claim is correct, and irrelevant. The parametric MGE likelihood really does appear to have only measure-zero singularities — it never produced a non-finite value here. What is wrong is using that claim to justify resurrect=False, because the deaths come from the prior, which the claim never covered.
• This is not MGE-specific. Any model with a hard-box prior and unbounded stepping has it. The pixelized cells are very likely affected.
• The fix is not resurrect=True. Resurrection redraws a lane that has walked out of the box, which is why it appears to help, but the lane will walk out again. The real options are bounded stepping (projection/clipping onto the prior support) or soft-walled priors. That belongs in PyAutoFit and is out of scope here.
• "62% died to value-NaN" is better read as "62% of lane-steps were spent frozen outside the prior box".

The ell_comps plateau was masked, not cleared

This is the finding I did not expect and it corrects the framing in the issue description. The baseline reports n_constrained_lane_steps = 0, which read as "the ell_comps plateau is not involved". With the prior deaths removed it is 667 (27.79%).

Lanes were dying of prior-exit before they could reach the saturation plateau. The zero was real and correctly measured — the positive control was sound — but it meant "nothing got that far", not "nothing gets trapped there". So #1475's trapped-lane counter is measuring a live failure mode on this cell; it was hidden behind a larger one. Once lanes stop dying they stop being dead and start being stuck, which is the harder problem.

Caveats

• One of 14 baseline deaths (lane 9, step 39) re-evaluates finite in every term, with all parameters well inside their boxes. Something in the jitted/vmapped float32 path differs from the eager recompute there and I have not explained it.
• Single seed per arm, CPU, x64 not enabled. The three arms are internally consistent and the effect is far larger than run-to-run variation (1446 vs the filed 1498), but the GPU/float64 confirmation across seeds is still owed.
• The prior-neutering arm is a diagnostic, not a valid inference — it removes the prior from the posterior.

Reproducing

Environment: Python 3.12, autolens 2026.8.15.1 from PyPI with autofit/autogalaxy main overlaid --no-deps (the released wheels predate #1475, so they lack _constrained_lane_count), plus jaxnnls and optax installed explicitly. build_for_cell gives 15 free parameters. Baseline run is ~250s at 16x150; scripts are the four diagnostics described above and attach to the search without editing autofit.

---

Generated by Claude Code

[Jammy2211]

Correction: the dead lanes are not frozen — they keep walking

My earlier comment said lanes were "frozen outside the prior box" and that the 62% is time spent "frozen". That is wrong, and the correction matters because it changes what the waste actually is.

Recorded the full per-lane parameter trajectory (150 x 16 x 15) rather than just the death vector.

What actually happens

The overshoot at the fatal step is tiny. Median 3.1% of the box width, minimum 0.21%, maximum 8.9%. gamma_2 = +0.301278 against a bound of 0.3 is enough to kill a lane. Nothing walks "well outside" — they trip over the edge.

Then the lane keeps stepping, forever. Outside the box log_prior = -inf is a constant, so its derivative is zero and the total gradient is just the finite likelihood gradient. optax.apply_if_finite therefore never fires. Every prior-exit lane moved on 100% of its remaining steps (113/113, 108/108, 110/110, ...) and drifted a long way — lane 14's gamma_2 went +0.301 -> -1.756, lane 0's gamma_1 went +0.316 -> +1.442.

0 of 16 lanes ever revive.

So the lane is a zombie, not a corpse: it pays the full likelihood-and-gradient cost every step, navigates on likelihood alone with the prior contributing nothing, and has every result discarded because alive=False. The 62% is compute actively spent and thrown away, which is worse than idle.

Two populations, which I had conflated

population	lanes	behaviour
zombie — crossed a prior edge, gradient finite	11/14	keeps stepping, wanders far, never revives
latched — gradient non-finite, apply_if_finite zeroes the update	3/14	frozen at the cliff edge, zero movement after

The resurrect docstring's "apply_if_finite alone latches the start at the cliff edge" is accurate — but only for the 3-lane minority. The dominant mode does not latch.

This also resolves the lane-9 anomaly

I flagged lane 9 (step 39) as unexplained because every term re-evaluated finite eagerly. Its grad_finite was False at death. So its likelihood genuinely was NaN in the jitted path — the Fitness guard maps NaN to -inf, and the where makes the gradient NaN — while my eager recompute gives finite. A real jit-vs-eager numerical difference. It belongs in the latched population, not the prior population, and the 11/2/1 split above reflects that.

Prototype clipping arm

Projecting parameters back into the prior box after optax.apply_updates:

arm (16x150)	value-NaN	lanes dead	alive at end	pinned to a bound
baseline	1446 (60.25%)	14/16	2	—
shear box widened	1422 (59.25%)	15/16	1	—
prior term neutered (diagnostic)	215 (8.96%)	3/16	13	—
clip to prior box	425 (17.71%)	5/16	11	8 mid-run, 5 at end

Survivors go 2 -> 11. The residual 5 deaths are the NaN-gradient population, which clipping cannot touch. Pinning is real but not fatal — roughly 6 of the 11 survivors still explore freely.

Because the overshoot is only ~3% of a box width, clipping is close to a no-op on the trajectory rather than a violent projection — which is a better argument for it than the one I gave before.

Two incidental bugs, both only visible when lanes survive

Neither fires on the baseline, where 14/16 lanes are dead — i.e. they sit on a code path this cell has apparently never taken:

1. float32 is not JSON serializable. autofit/non_linear/paths/directory.py:80 save_json raises TypeError: Object of type float32 is not JSON serializable at the end of a successful clipped run.
2. A crashed run poisons the next run of the same name. The half-written output from (1) made the next same-named search fail with JSONDecodeError while trying to resume — a 4-second no-op that read as a clean result (zero deaths, because zero steps). A new form of the cached-result hazard already recorded in multistart-nan-step-diagnostics.

Fix now scoped

Filed as a search-agnostic Clipper in PyAutoMind draft/feature/autofit/prior_support_clipper.md. Scoped that way because the exposure is not MultiStart-only: LBFGS builds the same log-posterior Fitness, steps in physical space, and calls optimize.minimize with no bounds= despite L-BFGS-B supporting them natively. NUTS also targets the log posterior from a physical start, but diverges rather than dying — its own investigation.

---

Generated by Claude Code