Replace hardcoded NS coefficients with analytically optimal ones [skip-build]

Replaces the five fixed (a, b, c) quintic coefficients in
`_zeropower_via_newtonschulz5` with per-step optimal coefficients computed by
`_optimal_composition` / `_optimal_quintic` via the equioscillation (Chebyshev)
criterion. The precomputed `_coeffs_list` (10 iters, l=1e-3, safety_factor_eps=1e-2,
cushion=0.02) minimizes the worst-case deviation from 1 over the singular-value
interval, and the function now supports arbitrary `steps` by repeating the last
coefficient set. Docstrings updated accordingly.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

torch-ext/optimizer/newton_schulz.py

@@ -1,26 +1,106 @@
 import torch
+from itertools import repeat
+from math import inf, sqrt
+import numpy as np
 
 from .matmul_transpose_triton import matmul_transpose_assign
+from loguru import logger
 
 COMM_DTYPE = torch.bfloat16
 DEFAULT_CHUNK_SIZE_RATIO = 4
 
+def _optimal_quintic(l, u):
+    """
+    Use the simplified Remez algorithm to find the optimal odd quintic approximant
+    to the constant function x -> 1 over the interval [l, u].
+
+    Returns (a, b, c) for p(x) = ax + bx^3 + cx^5 that minimizes the maximum
+    approximation error max_{x in [l,u]} |p(x) - 1|. Iterates by updating the
+    two interior equioscillation nodes q, r until convergence. Returns the
+    closed-form equioscillating solution when l ≈ u.
+    """
+    assert 0 <= l <= u
+    if 1 - 5e-6 <= l / u:
+        return (15/8)/u, (-10/8)/(u**3), (3/8)/(u**5)
+    q = (3*l + u) / 4
+    r = (l + 3*u) / 4
+    E, old_E = inf, None
+    while not old_E or abs(old_E - E) > 1e-15:
+        old_E = E
+        LHS = np.array([
+            [l, l**3, l**5, 1],
+            [q, q**3, q**5, -1],
+            [r, r**3, r**5, 1],
+            [u, u**3, u**5, -1],
+        ])
+        a, b, c, E = np.linalg.solve(LHS, np.ones(4))
+        q, r = np.sqrt((-3*b + np.array([-1, 1]) *
+                        sqrt(9*b**2 - 20*a*c)) / (10*c))
+    return float(a), float(b), float(c)
 
-# This code snippet is a modified version adapted from the following GitHub repositories:
-# https://github.com/KellerJordan/Muon/blob/master/muon.py
-# Muon's Newton–Schulz iteration causes high variance in singular values
-# Idea: give each iteration its own 3 coefficients and optimize them via gradient descent.
+
+def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0):
+    """
+    Compute the Polar Express coefficient series for `num_iters` quintic iterations.
+
+    Builds a sequence of per-step optimal odd quintic coefficients (a, b, c) that
+    compose to map singular values from [l, 1] toward 1. At each step:
+      1. Solves `_optimal_quintic` on [max(l, cushion*u), u]. The `cushion`
+         prevents near-zero singular values from stalling by raising the effective
+         lower bound; if it is active (cushion*u > l), the coefficients are
+         rescaled so that p(l) and p(u) are centered around 1 w.r.t. the true [l, u].
+      2. Deflates the coefficients by (1 + safety_factor_eps)^degree for all but the
+         last iteration, providing numerical headroom at the cost of a slightly slower
+         final convergence step.
+      3. Advances the interval: l <- p(l), u <- 2 - p(l) (by symmetry of p around 1).
+
+    Returns a list of (a, b, c) tuples, one per iteration.
+
+    Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and
+    Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932
+    """
+    u = 1
+    assert 0 <= l <= u
+    safety_factor = 1 + safety_factor_eps
+    coefficients = []
+    for iter in range(num_iters):
+        a, b, c = _optimal_quintic(max(l, cushion*u), u)
+        if cushion*u > l:
+            pl = a*l + b*l**3 + c*l**5
+            pu = a*u + b*u**3 + c*u**5
+            rescaler = 2/(pl + pu)
+            a *= rescaler; b *= rescaler; c *= rescaler
+        if iter < num_iters - 1:
+            a /= safety_factor; b /= safety_factor**3; c /= safety_factor**5
+        coefficients.append((a, b, c))
+        l = a*l + b*l**3 + c*l**5
+        u = 2 - l
+    return coefficients
+
+
+_coeffs_list = _optimal_composition(l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02)
+
+# This code is adapted from:
+#   KellerJordan/Muon (https://github.com/KellerJordan/Muon/blob/master/muon.py)
+#   NoahAmsel/PolarExpress (https://github.com/NoahAmsel/PolarExpress)
+#   matmul_transpose_assign kernel from nil0x9/flash-muon (https://github.com/nil0x9/flash-muon)
 @torch.no_grad()
-# matmul_transpose_assign from : https://github.com/nil0x9/flash-muon
 def _zeropower_via_newtonschulz5(G, steps):
     """
-    Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
-    quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
-    of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
-    zero even beyond the point where the iteration no longer converges all the way to one everywhere
-    on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
-    where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
-    performance at all relative to UV^T, where USV^T = G is the SVD.
+    Compute the polar factor of G via the Polar Express method.
+
+    Applies `steps` quintic iterations X <- aX + bX^3 + cX^5, where (a, b, c)
+    are the Polar Express coefficients from `_coeffs_list`. Each step is the
+    optimal odd quintic approximant to x -> 1 over the current singular-value
+    interval, minimizing the maximum approximation error (Remez / minimax criterion).
+    The composition maps singular values from [l, 1] to near 1, producing the
+    polar factor (orthogonal factor in the polar decomposition G = UP).
+
+    `_coeffs_list` is precomputed for 10 iterations (l=1e-3, safety_factor_eps=1e-2,
+    cushion=0.02). If `steps` exceeds 10, the final coefficient set is repeated.
+
+    Reference: Amsel et al., "The Polar Express: Optimal Matrix Sign Methods and
+    Their Application to the Muon Algorithm", https://arxiv.org/abs/2505.16932
     """
     assert len(G.shape) == 2
     assert G.dtype == COMM_DTYPE
@@ -28,18 +108,14 @@ def _zeropower_via_newtonschulz5(G, steps):
 
     if G.size(0) > G.size(1):
         X = X.T
-    # Ensure spectral norm is at most 1
+
     X = X / (X.norm() + 1e-7)
+    hs = _coeffs_list[:steps] + list(
+        repeat(_coeffs_list[-1], steps - len(_coeffs_list)))
     buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
     buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
     # Perform the NS iterations
-    for a, b, c in [
-        (4.0848, -6.8946, 2.9270),
-        (3.9505, -6.3029, 2.6377),
-        (3.7418, -5.5913, 2.3037),
-        (2.8769, -3.1427, 1.2046),
-        (2.8366, -3.0525, 1.2012),
-    ]:
+    for a, b, c in hs:
         matmul_transpose_assign(X, buf1)
         matmul_transpose_assign(buf1, buf2)
         buf1.mul_(b).add_(buf2, alpha=c)
@@ -47,4 +123,5 @@ def _zeropower_via_newtonschulz5(G, steps):
 
     if G.size(0) > G.size(1):
         X = X.T
+
     return X