EchoMuon

Reference implementation of the preprint EchoMuon: Cross-Timescale Gating in Muon, and the Learning-Rate Confound in Gated Optimizers (Preprints.org, 2026).

Experimental status — read this first

This is a research prototype at an experimental stage, not a production optimizer.

The claims in earlier versions of this package were substantially overstated. A systematic audit of our own results (1,815 runs, all released) found that a large part of the originally reported advantage was an artifact of learning-rate selection, and cut the vision numbers by 42–80%. The mechanism is real and survives norm-matched controls on one benchmark; its useful regime is narrow, and both facts are documented below. Do not adopt this in place of a properly tuned Muon or AdamW baseline without measuring it on your own task. It is also single-GPU only.

Muon orthogonalizes the momentum of 2-D hidden weight matrices, giving every singular direction of the update equal trust. EchoMuon prices that trust by each direction's echo — its support in a second, slower momentum buffer (β=0.99 next to Muon's 0.95). Persistent signal appears in both buffers; noise flickers in the fast one and leaves little trace in the slow one.

The gate scores each singular direction by cross-timescale agreement, c_i = uᵢᵀ M₂M₁ᵀ uᵢ / σᵢ² — not by magnitude, so a large direction sustained by a few noisy batches is damped while a small persistent one is trusted. Scores are referenced to the layer median and clipped to [0.1, 1.0], then engaged in proportion to a measured retention gap λ ∈ [0, 1].

What the gate actually does to the step

Earlier documentation described the gate as mean-1, and therefore unable to act as a disguised learning-rate schedule. The logs say otherwise: the multiplier is clipped at 1.0 and never exceeds it (0 of 5,760 logged layer-steps), with mean ≈ 0.957. It is a contraction, not a reallocation.

The measurable consequence is that EchoMuon sits higher on the learning-rate axis than Muon. Normalising each arm to its own optimum, across four cells with 3-seed sweeps:

This holds in all four cells and in both directions. The window has shifted up, not widened. Sweep EchoMuon's learning rate separately from Muon's: a grid that sits too high for both arms flatters EchoMuon by roughly 0.046 in validation loss, which is comparable to the margins being measured.

(An earlier version of this file claimed the optimum was consistently exactly 2× Muon's. Our learning-rate grids are geometric with ratio 2, so any difference in the discrete pick is necessarily "2×"; that number measured grid resolution rather than the shift. The table above is the continuous measurement.)

What the audit established

The margin follows a measured law. Across ten vision cells spanning three datasets, four label-noise levels, two budgets and two profiles, the advantage scales with how far the model is from the quality it can reach:

margin (pp)  ~  0.027 x (baseline error rate %)  -  0.28

It crosses zero at roughly 90% accuracy. Above that, expect nothing. The LM budget ladder points the same way. This is the single most useful thing to know before trying it.

Install

pip install echomuon

Usage

EchoMuon handles only the 2-D hidden matrices, exactly like Muon; route embeddings, heads, gains and biases to AdamW:

import torch
from echomuon import EchoMuon, MemorizationGapController

hidden = [p for n, p in model.named_parameters()
          if p.ndim == 2 and "embed" not in n and "lm_head" not in n]
others = [p for n, p in model.named_parameters()
          if not any(p is h for h in hidden)]

opt = EchoMuon(hidden, lr=0.02)   # sweep separately from Muon's; see note above
aux = torch.optim.AdamW(others, lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)

Training loop with the retention-gap controller (optional; without it, set gate_lambda yourself — gate_lambda=0 is plain Muon):

import collections

ctl = MemorizationGapController(opt)     # probes every 200 steps by default
ring = collections.deque(maxlen=8)       # the last 8 training batches

for step, batch in enumerate(loader):
    loss = model(**batch).loss
    opt.zero_grad(); aux.zero_grad()
    loss.backward()
    opt.step(); aux.step()
    ring.append(batch)

    if ctl.due(step):
        with torch.no_grad():
            reseen = mean_loss(model, list(ring)[:4])   # ~4-8 steps old, see note
            fresh  = mean_loss(model, take_fresh(4))    # 4 held-back fresh batches
        ctl.update(fresh_loss=fresh, reseen_loss=reseen)

Note on the retention horizon. Appending every step, as above, makes the re-seen batches only ~4–8 steps old — not the few hundred that earlier documentation claimed. We tested the intended long horizon (append once per probe, giving 400–700 steps) on five cells: it changes nothing. The controller is insensitive to this horizon over two orders of magnitude, so the short ring is fine — but the "memorization over hundreds of steps" story that motivated it is not supported.

All constants (slow β=0.99, floor 0.1, median reference, the 2% normalizer, the 0.7 EMA) were fixed once and used unchanged throughout. The knob you tune is the learning rate — separately from Muon's.

When it might help

Limitations

• Single-GPU only. The gate needs the unsharded matrix to form its Gram product. FSDP or tensor-parallel execution would need an all-gather per refresh, or per-shard gates with different semantics. Untested.
• Memory: 9–12 B per 2-D parameter (M₁, M₂, cached basis U, gate g, all fp32), versus 4 B for Muon and 8 B for AdamW.
• Wall-clock: +19–22% per step at LM scale with the standard profile (median +21%), +13% with the fast profile. The vision overhead figures are unreliable: those models are small enough to be CPU-launch-bound rather than GPU-bound.
• No head-to-head comparison against other gated-Muon variants (DynMuon, Pion, MGUP, MAGMA, Bi-Maxwell). Everything here is measured against Muon and AdamW only. Any claim of superiority over those methods would be unsupported.
• Learning-rate selection is single-seed on most cells. Three-seed selection was used only where one seed proved insufficient; we demonstrated one case where a single-seed pick flipped with a 1.45 pp consequence on the test metric. Report your own selection margins.
• Largest model tested: 162M parameters. Nothing here has been tested at production scale.

Paper

S. Mastromichalakis, EchoMuon: Cross-Timescale Gating in Muon, and the Learning-Rate Confound in Gated Optimizers, Preprints.org, 2026. doi:10.20944/preprints202608.2150.v1

The code, all 1,815 runs and the audit behind every number are at github.com/MStamatis/echomuon.

Package/paper parity

Before 0.3.0 the installed package and the research code that produced the paper differed in one detail: the package refreshed the gate basis after folding in the current gradient, the research code did so before. Nothing tested it. The measured effect was smaller than seed-to-seed noise and not statistically distinguishable at n=8, so no published number changes, but 0.3.0 moves the package onto the paper's order and a parity test in the repository now runs both on the same gradients and requires the updates to match bit for bit.

Citation

@misc{mastromichalakis2026echomuon,
  title     = {EchoMuon: Cross-Timescale Gating in Muon, and the Learning-Rate
               Confound in Gated Optimizers},
  author    = {Mastromichalakis, Stamatis},
  year      = {2026},
  publisher = {Preprints.org},
  doi       = {10.20944/preprints202608.2150.v1},
  url       = {https://doi.org/10.20944/preprints202608.2150.v1},
  note      = {Preprint}
}

License

MIT