improve tests

Signed-off-by: Hao Wu <skyw@nvidia.com>

Author: skyw

emerging_optimizers/orthogonalized_optimizers/muown.py

@@ -40,37 +40,6 @@ class Muown(Muon):
     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:
 
-    - **Direction** ``v``: the Muon update (EMA momentum + Newton-Schulz orthogonalization), reusing the
-      parent's ``scaled_orthogonalize_fn``.
-    - **Magnitude** ``g`` (one entry per output row): Adam, which realizes the :math:`\\ell_\\infty` duality
-      map for the diagonal per-neuron gain.
-
-    The reparameterization ``W = Diag(g / ||v||_row) v`` is held implicitly inside the optimizer, so the
-    forward pass is unchanged. At init ``g``, ``v`` are seeded from ``W`` so Muown starts from the same
-    point as Muon. The row magnitude is the empirical driver of spectral-norm drift under Muon; making it an
-    explicit, separately-optimized variable controls that drift without the indiscriminate shrinkage of
-    plain weight decay.
-
-    A single ``lr`` drives both halves: the direction step carries the ``0.2 * sqrt(max(m, n))`` scaling
-    (via ``extra_scale_factor``) that matches Adam's update RMS norm, so no separate magnitude lr is needed.
-
-    State per parameter:
-      - ``step``
-      - ``g``: per-row magnitude, shape ``[out, 1]``.
-      - ``v_norm``: cached row norms ``||v||_row`` of the direction, shape ``[out, 1]``.
-      - ``momentum_buffer``: EMA momentum of the direction gradient (the Muon momentum on ``v``).
-      - ``m_g``, ``v_g``: Adam first/second moments of the magnitude gradient.
-
-    Note:
-        Weight decay is decoupled and applied to the magnitude ``g`` (which is exactly the spectral-norm
-        driver). Because ``W`` is recomposed from ``g`` after the decay, the invariant ``||W_row|| == g``
-        holds without a separate resync.
-
-    Warning:
-        - This optimizer requires that all parameters passed in are 2D.
-        - It should not be used for the embedding layer, the final fully connected layer, or any 1-D
-          parameters; those can all be optimized by a standard method (e.g., AdamW).
-        - This optimizer is experimental and may change in future versions.
 
     Args:
         params: Iterable of parameters to optimize or dicts defining parameter groups.
@@ -182,20 +151,16 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                 grad = p.grad.to(torch.float32)
                 weight = p.to(torch.float32)
 
-                # Decompose grad_W into magnitude and direction gradients via the weight-norm Jacobian.
-                # u = v / ||v||_row is O(1) per element; reconstruct the direction v from the live weight.
                 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)
 
-                # Direction: Muon update on v (EMA momentum + Newton-Schulz), reusing the parent LMO.
                 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: Adam on g (the l-infinity duality map for the diagonal gain), same lr.
                 magnitude_update = update_functions.calculate_adam_update(
                     grad_g,
                     state["m_g"],
@@ -208,12 +173,8 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                 )
                 g.add_(magnitude_update, alpha=-lr)
 
-                # Decoupled weight decay on the magnitude (the spectral-norm driver).
-                if weight_decay != 0.0:
-                    g.add_(g, alpha=-weight_decay * lr)
+                g.add_(g, alpha=-weight_decay * lr)
 
-                # Recompose W = g * v_new / ||v_new||_row and refresh the cached direction norm. Recomposing
-                # from the decayed g keeps the invariant ||W_row|| == g without a separate resync.
                 v_norm_new = v_new.norm(dim=1, keepdim=True)
                 p.copy_(g * (v_new / v_norm_new))
                 state["v_norm"] = v_norm_new

tests/test_muown.py

@@ -31,9 +31,6 @@ def setUpModule() -> None:
         torch.manual_seed(FLAGS.seed)
         if torch.cuda.is_available():
             torch.cuda.manual_seed_all(FLAGS.seed)
-    # Use the most precise fp32 matmul path so the reference (which calls the optimizer's
-    # orthogonalization outside the fp32_matmul_precision context) computes bit-comparable results.
-    torch.set_float32_matmul_precision("highest")
 
 
 class MuownTest(parameterized.TestCase):
@@ -43,7 +40,6 @@ def setUp(self):
 
     @parameterized.product(shape=[(8, 16), (16, 8), (33, 65)])
     def test_smoke(self, shape):
-        """A few steps run and keep the weight finite."""
         p = torch.nn.Parameter(torch.randn(shape, device=self.device))
         opt = Muown([p], lr=1e-2, weight_decay=0.01)
         for _ in range(3):
