magma-optimizer 0.1.0__tar.gz

magma_optimizer-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,63 @@
+name: Publish to PyPI
+
+on:
+  push:
+    tags:
+      - "v*"
+
+permissions:
+  contents: read
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install -e ".[dev]"
+      - name: Run tests
+        run: python -m pytest tests/ -v
+
+  build:
+    needs: test
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - name: Install build tools
+        run: python -m pip install --upgrade pip build
+      - name: Build sdist and wheel
+        run: python -m build
+      - name: Verify version matches tag
+        run: |
+          TAG="${GITHUB_REF#refs/tags/v}"
+          PKG_VERSION=$(python -c "import re; print(re.search(r'__version__ = \"(.*?)\"', open('magma/__init__.py').read()).group(1))")
+          if [ "$TAG" != "$PKG_VERSION" ]; then
+            echo "ERROR: Tag v$TAG does not match package version $PKG_VERSION"
+            exit 1
+          fi
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish:
+    needs: build
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

magma_optimizer-0.1.0/.gitignore

@@ -0,0 +1,17 @@
+__pycache__/
+*.py[cod]
+*$py.class
+*.egg-info/
+dist/
+build/
+*.egg
+.eggs/
+*.so
+.venv/
+venv/
+env/
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.vscode
+data/

magma_optimizer-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Andrij David
+
+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.

magma_optimizer-0.1.0/PKG-INFO

@@ -0,0 +1,17 @@
+Metadata-Version: 2.4
+Name: magma-optimizer
+Version: 0.1.0
+Summary: Momentum-Aligned Gradient Masking — block-wise stochastic masking wrapper for PyTorch optimizers
+Project-URL: Homepage, https://github.com/andrijdavid/magma-optimizer
+Project-URL: Repository, https://github.com/andrijdavid/magma-optimizer
+License: MIT
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+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
+Requires-Dist: torch>=2.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == 'dev'

magma_optimizer-0.1.0/README.md

