Which Geometry on Which Layer? Measured Sign-Spectral Assignment in Mixed-Optimizer Training

Hiroki Naganuma; Laura Gomezjurado; Mahdi Ghaznavi; Ganesh Talluri; Atsushi Nitanda; Seng Pei Liew; Ryuichiro Hataya; Ioannis Mitliagkas

cs.LGOct 12, 2026

Which Geometry on Which Layer? Measured Sign-Spectral Assignment in Mixed-Optimizer Training

Authors: Hiroki Naganuma, Laura Gomezjurado, Mahdi Ghaznavi, Ganesh Talluri, Atsushi Nitanda, Seng Pei Liew, Ryuichiro Hataya, Ioannis Mitliagkas

Organizations: Mila · Université de Montréal · Stanford University · Arizona State University ·

Abstract

Modern neural-network training uses mixed recipes, assigning AdamW-like or sign-based updates to embedding and output layers and spectral updates to hidden matrices. These assignments are usually fixed by layer type, without a clear criterion for selecting the appropriate geometry. We turn sign-versus-spectral assignment into a measurable, layerwise decision problem. For ideal memoryless steps, we derive a finite-batch ratio of conservative descent-lemma lower bounds for spectral and sign updates. Its closed form separates signal efficiency from a geometry-dependent global curvature factor. In the positive-signal regime, the same smoothness assumption gives two-sided bounds that determine which step has the larger conditional expected decrease when the ratio crosses explicit noise-dependent thresholds. We then derive a local directional quadratic score from gradient alignments and Hessian-vector products and use it in Generalized-Scion (G-Scion) to control layerwise routing. In a controlled input-rank experiment, including directional curvature matches the measured same-iterate ordering in all seven aggregate settings, whereas the curvature-free proxy selects the opposite geometry. On 31M GPT-2, the output score routes the output layer from spectral warmup to sign, and G-Scion improves validation loss over fixed Scion in memoryless and matched-momentum comparisons. ViT-B/16 results further show that a spectral edge assignment can outperform fixed Scion outside language modelling. These results suggest that optimizer geometry need not be hard-coded by layer type and are a first step toward understanding how it should be assigned.

Figures & tables

Method	Primal norm $\left\\|\cdot\right\\|$	Dual norm $\left\\|\cdot\right\\|_{*}$	Direction $D^{*}(G)$
(S)GD	Frobenius $\left\\|\cdot\right\\|_{\mathrm{F}}$	Frobenius $\left\\|\cdot\right\\|_{\mathrm{F}}$	$G/\left\\|G\right\\|_{\mathrm{F}}$
SignSGD	Entrywise $\ell_{\infty}$	Entrywise $\ell_{1}$	$\sign(G)$
SpecSGD	Spectral $\\|\cdot\\|_{\mathrm{op}}$	Nuclear $\\|\cdot\\|_{\mathrm{nuc}}$	$\mathrm{polar}(G)=UV^{\top}$

Table 1: Three optimizer geometries as norm-steepest descent.

Figure 1: Per-layer

\log{\color[rgb]{0,0,0}\widehat{\mathcal{R}}_{\mathrm{dir}}}

diagnostics (OWT,

31

2{,}000

steps; mean

\pm

std across three activations per method). Left: embedding diagnostics are descriptive because eligible embeddings were already forced spectral during warmup and their first readings differ by variant. Right: output-gated variants start spectral with a clearly negative score, satisfy the two-confirmation rule at

t{=}150

(dashed), and route to sign. Scion is diagnostic only and does not gate.

Assignment or routing variant	Final validation loss
PP: spectral on both edge groups	4.92
Frozen-route replay	5.23
G-Scion with warm-up	5.23
SP: spectral on output only	5.35
PS: spectral on embedding only	5.40
G-Scion	5.45

Table 3: Reported

31

M GPT-2/OWT follow-up. S denotes sign and P spectral; the embedding assignment is listed first. Lower validation loss is better. Values are rounded figure labels; per-seed values and the plotted error-bar definition are not available in the accompanying summary, so no uncertainty estimates or significance tests are inferred.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning
$W\in\mathbb{R}^{m\times n}$	weight matrix of a given layer
$m,n$	output / input dimension of the layer
$p:=\min(m,n)$	smaller side of the matrix
$f:\mathbb{R}^{m\times n}\to\mathbb{R}$	(possibly regularised) training loss at parameter $W$
$G:=\nabla f(W)$	full-batch (population) gradient at $W$
$G=U\Sigma V^{\top}$	thin SVD of $G$ , with $\Sigma\in\mathbb{R}^{p\times p}$