@@ -52,7 +48,6 @@ def test_smoke(self, shape):
         self.assertTrue(torch.isfinite(p).all())
 
     def test_raises_on_non_2d(self):
-        """Muown supports 2D parameters only."""
         for shape in [(8,), (2, 3, 4)]:
             p = torch.nn.Parameter(torch.randn(shape, device=self.device))
             p.grad = torch.randn_like(p)
@@ -61,7 +56,6 @@ def test_raises_on_non_2d(self):
                 opt.step()
 
     def test_raises_on_closure(self):
-        """Closures are not supported."""
         p = torch.nn.Parameter(torch.randn((8, 16), device=self.device))
         p.grad = torch.randn_like(p)
         opt = Muown([p], lr=1e-2)
@@ -70,7 +64,6 @@ def test_raises_on_closure(self):
 
     @parameterized.product(shape=[(8, 16), (16, 8)], weight_decay=[0.0, 0.1])
     def test_row_norm_equals_magnitude_state(self, shape, weight_decay):
-        """The reparameterization invariant ||W_row|| == g holds after each step."""
         p = torch.nn.Parameter(torch.randn(shape, device=self.device))
         opt = Muown([p], lr=1e-2, weight_decay=weight_decay)
         for _ in range(4):
@@ -80,50 +73,28 @@ def test_row_norm_equals_magnitude_state(self, shape, weight_decay):
             torch.testing.assert_close(
                 row_norm,
                 opt.state[p]["g"],
-                atol=1e-5,
+                atol=0,
                 rtol=1e-5,
-                msg=lambda m: f"Row norms of W must equal the magnitude state g.\n\n{m}",
             )
 
-    def test_weight_decay_shrinks_magnitude(self):
-        """Decoupled weight decay shrinks the magnitude g relative to the no-decay run."""
-        p_wd = torch.nn.Parameter(torch.randn((16, 32), device=self.device))
-        p_no = torch.nn.Parameter(p_wd.detach().clone())
-        opt_wd = Muown([p_wd], lr=1e-2, weight_decay=0.1)
-        opt_no = Muown([p_no], lr=1e-2, weight_decay=0.0)
-        for _ in range(5):
-            grad = torch.randn_like(p_wd)
-            p_wd.grad = grad.clone()
-            p_no.grad = grad.clone()
-            opt_wd.step()
-            opt_no.step()
-        self.assertLess(opt_wd.state[p_wd]["g"].sum().item(), opt_no.state[p_no]["g"].sum().item())
-
     @parameterized.product(shape=[(8, 16), (16, 8), (33, 65)], momentum=[0.0, 0.95])
-    def test_agrees_with_reference(self, shape, momentum):
-        """Muown matches the authors' reference implementation (no weight decay).
-
-        The reference uses classic (heavy-ball) momentum while Muown uses EMA momentum; the two differ by
-        a constant factor that the scale-invariant Newton-Schulz orthogonalization removes, so the updates
-        agree. Both use the same injected orthogonalization, so only float rounding (lerp vs mul/add,
-        Newton-Schulz normalization) separates them — hence a tight tolerance rather than bit-identity.
-        """
+    def test_close_reference(self, shape, momentum):
         p = torch.nn.Parameter(torch.randn(shape, device=self.device))
         p_ref = torch.nn.Parameter(p.detach().clone())
 
         opt = Muown(
             [p],
-            lr=1e-2,
+            lr=0.125,
             momentum=momentum,
             weight_decay=0.0,
-            extra_scale_factor=0.2,
+            extra_scale_factor=0.25,
             fp32_matmul_prec="highest",
         )
         # Hold orthogonalization identical: feed the reference Muown's own scaled_orthogonalize_fn.
         opt_ref = MuownReference(
             [p_ref],
             orthogonalize_fn=opt.scaled_orthogonalize_fn,
-            lr=1e-2,
+            lr=0.125,
             momentum=momentum,
             nesterov=False,
             weight_decay=0.0,
@@ -139,16 +110,14 @@ def test_agrees_with_reference(self, shape, momentum):
         torch.testing.assert_close(
             p.detach(),
             p_ref.detach(),
-            atol=1e-5,
-            rtol=1e-4,
-            msg=lambda m: f"Muown weight diverged from the reference implementation.\n\n{m}",
+            atol=1e-6,
+            rtol=1e-5,
         )
         torch.testing.assert_close(
             opt.state[p]["g"],
             opt_ref.state[p_ref]["g"],
-            atol=1e-5,
-            rtol=1e-4,
-            msg=lambda m: f"Muown magnitude g diverged from the reference implementation.\n\n{m}",
+            atol=1e-7,
+            rtol=1e-5,
         )