@@ -0,0 +1,78 @@
+# Magma
+
+**Momentum-Aligned Gradient Masking for Adaptive Optimizers**
+
+Magma is a lightweight wrapper that applies block-wise stochastic masking to any PyTorch optimizer, modulated by the alignment between gradient momentum and the current gradient. It is an implementation of the algorithm described in *"On Surprising Effectiveness of Masking Updates in Adaptive Optimizers"*[(arXiv 2602.15322)](https://arxiv.org/pdf/2602.15322).
+
+The core insight is deceptively simple. At each step, a per-parameter Bernoulli coin flip decides whether to keep or discard the update. Updates that survive are further scaled by a smoothed cosine similarity score between the gradient and its exponential moving average. The base optimizer's internal states i.e Adam's running means or RMSProp's squared gradients are always updated. Only the parameter itself is masked.
+
+This acts as a form of implicit regularization, particularly effective under the heterogeneous curvature and heavy-tailed gradient noise characteristic of transformer training.
+
+## Installation
+
+```bash
+pip install magma-optimizer
+```
+
+Or directly from source:
+
+```bash
+pip install git+https://github.com/andrijdavid/magma-optimizer.git
+```
+
+## Usage
+
+Magma wraps any instantiated PyTorch optimizer. The interface mirrors what you already know.
+
+```python
+from magma import Magma
+import torch
+
+model = ...  # your model
+base = torch.optim.Adam(model.parameters(), lr=1e-3)
+
+optimizer = Magma(
+    base,
+    mask_prob=0.5,        # prob of keeping an update
+    tau=2.0,              # temperature for the alignment sigmoid
+    momentum_beta=0.9,    # EMA coefficient for gradient momentum
+    alignment_ema=0.9,    # EMA coefficient for smoothing the alignment score
+    exclude=set(model.embed.parameters()),  # skip masking on embeddings
+)
+
+for x, y in dataloader:
+    optimizer.zero_grad()
+    loss = criterion(model(x), y)
+    loss.backward()
+    optimizer.step()
+```
+
+The `exclude` parameter accepts a set of tensors that should bypass masking entirely. The paper recommends excluding embedding layers, as their update dynamics differ from attention and MLP blocks.
+
+## Algorithm
+
+The procedure, applied at each step for each non-excluded parameter:
+
+1. Update momentum EMA: `μ = β·μ + (1−β)·g`
+2. Compute alignment: `s̃ = sigmoid(cosine_similarity(μ, g) / τ)`
+3. Smooth alignment: `s = 0.9·s_prev + 0.1·s̃`
+4. Run the base optimizer step (all internal states update normally)
+5. Sample mask: `m ~ Bernoulli(p)`
+6. Apply: `θ = (s·m)·θ_new + (1 − s·m)·θ_old`
+
+When the mask is zero, the parameter reverts to its pre-step value. When the mask is one, the update is scaled by the alignment score. The base optimizer sees every gradient regardless.
+
+## Citation
+
+```bibtex
+@article{joo2026magma,
+  title={On Surprising Effectiveness of Masking Updates in Adaptive Optimizers},
+  author={Joo, Taejong and Xia, Wenhan and Kim, Cheolmin and Zhang, Ming and Ie, Eugene},
+  journal={arXiv preprint arXiv:2602.15322},
+  year={2026}
+}
+```
+
+## License
+
+MIT

magma_optimizer-0.1.0/magma/__init__.py

@@ -0,0 +1,4 @@
+from magma.magma import Magma
+
+__version__ = "0.1.0"
+__all__ = ["Magma"]

magma_optimizer-0.1.0/magma/magma.py

@@ -0,0 +1,184 @@
+"""Magma: Momentum-Aligned Gradient Masking wrapper for PyTorch optimizers."""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+from torch.optim import Optimizer
+
+
+class Magma:
+    """
+    A Pytorch optimizer wraper with block-wise stochastic masking
+    modulated by momentum-gradient alignment.
+    As explained here https://arxiv.org/abs/2602.15322
+    """
+
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        mask_prob: float = 0.5,
+        tau: float = 2.0,
+        momentum_beta: float = 0.9,
+        alignment_ema: float = 0.9,
+        exclude: set[Tensor] | None = None,
+    ) -> None:
+        if not 0.0 <= mask_prob <= 1.0:
+            raise ValueError(f"mask_prob must be in [0, 1], got {mask_prob}")
+        if tau <= 0.0:
+            raise ValueError(f"tau must be positive, got {tau}")
+
+        self.optimizer = optimizer
+        self.mask_prob = mask_prob
+        self.tau = tau
+        self.momentum_beta = momentum_beta
+        self.alignment_ema = alignment_ema
+        self._exclude_ids: set[int] = {id(t) for t in (exclude or ())}
+
+        # Magma-specific per-parameter state keyed by param id for momentum and alignment
+        self._state: dict[int, dict[str, Tensor]] = {}
+
+    def step(self, closure=None):
+        """Perform a single Magma-wrapped optimisation step."""
+
+        # --- Phase 1: pre-step (momentum, alignment, snapshot) --------
+        snapshots: dict[int, Tensor] = {}  # id(p) -> p.data clone
+        alignment_scores: dict[int, float] = {}
+
+        for group in self.optimizer.param_groups:
+            for p in group["params"]:
+                if p.grad is None or id(p) in self._exclude_ids:
+                    continue
+
+                pid = id(p)
+                grad = p.grad.detach()
+
+                # Lazily initialise per-param Magma state
+                if pid not in self._state:
+                    self._state[pid] = {
+                        "momentum": torch.zeros_like(p.data),
+                        "alignment": torch.tensor(1.0, device=p.device),
+                    }
+
+                st = self._state[pid]
+
+                #  μ_t = β μ_{t-1} + (1-β) g_t
+                st["momentum"].mul_(self.momentum_beta).add_(
+                    grad, alpha=1.0 - self.momentum_beta
+                )
+
+                #  cossim(μ_t, g_t)
+                cos = torch.nn.functional.cosine_similarity(
+                    st["momentum"].flatten().unsqueeze(0),
+                    grad.flatten().unsqueeze(0),
+                ).item()
+                
+                # sigmoid(cos/τ) 
+                s_tilde = torch.sigmoid(
+                    torch.tensor(cos / self.tau, device=p.device)
+                ).item()
+
+                # EMA smoothing
+                s_prev = st["alignment"].item()
+                s = self.alignment_ema * s_prev + (1.0 - self.alignment_ema) * s_tilde
+                st["alignment"].fill_(s)
+
+                alignment_scores[pid] = s
+
+                # Save current params
+                snapshots[pid] = p.data.clone()
+
+        loss = self.optimizer.step(closure)
+
+        for group in self.optimizer.param_groups:
+            for p in group["params"]:
+                pid = id(p)
+                if pid not in snapshots:
+                    continue
+
+                # m_t ~ Bernoulli(mask_prob)
+                mask = float(torch.bernoulli(torch.tensor(self.mask_prob)).item())
+                s = alignment_scores[pid]
+                # blend = s_t * m_t
+                blend = s * mask
+
+                if blend == 0.0:
+                    # Fully revert parameter to pre-step value
+                    p.data.copy_(snapshots[pid])
+                elif blend != 1.0:
+                    # θ = blend*θ_new + (1-blend)*θ_old
+                    p.data.mul_(blend).add_(snapshots[pid], alpha=1.0 - blend)
+                # blend == 1.0, keep as it is
+
+        return loss
+
+    
+    def zero_grad(self, set_to_none: bool = True):
+        self.optimizer.zero_grad(set_to_none=set_to_none)
+
+    @property
+    def param_groups(self):
+        return self.optimizer.param_groups
+
+    def add_param_group(self, param_group: dict):
+        self.optimizer.add_param_group(param_group)
+
+    def state_dict(self):
+        id_to_key: dict[int, tuple[int, int]] = {}
+        for gi, group in enumerate(self.optimizer.param_groups):
+            for pi, p in enumerate(group["params"]):
+                id_to_key[id(p)] = (gi, pi)
+
+        magma_state = {}
+        for pid, st in self._state.items():
+            key = id_to_key.get(pid)
+            if key is not None:
+                magma_state[key] = {k: v.clone() for k, v in st.items()}
+
+        return {
+            "base": self.optimizer.state_dict(),
+            "magma_state": magma_state,
+            "mask_prob": self.mask_prob,
+            "tau": self.tau,
+            "momentum_beta": self.momentum_beta,
+            "alignment_ema": self.alignment_ema,
+        }
+
+    def load_state_dict(self, state_dict: dict):
+        self.optimizer.load_state_dict(state_dict["base"])
+        self.mask_prob = state_dict["mask_prob"]
+        self.tau = state_dict["tau"]
+        self.momentum_beta = state_dict["momentum_beta"]
+        self.alignment_ema = state_dict["alignment_ema"]
+
+        key_to_id: dict[tuple[int, int], int] = {}
+        for gi, group in enumerate(self.optimizer.param_groups):
+            for pi, p in enumerate(group["params"]):
+                key_to_id[(gi, pi)] = id(p)
+
+        self._state = {}
+        for key, st in state_dict["magma_state"].items():
+            key = tuple(key) if isinstance(key, list) else key
+            pid = key_to_id.get(key)
+            if pid is not None:
+                self._state[pid] = {k: v.clone() for k, v in st.items()}
+
+    def __getattr__(self, name: str):
+        # Fallback to base optimizer for anything not on Magma itself
+        try:
+            return getattr(self.optimizer, name)
+        except AttributeError:
+            raise AttributeError(
+                f"Neither 'Magma' nor the base optimizer have attribute '{name}'"
+            )
+
+    def __repr__(self) -> str:
+        return (
+            f"Magma(\n"
+            f"  mask_prob={self.mask_prob},\n"
+            f"  tau={self.tau},\n"
+            f"  momentum_beta={self.momentum_beta},\n"
+            f"  alignment_ema={self.alignment_ema},\n"
+            f"  base={self.optimizer}\n"
+            f")"
+        )