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

This PR adds a Newton-Schulz approximation of the polar retraction to the Iso optimizer.

• Extracts Stiefel retractions into a dedicated module.
• Adds newton_schulz as a supported Iso retraction with eight Polar Express iterations by default.
• Extends singular-value preservation coverage to tall and wide matrices.
• Adds a benchmark comparing latency and numerical accuracy across retraction methods.

Confidence Score: 5/5

The PR appears safe to merge, with no outstanding correctness or repository-rule issues.

No new actionable failures were identified in the changes since the previous review. All previous findings were manually resolved, including the CPU SYRK concern and the benchmark synchronization policy threads.

Important Files Changed

Filename	Overview
emerging_optimizers/riemannian_optimizers/isospectral.py	Routes Iso factor updates through the extracted retraction functions and accepts the new Newton-Schulz option.
emerging_optimizers/riemannian_optimizers/retractions/stiefel.py	Defines the existing Stiefel retractions and the new eight-step Polar Express Newton-Schulz retraction.
tests/test_isospectral.py	Covers singular-value preservation for Newton-Schulz retraction on tall and wide matrices.
benchmarks/benchmark_stiefel_retraction.py	Benchmarks retraction latency and reports orthogonality and polar-approximation metrics.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    Iso[Iso optimizer step] --> Factors[Update Stiefel factors U and V]
    Factors --> Choice{Retraction}
    Choice --> QR[QR]
    Choice --> Polar[SVD polar]
    Choice --> Cayley[Cayley transform]
    Choice --> NS[Newton-Schulz]
    NS --> PE[Polar Express polynomial<br/>8 iterations]
    QR --> Result[Retracted factors]
    Polar --> Result
    Cayley --> Result
    PE --> Result
    Result --> Recompose[Recompose parameter<br/>U S Vᵀ]

Loading

_{Reviews (7): Last reviewed commit: "Add retraction dosctrings" | 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.

[kogolobo]

Fixed, thanks for noticing this!

[mkhona-nvidia]

/claude review

[skyw]

typo is not nit and should be fixed.

[mkhona-nvidia]

/ok to test fea8a69

[mkhona-nvidia]

/claude review