Appendix

Table 4

Symbol	Meaning
$\left\langle A,\,B\right\rangle:=\mathrm{tr}(A^{\top}B)$	Frobenius inner product
$\left\\|\cdot\right\\|_{\mathrm{F}}$	Frobenius norm, $\left\\|A\right\\|_{\mathrm{F}}=\sqrt{\left\langle A,\,A\right\rangle}$
$\\|\cdot\\|_{\mathrm{op}}=\left\\|\cdot\right\\|_{\mathrm{op}}$	spectral (operator) norm, largest singular value
$\\|\cdot\\|_{\mathrm{nuc}}=\left\\|\cdot\right\\|_{*}$	nuclear (trace) norm, sum of singular values
$\\|\cdot\\|_{1}=\left\\|\cdot\right\\|_{1}$	entrywise $\ell_{1}$ norm, $\sum_{ij}\|A_{ij}\|$
$\left\\|\cdot\right\\|_{\infty}$	entrywise $\ell_{\infty}$ norm, $\max_{ij}\|A_{ij}\|$

Appendix

Table 5

Symbol	Meaning
$L_{\left\\|\cdot\right\\|}$	smoothness constant of $f$ w.r.t. a primal norm $\left\\|\cdot\right\\|$
$L_{\mathrm{F}}$	Frobenius smoothness constant
$L_{\mathrm{op}}$	operator (spectral) smoothness constant
$L_{\infty}$	entrywise $\ell_{\infty}$ smoothness constant
$\kappa:=L_{\infty}/L_{\mathrm{op}}$	curvature ratio (central configuration-side quantity)
$C(D):=\left\langle D,\,\nabla^{2}f(W)[D]\right\rangle$	directional curvature along unit direction $D$

Appendix

Table 6

Symbol	Meaning
$D^{*}(G)\in\argmax_{\left\\|D\right\\|\leq 1}\left\langle G,\,D\right\rangle$	norm-steepest-descent direction for primal norm $\left\\|\cdot\right\\|$
$D_{\mathrm{sign}}:=\sign(\widehat{G})$	SignSGD update direction
$D_{\mathrm{spec}}:=\mathrm{polar}(\widehat{G})$	SpecSGD / Muon-like update direction
$W^{+}$	next iterate after an ideal scaled step from $W$
$\Delta_{g}(W;\eta)$	conditional expected one-step decrease under geometry $g$ at step $\eta$
$B_{\mathrm{sign}}(W)$	SignSGD conservative lower bound (Eq. ( 3 ))

Appendix

Table 7

Symbol	Meaning
$E:=\widehat{G}-G$	stochastic gradient noise matrix
$\nu_{1}(W):=\mathbb{E}[\\|\widehat{G}-G\\|_{1}\mid W]$	$\ell_{1}$ (sign-geometry) noise scale
$\nu_{*}(W):=\mathbb{E}[\\|E\\|_{\mathrm{nuc}}\mid W]$	nuclear-norm (spectral-geometry) noise scale
$c_{1}:=\sqrt{B}\cdot\nu_{1}$	batch-size-normalised $\ell_{1}$ noise scale
$c_{}:=\sqrt{B}\cdot\nu_{}$	batch-size-normalised nuclear noise scale
$\mathrm{nr}(G):=\\|G\\|_{\mathrm{nuc}}^{2}/\left\\|G\right\\|_{\mathrm{F}}^{2}$	nuclear rank (effective rank in the $*$ -vs- $\mathrm{F}$ sense)

Appendix

Table 8

