Return best iterate from AMICATorchNG.fit (#51)

Multi-model NG log-likelihood was ~0.02 lower and ~13x more variable than
Fortran because fit() returned the LAST EM iterate under the deliberately
non-monotone lrate schedule (both NG and Fortran anneal only after an LL
decrease). The variance was driven by late Newton-fallback overshoots: on
the sample EEG one seed peaked at LL -3.357 then crashed to -3.545 in its
final iterations. It was return-last, not a bad basin.

fit() now tracks and returns the highest-LL iterate:
- keep_best (default True) restores the best iterate when the run ends more
  than _KEEP_BEST_TOL (1e-9) below its peak; a monotone single-model run has
  best == last, so no restore fires and issue #24 parity stays byte-for-byte
  identical (verified: max param diff 0.0). Inactive under do_reject (the
  good-sample set, hence the LL normalization, changes across iterations).
- final_ll_ reports the returned iterate's LL; ll_history stays the true
  trajectory. The AMICA wrapper, validate_implementations, and the ensemble
  scripts read final_ll_.
- state_dict format bumped 2 -> 3 (adds keep_best, final_ll).

At matched 100-iter budget this cuts the LL sd from 12.7x to 2.0x Fortran's.
The residual ~0.009 mean gap is convergence speed, not a worse optimum: NG
reaches Fortran's exact mean (-3.3541) at 200 iters and exceeds it at 300.

Tested: 69 torch tests pass (+3 keep_best tests); ruff clean; ty no new
diagnostics. ADR 0003, .context/issue-51/.

Author: neuromechanist

.context/decisions/0003-best-iterate-safeguard.md

@@ -0,0 +1,77 @@
+# ADR 0003: Return the best-log-likelihood iterate from AMICATorchNG.fit
+
+**Status:** accepted
+**Date:** 2026-07-06
+**Owner:** neuromechanist
+
+## Context
+
+Issue #51: on the real sample EEG, multi-model `AMICATorchNG` (`n_models=2`,
+100 iterations) reaches a final log-likelihood distribution that is ~0.02 lower
+and ~13x more variable than the Fortran reference, even though the per-block
+sufficient statistics and one M-step are bit-exact vs Fortran (~1e-15). The
+partition distribution is already statistically equivalent to Fortran's (ADR-era
+work on #27), so this is an optimizer-quality residual, not a correctness bug.
+
+Diagnosing the variance (NG-only sweep, 20 seeds): the spread is driven almost
+entirely by a *late overshoot*, not a bad basin. One seed climbed to LL -3.3573
+(dead in the pack) then crashed to -3.5452 in its final two iterations after
+Newton went non-positive-definite and fell back to the natural gradient; the
+lrate ramp re-inflated the step and the run ended mid-crash. Nine of twenty
+"good" runs also ended a small amount below their own peak. The learning-rate
+schedule is deliberately non-monotone (both NG and Fortran anneal the rate only
+*after* an LL decrease, `amica15.f90:1038-1058`), so the *last* EM iterate is not
+guaranteed to be the best one. `fit` returned the last iterate.
+
+## Decision
+
+`AMICATorchNG.fit` tracks the highest-log-likelihood iterate and restores it when
+the run ends more than `_KEEP_BEST_TOL` (1e-9, per sample-channel LL) below that
+peak. A new `keep_best: bool = True` constructor flag controls it; `final_ll_`
+exposes the log-likelihood of the *returned* parameters (while `ll_history` stays
+the true per-iteration trajectory, overshoot and all). The safeguard is inactive
+under `do_reject`, where the good-sample set (and thus the LL normalization)
+changes across iterations and per-iteration LLs are not comparable.
+
+## Consequences
+
+- The pathological low-LL tail is removed and the LL variance collapses toward
+  Fortran's: on the 20-seed sample ensemble NG goes from mean -3.3738 (sd 0.040)
+  to mean -3.363 (sd ~0.007). See `.context/issue-51/`.
+- **Single-model issue #24 parity stays bit-exact.** A monotone fit has its best
+  iterate == its last iterate, the gap is 0 < tol, no restore fires, and the
+  returned parameters are byte-for-byte identical to `keep_best=False` (verified:
+  max parameter difference 0.0 across A/W/mu/beta/alpha/rho/c/gm at 100 iters).
+- New obligation: `ll_history[-1]` is no longer the fitted model's LL when a
+  restore fired. Consumers must read `final_ll_` (or `max(ll_history)`). The
+  validation harness and ensemble scripts were updated accordingly.
+- A residual mean gap (~0.009) remains at 100 iterations, but it is *convergence
+  speed*, not a worse optimum: at 200 iterations NG reaches Fortran's exact mean
+  LL (-3.3541) and by 300 slightly exceeds it. NG's per-iteration progress is ~2x
+  slower than Fortran's; the reachable solution is identical (the M-step is
+  bit-exact vs Fortran, #27). Not a correctness issue.
+- `state_dict` format bumped 2 -> 3 (adds `keep_best`, `final_ll`).
+
+## Alternatives considered
+
+- **Do nothing / raise `max_iter`:** more iterations let a crashed run's annealing
+  recover, but waste compute on every run and do not help a fit that already
+  peaked and then overshot. Rejected as the primary fix (kept as a knob).
+- **Trust-region / reject-the-update in-loop:** clamp any LL-decreasing step. This
+  changes the optimization trajectory and risks perturbing the bit-exact
+  single-model path; return-best is a pure post-hoc selection that leaves the
+  trajectory (and `ll_history`) untouched. Rejected.
+- **Rewrite `ll_history[-1]` to the best value:** dishonest (hides the overshoot)
+  and breaks the fixed-length trajectory contract. Rejected in favor of a
+  separate `final_ll_`.
+- **Restore best even on a degenerate (nan_ll) stop:** would salvage finite
+  parameters from a diverged run, but entangles with issue #50's degenerate-fit
+  contract; the safeguard is skipped for degenerate stops and left to #50.
+
+## Receipts
+
+- `pyAMICA/torch_impl/amica_torch_ng.py` (`keep_best`, `_snapshot_params`/
+  `_restore_params`, `final_ll_`, `_KEEP_BEST_TOL`).
+- `pyAMICA/tests/torch_tests/test_ng_backend.py::test_keep_best_*`.
+- `.context/issue-51/ensemble_ll.py` (Fortran-vs-NG LL ensemble, real data).
+- Fortran schedule: `pyAMICA/amica15.f90:1038-1058` (anneal-on-decrease).

.context/issue-27/multimodel_ensemble.py

@@ -93,7 +93,7 @@ def run_ng(data, seed):
     ng.fit(data, max_iter=MAX_ITER, verbose=False)
     return (
         np.vstack([ng.get_unmixing_matrix(0), ng.get_unmixing_matrix(1)]),
-        ng.ll_history[-1],
+        ng.final_ll_,  # LL of the returned iterate (issue #51 best-iterate safeguard)
     )
 
 

.context/issue-51/best_iterate_safeguard.md

@@ -0,0 +1,78 @@
+# Multi-model NG log-likelihood: best-iterate safeguard (issue #51)
+
+**Bottom line.** The multi-model NG log-likelihood was ~0.02 lower and ~13x more
+variable than Fortran because `AMICATorchNG.fit` returned the *last* EM iterate
+under a deliberately non-monotone learning-rate schedule. Returning the *best*
+iterate (`keep_best`, default on) removes the variance pathology; the remaining
+mean gap is convergence speed, not a worse optimum -- with more iterations NG
+reaches Fortran's exact solution.
+
+## Root cause: return-last, not a bad basin
+
+An NG-only sweep (20 seeds, `n_models=2`, 100 iters, real sample EEG) reproduced
+the #51 finding (mean -3.3738, sd 0.040) and localized the variance to a **late
+overshoot**, not a wrong basin:
+
+- The single variance-driving seed (#3) climbed to LL **-3.3573** (dead in the
+  pack) by iter 97, then **crashed to -3.5452** in its final two iterations after
+  Newton went non-positive-definite and fell back to the natural gradient; the
+  lrate ramp re-inflated the step and the run ended mid-crash.
+- 9 of 20 "good" seeds also ended a small amount below their own peak.
+
+The lrate schedule anneals only *after* an LL decrease (`amica15.f90:1038-1058`,
+mirrored in NG), so the last iterate is not guaranteed to be the best. `fit`
+returned the last iterate.
+
+## Fix: return the best iterate
+
+`AMICATorchNG` tracks the highest-LL iterate and restores it when the run ends
+more than `_KEEP_BEST_TOL` (1e-9) below that peak. `keep_best=True` by default;
+`final_ll_` reports the returned iterate's LL, while `ll_history` stays the true
+trajectory. Inactive under `do_reject` (the good-sample set, hence the LL
+normalization, changes across iterations). A monotone single-model fit has
+best == last, so no restore fires and **issue #24 parity stays bit-exact**
+(verified: max parameter difference 0.0 with keep_best on vs off).
+
+## Results (real sample EEG, Fortran binary + NG, NO MOCK)
+
+`ensemble_ll.py`, N=20 each, `n_models=2`, matched schedule:
+
+| max_iter=100 | mean LL | sd | mean gap | sd ratio | KS p | TOST(±0.01) |
+|---|---:|---:|---:|---:|---:|---|
+| Fortran | -3.3541 | 0.0031 | -- | -- | -- | -- |
+| NG return-last | -3.3738 | 0.0399 | -0.0197 | 12.7x | 9.5e-6 | inconclusive |
+| **NG keep_best** | **-3.3634** | **0.0064** | **-0.0093** | **2.0x** | 5.6e-5 | inconclusive |
+
+keep_best cuts the variance from **12.7x -> 2.0x** Fortran's sd (the headline
+"~13x more variable" defect) and halves the mean gap. See `ll_before_after.png`.
+
+## The residual is convergence speed, not a worse optimum
+
+The ~0.009 mean gap that remains at 100 iters is *iteration budget*, not a wrong
+term. NG keep_best mean LL vs iteration budget (8 seeds):
+
+| max_iter | NG mean | NG sd | Fortran mean (100 it) |
+|---:|---:|---:|---:|
+| 100 | -3.3639 | 0.0076 | -3.3541 |
+| **200** | **-3.3541** | 0.0050 | -3.3541 |
+| 300 | -3.3523 | 0.0040 | -3.3541 |
+
+At 200 iterations NG reaches Fortran's **exact** mean (-3.3541); by 300 it slightly
+exceeds it. NG's per-iteration progress is ~2x slower than Fortran's, but it
+converges to the same optimum -- as expected from the M-step being bit-exact vs
+Fortran (#27). This is optimizer efficiency, not correctness: the reachable
+solution is identical.
+
+## Acceptance
+
+1. **Variance pathology removed** at matched budget (12.7x -> 2.0x). *(Held.)*
+2. **Same optimum**: NG reaches Fortran's exact mean LL with adequate iterations
+   (200-iter mean == -3.3541). *(Held.)*
+3. **Single-model #24 parity** stays bit-exact under the safeguard. *(Held --
+   `test_keep_best_single_model_is_bit_exact`.)*
+
+## Reproduction
+
+- `ensemble_ll.py [N] [MAX_ITER]` -> `ensemble_ll.npz` (Fortran + NG, return-last
+  and keep_best final LLs). `plot_ll.py` renders `ll_before_after.png`.
+- Real sample data + the macOS Fortran binary (x86_64, Rosetta) only.

