[Bug] chunk_momentum_delta_rule diverges from the recurrent reference when log_mu is moderately negative #1578

Describe the bug

chunk_momentum_delta_rule disagrees with fused_recurrent_momentum_delta_rule / recurrent_momentum_delta_rule_ref once the momentum decay μ = exp(log_mu) is moderately small. chunk_momentum_delta_rule is the PyTorch chunk_momentum_delta_rule_ref path, and it is the default MomentumDeltaNet training mode.

The error grows with the position inside a chunk. The first positions of each chunk match to ~1e-6, and by the end of a 64-token chunk the output is essentially unrelated to the recurrent reference. Shrinking the chunk makes the disagreement go away (e.g. chunk 8 → 3e-6). It looks like float32 overflow / cancellation of the in-chunk log-space terms. exp(log_a_cum - log_m_cum + log_ct) grows like μ^-c within a chunk; with μ = 0.5 and c = 63, -log_m_cum reaches ~44, and temp = 1 - exp(x) then cancels.

Steps to reproduce

This runs on CPU and only loads fla/ops/momentum_delta_rule/naive.py. The settings are α = 0.99 and constant μ, with normalized k/p, beta ~ U(0,1), eta = 1, T = 128 and the default chunk size 64.

import importlib.util, torch, torch.nn.functional as F
spec = importlib.util.spec_from_file_location("m", "fla/ops/momentum_delta_rule/naive.py")
m = importlib.util.module_from_spec(spec); spec.loader.exec_module(m)
rel = lambda a, b: float((a - b).norm() / b.norm())
for mu in (0.95, 0.9, 0.8, 0.7, 0.5):
    torch.manual_seed(0)
    B, T, H, K, V = 2, 128, 2, 16, 16
    q, v = torch.randn(B, T, H, K), torch.randn(B, T, H, V)
    k, p = F.normalize(torch.randn(B, T, H, K), dim=-1), F.normalize(torch.randn(B, T, H, K), dim=-1)
    la = torch.full((B, T, H), torch.tensor(0.99).log().item())
    lm = torch.full((B, T, H), torch.tensor(mu).log().item())
    beta, eta = torch.rand(B, T, H), torch.ones(B, T, H)
    r, _ = m.recurrent_momentum_delta_rule_ref(q=q, k=k, v=v, p=p, log_alpha=la, log_mu=lm, beta=beta, eta=eta)
    c, _ = m.chunk_momentum_delta_rule_ref(q=q, k=k, v=v, p=p, log_alpha=la, log_mu=lm, beta=beta, eta=eta, chunk_size=64)
    print(mu, rel(c, r))

The same happens with per-token gates. For example, with log_mu = -softplus(N(0,1)) and log_alpha = -0.1·softplus(N(0,1)):

• T = 64, chunk 64: 0.89 relative error on o and ~1.0 on the final S.
• T = 64, chunk 16: 2e-3.
• T = 8, chunk 8: 3e-6.

Expected behavior

The chunk path should match the recurrent definition for every log_mu that MomentumDeltaNet can produce, within float32 error.

Why this matters in practice

In MomentumDeltaNet, log_mu = -Mu · softplus(m + mu_bias), with Mu ~ U(0, 16) and mu_bias initialized from dt ~ logU(0.001, 0.1). It is clamped only at min_log_mu = -2, so μ can go down to ~0.135. With the default init and m = 0, about 24% of heads start with log_mu <= -0.22 (μ ≤ 0.8) and about 16% with log_mu <= -0.36 (μ ≤ 0.7), before any input-dependent spread. Training uses mode='chunk', while eval with q_len <= 64 switches to fused_recurrent, so those heads would also see different outputs between training and short-sequence inference.

The existing parity test test_full_momentum_chunk_recurrent_backward_parity doesn't catch this. It uses T = 33 with chunk 16, and log_alpha ~ N(0,1), i.e. α up to e^3, under which the state grows to ~1e16 rather than decaying.

Possible directions

I haven't attempted a fix because it touches the chunk algorithm. Two options:

• Keep every exponent in the in-chunk masks as a difference that is ≤ 0. The other FLA chunk references do this with exp(decay_i - decay_j) for j ≤ i, instead of factoring out exp(-log_m_cum) terms that grow within the chunk.
• As a stopgap, sub-chunk or reduce the chunk size when log_mu is small.

Happy to help test a fix.

Environment

CPU, torch 2.9.1. Only the PyTorch reference path is involved; no Triton kernel is used.