Symbol	Meaning
$\Delta_{g}(W;\eta)$	actual expected decrease for geometry $g$ at step size $\eta$
$B_{g}(W,B)$	lower-bound quantity for the stated conservative step
${\color[rgb]{0,0,0}\mathcal{R}_{\mathrm{LB}}}(W,B)$	finite-batch geometry ratio $B_{\mathrm{spec}}/B_{\mathrm{sign}}$ (Eq. ( 5 ))
$\mathcal{R}_{\infty}(W)$	high-SNR form of the global-smoothness ratio (Eq. ( 9 ))
$\mathcal{R}_{0}(W):=\\|G\\|_{\mathrm{nuc}}^{2}/\\|G\\|_{1}^{2}$	curvature-free proxy (Eq. ( 10 ))
${\color[rgb]{0,0,0}\widehat{\mathcal{R}}_{\mathrm{dir}}}(W)$	empirical directional quadratic score (Eq. ( 11 ))

Appendix

Table 9

Symbol	Meaning
$\sigma(\cdot)$	pointwise activation function
$X_{\ell}$	pre-activation at layer $\ell$ : $X_{\ell}=W_{\ell}A_{\ell-1}$
$A_{\ell}:=\sigma(X_{\ell})$	post-activation at layer $\ell$
$\mathrm{st}(A):=\left\\|A\right\\|_{\mathrm{F}}^{2}/\\|A\\|_{\mathrm{op}}^{2}$	stable rank of a matrix $A$
$\mathrm{er}(A)$	effective rank (alternative rank surrogate; see Appendix D )
$s_{\sigma}:=\mathbb{E}[\sigma(\gamma)^{2}]/(\mathbb{E}[\sigma(\gamma)])^{2}$	activation noise-to-signal ratio, $\gamma\sim\mathcal{N}(0,1)$

Appendix

Table 10

Name	Meaning
(S)GD	(stochastic) gradient descent; Frobenius-norm steepest descent
SignSGD, SignSGD	sign-based steepest descent under entrywise $\ell_{\infty}$ geometry
SpecSGD, SpecSGD	spectral steepest descent under operator-norm geometry; direction $\mathrm{polar}(\widehat{G})$
Muon	practical spectral-geometry optimizer used as our SpecSGD instantiation
AdamW	Adam with decoupled weight decay (standard baseline)
Scion	fixed mixed recipe: SignSGD on embed / head, Muon on hidden layers

Appendix

Table 11

Symbol	Meaning
$\mathbb{R},\mathbb{N},\mathbb{Z}$	real numbers, natural numbers, integers
$\mathbb{E}[\cdot]$	expectation (over minibatch sampling unless specified otherwise)
$\mathrm{Var}[\cdot],\mathrm{Cov}[\cdot,\cdot]$	variance, covariance
$\mathcal{N}(\mu,\Sigma)$	Gaussian distribution with mean $\mu$ and covariance $\Sigma$
$\mathbb{1}[\cdot]$	indicator function
$(x)_{+}:=\max(x,0)$	non-negative part

Appendix

Table 12

$r$	$\hat{\kappa}_{\mathrm{dir}}$	$\mathcal{R}_{0}$	$R_{\mathrm{quad}}$	$R_{\mathrm{emp}}$
$2$	$20{,}100$	$5.7\!\cdot\!10^{-5}$	$1.14$	$1.13$
$4$	$14{,}100$	$9.6\!\cdot\!10^{-5}$	$1.36$	$1.33$
$8$	$13{,}700$	$1.5\!\cdot\!10^{-4}$	$2.04$	$2.03$
$16$	$\phantom{1}8{,}100$	$2.9\!\cdot\!10^{-4}$	$2.31$	$2.24$
$32$	$\phantom{1}7{,}000$	$4.7\!\cdot\!10^{-4}$	$3.27$	$3.09$
$64$	$\phantom{1}5{,}200$	$7.9\!\cdot\!10^{-4}$	$4.09$	$3.54$

Appendix

Table 4: Toy directional-curvature sweep at block0.mlp.c_fc , step

200

, GELU,

24

trials per row (geometric means).

R_{\mathrm{quad}}/R_{\mathrm{emp}}\in[1.01,1.31]

across the aggregate rows ;

\mathcal{R}_{0}

disagrees in sign at every

r

Figure 2: Toy directional-curvature sweep at block0.mlp.c_fc (step

200

, GELU,

24

trials per point). As input rank

r

varies

2\to 128

, the directional curvature ratio

\hat{\kappa}_{\mathrm{dir}}

(gray, right axis) moves over

\sim\!5\times

. The deterministic score

R_{\mathrm{quad}}

(blue) and empirical one-step ratio

R_{\mathrm{emp}}

(green) move together. The curvature-free proxy

\mathcal{R}_{0}

(red) predicts the wrong side throughout. The intervention also changes

\mathcal{R}_{0}

, so the figure demonstrates the empirical value of retaining directional curvature rather than varying curvature alone .

Activation $\sigma(t)$	Unit-scale $s_{\sigma}$	Scope
ReLU $\max(0,t)$	$\pi\approx 3.14$	homogeneous corollary
SiLU $t(1+e^{-t})^{-1}$	$\approx 8.33$	descriptive; non-homogeneous
GELU $t\,\Phi(t)$	$\approx 5.34$	descriptive; non-homogeneous
ReLU 2 $(\max(0,t))^{2}$	$6$	homogeneous corollary

Appendix

Table 5: Unit-scale activation moment ratios. For non-homogeneous SiLU and GELU these values are descriptive and are not the scale-uniform constants required by the full theorem .

Figure 3: Batch-size sweep (Tiny Shakespeare, G-Scion variants). Left: Training loss curves at effective batch

256

(

1{,}600

steps). G-Scion (embedding, output) leads throughout; G-Scion (embedding) tracks Scion closely, while G-Scion (output) and G-Scion (embedding, output) pull away via output-layer adaptation. Right: Mean last-

20

training loss vs. effective batch size (

\log_{2}

scale). All G-Scion variants beat Scion at every reported batch size; the observed margin is largest at small batch and shrinks as the effective batch grows.

Activation	Scion	G-Scion (embed.)	G-Scion (out.)	G-Scion (embed., out.)	Best–Worst
Single seed (seed 1337)
GELU	$5.526$	$5.278$	$5.459$	$\mathbf{5.161}$	$0.365$
ReLU 2	$5.490$	$5.293$	$5.421$	$\mathbf{5.196}$	$0.294$
SiLU	$5.531$	$5.309$	$5.462$	$\mathbf{5.205}$	$0.326$
Mean	$5.516$	$5.293$	$5.447$	$\mathbf{5.187}$
$\Delta$ vs. Scion	—	$+0.22$	$+0.07$	$\mathbf{+0.33}$

Appendix

Table 6: Final validation loss, OWT

31

{\color[rgb]{0,0,0}\widehat{\mathcal{R}}_{\mathrm{dir}}}

. Top: single-seed runs (

12

runs, seed 1337). Bottom: two-seed replication (

24

runs, seeds 42 and 2024; mean

\pm

population std). The validation-loss ordering (lower is better) G-Scion (embedding, output)

<

G-Scion (embedding)

<

G-Scion (output)

<

Scion is uniform across all activations and both seed sets.

Figure 4: Activation $\times$ method validation loss (OWT, 31M, 2 000 steps, ${\color[rgb]{0,0,0}\widehat{\mathcal{R}}_{\mathrm{dir}}}$ ; mean of last-5 checkpoints). Lower numerical values indicate better validation loss. The validation-loss ordering (lower is better) G-Scion (embedding, output) $<$ G-Scion (embedding) $<$ G-Scion (output) $<$ Scion holds uniformly across all three activations. Activation function is a second-order effect (within-column spread $<0.05$ in validation loss ).

Figure 5: Validation loss curves (OWT, 31M, 2 000 steps,

{\color[rgb]{0,0,0}\widehat{\mathcal{R}}_{\mathrm{dir}}}

) for all three MLP activations. Colours: G-Scion (embedding, output) (blue), G-Scion (embedding) (green), G-Scion (output) (orange), Scion (vermilion). Ordering is preserved across activations; G-Scion (embedding, output) leads in every panel.

Method	Final val loss	$\Delta$ vs. Muon
Muon (AdamW embed/head + Muon hidden)	$4.755\pm 0.035$	—
G-Muon (AdamW $\to$ Spectral embed/head + Muon hidden)	$4.851\pm 0.016$	$-0.096$

Appendix

Table 7: AdamW-side reference at

31

M (same setting as Table 2 ). Muon : AdamW on embed/head + Muon on hidden, the standard practical recipe. G-Muon : one-way AdamW

\to

