[kogolobo]

Summary

This PR fixes the mathematical implementation of the SPEL (Spel) optimizer (emerging_optimizers/orthogonalized_optimizers/spel.py) to align with Manifold Constrained Steepest Descent (arXiv:2601.21487) and adds configurable manifold scaling (scale_mode).

---

Key Improvements

1. Tangent Space Projection on the Stiefel Manifold

In the previous implementation, Spel did not override orthogonalize(). It passed the unconstrained ambient gradient directly to the inner Newton–Schulz step, effectively running Euclidean Muon followed by a post-hoc weight retraction.

This PR implements the true Stiefel tangent space projection in orthogonalize() before applying the spectral LMO:
$$P_{T_W}(G) = G - \frac{1}{\sigma^2} W , \text{sym}(W^\top G) = G - \frac{1}{2\sigma^2} W (W^\top G + G^\top W)$$

This ensures the descent direction is strictly tangential to the manifold at the current iterate $W_t$, fulfilling the Manifold Constrained Steepest Descent formulation.

2. Scaled Stiefel Manifold Support (scale_mode)

Added scale_mode: SpelScaleT ("unit_spectral_norm" | "unit_rms_to_rms_norm"):

• unit_spectral_norm (default): Standard Stiefel manifold constraint ($\sigma = 1.0$, $W^\top W = I$).
• unit_rms_to_rms_norm: Scaled Stiefel manifold where $\sigma = \sqrt{m / n}$ ($W^\top W = \frac{m}{n} I$). This ensures the linear layer preserves the RMS norm of activations across network depth without requiring separate normalization layers.

Both the tangent projection and the post-update retraction in post_weight_update_fn_inplace correctly account for $\sigma$.

---

Testing & Verification

1. Tangent Space Verification (test_tangent_space_projection):

   • Verifies that for any parameter matrix $W$ on the (scaled) manifold, the projected gradient $V = P_{T_W}(G)$ satisfies the Stiefel tangent condition: $\text{sym}(W^\top V) = \frac{1}{2}(W^\top V + V^\top W) = 0$.
   • Parameterized over multiple rectangular and square shapes and both scale_mode options.

2. Manifold Retraction Scaling (test_post_update_produces_appropriately_scaled_orthogonal_weights):

   • Verifies that after multiple optimization steps, the Gram matrix satisfies $W^\top W \approx \sigma^2 I$ (or $W W^\top \approx \sigma^2 I$) within numerical tolerance for both scaling modes.

3. CI & Code Quality:

   • Passed L0 test suites: bash tests/ci/L0_Tests_CPU.sh and bash tests/ci/L0_Tests_GPU.sh.
   • Passed pre-commit hooks (uv run ruff check --fix ., uv run ruff format .).

[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]

Confidence Score: 5/5

The PR appears safe to merge; no new actionable issues remain after the previous findings were resolved.

Summary

This PR corrects SPEL’s Riemannian gradient calculation and introduces configurable scaled-Stiefel behavior.

• Projects gradients onto the appropriate column- or row-Stiefel tangent space before Newton–Schulz orthogonalization.
• Adds spectral and RMS-preserving scale modes used consistently by tangent projection and post-update retraction.
• Expands tests across square, tall, and wide matrices and both scaling modes.
• The changes since the previous review only add the repository’s standard license header.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    G[Gradient or momentum] --> T[Tangent-space projection]
    T --> L[Newton–Schulz spectral LMO]
    L --> U[Learning-rate update]
    U --> R[Scaled Stiefel retraction]
    R --> W[Constrained weight]

Loading

_{Reviews (6) · Last reviewed commit: "Add copyright notice back"}

[mkhona-nvidia]

@kogolobo can you update the copyright check

[mkhona-nvidia]

/ok to test bb2dd17