Optimization on Manifolds

Manifold optimization adapts gradient steps to curved constraint sets using tangent projections and retractions, keeping iterates exactly feasible.

Project the gradient to the tangent, step flat, then retract onto the surface

Why Does This Exist?

Orthogonal RNNs, low-rank adapters, and subspace models need constraints to hold at every step, not approximately at the end. Penalty terms drift and need tuning; Euclidean steps with hard projection zigzag. Riemannian optimization moves natively on the set: each iterate satisfies constraints by construction and each step follows the surface.

It powers orthogonal initialization, Grassmannian subspace tracking, and rotation averaging. Prerequisites: manifolds, geodesics.

The one idea you need first

A retraction $R_p(v)$ maps a tangent step back to the manifold with $R_p(0) = p$ and local rigidity (it matches the exponential to first order). QR-based retraction on Stiefel ( $qf$ of $X + V$ ) is the workhorse example.

Think of It Like This

Walking while glued to a globe

You must stay on the globe's surface with every footstep. Each move has three beats: face a compass direction (tangent), stride onto a flat local map (linear step), then snap your boots back to the soil (retraction).

Short strides make the snap tiny; giant leaps teleport you somewhere unintended. Good manifold code keeps strides chart-sized.

Where the analogy stops: retractions are first-order approximations, cheaper than true exponentials, and the savings fund larger models.

How It Actually Works

Riemannian gradient descent: $X_{k+1} = R_{X_k}(-\eta \,\text{grad}\,f(X_k))$ where $\text{grad}\,f$ is the tangent-projected gradient. On Stiefel ( $X^TX = I$ ): $\text{grad}\,f = G - X\,\text{sym}(X^TG)$ with $G$ the Euclidean gradient. Convergence mirrors Euclidean rates with curvature-dependent constants; trust-region variants use Hessian actions in tangent space.

Worked example

Minimize $f(x) = -x_1$ on unit circle $S^1$ from $p_0 = (0, 1)$ with $\eta = 0.5$ :

Euclidean gradient: $(-1, 0)$ . Tangent at $(0,1)$ : $v_2 = 0$ , so projected gradient $(-1, 0)$ is already tangent.
Linear step: $(0,1) - 0.5(-1,0) = (0.5, 1)$ . Norm $\sqrt{1.25} \approx 1.118$ .
Retract (normalize): $(0.447, 0.894)$ . Loss $-0.447$ vs start $0$ . Improved and feasible.
Repeat: gradient $(-1,0)$ , tangent projection removes radial part $-0.447 \times (0.447,0.894) \approx (-0.2,-0.4)$ , tangent $\approx (-0.8, 0.4)$ . Steps march toward $(1, 0)$ , the true minimizer with loss $-1$ .

Code

1import numpy as np2
3
4def rgrad_circle(x: np.ndarray, g: np.ndarray) -> np.ndarray:5    return g - (g @ x) * x  # remove radial component6
7
8x = np.array([0.0, 1.0])9for _ in range(20):10    g = np.array([-1.0, 0.0])11    x = x - 0.5 * rgrad_circle(x, g)12    x = x / np.linalg.norm(x)  # retraction by normalization13
14print(np.round(x, 3))  # -> about [1. 0.], the minimizer

Watch Out For

Retracting after oversized steps

Large tangent steps leave the chart where retraction approximates the exponential. QR and polar retractions can jump to distant representatives and loss spikes.

Scale steps by the manifold's injectivity radius proxy (curvature-aware caps) and shrink on retraction failure.

Forgetting vector transport for momentum

Tangent spaces differ per point, so yesterday's momentum vector lives in the wrong space today. Adding it raw breaks feasibility and direction.

Transport momentum ( $P_{x \to y}(m)$ ) before combining with the new gradient.

The Quick Version

Riemannian steps project gradients, move in tangents, and retract to feasibility.
Stiefel and sphere formulas make projection plus retraction cheap.
Rates mirror Euclidean theory with curvature constants added.
Momentum needs vector transport across differing tangent planes.
Step sizes must respect chart validity, not just loss decrease.