spectral gate on embed/head + Muon on hidden. AdamW adds momentum and per-parameter adaptivity outside the one-step bound, so this is a reference rather than a like-for-like sign-vs-spectral comparison; the gate does not improve over Muon in this regime. The scope of the theorem does not explain this performance difference.

LR	Out $\log{\color[rgb]{0,0,0}\widehat{\mathcal{R}}_{\mathrm{dir}}}$ , spectral phase	Out $\log{\color[rgb]{0,0,0}\widehat{\mathcal{R}}_{\mathrm{dir}}}$ , sign phase
$10^{-4}$	$-1.83$ ( $100\%$ sign-side)	$-1.76$ ( $100\%$ sign-side)
$3{\times}10^{-4}$	$-2.46$ ( $100\%$ sign-side)	$-1.94$ ( $100\%$ sign-side)
$10^{-3}$	$-1.12$ ( $100\%$ sign-side)	$-0.29$ ( $72\%$ sign-side)
$3{\times}10^{-3}$	$-0.27$ ( $100\%$ sign-side)	$\mathbf{+0.70}$ ( $19\%$ sign-side)

Appendix

Table 8: Output

\log{\color[rgb]{0,0,0}\widehat{\mathcal{R}}_{\mathrm{dir}}}

by learning rate and phase (OWT

31

M, 80-run sweep,

{\color[rgb]{0,0,0}\widehat{\mathcal{R}}_{\mathrm{dir}}}

). “Spectral phase”: steps

0

–

100

(before gate). “Sign phase”: steps

150

–

1950

(after gate). All gated-output runs switch at step

150

regardless of LR.

Effective batch	Scion	G-Scion (embedding)	G-Scion (output)	G-Scion (embedding, output)	$\Delta$
$128$	$30.9\pm 1.2$	$32.5\pm 2.1$	$43.8\pm 2.8$	$\mathbf{42.8\pm 0.4}$	$+11.9$
$256$	$39.5\pm 0.7$	$41.0\pm 0.4$	$45.3\pm 0.4$	$\mathbf{51.0\pm 1.1}$	$+11.5$
$512$	$40.0\pm 0.7$	$43.9\pm 2.7$	$44.1\pm 1.6$	$\mathbf{47.9\pm 1.2}$	$+\phantom{0}7.9$

Appendix

Table 9: Final validation accuracy on ViT-B/16 CIFAR-100 (

1{,}000

steps,

2

seeds per cell, lr

{=}3{\times}10^{-3}

). The joint spectral-edge assignment trajectory exceeds fixed Scion at every reported effective batch size.

\Delta

is its accuracy minus Scion accuracy in percentage points.

Selection rule	Correct / 288	Mean regret (nats/token)
Full-group directional diagnostic + fallback	219	0.005037
Always sign	24	0.024432
Always spectral	264	0.003471
$Q_{\mathrm{actual}}$ (development only)	288	0.000000

Appendix

Table 10: Development scheduled-update ranking. Regret is selected held-out loss minus the better sign/spectral loss, averaged over the 288 tuples. The four undefined directional ratios retain the source-route fallback.

Spectral Optimization

See all 22 papers

May 22, 2026

Spectral Asymptotics of Neural Network Loss Landscapes: An Exact Decomposition of the Curvature Exponent

The curvature exponent

α

h_k \propto σ_k^α

-- governing how Hessian eigenvalues scale with gradient singular values -- varies systematically across layer types (

α\approx 2

for convolutions,

\approx 1

for transformer attention,

< 1

for MLP up-projections). Why? We prove the Spectral Alignment Decomposition:

α= 2 + d\logΦ_k / d\logσ_k

, where

Φ_k

measures alignment between Kronecker factor eigenbases and gradient singular directions. This reduces "why does

α

vary?" to a geometric question we answer for LayerNorm, residual connections, and softmax heads. The decomposition implies a spectral transfer identity

s = αγ

linking curvature exponent, effective gradient rank-decay

γ

, and Hessian decay exponent

s

. The identity is algebraic; its empirical content is that

α

and

γ

, fit on independent data (HVPs vs. SVD), recover

s

to ~2% median error across 93 layers, five architectures, and three datasets -- with no free parameters. A zeta-function bound on participation ratio shows curvature concentrates onto effectively one direction per layer. As a proof of concept, we derive the architecture-adaptive preconditioner

