Optimization Function Classes

Riemannian Optimization and Manifold Constraints

Optimization on curved spaces: the Stiefel manifold for orthogonal matrices, symmetric positive definite matrices, Riemannian gradient descent, retractions, and why flat-space intuitions break on manifolds. The geometric backbone of Shampoo, Muon, and constrained neural network training.

AdvancedTier 2Current~55 min

Prerequisites

Convex Optimization Basics The Hessian Matrix Eigenvalues and Eigenvectors

Why This Matters

Standard gradient descent assumes parameters live in flat Euclidean space $\mathbb{R}^d$ . But many ML problems have constraints that define curved spaces: weight matrices that must stay orthogonal, covariance matrices that must stay positive definite, probability distributions that must stay on the simplex, or embeddings that must stay on a sphere.

If you just project back to the constraint set after each gradient step, you waste computation and can get poor convergence. Riemannian optimization does better: it works directly on the curved surface (manifold), computing gradients that are tangent to the surface and taking steps that stay on it.

This is not esoteric geometry. The Muon optimizer uses Newton-Schulz iteration to approximate projection onto the Stiefel manifold. Shampoo preconditions gradients using the space of symmetric positive definite matrices. Understanding the manifold geometry explains why these optimizers work and where they fail.

Core Concepts

Definition

Smooth Manifold

A smooth manifold $\mathcal{M}$ is a space that locally looks like $\mathbb{R}^d$ but may be globally curved. Each point $p \in \mathcal{M}$ has a tangent space $T_p\mathcal{M}$ , a vector space of directions you can move from $p$ while staying on $\mathcal{M}$ (infinitesimally).

For ML, the key manifolds are:

The Stiefel manifold $\text{St}(n, p) = \{W \in \mathbb{R}^{n \times p} : W^\top W = I_p\}$ : matrices with orthonormal columns
The SPD manifold $\mathcal{S}_{++}^n$ : $n \times n$ symmetric positive definite matrices
The sphere $S^{d-1} = \{x \in \mathbb{R}^d : \|x\| = 1\}$ : unit vectors
The Grassmannian $\text{Gr}(n, p)$ : $p$ -dimensional subspaces of $\mathbb{R}^n$

Definition

Riemannian Metric

A Riemannian metric assigns an inner product $g_p: T_p\mathcal{M} \times T_p\mathcal{M} \to \mathbb{R}$ to each tangent space, varying smoothly with $p$ . This defines distances, angles, and curvature on the manifold.

For the Stiefel manifold embedded in $\mathbb{R}^{n \times p}$ , the simplest metric is the inherited Euclidean one: $g_W(\xi, \eta) = \text{tr}(\xi^\top \eta)$ for tangent vectors $\xi, \eta \in T_W \text{St}(n, p)$ .

Definition

Riemannian Gradient

The Riemannian gradient of $f: \mathcal{M} \to \mathbb{R}$ at $p$ is the tangent vector $\text{grad } f(p) \in T_p\mathcal{M}$ such that:

$g_p(\text{grad } f(p), \xi) = Df(p)[\xi] \quad \forall \xi \in T_p\mathcal{M}$

For a submanifold of $\mathbb{R}^d$ , this is the projection of the Euclidean gradient onto the tangent space:

$\text{grad } f(p) = \Pi_{T_p\mathcal{M}}(\nabla f(p))$

The Riemannian gradient points in the steepest descent direction on the manifold, ignoring directions that would leave the constraint surface.

The Stiefel Manifold

The Stiefel manifold is the most important manifold in ML optimization because orthogonality constraints appear in neural network weight matrices, embedding layers, and preconditioner updates.

Proposition

Tangent Space and Retraction on the Stiefel Manifold

Statement

The tangent space of the Stiefel manifold at $W$ is:

$T_W \text{St}(n, p) = \{\xi \in \mathbb{R}^{n \times p} : W^\top \xi + \xi^\top W = 0\}$

The Riemannian gradient of $f$ at $W$ is obtained by projecting the Euclidean gradient $G = \nabla f(W)$ :

$\text{grad } f(W) = G - W \cdot \text{sym}(W^\top G)$

where $\text{sym}(A) = (A + A^\top)/2$ . A retraction (first-order approximation to the exponential map) is the QR-based retraction:

$R_W(\xi) = \text{qf}(W + \xi)$

where $\text{qf}$ extracts the $Q$ factor from QR decomposition.

Intuition

The tangent space condition $W^\top \xi + \xi^\top W = 0$ says that tangent directions must be "skew" relative to $W$ . Moving in direction $\xi$ from $W$ should keep $W^\top W$ close to $I$ to first order. The projection $G - W \cdot \text{sym}(W^\top G)$ removes the component of the Euclidean gradient that would move $W$ off the manifold. The retraction maps back to the manifold after taking a step.

Why It Matters

This is exactly what Muon's Newton-Schulz iteration approximates. Instead of the expensive QR decomposition, Newton-Schulz uses polynomial iterations to approximate the orthogonal polar factor, which is another valid retraction. The choice of retraction affects computational cost but not the direction of descent (to first order).

Shampoo also uses the Stiefel manifold implicitly: its preconditioner updates maintain structure on the space of positive definite matrices, and the Kronecker factorization exploits the product structure of weight matrices.

Failure Mode

QR retraction costs $O(np^2)$ , which is expensive for large matrices. Cayley retraction and polar retraction are alternatives with different cost-accuracy tradeoffs. For very large $n$ (e.g., transformer weight matrices with $n = 4096$ ), even the retraction cost can dominate training time, which is why approximations like Newton-Schulz are used in practice.

Riemannian Gradient Descent

Theorem

Riemannian Gradient Descent

Statement

Riemannian gradient descent updates:

