echomuon 0.1.0__tar.gz

echomuon-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Stamatis Mastromichalakis
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

echomuon-0.1.0/PKG-INFO

@@ -0,0 +1,125 @@
+Metadata-Version: 2.4
+Name: echomuon
+Version: 0.1.0
+Summary: EchoMuon: Muon with a per-direction temporal trust gate and a memorization-gap controller. Better than scheduled Muon wherever data are imperfect.
+Author-email: Stamatis Mastromichalakis <stamatis@tmnetworks.gr>
+License: MIT
+Project-URL: Homepage, https://github.com/MStamatis/echomuon
+Project-URL: Paper, https://arxiv.org/abs/ARXIV-PLACEHOLDER
+Keywords: optimizer,muon,deep-learning,pytorch,label-noise,orthogonalization
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Dynamic: license-file
+
+# EchoMuon
+
+**Muon with a per-direction temporal trust gate and a memorization-gap controller.
+Better than scheduled Muon wherever data are imperfect; ties it everywhere else.**
+
+[Muon](https://kellerjordan.github.io/posts/muon/) orthogonalizes the momentum of
+2-D hidden weight matrices, giving every singular direction of the update exactly
+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 no 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;
+- is **median-normalized per layer** (mean ≈ 1): trust is reallocated across
+  directions at constant total step, so the gate cannot act as a disguised
+  learning-rate schedule;
+- is engaged in proportion to a **measured memorization gap** λ ∈ [0, 1]: the loss
+  on fresh batches minus the loss on batches seen a few hundred steps ago, under
+  the same current weights. At λ=0 EchoMuon *is* Muon, structurally — never worse
+  by construction on clean data.
+
+Headline results (paired seeds, schedule parity, per-arm lr sweeps): **+0.9 to
++1.8pp** over scheduled Muon on six vision cells (CIFAR-10/100, Tiny ImageNet,
+clean and 20% label noise), **−0.031 nats (t=−6.3)** on a LLaMA-style 162M
+transformer on FineWeb-Edu, Muon-tier quality in 80–85% of Muon's steps in every
+seed, ~13% step overhead at 162M with the fast profile. See the paper for the
+boundaries (byte-level LMs, SSMs, strong augmentation recipes) — they are reported,
+measured, and part of the result.
+
+## Install
+
+```bash
+pip install echomuon
+```
+
+## Usage
+
+EchoMuon handles only the 2-D hidden matrices, exactly like Muon; route
+embeddings, heads, gains and biases to AdamW:
+
+```python
+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 lr as you would for Muon
+aux = torch.optim.AdamW(others, lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)
+```
+
+Training loop with the memorization-gap controller (optional but recommended —
+without it, set `gate_lambda` yourself; `gate_lambda=0` is plain Muon):
+
+```python
+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])   # oldest ~400-800 steps ago
+            fresh  = mean_loss(model, take_fresh(4))    # 4 held-back fresh batches
+        ctl.update(fresh_loss=fresh, reseen_loss=reseen)
+```
+
+All constants (slow β=0.99, floor 0.1, median reference, the 2% normalizer, the
+0.7 EMA) were fixed once and used unchanged in every experiment of the paper —
+the only knob you tune is Muon's own learning rate.
+
+## When to use it
+
+| Your setting | Recommendation |
+|---|---|
+| Web-scale corpora, label noise, ambiguous labels, light augmentation | **EchoMuon** — this is where the margins live |
+| Clean data / strong augmentation (RandAugment + mixup) | Tie with Muon; λ backs the gate off automatically |
+| Byte-level LMs at scale | Scheduled Muon (measured boundary; EchoMuon concedes ≤0.5%) |
+| Mamba-style SSMs | AdamW beats the whole Muon family there (measured boundary) |
+
+## Citation
+
+```bibtex
+@article{mastromichalakis2026echomuon,
+  title  = {EchoMuon: Better Than Scheduled Muon Wherever Data Are Imperfect},
+  author = {Mastromichalakis, Stamatis},
+  year   = {2026},
+  note   = {arXiv preprint}
+}
+```
+
+## License
+
+MIT

echomuon-0.1.0/README.md

@@ -0,0 +1,105 @@
+# EchoMuon
+
+**Muon with a per-direction temporal trust gate and a memorization-gap controller.
+Better than scheduled Muon wherever data are imperfect; ties it everywhere else.**
+
+[Muon](https://kellerjordan.github.io/posts/muon/) orthogonalizes the momentum of
+2-D hidden weight matrices, giving every singular direction of the update exactly
+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 no 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;
+- is **median-normalized per layer** (mean ≈ 1): trust is reallocated across
+  directions at constant total step, so the gate cannot act as a disguised
+  learning-rate schedule;
+- is engaged in proportion to a **measured memorization gap** λ ∈ [0, 1]: the loss
+  on fresh batches minus the loss on batches seen a few hundred steps ago, under
+  the same current weights. At λ=0 EchoMuon *is* Muon, structurally — never worse
+  by construction on clean data.
+
+Headline results (paired seeds, schedule parity, per-arm lr sweeps): **+0.9 to
++1.8pp** over scheduled Muon on six vision cells (CIFAR-10/100, Tiny ImageNet,
+clean and 20% label noise), **−0.031 nats (t=−6.3)** on a LLaMA-style 162M
+transformer on FineWeb-Edu, Muon-tier quality in 80–85% of Muon's steps in every
+seed, ~13% step overhead at 162M with the fast profile. See the paper for the
+boundaries (byte-level LMs, SSMs, strong augmentation recipes) — they are reported,
+measured, and part of the result.
+
+## Install
+
+```bash
+pip install echomuon
+```
+
+## Usage
+
+EchoMuon handles only the 2-D hidden matrices, exactly like Muon; route
+embeddings, heads, gains and biases to AdamW:
+
+```python
+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 lr as you would for Muon
+aux = torch.optim.AdamW(others, lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)
+```
+
+Training loop with the memorization-gap controller (optional but recommended —
+without it, set `gate_lambda` yourself; `gate_lambda=0` is plain Muon):
+
+```python
+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])   # oldest ~400-800 steps ago
+            fresh  = mean_loss(model, take_fresh(4))    # 4 held-back fresh batches
+        ctl.update(fresh_loss=fresh, reseen_loss=reseen)
+```
+
+All constants (slow β=0.99, floor 0.1, median reference, the 2% normalizer, the
+0.7 EMA) were fixed once and used unchanged in every experiment of the paper —
+the only knob you tune is Muon's own learning rate.
+
+## When to use it
+
+| Your setting | Recommendation |
+|---|---|
+| Web-scale corpora, label noise, ambiguous labels, light augmentation | **EchoMuon** — this is where the margins live |
+| Clean data / strong augmentation (RandAugment + mixup) | Tie with Muon; λ backs the gate off automatically |
+| Byte-level LMs at scale | Scheduled Muon (measured boundary; EchoMuon concedes ≤0.5%) |
+| Mamba-style SSMs | AdamW beats the whole Muon family there (measured boundary) |
+
+## Citation
+
+```bibtex
+@article{mastromichalakis2026echomuon,
+  title  = {EchoMuon: Better Than Scheduled Muon Wherever Data Are Imperfect},
+  author = {Mastromichalakis, Stamatis},
+  year   = {2026},
+  note   = {arXiv preprint}
+}
+```
+
+## License
+
+MIT

echomuon-0.1.0/pyproject.toml

@@ -0,0 +1,28 @@
+[build-system]
+requires = ["setuptools>=64"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "echomuon"
+version = "0.1.0"
+description = "EchoMuon: Muon with a per-direction temporal trust gate and a memorization-gap controller. Better than scheduled Muon wherever data are imperfect."
+readme = "README.md"
+requires-python = ">=3.9"
+license = { text = "MIT" }
+authors = [{ name = "Stamatis Mastromichalakis", email = "stamatis@tmnetworks.gr" }]
+keywords = ["optimizer", "muon", "deep-learning", "pytorch", "label-noise", "orthogonalization"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = ["torch>=2.0"]
+
+[project.urls]
+Homepage = "https://github.com/MStamatis/echomuon"
+Paper = "https://arxiv.org/abs/ARXIV-PLACEHOLDER"
+
+[tool.setuptools.packages.find]
+where = ["src"]

echomuon-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

echomuon-0.1.0/src/echomuon/__init__.py

@@ -0,0 +1,7 @@
+"""EchoMuon: Muon with a per-direction temporal trust gate and a
+memorization-gap controller. Better than scheduled Muon wherever data are
+imperfect; ties it everywhere else."""
+from .optimizer import EchoMuon, MemorizationGapController, newton_schulz5
+
+__version__ = "0.1.0"
+__all__ = ["EchoMuon", "MemorizationGapController", "newton_schulz5", "__version__"]

echomuon-0.1.0/src/echomuon/optimizer.py

@@ -0,0 +1,228 @@
+"""EchoMuon: Muon with a per-direction temporal trust gate and an optional
+memorization-gap controller.
+
+Muon (Jordan et al., 2024) 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. Directions whose fast/slow buffers agree keep their step;
+directions with no echo are damped. Gates are median-normalized per layer
+(mean ~1), so trust is reallocated across directions at constant total step --
+the gate cannot act as a disguised learning-rate schedule -- and a floor keeps
+every direction alive. A scalar ``gate_lambda`` in [0, 1] interpolates the gate
+between OFF (exactly plain Muon) and fully ON; the MemorizationGapController
+sets it from a measured old-vs-fresh loss gap.
+
+Reference: S. Mastromichalakis, "EchoMuon: Better Than Scheduled Muon Wherever
+Data Are Imperfect", 2026.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+
+
+__all__ = ["EchoMuon", "MemorizationGapController", "newton_schulz5"]
+
+
+def newton_schulz5(G: torch.Tensor, steps: int = 5, eps: float = 1e-7) -> torch.Tensor:
+    """Quintic Newton-Schulz orthogonalization (Keller Jordan's coefficients)."""
+    a, b, c = 3.4445, -4.7750, 2.0315
+    X = G.to(torch.bfloat16) if G.is_cuda else G.to(torch.float32)
+    transposed = X.size(0) > X.size(1)
+    if transposed:
+        X = X.T
+    X = X / (X.norm() + eps)
+    for _ in range(steps):
+        A = X @ X.T
+        B = b * A + c * (A @ A)
+        X = a * X + B @ X
+    if transposed:
+        X = X.T
+    return X.to(G.dtype)
+
+
+class EchoMuon(torch.optim.Optimizer):
+    """EchoMuon optimizer for 2-D hidden weight matrices.
+
+    Like Muon, this optimizer is meant ONLY for the hidden matrix parameters of
+    a network (attention/MLP weights). Embeddings, output heads, gains, and
+    biases should be handled by a separate AdamW -- exactly as with Muon.
+
+    Arguments:
+        params: 2-D parameters to optimize (an iterable of tensors or dicts).
+        lr: learning rate (Muon convention; sweep it as you would for Muon).
+        momentum: fast EMA-sum decay (Muon's own buffer), default 0.95.
+        slow_beta: slow buffer decay, default 0.99. The gate scores each
+            singular direction of the fast buffer by its normalized support in
+            this slower buffer.
+        nesterov: use the Nesterov form of the fast buffer (default True).
+        weight_decay: decoupled weight decay (default 0).
+        ns_steps: Newton-Schulz iterations (default 5).
+        gate_every: refresh the gate basis every this many steps (default 100,
+            the paper's "fast profile"; 25 is the standard profile -- same
+            quality within noise, ~2x the gate overhead).
+        gate_floor: minimum gate value so no direction is fully silenced
+            (default 0.1).
+        gate_lambda: gate strength in [0, 1]. 0 = exactly plain Muon; 1 = full
+            gate. Set it directly, or let MemorizationGapController drive it
+            from a measured memorization gap.
+
+    Example::
+
+        hidden = [p for n, p in model.named_parameters()
+                  if p.ndim == 2 and "embed" not in n and "head" not in n]
+        others = [p for n, p in model.named_parameters() if p not in set(hidden)]
+        opt = EchoMuon(hidden, lr=0.02)
+        aux = torch.optim.AdamW(others, lr=3e-4, weight_decay=0.1)
+    """
+
+    def __init__(self, params, lr: float = 0.02, momentum: float = 0.95,
+                 slow_beta: float = 0.99, nesterov: bool = True,
+                 weight_decay: float = 0.0, ns_steps: int = 5,
+                 gate_every: int = 100, gate_floor: float = 0.1,
+                 gate_lambda: float = 1.0):
+        if not 0.0 <= gate_lambda <= 1.0:
+            raise ValueError("gate_lambda must be in [0, 1]")
+        defaults = dict(lr=lr, momentum=momentum, slow_beta=slow_beta,
+                        nesterov=nesterov, weight_decay=weight_decay,
+                        ns_steps=ns_steps, gate_every=max(1, gate_every),
+                        gate_floor=gate_floor)
+        super().__init__(params, defaults)
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.ndim != 2:
+                    raise ValueError(
+                        "EchoMuon only accepts 2-D parameters; route "
+                        f"{tuple(p.shape)} tensors to AdamW instead.")
+        self.gate_lambda = float(gate_lambda)
+        self._t = 0
+
+    @torch.no_grad()
+    def _refresh_gate(self, p, group):
+        """Recompute the temporal-consistency gate for one parameter.
+
+        For each singular direction u_i of the fast buffer M1 (via the Gram
+        eigendecomposition on the small side), consistency
+        c_i = (u_i^T M2 M1^T u_i) / sigma_i^2 is the slow buffer's relative
+        support along that direction. Gates are median-normalized per layer so
+        they redistribute (mean ~1) rather than rescale.
+        """
+        st = self.state[p]
+        M1 = st["buf"].float()
+        M2 = st["buf2"].float()
+        if M1.size(0) > M1.size(1):
+            M1, M2 = M1.T, M2.T
+        G = M1 @ M1.T
+        evals, U = torch.linalg.eigh(G)
+        evals = evals.clamp_min(1e-12)
+        C = U.T @ (M2 @ M1.T) @ U
+        c = C.diagonal() / evals
+        med = torch.quantile(c, 0.5).clamp_min(1e-12)
+        g = (c / med).clamp(group["gate_floor"], 1.0)
+        st["gate_U"], st["gate_g"] = U, g
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+        self._t += 1
+        refresh = (self._t % self.param_groups[0]["gate_every"] == 0)
+        for group in self.param_groups:
+            beta = group["momentum"]
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                st = self.state[p]
+                if "buf" not in st:
+                    st["buf"] = torch.zeros_like(p)
+                    st["buf2"] = torch.zeros_like(p)
+                buf, buf2 = st["buf"], st["buf2"]
+                buf.mul_(beta).add_(p.grad)
+                buf2.mul_(group["slow_beta"]).add_(p.grad)
+                if refresh:
+                    self._refresh_gate(p, group)
+                u = p.grad.add(buf, alpha=beta) if group["nesterov"] else buf
+                d = newton_schulz5(u, group["ns_steps"])
+                if "gate_U" in st and self.gate_lambda > 0.0:
+                    U, g = st["gate_U"], st["gate_g"]
+                    g_eff = 1.0 - self.gate_lambda * (1.0 - g)
+                    transposed = d.size(0) > d.size(1)
+                    O = (d.T if transposed else d).float()
+                    O = O - U @ ((1.0 - g_eff).unsqueeze(1) * (U.T @ O))
+                    d = (O.T if transposed else O).to(d.dtype)
+                d = d * math.sqrt(max(1.0, p.size(0) / p.size(1)))
+                if group["weight_decay"]:
+                    p.mul_(1 - group["lr"] * group["weight_decay"])
+                p.add_(d, alpha=-group["lr"])
+        return loss
+
+    def state_dict(self):
+        sd = super().state_dict()
+        sd["echomuon_t"] = self._t
+        sd["echomuon_lambda"] = self.gate_lambda
+        return sd
+
+    def load_state_dict(self, state_dict):
+        self._t = state_dict.pop("echomuon_t", 0)
+        self.gate_lambda = state_dict.pop("echomuon_lambda", self.gate_lambda)
+        super().load_state_dict(state_dict)
+
+
+class MemorizationGapController:
+    """Drives ``EchoMuon.gate_lambda`` from a measured memorization gap.
+
+    The signal: re-evaluate, under the CURRENT weights, batches the model
+    trained on a few hundred steps ago, alongside fresh batches. The gap
+    (fresh loss - re-seen loss) is memorization, measured on the model itself;
+    lambda = clip(gap / (target_frac * fresh_loss), 0, 1), EMA-smoothed. On
+    clean data the gap vanishes and EchoMuon relaxes to exactly plain Muon.
+
+    The controller does not run the forward passes itself -- your training
+    loop supplies the two loss values (this keeps the package framework-free).
+
+    Example::
+
+        ctl = MemorizationGapController(opt)          # opt is an EchoMuon
+        ring = collections.deque(maxlen=8)            # your recent batches
+        for step, batch in enumerate(loader):
+            train_step(batch); ring.append(batch)
+            if ctl.due(step):
+                with torch.no_grad():
+                    reseen = mean_loss(model, list(ring)[:4])   # oldest 4
+                    fresh  = mean_loss(model, next_fresh_batches(4))
+                ctl.update(fresh_loss=fresh, reseen_loss=reseen)
+
+    Arguments:
+        optimizer: the EchoMuon instance whose gate_lambda is controlled.
+        probe_every: probe cadence in steps (default 200, the fast profile;
+            100 is the standard profile).
+        target_frac: gap normalizer as a fraction of the fresh loss
+            (default 0.02 -- fixed a priori in the paper, never retuned).
+        ema: smoothing of lambda across probes (default 0.7).
+    """
+
+    def __init__(self, optimizer: EchoMuon, probe_every: int = 200,
+                 target_frac: float = 0.02, ema: float = 0.7):
+        self.opt = optimizer
+        self.probe_every = max(1, probe_every)
+        self.target_frac = target_frac
+        self.ema = ema
+        self._lam = optimizer.gate_lambda
+
+    def due(self, step: int) -> bool:
+        return step > 0 and step % self.probe_every == 0
+
+    def update(self, fresh_loss: float, reseen_loss: float) -> float:
+        gap = float(fresh_loss) - float(reseen_loss)
+        target = self.target_frac * max(float(fresh_loss), 1e-12)
+        lam_now = min(max(gap / target, 0.0), 1.0)
+        self._lam = self.ema * self._lam + (1.0 - self.ema) * lam_now
+        self.opt.gate_lambda = self._lam
+        return self._lam
+
+    @property
+    def gate_lambda(self) -> float:
+        return self._lam

echomuon-0.1.0/src/echomuon.egg-info/PKG-INFO

@@ -0,0 +1,125 @@
+Metadata-Version: 2.4
+Name: echomuon
+Version: 0.1.0
+Summary: EchoMuon: Muon with a per-direction temporal trust gate and a memorization-gap controller. Better than scheduled Muon wherever data are imperfect.
+Author-email: Stamatis Mastromichalakis <stamatis@tmnetworks.gr>
+License: MIT
+Project-URL: Homepage, https://github.com/MStamatis/echomuon
+Project-URL: Paper, https://arxiv.org/abs/ARXIV-PLACEHOLDER
+Keywords: optimizer,muon,deep-learning,pytorch,label-noise,orthogonalization
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Dynamic: license-file
+
+# EchoMuon
+
+**Muon with a per-direction temporal trust gate and a memorization-gap controller.
+Better than scheduled Muon wherever data are imperfect; ties it everywhere else.**
+
+[Muon](https://kellerjordan.github.io/posts/muon/) orthogonalizes the momentum of
+2-D hidden weight matrices, giving every singular direction of the update exactly
+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 no 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;
+- is **median-normalized per layer** (mean ≈ 1): trust is reallocated across
+  directions at constant total step, so the gate cannot act as a disguised
+  learning-rate schedule;
+- is engaged in proportion to a **measured memorization gap** λ ∈ [0, 1]: the loss
+  on fresh batches minus the loss on batches seen a few hundred steps ago, under
+  the same current weights. At λ=0 EchoMuon *is* Muon, structurally — never worse
+  by construction on clean data.
+
+Headline results (paired seeds, schedule parity, per-arm lr sweeps): **+0.9 to
++1.8pp** over scheduled Muon on six vision cells (CIFAR-10/100, Tiny ImageNet,
+clean and 20% label noise), **−0.031 nats (t=−6.3)** on a LLaMA-style 162M
+transformer on FineWeb-Edu, Muon-tier quality in 80–85% of Muon's steps in every
+seed, ~13% step overhead at 162M with the fast profile. See the paper for the
+boundaries (byte-level LMs, SSMs, strong augmentation recipes) — they are reported,
+measured, and part of the result.
+
+## Install
+
+```bash
+pip install echomuon
+```
+
+## Usage
+
+EchoMuon handles only the 2-D hidden matrices, exactly like Muon; route
+embeddings, heads, gains and biases to AdamW:
+
+```python
+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 lr as you would for Muon
+aux = torch.optim.AdamW(others, lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)
+```
+
+Training loop with the memorization-gap controller (optional but recommended —
+without it, set `gate_lambda` yourself; `gate_lambda=0` is plain Muon):
+
+```python
+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])   # oldest ~400-800 steps ago
+            fresh  = mean_loss(model, take_fresh(4))    # 4 held-back fresh batches
+        ctl.update(fresh_loss=fresh, reseen_loss=reseen)
+```
+
+All constants (slow β=0.99, floor 0.1, median reference, the 2% normalizer, the
+0.7 EMA) were fixed once and used unchanged in every experiment of the paper —
+the only knob you tune is Muon's own learning rate.
+
+## When to use it
+
+| Your setting | Recommendation |
+|---|---|
+| Web-scale corpora, label noise, ambiguous labels, light augmentation | **EchoMuon** — this is where the margins live |
+| Clean data / strong augmentation (RandAugment + mixup) | Tie with Muon; λ backs the gate off automatically |
+| Byte-level LMs at scale | Scheduled Muon (measured boundary; EchoMuon concedes ≤0.5%) |
+| Mamba-style SSMs | AdamW beats the whole Muon family there (measured boundary) |
+
+## Citation
+
+```bibtex
+@article{mastromichalakis2026echomuon,
+  title  = {EchoMuon: Better Than Scheduled Muon Wherever Data Are Imperfect},
+  author = {Mastromichalakis, Stamatis},
+  year   = {2026},
+  note   = {arXiv preprint}
+}
+```
+
+## License
+
+MIT

echomuon-0.1.0/src/echomuon.egg-info/SOURCES.txt

@@ -0,0 +1,11 @@
+LICENSE
+README.md
+pyproject.toml
+src/echomuon/__init__.py
+src/echomuon/optimizer.py
+src/echomuon.egg-info/PKG-INFO
+src/echomuon.egg-info/SOURCES.txt
+src/echomuon.egg-info/dependency_links.txt
+src/echomuon.egg-info/requires.txt
+src/echomuon.egg-info/top_level.txt
+tests/test_echomuon.py

echomuon-0.1.0/src/echomuon.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

echomuon-0.1.0/src/echomuon.egg-info/requires.txt

@@ -0,0 +1 @@
+torch>=2.0

echomuon-0.1.0/src/echomuon.egg-info/top_level.txt

@@ -0,0 +1 @@
+echomuon

echomuon-0.1.0/tests/test_echomuon.py

@@ -0,0 +1,133 @@
+"""CPU tests for the echomuon package. Run: python tests/test_echomuon.py (or pytest)."""
+import copy
+import sys
+
+import torch
+
+from echomuon import EchoMuon, MemorizationGapController
+
+
+def make_model(seed=0):
+    torch.manual_seed(seed)
+    return torch.nn.Sequential(
+        torch.nn.Linear(32, 64, bias=False),
+        torch.nn.ReLU(),
+        torch.nn.Linear(64, 10, bias=False),
+    )
+
+
+def data(seed=1, n=60):
+    g = torch.Generator().manual_seed(seed)
+    for _ in range(n):
+        x = torch.randn(16, 32, generator=g)
+        y = torch.randint(0, 10, (16,), generator=g)
+        yield x, y
+
+
+def run(model, opt, n=60, lam=None):
+    losses = []
+    for x, y in data(n=n):
+        loss = torch.nn.functional.cross_entropy(model(x), y)
+        opt.zero_grad()
+        loss.backward()
+        if lam is not None:
+            opt.gate_lambda = lam
+        opt.step()
+        losses.append(float(loss.detach()))
+    return losses
+
+
+def flat(model):
+    return torch.cat([p.detach().flatten() for p in model.parameters()])
+
+
+def test_trains():
+    # a memorizable task: cycle over the same 4 fixed batches
+    m = make_model()
+    opt = EchoMuon(m.parameters(), lr=0.02, gate_every=10)
+    batches = list(data(n=4))
+    losses = []
+    for i in range(80):
+        x, y = batches[i % 4]
+        loss = torch.nn.functional.cross_entropy(m(x), y)
+        opt.zero_grad(); loss.backward(); opt.step()
+        losses.append(float(loss.detach()))
+    assert all(torch.isfinite(torch.tensor(losses)))
+    assert sum(losses[-8:]) < sum(losses[:8]), "loss did not decrease"
+
+
+def test_lambda_zero_is_plain_muon():
+    # gate_lambda=0 must be bit-identical to never refreshing the gate at all
+    m0 = make_model()
+    o0 = EchoMuon(m0.parameters(), lr=0.02, gate_every=10, gate_lambda=0.0)
+    m1 = make_model()
+    o1 = EchoMuon(m1.parameters(), lr=0.02, gate_every=10 ** 9)  # gate never built
+    run(m0, o0, n=30)
+    run(m1, o1, n=30)
+    assert torch.allclose(flat(m0), flat(m1), atol=0), "lambda=0 != plain Muon"
+
+
+def test_gate_changes_trajectory():
+    m0 = make_model()
+    o0 = EchoMuon(m0.parameters(), lr=0.02, gate_every=10, gate_lambda=1.0)
+    m1 = make_model()
+    o1 = EchoMuon(m1.parameters(), lr=0.02, gate_every=10, gate_lambda=0.0)
+    run(m0, o0, n=30)
+    run(m1, o1, n=30)
+    assert not torch.allclose(flat(m0), flat(m1)), "gate had no effect"
+
+
+def test_rejects_non_2d():
+    lin = torch.nn.Linear(8, 8, bias=True)  # bias is 1-D
+    try:
+        EchoMuon(lin.parameters(), lr=0.02)
+    except ValueError:
+        return
+    raise AssertionError("non-2D parameter was accepted")
+
+
+def test_state_dict_roundtrip():
+    m = make_model()
+    opt = EchoMuon(m.parameters(), lr=0.02, gate_every=10)
+    run(m, opt, n=25)
+    sd_model = copy.deepcopy(m.state_dict())
+    sd_opt = copy.deepcopy(opt.state_dict())
+    # continue 10 more steps -> reference
+    ref = copy.deepcopy(m)
+    ref_opt = EchoMuon(ref.parameters(), lr=0.02, gate_every=10)
+    ref.load_state_dict(sd_model)
+    ref_opt.load_state_dict(copy.deepcopy(sd_opt))
+    # note: generators differ, so drive both with the SAME fresh data
+    g = torch.Generator().manual_seed(99)
+    batches = [(torch.randn(16, 32, generator=g), torch.randint(0, 10, (16,), generator=g))
+               for _ in range(10)]
+    for x, y in batches:
+        loss = torch.nn.functional.cross_entropy(m(x), y)
+        opt.zero_grad(); loss.backward(); opt.step()
+    m2 = ref
+    for x, y in batches:
+        loss = torch.nn.functional.cross_entropy(m2(x), y)
+        ref_opt.zero_grad(); loss.backward(); ref_opt.step()
+    assert torch.allclose(flat(m), flat(m2), atol=1e-6), "resume diverged"
+
+
+def test_controller():
+    m = make_model()
+    opt = EchoMuon(m.parameters(), lr=0.02, gate_every=10)
+    ctl = MemorizationGapController(opt, probe_every=5, ema=0.0)
+    assert not ctl.due(3) and ctl.due(5)
+    lam = ctl.update(fresh_loss=2.0, reseen_loss=2.0)   # no gap -> 0
+    assert lam == 0.0 and opt.gate_lambda == 0.0
+    lam = ctl.update(fresh_loss=2.0, reseen_loss=1.0)   # huge gap -> clip 1
+    assert lam == 1.0 and opt.gate_lambda == 1.0
+    lam = ctl.update(fresh_loss=2.0, reseen_loss=2.0 - 0.02)  # gap = 0.5*target
+    assert abs(lam - 0.5) < 1e-6
+
+
+if __name__ == "__main__":
+    fns = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
+    for fn in fns:
+        fn()
+        print(f"  {fn.__name__} ... ok")
+    print(f"ALL {len(fns)} TESTS PASS")
+    sys.exit(0)