Modular Duality in Deep Learning

TL;DR

Bernstein & Newhouse construct a duality map for general neural architectures from first principles: gradients live in the dual weight space and cannot be subtracted from weights until they’ve been mapped back through a duality map. The construction is recursive — assign an operator norm to each atomic layer (Linear, Embed, Conv2D), derive the corresponding per-layer duality map from steepest-descent theory, then compose via the modular norm (max-weighted composition) to lift the map to the full network. The paper shows that maximal update parameterization (µP) and Shampoo both emerge as partial approximations to the single RMS→RMS operator-norm duality map for Linear layers, and proposes a new rectangular Newton-Schulz iteration for computing it — the same iteration that would later become the algorithmic core of Muon.

Key claims

Gradients belong to the dual weight space and should be passed through a duality map before being subtracted from weights; without this “type check” the update ignores that loss curvature is non-isotropic [§1, Eqs. 1–2].
Steepest descent under a norm ‖·‖ with sharpness λ minimizes ⟨g, Δw⟩ + (λ/2)‖Δw‖² at Δw* = −(‖g‖_†/λ) · duality_map(g), where ‖·‖_† is the dual norm and the duality map returns the argmax of ⟨g,t⟩ over the unit ball — this is the entire recipe [§3.2, Prop. 1].
For a Linear layer with input and output both RMS-normed, the induced operator norm is a rescaled spectral norm ‖W‖_{RMS→RMS} = √(d_in/d_out)·‖W‖_spectral, and the duality map is √(d_out/d_in)·UVᵀ where G = UΣVᵀ is the gradient’s reduced SVD [§4.1, Example 5, Table 1].
This single duality map recovers essential features of both µP and Shampoo: the √(d_out/d_in) scaling factor implements µP’s spectral update scaling, and the UVᵀ projection is Shampoo without accumulation [§4.1, Example 5, remark].
For Embed layers the input space is one-hot (∞-normed) not RMS-normed, so the induced norm is a column-max and the duality map normalizes each column of G to unit RMS — a different map than for Linear, despite the identical matrix shape [§4.1, Example 6, Table 1].
For Conv2D the max-over-kernel-indices norm decouples the duality map into one SVD-projection per kernel spatial index [§4.1, Example 7, Table 1].
Duality maps compose recursively: for composite f₂ ∘ f₁ the composite duality map applies each sub-module’s map independently, weighted by the modular-norm coefficients that also control per-module feature-learning contribution [§4.3, Eqs. 11–12].
The paper introduces a novel rectangular Newton-Schulz iteration for computing the UVᵀ projection efficiently: after normalizing X so its singular values are ≤ √3, iterate X ← (3/2)X − (1/2)XXᵀX; the cubic f(σ) = (3/2)σ − (1/2)σ³ drives every singular value to 1 [§5.3, Eq. 15]. A degree-2k+1 family generalizes this with tunable coefficients [§5.3, Eq. 16].
The rectangular Newton-Schulz iteration was communicated to Keller Jordan via the workshop paper and used to set a new NanoGPT training-speed record [§6.2].

Method

The construction has three layers. Layer 1 (theoretical preliminaries): the dual space of a weight space $W$ is defined as the space of linear functionals on $W$ ; a duality map $d: W^\* \to W$ identifies dual vectors with primal ones. Given a norm $\lVert \cdot \rVert$ on $W$ , the canonical duality map is $d(g) = \arg\max_{\lVert t\rVert = 1} \langle g, t\rangle$ , and steepest descent under this norm is $\Delta w = -(\lVert g\rVert_\dagger / \lambda) \cdot d(g)$ (Prop. 1). Layer 2 (per-layer maps): for each atomic layer type (Linear, Embed, Conv2D), assign RMS or ∞ norms to input/output spaces based on semantics (e.g. one-hot inputs to Embed → ∞-norm), compute the induced operator norm on the weight space, and derive the corresponding duality map via SVD (Examples 5–7). The Linear case gives d(G) = √(d_out/d_in) · UVᵀ for G = UΣVᵀ. Layer 3 (composition): the modular norm from Large et al. (2024) composes per-module norms via ‖(w₁, w₂)‖ = max(s₁ · ‖w₁‖₁, s₂ · ‖w₂‖₂) where sᵢ are module-mass hyperparameters that control per-module feature learning. Because the composition is a max (not a sum), each sub-module’s duality map is invoked independently, giving straightforward all-layer updates — in contrast to prior work (Flynn 2017) that used a sum-composition and produced one-layer-at-a-time updates that had to be heuristically extended.

For fast computation of the SVD-projection UVᵀ, the paper compares three approaches: sketching (fast on well-conditioned matrices, unreliable on small singular values), Newton iteration for inverse matrix roots via (XXᵀ)^{-1/2}X = UVᵀ (well-known but breaks on low-rank matrices), and the new rectangular Newton-Schulz iteration (§5.3). The Newton-Schulz iteration is derived by observing that under spectral normalization the cubic f(σ) = aσ + bσ³ with the right coefficients drives all singular values to 1 simultaneously — mathematically identical to computing UVᵀ, but expressible entirely in matrix-multiply primitives with no SVD.

Results

The paper is theoretical, not empirical — no benchmark numbers, no ablations. The key deliverables are (i) a unified framework in which µP (Yang & Hu 2021) and Shampoo (Gupta et al. 2018) both fall out as partial approximations of the same object, and (ii) the rectangular Newton-Schulz iteration itself, whose downstream empirical impact is documented in later papers (Muon and Modded-NanoGPT speedruns). The paper reports (§6.2) that Keller Jordan used the iteration communicated in the companion workshop paper to set a new NanoGPT training-speed record; the wiki’s This Simple Optimizer Is Revolutionizing How We Train AI [Muon] and downstream Muon papers document the empirical trajectory from there.

Why it’s interesting

This is the theoretical origin document for the entire Muon lineage on the wiki. The rectangular Newton-Schulz iteration in §5.3 is exactly the primitive Muon uses to orthogonalize its momentum matrix; Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon speeds this same primitive up 2× via Gram-form iteration + IO-aware kernels, and The Newton-Muon Optimizer derives a Muon variant via a distinct Newton-analysis path that converges on the same msign family. The paper’s second big claim — that µP and Shampoo are partial approximations to one duality map — reframes the entire hyperparameter-scaling-laws vs training-stability-at-scale debate: Spectral Condition for μP under Width-Depth Scaling and Beyond MuP 3: Special Cases, Special Treatment (Embedding, LM Head, RMS Norm) can now be read as filling in which per-layer norms/masses reproduce µP’s transfer properties, while Controlled LLM Training on Spectral Sphere‘s claim that Muon is “half-aligned with µP” (controls update spectra but not weight spectra) is a claim about which pieces of the duality-map framework Muon actually implements. Finally, the modular structure — recursive composition, one duality map per module type — is what Modular Manifolds later extends from norms to full weight-space manifolds (Stiefel + spectral norm), operationalized in Training Transformers with Enforced Lipschitz Constants. For a Luma team debugging optimizer-side pretraining instabilities, this is the paper that says why per-layer learning rates and the choice of orthogonalization primitive are not independent knobs but two views of the same object.