From 58d53151ba1a3a4a8548f73963fd526304d73e4f Mon Sep 17 00:00:00 2001 From: Seyed Yahya Shirazi Date: Sun, 19 Jul 2026 06:46:06 -0700 Subject: [PATCH 1/2] Fix MLX rholrate to ratchet at maxdecs (#195) --- pamica/mlx_impl/core.py | 22 +++++++- pamica/tests/mlx_tests/test_mlx_backend.py | 62 ++++++++++++++++++++++ 2 files changed, 82 insertions(+), 2 deletions(-) diff --git a/pamica/mlx_impl/core.py b/pamica/mlx_impl/core.py index 0d04289..32a19ce 100644 --- a/pamica/mlx_impl/core.py +++ b/pamica/mlx_impl/core.py @@ -93,6 +93,7 @@ def __init__( lratefact: float = 0.5, maxdecs: int = 5, newt_ramp: int = 10, + newt_start: int = 50, do_newton: bool = False, rho0: float = 1.5, minrho: float = 1.0, @@ -134,6 +135,11 @@ def __init__( self.lratefact = lratefact self.maxdecs = maxdecs self.newt_ramp = newt_ramp + # Schedule threshold for the rholrate ceiling ratchet (Fortran + # amica15.f90:1049 gates it on iter > newt_start, independent of + # do_newton). MLX does no Newton, but keeps newt_start as this gate so the + # rho-rate schedule matches Fortran/torch. + self.newt_start = newt_start self.do_newton = False self.rho0 = rho0 @@ -602,7 +608,18 @@ def fit( self.ll_history.append(ll) # Learning-rate control (Fortran amica17.f90:1062-1108): anneal on an - # LL decrease; ratchet the ceiling after maxdecs persistent decreases. + # LL decrease; ratchet the ceilings after maxdecs persistent decreases. + # + # rholrate is a maxdecs-ratcheted CEILING, not a per-decrease-annealed + # working rate. Fortran resets rholrate=rholrate0 every iteration before + # the rho update (amica15.f90:1788) and only tightens the rholrate0 + # ceiling at maxdecs (amica15.f90:1050, gated on iter > newt_start), so + # its per-decrease rholrate*=rholratefact (:1045) never reaches the rho + # update. rho has no ramp, so self.rholrate carries that ceiling directly + # (reset to rholrate0 at fit start, nothing re-inflates it) and must + # ratchet ONLY at maxdecs. The previous per-decrease decay collapsed the + # rho rate to ~1e-5 within a few hundred iterations and froze rho at a + # stale shape (issue #195, mirroring the torch/numpy fix in #193/#194). if len(self.ll_history) > 1 and ll < self.ll_history[-2]: if self.lrate <= self.minlrate: logger.warning( @@ -613,10 +630,11 @@ def fit( self.stop_reason = "lrate_floor" break self.lrate *= self.lratefact - self.rholrate *= self.rholratefact numdecs += 1 if numdecs >= self.maxdecs: self.lrate_cap *= self.lratefact + if it > self.newt_start: + self.rholrate *= self.rholratefact numdecs = 0 if self.stop_reason in self._DEGENERATE_STOP_REASONS: diff --git a/pamica/tests/mlx_tests/test_mlx_backend.py b/pamica/tests/mlx_tests/test_mlx_backend.py index 56c11ed..9a8909b 100644 --- a/pamica/tests/mlx_tests/test_mlx_backend.py +++ b/pamica/tests/mlx_tests/test_mlx_backend.py @@ -310,3 +310,65 @@ def test_multimodel_dead_model_keeps_prior_c(): m.c = mx.array(np.stack([np.zeros(NW, np.float32), marker], axis=1)) m._update_parameters(acc, xt.shape[1]) assert np.allclose(np.array(m.c)[:, 1], marker) # dead model kept its prior c + + +def test_rholrate_ratchets_at_maxdecs_not_per_decrease(): + """Issue #195 (mirrors #193/#194 for torch/numpy): the MLX rho learning rate is + a maxdecs-ratcheted CEILING, not a per-LL-decrease monotone decay. + + Fortran resets ``rholrate = rholrate0`` each iteration before the rho update + (amica15.f90:1788) and only tightens the ceiling at ``maxdecs`` + (amica15.f90:1050, gated on ``iter > newt_start``). MLX previously decayed + ``rholrate`` on EVERY LL decrease with no reset, collapsing it toward ~1e-5 + within a few hundred iterations and freezing rho at a stale shape. + + An aggressive-lrate natural-gradient run on the real sample overshoots and + triggers several LL decreases; with ``newt_start=0`` (gate always open) the + ceiling ratchets once per ``maxdecs`` decreases, NOT once per decrease. + """ + from pamica.mlx_impl import AMICAMLXNG + + data = _load_real_data() + m = AMICAMLXNG( + n_channels=NW, n_mix=NMIX, seed=SEED, block_size=256, + lrate=0.5, lratefact=0.5, rholrate=0.05, rholratefact=0.5, + maxdecs=3, newt_start=0, + ) # fmt: skip + m.fit(data, max_iter=400, verbose=False) + + ll = m.ll_history + n_dec = sum(1 for i in range(1, len(ll)) if ll[i] < ll[i - 1]) + assert m.lrate > m.minlrate, "run hit the lrate floor; use a gentler config" + assert n_dec >= m.maxdecs, "too few LL decreases to exercise the ratchet" + + # The ceiling ratcheted a whole number of times at the maxdecs cadence + # (rholrate0 * rholratefact**k), far fewer steps than one-per-decrease. + assert m.rholrate < m.rholrate0 + k = round(np.log(m.rholrate / m.rholrate0) / np.log(m.rholratefact)) + assert m.rholrate == pytest.approx(m.rholrate0 * m.rholratefact**k) + assert k <= n_dec // m.maxdecs + 1 + # Guard against a regression to the old per-decrease decay (orders below). + buggy = m.rholrate0 * (m.rholratefact**n_dec) + assert m.rholrate > buggy * 10 + + +def test_rholrate_ceiling_ratchet_gated_on_newt_start(): + """The rholrate ceiling ratchet is gated on ``iter > newt_start`` (Fortran + amica15.f90:1049, independent of do_newton). With ``newt_start`` past the whole + budget, LL decreases must NOT tighten the rho ceiling -- it stays at + ``rholrate0`` exactly (only ``lrate_cap`` ratchets, which is not gated).""" + from pamica.mlx_impl import AMICAMLXNG + + data = _load_real_data() + m = AMICAMLXNG( + n_channels=NW, n_mix=NMIX, seed=SEED, block_size=256, + lrate=0.5, lratefact=0.5, rholrate=0.05, rholratefact=0.5, + maxdecs=3, newt_start=100_000, + ) # fmt: skip + m.fit(data, max_iter=400, verbose=False) + + ll = m.ll_history + n_dec = sum(1 for i in range(1, len(ll)) if ll[i] < ll[i - 1]) + assert n_dec >= m.maxdecs, "config did not exercise the decrease path" + # Gated off for the whole run: the rho ceiling never moved. + assert m.rholrate == m.rholrate0 From 840db74ef9f292793c9913d50173f5ca080d89e1 Mon Sep 17 00:00:00 2001 From: Seyed Yahya Shirazi Date: Sun, 19 Jul 2026 09:19:22 -0700 Subject: [PATCH 2/2] Match newt_start default to torch/numpy; add reset test --- pamica/mlx_impl/core.py | 2 +- pamica/tests/mlx_tests/test_mlx_backend.py | 16 ++++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/pamica/mlx_impl/core.py b/pamica/mlx_impl/core.py index 32a19ce..b46c0a8 100644 --- a/pamica/mlx_impl/core.py +++ b/pamica/mlx_impl/core.py @@ -93,7 +93,7 @@ def __init__( lratefact: float = 0.5, maxdecs: int = 5, newt_ramp: int = 10, - newt_start: int = 50, + newt_start: int = 20, do_newton: bool = False, rho0: float = 1.5, minrho: float = 1.0, diff --git a/pamica/tests/mlx_tests/test_mlx_backend.py b/pamica/tests/mlx_tests/test_mlx_backend.py index 9a8909b..559fce4 100644 --- a/pamica/tests/mlx_tests/test_mlx_backend.py +++ b/pamica/tests/mlx_tests/test_mlx_backend.py @@ -372,3 +372,19 @@ def test_rholrate_ceiling_ratchet_gated_on_newt_start(): assert n_dec >= m.maxdecs, "config did not exercise the decrease path" # Gated off for the whole run: the rho ceiling never moved. assert m.rholrate == m.rholrate0 + + +def test_rholrate_ceiling_resets_at_fit_start(): + """Issue #195: the rho-rate ceiling is reset to rholrate0 at fit start + (``_initialize_parameters``, the same call ``fit`` makes), so a previously + ratcheted ``rholrate`` does not carry across a re-fit/restart -- parity with + the numpy backend's ``test_reinitialize_for_restart_resets_rho_ceiling``.""" + from pamica.mlx_impl import AMICAMLXNG + + m = AMICAMLXNG(n_channels=NW, n_mix=NMIX, seed=SEED) + # Simulate a prior fit that ratcheted both ceilings down at maxdecs. + m.rholrate = m.rholrate0 * 0.25 + m.lrate_cap = m.lrate0 * 0.25 + m._initialize_parameters() # fit-start reset + assert m.rholrate == m.rholrate0 + assert m.lrate_cap == m.lrate0