Skip to content

Reset rholrate per iteration to match Fortran - #194

Merged
neuromechanist merged 3 commits into
mainfrom
fix/issue-193-rholrate
Jul 19, 2026
Merged

Reset rholrate per iteration to match Fortran#194
neuromechanist merged 3 commits into
mainfrom
fix/issue-193-rholrate

Conversation

@neuromechanist

Copy link
Copy Markdown
Member

Summary

AMICATorchNG (and the legacy numpy_impl backend) diverged from the Fortran reference in the rho (generalized-Gaussian shape) learning-rate schedule: torch permanently decayed rholrate on every log-likelihood decrease, whereas Fortran resets rholrate = rholrate0 each iteration and only ratchets the ceiling at maxdecs. Once an LL decrease shrank rholrate it stayed shrunk (monotone decay), so on a long run the rho update was progressively under-relaxed and eventually frozen (rate collapsed to ~1e-5 within a few hundred iterations). Refs #193.

Fortran expressions

  • amica15.f90:1788/:1795update_params sets rholrate = rholrate0 every iteration, before the rho update at :1814. So Fortran's per-decrease rholrate = rholrate*rholratefact (:1045) is overwritten and never reaches the rho update; the effective rate is always the rholrate0 ceiling.
  • amica15.f90:1050 — the ceiling rholrate0 ratchets *= rholratefact only at maxdecs, gated on iter > newt_start (unlike newtrate, which is additionally gated on do_newton).

Fix

rho has no per-iteration ramp, so the effective rho rate is exactly its ceiling. Both backends now carry that ceiling directly:

  • torch (pamica/torch_impl/core.py): self.rholrate is the ceiling — reset to rholrate0 at fit start (existing), ratcheted only in the maxdecs block (gated it > newt_start), never per-decrease. The buggy per-decrease self.rholrate *= self.rholratefact is removed.
  • numpy (pamica/numpy_impl/core.py): same — self.rholrate is the ceiling, reset to rholrate0 at fit start, ratcheted at maxdecs (the previously-dead rholrate0 ratchet is retargeted to rholrate); the per-decrease decay is removed.

rholrate0 stays the pristine constructor value in both. No state-dict format change (the persisted rholrate field now carries the ceiling; round-trip is unchanged).

Parity validation (vs amica15mac, bundled 32ch/30504 data, 2000 iters, do_newton)

Converged rho distribution (shape param, bounded [1,2]) — sorted, permutation-free:

min q25 med q75 max
torch (fixed) 1.000 1.606 1.822 1.994 2.000
Fortran run #1 1.000 1.614 1.829 2.000 2.000
Fortran run #2 1.000 1.636 1.845 2.000 2.000

Sorted-rho L1 distance: torch-vs-Fortran = 0.032, Fortran-vs-Fortran = 0.018 — torch's rho distribution sits within Fortran's own run-to-run spread on this data. Post-fit rholrate ends at 0.00039 (= 0.05 * 0.5**7, seven maxdecs ratchets, matching newtrate), not the buggy per-decrease collapse (0.05 * 0.5**19 ~ 1e-7).

Note on #145

This is a correctness/parity fix on its own merits; it is evaluated against the Fortran binary, NOT the #145 long-budget Newton metric. As documented in #193, matching Fortran here worsens the #145 min-correlation on the (data-inadequate, k=29.8) bundled recording, because the accidental decay was freezing rho and removing a Newton noise source. That is expected and is a red herring for #145 (see the #145 investigation notes).

Test plan

  • New real-data regression test test_rholrate_ratchets_at_maxdecs_not_per_decrease (marked slow): on the sample EEG a 300-iter Newton run overshoots and triggers several LL decreases; asserts the surviving rholrate ratcheted exactly as often as newtrate (both gated on iter > newt_start at maxdecs), and is orders of magnitude above the old per-decrease-decay value.
  • Existing gates green: test_full_fit_parity_numpy_vs_ng, test_end_to_end_correlation_vs_fortran, test_sample_data.py (numpy-vs-Fortran), torch/numpy non-slow suites, ruff, ty.

Refs #193

@neuromechanist

Copy link
Copy Markdown
Member Author

Review (/review-pr, reviewers = Sonnet)

Ran the pr-review-toolkit code-reviewer and pr-test-analyzer on the diff. Correctness checks were clean (ceiling-ratchet gating iter > newt_start correct and distinct from the do_newton-gated newtrate ratchet; rholrate0 stays pristine; torch state_dict round-trip unaffected; cross-backend consistency confirmed). Findings and resolutions:

1. numpy fix had no test that would fail without it (test-analyzer, sev 8) — ADDRESSED. Extended the existing fast, deterministic unit test test_check_convergence_ratchets_lrate_on_decrease to assert the rho ceiling: it stays untouched on a decrease below max_decs (pre-#193 it decayed on the first decrease) and ratchets by rholratefact when numdecs hits max_decs. Fails on pre-fix numpy, no real data needed. (commit 7bae7f0)

2. numpy _reinitialize_for_restart rho reset was new + untested (test-analyzer, sev 7) — ADDRESSED. Added test_reinitialize_for_restart_resets_rho_ceiling: a direct unit test that a ratcheted rholrate is restored to rholrate0 across a restart, mirroring the existing lrate reset. (commit 9dd3141)

3. MLX backend has the same bug (code-reviewer) — TRACKED as #195, not folded in. mlx_impl/core.py:616 still decays rholrate per-decrease with no reset/ceiling. Out of #193's stated torch+numpy scope, and MLX is natural-gradient-only with no newt_start/Newton concept, so the Fortran ceiling-ratchet gate (iter > newt_start) needs a design decision. Filed #195 with the fix sketch; MLX is explicitly not a float64-parity backend, so it does not gate this PR.

4. torch fit-start refit reset has no direct test (test-analyzer, sev 4) — pre-existing, out of scope. torch_impl/core.py:782 (self.rholrate = self.rholrate0) is unchanged by this PR and mirrors the existing lrate/newtrate/lrate_cap fit-start resets; noting as a possible follow-up rather than expanding this diff.

All added tests + existing parity gates (test_full_fit_parity_numpy_vs_ng, test_end_to_end_correlation_vs_fortran, test_sample_data numpy-vs-Fortran, torch/numpy suites, ruff, ty) are green.

@neuromechanist
neuromechanist merged commit a6b967f into main Jul 19, 2026
7 checks passed
@neuromechanist
neuromechanist deleted the fix/issue-193-rholrate branch July 19, 2026 03:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant