Polar-Muon Stiefel Optimizer

(1) Integration point: replace the optimizer update for each rectangular neural-network weight matrix $X\in\mathbb{R}^{m\times n}$ whose columns are intended to remain orthonormal. Suitable targets include tall linear layers, orthogonal projection layers, recurrent transition matrices, and low-rank adapter factors. Do not apply the update to biases or unconstrained parameters. Store a full-precision momentum matrix $M$ with the same shape as $X$.

(2) Pseudocode:

initialize X with orthonormal columns and M = 0
for each minibatch:
    loss = model(batch, X)
    G = autograd(loss, X)
    M = beta*M + (1-beta)*G
    S = 0.5*(X.T @ M + M.T @ X)
    T = M - X @ S
    U, Sigma, Vt = thin_svd(T)
    D = -(U @ Vt)
    Z = X + eta*D
    U2, Sigma2, Vt2 = thin_svd(Z)
    X = U2 @ Vt2

The line defining $T$ is the tangent projection $P_X(M)$; the line defining $D$ is the paper's spectral-norm linear-minimization solution; the final SVD is the polar retraction. Use bias correction if desired, but keep the first comparison simple.

(3) Computed directly from the mathematics: the tangent projection, polar direction, and exact orthonormal retraction. Estimated empirically: the learning rate $\eta$, momentum coefficient $\beta$, whether tangent projection is necessary, and the tradeoff between SVD cost and optimization quality. Log $\|X^\top X-I\|_F$, the singular values of $X$, training loss, validation loss, gradient norms, wall-clock time, and SVD FLOPs.

(4) First cheap experiment: train a four-layer width-256 MLP on MNIST or CIFAR-10, making two dense layers Stiefel-constrained. Compare AdamW with an orthogonality penalty, ordinary Muon followed by QR, and this optimizer at matched batch size and approximately matched compute. Sweep $\eta\in\{10^{-3},3\times10^{-3},10^{-2}\}$ and $\beta\in\{0.9,0.95,0.99\}$. A positive result is lower validation loss at equal optimizer steps, near-machine-precision orthogonality without penalty tuning, and equal or better final accuracy. Also report wall-clock performance; repeated SVDs may erase any optimization gain.