T(σ;α)

and show that Spectral Newton -- implementing

T

in the gradient singular basis -- outperforms AdamW on vision benchmarks where

α\approx 2

Anherutowa Calvo

19D Labs

Neural Network Optimization Second-Order Optimization

Ask PDF

Mar 14, 2026

Spectral Edge Dynamics of Training Trajectories: Signal--Noise Geometry Across Scales

Despite hundreds of millions of parameters, transformer training trajectories evolve within only a few coherent directions. We introduce Spectral Edge Dynamics (SED) to quantify this structure: a rolling-window SVD of parameter updates reveals a sharp boundary -- the spectral edge -- between coherent optimization directions and stochastic noise, identified via the maximum consecutive singular value ratio

σ_k / σ_{k+1}

. Across a 51M-parameter TinyStories model (4 seeds) and GPT-2 124M under distribution shift, the spectral edge exhibits a universal three-phase pattern (rise, plateau, collapse). The effective signal rank adapts to task complexity (

k^* = 2

at 51M,

k^* = 3

at 124M), and the directional coupling between spectral geometry and validation loss reverses with window size -- a lag flip reflecting the timescale of trajectory integration. Johnson--Lindenstrauss projection to

d = 10W

dimensions (e.g.,

d = 100

for

W = 10

) preserves the spectral gap within

5.7\%

, making the framework applicable to models of arbitrary scale. In companion work, the same spectral geometry provides early-warning signals of grokking -- predicting generalization

600

1{,}700

steps before it occurs across modular arithmetic, Dyck languages, and the SCAN benchmark.

Yongzhong Xu

Transformer Training Dynamics Neural Network Training Dynamics Gradient Descent Dynamics Spectral Methods

Ask PDF

Jun 11, 2026

Different Layers, Different Manifolds: Module-Wise Weight-Space Geometry in Transformer Optimization

Weight-space geometry plays a central role in neural network optimization, yet manifold constraints are often applied uniformly across all weight matrices. In this work, we ask whether different transformer modules prefer different manifold geometries. We study Manifold Muon for GPT-2 pretraining and compare layer-wise assignments of Stiefel and DGram constraints across attention and MLP blocks. Our results show a clear asymmetry: constraining attention layers with Stiefel geometry while assigning DGram geometry to MLP layers gives the best performance among the tested configurations, whereas the inverted assignment and all-DGram configuration become unstable under the shared hyperparameter setting. We trace this failure to singular value growth in DGram-constrained attention weights, which can amplify attention logits and induce softmax saturation. These findings suggest that symmetry-aware and geometry-aware optimization for transformers should be module-specific rather than uniform.

Kirato Yoshihara

School of Engineering Science, The University of Osaka, Osaka, Japan.

Stiefel Manifold Optimization Neural Network Optimization Weight-Space Analysis Constrained Optimization Neural Network Training Dynamics

Ask PDF

Symbol	Meaning
$\left\langle A,\,B\right\rangle:=\mathrm{tr}(A^{\top}B)$	Frobenius inner product
$\left\\|\cdot\right\\|_{\mathrm{F}}$	Frobenius norm, $\left\\|A\right\\|_{\mathrm{F}}=\sqrt{\left\langle A,\,A\right\rangle}$
$\\|\cdot\\|_{\mathrm{op}}=\left\\|\cdot\right\\|_{\mathrm{op}}$	spectral (operator) norm, largest singular value
$\\|\cdot\\|_{\mathrm{nuc}}=\left\\|\cdot\right\\|_{*}$	nuclear (trace) norm, sum of singular values
$\\|\cdot\\|_{1}=\left\\|\cdot\right\\|_{1}$	entrywise $\ell_{1}$ norm, $\sum_{ij}\|A_{ij}\|$
$\left\\|\cdot\right\\|_{\infty}$	entrywise $\ell_{\infty}$ norm, $\max_{ij}\|A_{ij}\|$

Abstract

Spectral Optimization

Explore similar work

Spectral Asymptotics of Neural Network Loss Landscapes: An Exact Decomposition of the Curvature Exponent

Spectral Edge Dynamics of Training Trajectories: Signal--Noise Geometry Across Scales

Different Layers, Different Manifolds: Module-Wise Weight-Space Geometry in Transformer Optimization