Modular Manifolds

TL;DR

A pedagogical Thinking Machines post by Jeremy Bernstein arguing that weight matrices should be kept on submanifolds (e.g. the Stiefel manifold of matrices with unit condition number), with the optimizer co-designed against that constraint rather than projecting back after the fact. Derives manifold Muon — gradient descent under a spectral-norm budget on the Stiefel manifold — via dual ascent on a Lagrangian, producing the update A_opt = -η · msign(G + 2W(Λ+Λᵀ)) with a retraction W ← msign(W). Generalizes to modular manifolds: composing per-layer manifold + norm pairs into a network-wide norm via Cartesian-product manifolds and max-weighted norm composition, so per-layer learning rates can be budgeted from a single network-Lipschitz statement. Companion to Training Transformers with Enforced Lipschitz Constants (same lab, same recipe of “constrain weights, not just updates”), and pitched explicitly as a research-area introduction rather than a finished method.

Key claims

A first-order manifold optimizer has three steps: (i) find the unit-length tangent vector with maximum gradient projection, (ii) take a step of length η in that direction, (iii) retract back to the manifold; the choice of manifold + tangent-space norm together determine the algorithm [§“The shape of a manifold optimizer”].
The same recipe with different (manifold, norm) pairs reproduces vanilla GD (ℝⁿ + Euclidean), sign-GD (ℝⁿ + ∞-norm), hyperspherical descent (Sⁿ + Euclidean), Muon (matrix space + spectral), and manifold Muon (Stiefel + spectral) [Table 1, §“The shape of a manifold optimizer”].
Choosing the Stiefel manifold (matrices with WᵀW = I) plus the spectral norm gives a co-designed constraint+norm pair: the manifold pins all singular values to 1 (unit condition number, no input-vector blow-up or collapse), and the spectral norm budget caps how much any update can stretch any input direction [§“Manifold Muon”].
Manifold Muon’s update problem is convex and can be solved by dual ascent on a matrix Lagrange multiplier Λ ∈ ℝⁿˣⁿ, with inner minimizer A_opt(Λ) = -η · msign(G + 2W(Λ+Λᵀ)) and dual gradient H(Λ) = -2η · ‖G + 2W(Λ+Λᵀ)‖_nuclear-derivative; final update applies msign retraction W ← msign(W) [§“Manifold Muon”].
The dual-ascent solution is closely related to but distinct from Jianlin Su’s fixed-point solution to the same problem; both build on Franz Louis Cesista’s earlier heuristic, and Bernstein had posed the rectangular case as an open problem on the Modula docs [§“Manifold Muon”, footnote on prior work].
Whole networks compose by Cartesian-product manifolds (M₃ = M₁ × M₂) and max-weighted norms (‖(w₁,w₂)‖₃ = max(s₁·‖w₁‖₁, s₂·‖w₂‖₂)); the scalar coefficients sᵢ are exactly the per-layer learning-rate budget, derived from network-Lipschitz analysis rather than tuned [§“Modular manifolds”].
A StiefelLinear module fed unit-norm inputs is Lipschitz-with-constant-one in the spectral norm: ‖(W+ΔW)x − Wx‖₂ ≤ ‖ΔW‖_spectral, providing a building block for the modular norm composition [§“Modular manifolds”].
Sanity-checked with a small CIFAR-10 experiment whose runs each finish in under a minute; code at github.com/thinking-machines-lab/manifolds/ [§“Manifold Muon”, Fig. CIFAR10].

Method

The post derives optimization-on-a-manifold from first principles in two stages. Stage one: a hypersphere parameter w ∈ ℝᵈ, ‖w‖₂=1 trained on a loss over ℝᵈ. Casting “best tangent step under a budget” as a constrained linear program gives min_a aᵀg s.t. ‖a‖₂=η, aᵀw=0, solved via Lagrange multipliers to a_opt = -η · (g - wwᵀg)/‖g - wwᵀg‖₂ — i.e. subtract the radial gradient component, normalize, scale by η. A retraction by Pythagoras (w_new ← w_new / √(1+η²)) projects back. Stage two: lifts this to matrix parameters on the Stiefel manifold {W : WᵀW = Iₙ} with the spectral norm. The tangent condition generalizes from aᵀw=0 to WᵀA + AᵀW = 0. The resulting constrained problem min_A ⟨A,G⟩ s.t. ‖A‖_spectral≤η, WᵀA+AᵀW=0 is convex; the post solves it via dual ascent on a matrix Lagrange multiplier Λ ∈ ℝⁿˣⁿ, deriving the inner minimizer in closed form using the matrix sign function msign (efficiently computed on GPU via Newton–Schulz or Polar Express). The full algorithm is: (1) ascend Λ ← Λ + α·H(Λ), (2) A_opt = -η · msign(G + 2W(Λ_opt + Λ_optᵀ)), (3) W ← W + A_opt, (4) retract W ← msign(W).

The third stage — modular manifolds — extends this beyond a single tensor. A module is defined as a triple (forward function f: W×X → Y, weight-manifold M ⊂ W, weight-norm ‖·‖: W → ℝ). Composing modules via function composition f₃ = f₂ ∘ f₁ yields a product manifold M₁ × M₂ and a composite norm ‖(w₁,w₂)‖₃ = max(s₁·‖w₁‖₁, s₂·‖w₂‖₂), where the coefficients sᵢ derive from per-module Lipschitz statements and chain together to give a network-level bound. The optimizer-derivation recipe is then applied to each layer with its own (manifold, norm) — the LRs are budgeted across layers by the sᵢ, not by hand.

Results

The post is pedagogical, not benchmark-driven. The only empirical artifact is a CIFAR-10 sanity-check (each run < 1 minute) confirming that manifold Muon trains successfully under the Stiefel constraint, with code released at github.com/thinking-machines-lab/manifolds/. No comparisons against AdamW / Muon / SSO / weight-decay baselines, no LLM-scale results, no quantitative claims about loss or stability. The post is explicit that “hard manifold constraints may not ultimately be the right way to constrain weight matrices” and frames the recipe as “an introduction to a research area” — directions for future work include architecture-optimizer co-design, non-Riemannian geometry (operator norms have sharp corners, no unique gradient flow), regularization via constraint radii, and practical GPU implementation of manifold operations.

Why it’s interesting

This post sits cleanly inside the optimizer-side stability thread on the wiki — it is the conceptual framework that Training Transformers with Enforced Lipschitz Constants (same lab; Bernstein is a co-author) operationalizes, and Controlled LLM Training on Spectral Sphere (SSO) independently arrives at via a very similar diagnosis: “Muon controls update spectra but not weight spectra; close the loop by constraining both.” Three handles for Luma. (1) The modular-norm machinery gives a derivation of per-layer learning-rate budgets from a single network-Lipschitz statement — directly relevant if pretraining instabilities are being chased by hand-tuning per-module LRs. (2) The Stiefel + spectral-norm recipe is a candidate manifold for attention Q/K/O projections, where unit condition number coincides with the QK-Norm intuition in A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training without needing a normalization layer. (3) The unsolved question Bernstein flags — “what manifolds should attention heads, embeddings, and unembeddings live on?” — is the next research lever, and complements the architecture-side answers in ViT-5: Vision Transformers for The Mid-2020s (QK-Norm, LayerScale) and mHC: Manifold-Constrained Hyper-Connections (Birkhoff-constrained residual mixing).