Add Newton Schultz via Polar Express as a retraction for Iso optimizer

[kogolobo]

Summary

Contributes to #28.

This PR adds a Newton-Schulz matrix-sign retraction (retraction="newton_schulz") to the Iso (isospectral) optimizer. This allows to approximate the polar retraction without using torch.linalg.svd. This is an incremental step toward #28, exploring iterative matrix polynomial retractions for Stiefel factors.

Key Changes

• Retraction implementation: Added _newton_schulz_retraction leveraging muon_utils.newton_schulz (defaults to steps=8 with polar_express coefficients).
• Unit test suite: Extended tests/riemannian_optimizers/test_isospectral.py with parameterization for newton_schulz, validating singular value preservation and Stiefel factor orthogonality ($|Q^T Q - I|_\infty \le 10^{-5}$).
• Benchmark: Added becnmarks/benchmark_isospectral_retraction.py which uses triton.testing.do_bench to measure combined retraction time for both Stiefel factors $(U, V)$ across typical NN layer sizes.

Benchmark: Retraction Latency & Manifold Precision

Benchmarked on NVIDIA GeForce RTX 4090 using with TF32 enabled (fp32_matmul_prec="high"):

Matrix Shape $(M \times N)$	Layer Analogue	polar (SVD)	qr	cayley	newton_schulz (steps=8)	Speedup vs SVD	Orthogonality Drift $|Q^T Q - I|_\infty$	Abs Error vs SVD Polar
1024 × 1024	Small Hidden Dim	125.55 ms	7.63 ms	4.37 ms	1.65 ms	76.2x	$4.77 \times 10^{-7}$	$3.94 \times 10^{-5}$
2048 × 2048	7B/8B Attention	422.14 ms	18.45 ms	12.33 ms	11.83 ms	35.7x	$9.54 \times 10^{-7}$	$5.67 \times 10^{-5}$
4096 × 4096	70B Attention	4619.68 ms	50.08 ms	42.94 ms	82.79 ms	55.8x	$3.10 \times 10^{-6}$	$7.69 \times 10^{-5}$
4096 × 2048	Intermediate MLP	370.79 ms	21.67 ms	22.18 ms	13.86 ms	26.8x	$7.15 \times 10^{-7}$	$2.39 \times 10^{-5}$
8192 × 2048	SwiGLU Projection	376.01 ms	26.81 ms	57.25 ms	21.53 ms	17.5x	$7.15 \times 10^{-7}$	$1.93 \times 10^{-5}$

Observations:

1. Mathematical Fidelity to Exact Polar: Across all tested configurations, 8 iterations of newton_schulz (polar_express) track the analytical SVD polar factor within $\le 7.69 \times 10^{-5}$ absolute error while holding manifold drift $|Q^T Q - I|_\infty$ to machine precision ($\approx 10^{-7}$ to $10^{-6}$), directly matching QR and Cayley.
2. Speedup over SVD: Delivers a 17.5x to 76.2x speedup over cuSOLVER SVD (e.g., dropping from 4.62 seconds down to 82.79 ms at $4096 \times 4096$)
3. Efficiency on Rectangular Projections: For asymmetric weights typical of transformer MLPs ($8192 \times 2048$), newton_schulz (21.53 ms) is 2.66x faster than Cayley (57.25 ms) and faster than QR (26.81 ms), avoiding the $M \times M$ ($8192 \times 8192$) skew-symmetric linear solve required by Cayley.
4. Step-Count Selection: While 5 steps achieves lower latency, it leaves substantial manifold drift ($\approx 0.09\text{--}0.22$), confirming num_ns_steps=8 as the necessary default for numerical stability.

Testing

• pre-commit run --all-files passes cleanly.
• pytest tests/riemannian_optimizers/test_isospectral.py passes on CPU and CUDA.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

The PR adds a Newton–Schulz Stiefel retraction to the Iso optimizer and extracts the retraction implementations into a reusable module.

• Adds newton_schulz as a validated Iso retraction option.
• Uses the existing Polar Express coefficients with eight iterations by default.
• Extends tall and wide matrix coverage and adds a retraction benchmark.

Confidence Score: 4/5

The PR should not merge until the outstanding device-to-host synchronization in the benchmark workflow is resolved.

The CPU SYRK failure has been fixed by selecting SYRK only for CUDA tensors, but the benchmark still converts stacked CUDA metrics to a Python list and synchronizes the host despite the repository rule prohibiting that operation outside tests.

Files Needing Attention: benchmarks/benchmark_stiefel_retraction.py

Important Files Changed