.context/issue-51/ensemble_ll.npz

@@ -0,0 +1,138 @@
+"""Issue #51 acceptance measurement: does the best-iterate safeguard (keep_best)
+make NG's multi-model log-likelihood distribution equivalent to Fortran's?
+
+Runs N Fortran + N NG fits on the real sample EEG (n_models=2, matched schedule),
+comparing the final-LL distributions three ways:
+
+  - Fortran vs NG *return-last*  (keep_best=False; reproduces the #51 defect)
+  - Fortran vs NG *keep_best*    (keep_best=True; the fix, reports final_ll_)
+
+Reports mean/sd, KS, and TOST (mean-equivalence within +/-DELTA) for each, and
+the sd ratio. Real sample data + the macOS Fortran binary only (NO MOCK).
+
+    uv run python .context/issue-51/ensemble_ll.py [N] [MAX_ITER]
+"""
+
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import numpy as np
+from scipy import stats
+
+from pyAMICA.torch_impl import AMICATorchNG
+from pyAMICA.torch_impl.utils import load_eeglab_data
+
+HERE = Path(__file__).resolve().parent
+REPO = HERE.parents[1]
+BIN = REPO / "pyAMICA/sample_data/amica15mac"
+FDT = REPO / "pyAMICA/sample_data/eeglab_data.fdt"
+FIXTURE = REPO / "pyAMICA/tests/torch_tests/_ng_e2e_tmp/fortran_run/input.param"
+NW, FIELD = 32, 30504
+DELTA_LL = 0.01  # TOST equivalence margin on the mean LL (per sample-channel)
+
+
+def load_data():
+    return load_eeglab_data(str(FDT), data_dim=NW, field_dim=FIELD).astype(np.float64)
+
+
+def run_fortran(work, tag, max_iter):
+    d = work / f"fort_{tag}"
+    (d / "fortran_output").mkdir(parents=True, exist_ok=True)
+    shutil.copy(FDT, d / "eeglab_data.fdt")
+    lines = []
+    for ln in FIXTURE.read_text().splitlines():
+        if ln.startswith("num_models"):
+            lines.append("num_models 2")
+        elif ln.startswith("max_iter"):
+            lines.append(f"max_iter {max_iter}")
+        else:
+            lines.append(ln)
+    (d / "input.param").write_text("\n".join(lines) + "\n")
+    orig = os.getcwd()
+    os.chdir(d)
+    try:
+        r = subprocess.run(
+            [str(BIN), "input.param"], capture_output=True, text=True, timeout=900
+        )
+    finally:
+        os.chdir(orig)
+    if r.returncode != 0:
+        raise RuntimeError(r.stderr[-400:])
+    return next(
+        (
+            float(ln.split("LL =")[1].split()[0])
+            for ln in reversed(r.stdout.splitlines())
+            if "LL =" in ln
+        ),
+        np.nan,
+    )
+
+
+def run_ng(data, seed, keep_best, max_iter):
+    ng = AMICATorchNG(
+        n_channels=NW, n_models=2, n_mix=3, block_size=512, lrate=0.05, minlrate=1e-8,
+        lratefact=0.5, maxdecs=3, do_newton=True, newt_start=50, newt_ramp=10,
+        newtrate=1.0, rho0=1.5, minrho=1.0, maxrho=2.0, rholrate=0.05,
+        rholratefact=0.5, invsigmin=1e-8, invsigmax=100.0, doscaling=True,
+        scalestep=1, seed=seed, device="cpu", keep_best=keep_best,
+    )  # fmt: skip
+    ng.fit(data, max_iter=max_iter, verbose=False)
+    return ng.final_ll_
+
+
+def compare(name, F, G):
+    F, G = np.asarray(F), np.asarray(G)
+    ks = stats.ks_2samp(G, F).pvalue
+    diff = G.mean() - F.mean()
+    se = np.sqrt(G.var(ddof=1) / G.size + F.var(ddof=1) / F.size)
+    p_tost = max(
+        stats.norm.sf((diff + DELTA_LL) / se), stats.norm.cdf((diff - DELTA_LL) / se)
+    )
+    print(
+        f"  {name:16s} F={F.mean():.4f}(sd {F.std():.4f})  "
+        f"NG={G.mean():.4f}(sd {G.std():.4f})  diff={diff:+.4f}  "
+        f"sd_ratio={G.std() / max(F.std(), 1e-9):.1f}x  "
+        f"KS p={ks:.1e}  TOST(+/-{DELTA_LL}) p={p_tost:.1e} "
+        f"{'EQUIV' if p_tost < 0.05 else 'inconclusive'}"
+    )
+
+
+def main():
+    n = int(sys.argv[1]) if len(sys.argv) > 1 else 20
+    max_iter = int(sys.argv[2]) if len(sys.argv) > 2 else 100
+    data = load_data()
+    work = Path(tempfile.mkdtemp(prefix="amica_ll51_"))
+    print(f"scratch: {work}  N={n}  max_iter={max_iter}")
+
+    F = []
+    for k in range(n):
+        F.append(run_fortran(work, str(k), max_iter))
+        print(f"  Fortran {k + 1}/{n}: LL={F[-1]:.4f}", flush=True)
+    G_last, G_best = [], []
+    for k in range(n):
+        G_last.append(run_ng(data, k, False, max_iter))
+        G_best.append(run_ng(data, k, True, max_iter))
+        print(
+            f"  NG {k + 1}/{n}: last={G_last[-1]:.4f} best={G_best[-1]:.4f}", flush=True
+        )
+
+    print(f"\n==== N={n} max_iter={max_iter} ====")
+    compare("return-last", F, G_last)
+    compare("keep_best", F, G_best)
+    np.savez(
+        HERE / "ensemble_ll.npz",
+        F=F,
+        G_last=G_last,
+        G_best=G_best,
+        n=n,
+        max_iter=max_iter,
+    )
+    print(f"saved -> {HERE / 'ensemble_ll.npz'}")
+
+
+if __name__ == "__main__":
+    main()

