echomuon 0.1.0__py3-none-any.whl

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/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.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,7 @@
+echomuon/__init__.py,sha256=5TsodsqwFiyrUuKNBVoivw1ahXvkOk3rzzILexuY5G8,361
+echomuon/optimizer.py,sha256=dRVTtQOl3Vv28r16CaqRwKNsrR7p9MKTPQmOiJ-bYjQ,9991
+echomuon-0.1.0.dist-info/licenses/LICENSE,sha256=g8bStZWYbDUnFBDn4xfMJTKECPPKPFnMYduSTpyFiXU,1082
+echomuon-0.1.0.dist-info/METADATA,sha256=3rLGzUmSK8BzzdaQLZmOA9bQBiX-dRNwKqGl02NGUyY,5096
+echomuon-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+echomuon-0.1.0.dist-info/top_level.txt,sha256=f1d36YUQTRbY7mg3oOXgxR5AxMXkE3FuzwUKL9_Bsxc,9
+echomuon-0.1.0.dist-info/RECORD,,

echomuon-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

echomuon-0.1.0.dist-info/licenses/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.dist-info/top_level.txt

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