Merge branch 'main' into skyw/stacked-soap

Author: skyw

docs/apidocs/orthogonalized-optimizers.md

@@ -46,6 +46,8 @@ emerging_optimizers.orthogonalized_optimizers
 .. autoclass:: PolarGrad
     :members:
 
+.. autofunction:: right_polargrad_orth_fn
+
 
 :hidden:`AdaptiveMuon`
 ~~~~~~~~~~~~~~~~~~~~~~~

emerging_optimizers/contrib/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

emerging_optimizers/contrib/muown.py

@@ -0,0 +1,200 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING, Callable, override
+
+
+if TYPE_CHECKING:
+    from typing import overload
+
+import torch
+from torch.optim.optimizer import ParamsT
+
+from emerging_optimizers import utils
+from emerging_optimizers.orthogonalized_optimizers.muon import Muon, MuonScaleT
+from emerging_optimizers.orthogonalized_optimizers.muon_utils import NSCoeffT
+from emerging_optimizers.scalar_optimizers import update_functions
+from emerging_optimizers.utils import FP32MatmulPrecT
+
+
+__all__ = ["Muown"]
+
+
+@torch.compile
+def _weight_norm_decompose(
+    weight: torch.Tensor,
+    grad: torch.Tensor,
+    g: torch.Tensor,
+    v_norm: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    r"""Reconstructs the direction and splits the gradient under the weight-norm reparameterization.
+
+    Args:
+        weight: The current 2D weight ``W``.
+        grad: The gradient ``grad_W`` with respect to ``W``.
+        g: Per-row magnitude, shape ``[rows, 1]``.
+        v_norm: Cached row norms ``||v||_row`` of the direction, shape ``[rows, 1]``.
+
+    Returns:
+        ``(v, grad_g, grad_v)``: the reconstructed direction and the magnitude and direction gradients.
+    """
+    u = weight / g
+    v = u * v_norm
+    grad_g = (grad * u).sum(dim=1, keepdim=True)
+    grad_v = (g / v_norm) * (grad - u * grad_g)
+    return v, grad_g, grad_v
+
+
+class Muown(Muon):
+    """Muown: Muon with internal weight normalization (row-norm control).
+
+    Muown (Lion et al., *Muown: Row-Norm Control for Muon Optimization*, arXiv:2605.10797) is a drop-in
+    replacement for :class:`~emerging_optimizers.orthogonalized_optimizers.muon.Muon` that splits each 2D
+    weight into a per-row magnitude and a direction, then optimizes them under their natural geometries:
+
+
+    Args:
+        params: Iterable of parameters to optimize or dicts defining parameter groups.
+        lr: Learning rate shared by the direction and magnitude updates.
+        momentum: EMA momentum for the direction (Muon) update.
+        betas: Adam ``(beta1, beta2)`` for the magnitude update.
+        adam_eps: Adam epsilon for the magnitude update.
+        weight_decay: Decoupled weight decay coefficient, applied to the magnitude ``g``.
+        fp32_matmul_prec: Precision for the orthogonalization GEMM operations.
+        coefficient_type: Newton-Schulz coefficient set (see :class:`Muon`).
+        num_ns_steps: Number of Newton-Schulz iteration steps.
+        scale_mode: Update scale mode (see :func:`~emerging_optimizers.orthogonalized_optimizers.muon.get_muon_scale_factor`).
+        extra_scale_factor: Extra scale on the direction update; ``0.2`` matches Adam's update RMS norm.
+        use_syrk: Whether to use the Triton SYRK kernel for Newton-Schulz.
+    """
+
+    def __init__(
+        self,
+        params: ParamsT,
+        lr: float = 3e-4,
+        momentum: float = 0.95,
+        weight_decay: float = 0.0,
+        *,
+        betas: tuple[float, float] = (0.9, 0.95),
+        adam_eps: float = 1e-8,
+        fp32_matmul_prec: FP32MatmulPrecT = "medium",
+        coefficient_type: NSCoeffT = "quintic",
+        num_ns_steps: int = 5,
+        scale_mode: MuonScaleT = "spectral",
+        extra_scale_factor: float = 1.0,
+        use_syrk: bool = False,
+    ) -> None:
+        self.betas = betas
+        self.adam_eps = adam_eps
+
+        super().__init__(
+            params,
+            lr,
+            momentum,
+            weight_decay,
+            nesterov=False,
+            weight_decay_method="decoupled",
+            fp32_matmul_prec=fp32_matmul_prec,
+            coefficient_type=coefficient_type,
+            num_ns_steps=num_ns_steps,
+            scale_mode=scale_mode,
+            extra_scale_factor=extra_scale_factor,
+            use_syrk=use_syrk,
+        )
+
+    @torch.no_grad()  # type: ignore[misc]
+    @override
+    def _init_group(self, group: dict, skip_non_grad_params: bool = True) -> None:
+        for p in group["params"]:
+            if skip_non_grad_params and p.grad is None:
+                continue  # pragma: no cover
+            if p.dim() != 2:
+                raise TypeError("Muown is only supported for 2D parameters")
+            state = self.state[p]
+            if len(state) == 0:
+                # Seed g, v from the current weight so Muown starts from the same point as Muon. Floor the
+                # row norm so an all-zero weight row does not give g=0, which would make u = weight / g a
+                # 0/0 NaN on the first step.
+                row_norm = p.norm(dim=1, keepdim=True).to(torch.float32).clamp_min(1e-12)
+                state["step"] = 0
+                state["g"] = row_norm.clone()
+                state["v_norm"] = row_norm.clone()
+                state["momentum_buffer"] = torch.zeros_like(p, dtype=torch.float32)
+                state["m_g"] = torch.zeros_like(row_norm)
+                state["v_g"] = torch.zeros_like(row_norm)
+
+    if TYPE_CHECKING:
+
+        @overload
+        def step(self, closure: None = ...) -> None: ...
+
+        @overload
+        def step(self, closure: Callable[[], float]) -> float: ...
+
+    @torch.no_grad()  # type: ignore[misc]
+    @override
+    def step(self, closure: Callable[[], float] | None = None) -> float | None:
+        """Performs a single optimization step.
+
+        Args:
+            closure: Unsupported; must be ``None``.
+        """
+        if closure is not None:
+            raise ValueError("closure is not supported")
+
+        for group in self.param_groups:
+            self._init_group(group)
+
+            lr = group["lr"]
+            momentum = group["momentum"]
+            weight_decay = group["weight_decay"]
+            for p in group["params"]:
+                if p.grad is None:
+                    continue  # pragma: no cover
+
+                state = self.state[p]
+                curr_iter_1_based = state["step"] + 1
+                g = state["g"]
+                v_norm = state["v_norm"]
+
+                v, grad_g, grad_v = _weight_norm_decompose(p.to(torch.float32), p.grad.to(torch.float32), g, v_norm)
+
+                state["momentum_buffer"].lerp_(grad_v, 1 - momentum)
+                with utils.fp32_matmul_precision(self.fp32_matmul_prec):
+                    direction_update = self.scaled_orthogonalize_fn(state["momentum_buffer"])
+                v_new = v.add(direction_update, alpha=-lr)
+
+                magnitude_update = update_functions.calculate_adam_update(
+                    grad_g,
+                    state["m_g"],
+                    state["v_g"],
+                    betas=self.betas,
+                    eps=self.adam_eps,
+                    correct_bias=True,
+                    nesterov=False,
+                    step=curr_iter_1_based,  # 1-based iteration index is used for bias correction
+                )
+                g.add_(magnitude_update, alpha=-lr)
+
+                # Decoupled weight decay on the magnitude (the spectral-norm driver).
+                self._apply_weight_decay_inplace(g, grad_g, lr, weight_decay)
+
+                v_norm_new = v_new.norm(dim=1, keepdim=True)
+                p.copy_(g * (v_new / v_norm_new))
+                state["v_norm"] = v_norm_new
+
+                state["step"] += 1
+
+        return None

emerging_optimizers/orthogonalized_optimizers/polargrad.py

@@ -23,9 +23,10 @@
 from emerging_optimizers.orthogonalized_optimizers.muon_utils import NSCoeffT
 from emerging_optimizers.orthogonalized_optimizers.orthogonalized_optimizer import OrthogonalizedOptimizer, _args_doc
 from emerging_optimizers.utils import FP32MatmulPrecT
+from emerging_optimizers.utils.eig import eigh_with_fallback
 
 
-__all__ = ["PolarGrad"]
+__all__ = ["PolarGrad", "right_polargrad_orth_fn"]
 
 
 @registry.register_optimizer("polargrad")
@@ -103,3 +104,75 @@ def scaled_orthogonalize_fn(grad: torch.Tensor) -> torch.Tensor:
 
 
 PolarGrad.__doc__ = PolarGrad.__doc__.format(_args_doc=_args_doc)  # type: ignore[union-attr]
+
+
+def right_polargrad_orth_fn(
+    grad: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    center_rows: bool = False,
+    eps: float = 1e-15,
+    extra_scale_factor: float = 1.0,
+) -> torch.Tensor:
+    r"""Right-spectral (one-sided polar) orthogonalization for tall matrices.
+
+    Orthogonalizes only the right factor of a tall matrix ``G`` (e.g. an embedding or LM-head weight,
+    ``vocab x hidden``):
+
+    .. math::
+        u = G \, (G^\top G)^{-1/2}, \qquad \text{update} = \lVert G \rVert_*^{\,\alpha} \, u
+
+    .. code-block:: python
+       :caption: Define a ``RightPolarGrad`` by partially applying this orthogonalization
+
+       class RightPolarGrad(OrthogonalizedOptimizer):
+           def __init__(
+               self,
+               params,
+               lr: float = 3e-4,
+               momentum: float = 0.95,
+               weight_decay: float = 0.01,
+               *,
+               ...
+               alpha: float = 1.0,
+               center_rows: bool = False,
+               eps: float = 1e-15,
+               extra_scale_factor: float = 1.0,
+           ) -> None:
+               scaled_orthogonalize_fn = functools.partial(
+                   right_polargrad_orth_fn,
+                   alpha=alpha,
+                   center_rows=center_rows,
+                   eps=eps,
+                   extra_scale_factor=extra_scale_factor,
+               )
+               super().__init__(
+                   ...
+               )
+
+    Args:
+        grad: The (momentum) tensor to orthogonalize.
+        alpha: Exponent applied to the nuclear-norm scale factor.
+        center_rows: If True, subtract the per-column mean (the average over the row / vocabulary axis,
+            ``dim=0``) before and after the update, so each column is zero-mean.
+        eps: Floor on the right-Gram eigenvalues for the inverse sqrt and nuclear-norm computation.
+        extra_scale_factor: Extra multiplier on the update.
+
+    Returns:
+        The scaled right-polar update, same shape and dtype as ``grad``.
+    """
+    m = grad.to(torch.float32)
+    if center_rows:
+        m = m - m.mean(dim=0, keepdim=True)
+
+    eigvals, eigvecs = eigh_with_fallback(m.transpose(-1, -2) @ m)
+    eigvals.clamp_min_(eps)
+    right_gram_inv_sqrt = (eigvecs * eigvals.rsqrt().unsqueeze(-2)) @ eigvecs.transpose(-1, -2)
+
+    u = m @ right_gram_inv_sqrt
+    nuclear_norm = eigvals.sqrt().sum()
+    update = u * nuclear_norm.pow(alpha) * extra_scale_factor
+
+    if center_rows:
+        update = update - update.mean(dim=0, keepdim=True)
+    return update.to(grad.dtype)