Demo: When Does Newton–Schulz Reach the Polar Factor?

Muon approximates a polar factor with repeated matrix multiplications. Each active singular value follows \(p(x)=(3x-x^3)/2\) toward one.

Question 1: Which Singular Direction Should Be Slowest?

We construct two matrices with the same singular vectors and largest singular value. One has all singular values at least \(0.4\). The other has a smallest singular value of \(0.01\).

import numpy as np
import matplotlib.pyplot as plt

plt.style.use("../figure_style.mplstyle")
rng = np.random.default_rng(18)
size = 12
U, _ = np.linalg.qr(rng.standard_normal((size, size)))
V, _ = np.linalg.qr(rng.standard_normal((size, size)))

spectra = {
    "Well conditioned": np.linspace(1.0, 0.4, size),
    "Poorly conditioned": np.geomspace(1.0, 0.01, size),
}

Answer. The poorly conditioned matrix begins with one direction much farther from one. That weakest direction should control how many iterations are needed.

Question 2: How Does Conditioning Change Convergence?

We apply the same Newton–Schulz update to both matrices and track the largest singular-value error. A second view shows every singular value after five steps.

steps = 12
singular_values = {}
errors = {}
for name, spectrum in spectra.items():
    X = U @ np.diag(spectrum) @ V.T
    values = [np.linalg.svd(X, compute_uv=False)]
    for step in range(steps):
        X = 1.5 * X - 0.5 * X @ X.T @ X
        values.append(np.linalg.svd(X, compute_uv=False))
    singular_values[name] = np.array(values)
    errors[name] = np.max(np.abs(singular_values[name] - 1), axis=1)

The left panel tracks the largest error across singular directions. The right panel compares both matrices after the same five updates.

fig, axes = plt.subplots(1, 2, figsize=(8, 3.5), layout="constrained")
for name, marker in [("Well conditioned", "o"),
                             ("Poorly conditioned", "s")]:
    axes[0].plot(range(steps + 1), errors[name], marker=marker, linewidth=1.5, markersize=3, label=name)
axes[0].set_yscale("log")
axes[0].set_xlabel("Newton–Schulz Step")
axes[0].set_ylabel("Largest Singular-Value Error")
axes[0].legend(frameon=False)

shown_step = 5
positions = np.arange(size)
width = 0.36
for offset, name in [(-width / 2, "Well conditioned"),
                            (width / 2, "Poorly conditioned")]:
    after = singular_values[name][shown_step]
    axes[1].bar(positions + offset, after, width=width, label=name)
axes[1].axhline(1.0, linewidth=1, linestyle=":")
axes[1].set_xlabel("Singular Direction")
axes[1].set_ylabel(f"Singular Value after {shown_step} steps")
axes[1].set_ylim(0, 1.35)
axes[1].legend(frameon=False, loc="upper center")
plt.show()

for name in spectra:
    print(f"{name}: error after {shown_step} steps = {errors[name][shown_step]:.3e}")

Well conditioned: error after 5 steps = 1.258e-04
Poorly conditioned: error after 5 steps = 9.242e-01

Answer. Singular values already near one converge quickly, while the smallest nonzero value needs more iterations to catch up. A fixed iteration count is therefore an approximation whose quality depends on conditioning. An exactly zero singular value would remain zero.