Filename	Overview
emerging_optimizers/riemannian_optimizers/isospectral.py	Adds selection and validation for the new Newton–Schulz retraction while delegating all Stiefel retractions to the extracted module.
emerging_optimizers/riemannian_optimizers/retractions/stiefel.py	Defines reusable QR, polar, Cayley, and Newton–Schulz Stiefel retractions with CPU-safe SYRK selection.
tests/test_isospectral.py	Extends singular-value preservation coverage to tall and wide matrices using the Newton–Schulz retraction.
benchmarks/benchmark_stiefel_retraction.py	Benchmarks the supported Stiefel retractions and reports their numerical accuracy.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    G[Factor gradient and momentum] --> S[Finite factor step]
    S --> N[Newton-Schulz retraction]
    N --> U[Retracted U]
    N --> V[Retracted V]
    U --> R[Rebuild U Sigma V.T]
    V --> R
    F[Fixed singular values Sigma] --> R
    R --> P[Updated isospectral parameter]

Loading

_{Reviews (6): Last reviewed commit: "Renmae the Stiefel retraction benchmark" | Re-trigger Greptile}

[mkhona-nvidia]

Thank you for contribution!

It looks like the retractions are becoming big. Could you move this entire block into a retractions/ subdirectory?:

RetractionT = Literal["qr", "polar", "cayley", "newton_schulz"]

def _qr_retraction(
    point: torch.Tensor,
    momentum: torch.Tensor,
    step_size: float,
) -> torch.Tensor:
    matrix = point - step_size * momentum
    q, r = torch.linalg.qr(matrix, mode="reduced")
    signs = torch.diagonal(r).sign()
    signs.masked_fill_(signs == 0, 1)
    return q * signs


def _polar_retraction(
    point: torch.Tensor,
    momentum: torch.Tensor,
    step_size: float,
) -> torch.Tensor:
    matrix = point - step_size * momentum
    u, _, vh = torch.linalg.svd(matrix, full_matrices=False)
    return u @ vh


def _cayley_retraction(
    point: torch.Tensor,
    momentum: torch.Tensor,
    step_size: float,
) -> torch.Tensor:
    direction = -momentum
    skew = direction @ point.mT - point @ direction.mT
    identity = torch.eye(point.shape[0], dtype=point.dtype, device=point.device)
    lhs = identity - 0.5 * step_size * skew
    rhs = (identity + 0.5 * step_size * skew) @ point
    return torch.linalg.solve(lhs, rhs)


def _newton_schulz_retraction(
    point: torch.Tensor,
    momentum: torch.Tensor,
    step_size: float,
    coefficient_type: NSCoeffT = "polar_express",
    num_ns_steps: int = 8,
) -> torch.Tensor:
    matrix = point - step_size * momentum
    return newton_schulz(
        matrix,
        steps=num_ns_steps,
        coefficient_type=coefficient_type,
        use_syrk=matrix.is_cuda,
    )

[kogolobo]

Thank you for contribution!

It looks like the retractions are becoming big. Could you move this entire block into a retractions/ subdirectory?:

RetractionT = Literal["qr", "polar", "cayley", "newton_schulz"]

def _qr_retraction(
    point: torch.Tensor,
    momentum: torch.Tensor,
    step_size: float,
) -> torch.Tensor:
    matrix = point - step_size * momentum
    q, r = torch.linalg.qr(matrix, mode="reduced")
    signs = torch.diagonal(r).sign()
    signs.masked_fill_(signs == 0, 1)
    return q * signs


def _polar_retraction(
    point: torch.Tensor,
    momentum: torch.Tensor,
    step_size: float,
) -> torch.Tensor:
    matrix = point - step_size * momentum
    u, _, vh = torch.linalg.svd(matrix, full_matrices=False)
    return u @ vh


def _cayley_retraction(
    point: torch.Tensor,
    momentum: torch.Tensor,
    step_size: float,
) -> torch.Tensor:
    direction = -momentum
    skew = direction @ point.mT - point @ direction.mT
    identity = torch.eye(point.shape[0], dtype=point.dtype, device=point.device)
    lhs = identity - 0.5 * step_size * skew
    rhs = (identity + 0.5 * step_size * skew) @ point
    return torch.linalg.solve(lhs, rhs)


def _newton_schulz_retraction(
    point: torch.Tensor,
    momentum: torch.Tensor,
    step_size: float,
    coefficient_type: NSCoeffT = "polar_express",
    num_ns_steps: int = 8,
) -> torch.Tensor:
    matrix = point - step_size * momentum
    return newton_schulz(
        matrix,
        steps=num_ns_steps,
        coefficient_type=coefficient_type,
        use_syrk=matrix.is_cuda,
    )

Done, thanks!

[mkhona-nvidia]

nit: typo in new schultz

[skyw]

typo is not nit and should be fixed.

[mkhona-nvidia]

/claude review

[skyw]

typo is not nit and should be fixed.