.context/issue-51/ensemble_ll.py

@@ -0,0 +1,70 @@
+"""Render the issue #51 before/after LL distribution figure from ensemble_ll.npz
+(written by ensemble_ll.py). Fortran vs NG return-last vs NG keep_best."""
+
+from pathlib import Path
+
+import numpy as np
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
+
+HERE = Path(__file__).resolve().parent
+C_FORT, C_LAST, C_BEST = "#0072B2", "#D55E00", "#009E73"  # Okabe-Ito
+
+
+def main():
+    d = np.load(HERE / "ensemble_ll.npz")
+    F, G_last, G_best = d["F"], d["G_last"], d["G_best"]
+    n, mi = int(d["n"]), int(d["max_iter"])
+
+    plt.rcParams.update(
+        {"font.size": 11, "axes.spines.top": False, "axes.spines.right": False}
+    )
+    fig, ax = plt.subplots(figsize=(8.2, 4.6))
+    lo = min(F.min(), G_last.min()) - 0.005
+    hi = max(F.max(), G_last.max()) + 0.005
+    bins = np.linspace(lo, hi, 40)
+    for arr, col, lab in [
+        (F, C_FORT, f"Fortran  ({F.mean():.4f}, sd {F.std():.4f})"),
+        (
+            G_last,
+            C_LAST,
+            f"NG return-last  ({G_last.mean():.4f}, sd {G_last.std():.4f})",
+        ),
+        (G_best, C_BEST, f"NG keep_best  ({G_best.mean():.4f}, sd {G_best.std():.4f})"),
+    ]:
+        ax.hist(arr, bins=bins, density=True, color=col, alpha=0.32)
+        ax.hist(
+            arr, bins=bins, density=True, histtype="step", color=col, lw=2, label=lab
+        )
+        ax.axvline(arr.mean(), color=col, ls="--", lw=1.2)
+    ax.set_xlabel("final log-likelihood (per sample-channel)")
+    ax.set_ylabel("density")
+    ax.set_title(
+        f"Multi-model AMICA (n_models=2, N={n}, {mi} iters): LL distributions",
+        loc="left",
+        fontweight="bold",
+    )
+    ax.legend(frameon=False, fontsize=9, loc="upper left")
+    ax.text(
+        0.98,
+        0.97,
+        "keep_best (issue #51): sd 12.7x -> 2.0x Fortran,\n"
+        "mean gap 0.020 -> 0.009. Removes the low-LL\n"
+        "overshoot tail; small residual is optimizer\n"
+        "efficiency (NG peaks sit ~0.009 below Fortran).",
+        transform=ax.transAxes,
+        va="top",
+        ha="right",
+        fontsize=8.5,
+        bbox=dict(boxstyle="round", fc="white", ec="0.7", alpha=0.9),
+    )
+    fig.tight_layout()
+    fig.savefig(HERE / "ll_before_after.png", bbox_inches="tight", dpi=200)
+    fig.savefig(HERE / "ll_before_after.pdf", bbox_inches="tight")
+    print(f"figure -> {HERE / 'll_before_after.png'}")
+
+
+if __name__ == "__main__":
+    main()

.context/issue-51/ll_before_after.pdf

@@ -30,6 +30,12 @@
 - [x] Multi-model AMICA per-model bias `c` update (issue #27): ported to both backends, guarded
       no-op for `n_models=1`; controlled A/B shows +0.011 cross-corr, gap is intrinsic partition
       ambiguity (see `.context/issue-27/multimodel_c_update.md`).
+- [x] Best-iterate safeguard (issue #51): `AMICATorchNG.fit` returns the highest-LL iterate
+      (`keep_best`, `final_ll_`), not the last, so a late Newton-fallback overshoot no longer leaves
+      the model below a peak it reached. Root cause was return-last, not a bad basin (the sole
+      variance-driving seed peaked at -3.357 then crashed to -3.545 in its final iterations).
+      Single-model #24 parity stays bit-exact (monotone => no restore). See ADR 0003,
+      `.context/issue-51/`.
 - [ ] Component sharing
 
 ### Priority 3: Testing & validation