$p_{t+1} = R_{p_t}(-\eta \cdot \text{grad } f(p_t))$

where $R$ is a retraction. Under geodesic $L$ -smoothness (the Riemannian analog of $L$ -smoothness), with step size $\eta = 1/L$ :

$f(p_T) - f(p^*) \leq \frac{L \cdot d^2(p_0, p^*)}{2T}$

where $d$ is the geodesic distance. This matches the $O(1/T)$ Euclidean rate.

Intuition

Riemannian gradient descent is the direct analog of Euclidean gradient descent: compute the steepest descent direction (now on the manifold), take a step (now using the retraction instead of addition), and repeat. The convergence theory parallels the Euclidean theory, with geodesic convexity and smoothness replacing their Euclidean counterparts.

Why It Matters

This guarantees that constrained optimization on manifolds is not harder in rate than unconstrained optimization. The same $O(1/T)$ rate holds, and with geodesic strong convexity you get the same $O(\kappa \log(1/\epsilon))$ linear rate. The manifold structure is exploited, not fought against.

Failure Mode

Geodesic convexity is a stronger requirement than Euclidean convexity restricted to the manifold. Some functions that are convex in $\mathbb{R}^d$ are not geodesically convex on curved manifolds, and vice versa. The curvature of the manifold introduces additional terms in the convergence analysis. On manifolds with high curvature (small injectivity radius), the step size must be correspondingly small.

Why Low-Rank Does Not Mean Low Complexity

Watch Out

Low rank does not mean simple or easy to optimize

Low-rank approximation via SVD truncation is a Euclidean operation: find the closest rank- $k$ matrix in Frobenius norm. But the set of rank- $k$ matrices is not a convex set and not even a smooth manifold (it has singularities where the rank drops further). The smooth part (matrices of exactly rank $k$ ) forms a manifold, but it has curvature that depends on the singular value gaps.

Optimizing over low-rank matrices is a non-convex problem. The landscape has saddle points corresponding to singular value crossings. Riemannian optimization on the fixed-rank manifold avoids the singularities and can converge provably, but it requires careful treatment of the tangent space (which involves both the column space and row space of the matrix).

Watch Out

The Stiefel manifold is not flat

The Stiefel manifold $\text{St}(n, p)$ is a curved space embedded in $\mathbb{R}^{n \times p}$ . Straight-line steps in the ambient space leave the manifold. This is not just a technical nuisance: the curvature means that parallel transport (moving vectors from one tangent space to another) is path-dependent, and the shortest path between two orthogonal matrices (the geodesic) is not a straight line but a matrix exponential curve.

For the special case $p = n$ (orthogonal group $O(n)$ ), the geodesics are $W(t) = W_0 \exp(t \Omega)$ where $\Omega$ is skew-symmetric. The curvature of $O(n)$ is positive and equals $1/4$ in appropriate normalization.

SPD Matrices and the Information Geometry Connection

The manifold of symmetric positive definite (SPD) matrices $\mathcal{S}_{++}^n$ appears in ML through covariance estimation, Fisher information matrices, kernel matrices, and preconditioner design.

The natural metric on $\mathcal{S}_{++}^n$ is the affine-invariant metric:

$g_P(\xi, \eta) = \text{tr}(P^{-1} \xi P^{-1} \eta)$

Under this metric, the geodesic distance between two SPD matrices $P, Q$ is:

$d(P, Q) = \|\log(P^{-1/2} Q P^{-1/2})\|_F$

This connects directly to information geometry: the Fisher information metric on statistical models induces the affine-invariant metric on the space of Gaussian covariance matrices. Shampoo's preconditioner updates operate on this space.

Exercises

ExerciseCore

Problem

Verify that the tangent space condition $W^\top \xi + \xi^\top W = 0$ is necessary for $\xi$ to be tangent to $\text{St}(n, p)$ at $W$ . Start from the constraint $W^\top W = I_p$ and differentiate.

ExerciseAdvanced

Problem

The polar retraction on the Stiefel manifold is $R_W(\xi) = (W + \xi)(I_p + \xi^\top \xi)^{-1/2}$ . Show that this maps back to $\text{St}(n, p)$ , i.e., verify $R_W(\xi)^\top R_W(\xi) = I_p$ . Why is this preferred over QR retraction in some settings?

References

Canonical:

Absil, Mahony, Sepulchre, Optimization Algorithms on Matrix Manifolds (2008). The definitive reference.
Edelman, Arias, Smith, "The Geometry of Algorithms with Orthogonality Constraints" (SIMAX, 1998). Foundational paper on Stiefel/Grassmann optimization.

Current:

Boumal, An Introduction to Optimization on Smooth Manifolds (2023). Modern textbook, freely available.
Hu et al., "Brief Introduction to Manifold Optimization" (JMLR, 2020)
Becigneul & Ganea, "Riemannian Adaptive Optimization Methods" (ICLR 2019). Riemannian Adam.
Boyd & Vandenberghe, Convex Optimization (2004), Chapters 2-5

Next Topics

Optimizer theory: SGD, Adam, Muon: how Muon uses Newton-Schulz to approximate Stiefel retraction
Second-order optimization: natural gradient and preconditioned methods that operate on manifold structure

Last reviewed: April 2026

Prerequisites

Foundations this topic depends on.

Convex Optimization BasicsLayer 1
Differentiation in RnLayer 0A
Sets, Functions, and RelationsLayer 0A
Basic Logic and Proof TechniquesLayer 0A
Matrix Operations and PropertiesLayer 0A
The Hessian MatrixLayer 0A
Eigenvalues and EigenvectorsLayer 0A

Next Topics

Optimizer Theory SGD Adam MuonContinue →Second Order Optimization MethodsContinue →