Can Representation Learning Decouple from Loss Minimization? Polar Updates Have an Answer

Akash Kumar

cs.LGSep 28, 2026

Can Representation Learning Decouple from Loss Minimization? Polar Updates Have an Answer

Authors: Akash Kumar

Organizations: Department of Computer Science and Engineering University of California San Diego

Abstract

Does representation learning stop when the training loss stops improving? We study this question for matrix Muon, whose polar-normalised updates have a step length set by the gradient's rank rather than its norm. Near the edge of stability, full-batch Muon on teacher-student problems enters approximately period-2 loss oscillations that persist for thousands of steps: the cycle-mean loss stays flat or rises, yet the weights keep moving and the learned features continue to align with the teacher subspace. For linear teacher-student learning toys, we derive explicit cycle and alignment formulas and conditional plateau and decay bounds. For a population mean-field ReLU model, we prove that, under stated dimension, initialisation and small-head conditions, the leading eigenspace of the average gradient outer product (AGOP) recovers the teacher subspace exactly during a loss plateau, before the loss later drops. In all 33 ReLU, GELU and SiLU teacher configurations we study, direction-only alignment metrics show the student AGOP aligned with, or still aligning to, the teacher subspace during the period-2 oscillations; projected head refitting on selected configurations shows that the learned directions are useful for prediction, and further measurements distinguish AGOP alignment from weight-mass concentration. In deep residual ReLU students, freezing the downstream layers while the first layer trains with full-batch exact polar updates recreates a nearly flat cycle-mean loss with improving input-AGOP alignment; freezing and unfreezing switch between this plateau and loss decrease, and the effect is sensitive to momentum and to the choice of orthogonaliser.

Figures & tables

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Teacher	$r_{t}$	$r_{s}$	$\eta$	$1-\rho_{2}$	$L/L_{V,\mathrm{opt}}$	$\Delta\cos^{2}_{\min}$	final $\cos^{2}_{\min}$	Outcome
ReLU	4	50	2	$4.6\!\times\!10^{-7}$	3045	$-0.028$	0.670	no late gain
ReLU	6	50	1.5	$9.2\!\times\!10^{-9}$	173	$+0.016$	0.955	late gain
ReLU	8	30	1.5	$1.0\!\times\!10^{-10}$	16.2	$+0.532$	0.831	late gain
ReLU	8	50	1	$1.0\!\times\!10^{-10}$	13.3	$+0.289$	0.857	late gain
ReLU	8	50	1.5	$1.2\!\times\!10^{-10}$	40.9	$+0.189$	0.907	late gain
ReLU	12	50	1.5	$1.0\!\times\!10^{-10}$	23.1	$+0.077$	0.820	late gain

Appendix

Table 2: All

33

one-hidden-layer teacher configurations (

p=100

n=15{,}000

8{,}000

steps,

5

seeds each:

5,11,23,37,41

; medians over seeds).

1-\rho_{2}

: second-half period-

2

signature (values below

10^{-10}

shown as

10^{-10}

);

L/L_{V,\mathrm{opt}}

: cycle-mean training loss over the last

30\%

of steps divided by the head oracle;

\Delta\cos^{2}_{\min}

: change from step

4{,}000

to the end; final

\cos^{2}_{\min}

at the last step. Outcome, from the medians: not recovered if final

\cos^{2}_{\min}<0.2

; otherwise aligned by midpoint if

\Delta\cos^{2}_{\min}<0.005

and final

\cos^{2}_{\min}\geq 0.95

; no late gain if

\Delta\cos^{2}_{\min}<0.005

otherwise; late gain if

\Delta\cos^{2}_{\min}\geq 0.005

. This five-seed grid is distinct from the single-seed search above and from the

53

-seed runs of Figures 2 and 3 .

Config $(\sigma_{t},r_{t},r_{s})$	Predicted	Measured	Predicted	Measured	$M$	Measured
	$V$ -mass		$V^{\perp}$ -mass		Total mass
ReLU, $r_{t}=4$ , $r_{s}=50$	4	3.79	46	46.21	50	50.00
ReLU, $r_{t}=8$ , $r_{s}=50$	8	7.86	42	42.14	50	50.00
ReLU, $r_{t}=12$ , $r_{s}=50$	12	11.76	38	38.24	50	50.00

Appendix

Table 3: Polar mass split on three long-run ReLU checkpoints. The total equals

M

to numerical precision (Proposition 3 (i)); the

V

-mass is approximately

2

–

5\%

below

r_{t}

. The gradient is recomputed from the saved weights with a unit output head, and the leakage scalars of Proposition 3 (iii) were not measured separately.

input dimension $p$	$64$	teacher rank $r_{t}$	$4$
student width $m$	$32$	teacher depth	$2$ Eq. ( 77 )
training inputs $n$	$4096$	test inputs $n_{\mathrm{test}}$	$8192$
AGOP inputs $n_{m}$	$2048$	steps $T$	$1000$
initial gain $g$	$1$	head scale $c_{a}$	$1$
residual scale $\alpha$	$1/(L-1)$	architecture	residual Eq. ( 78 )

Appendix

Table 7: Settings shared by every deep run unless an arm overrides them.

arm	$L$	$\eta_{1}$	$\eta_{\mathrm{deep}}$	$c_{a}$	$\alpha$
D6	6	0.66	$3.125\times 10^{-3}$	0.75	1/5
D8	8	0.66	$7.8125\times 10^{-4}$	1	1/7
D10	10	0.70	$3.90625\times 10^{-4}$	1	1/9
D12	12	0.70	$1.953125\times 10^{-4}$	1	1/11
D16	16	0.74	$1.953125\times 10^{-4}$	1	1/15

Appendix

Table 8: Depth arms. The fresh-D8 block reuses the D8 configuration with new seeds.

arm	strict runs	problems	$\overline{\cos^{2}}$ gain	$\cos^{2}_{\min}$ gain	cycle mean	held-out refit
D6	94/100	20/20	$+0.120$	$+0.384$	$+2.62\%$	$-6.15\%$
D8	97/100	20/20	$+0.206$	$+0.515$	$+0.62\%$	$-6.88\%$
D10	96/100	19/20	$+0.200$	$+0.441$	$+1.10\%$	$-5.38\%$
D12	95/100	20/20	$+0.193$	$+0.436$	$+1.55\%$	$-5.20\%$
D16	96/100	19/20	$+0.187$	$+0.418$	$+1.75\%$	$-4.60\%$
Fresh D8	91/100	20/20	$+0.153$	$+0.396$	$+1.24\%$	$-6.58\%$

Appendix

Table 9: Depth campaigns under the strict protocol (Appendix D.2 ): runs passing, problems succeeding, and, as medians over problems of the per-problem median over initialisations, the window gains of

\overline{\cos^{2}}

and

\cos^{2}_{\min}

and the relative change over the window of the cycle-mean training loss and of the held-out refit loss.

model	data	architecture	LN removed	$T_{0}$	$\delta$
transformer	Gaussian tokens	1 block, width 64, 4 heads	output	0	1.0
ViT, four $4\times 4$ patches	digits	1 block, width 64, 4 heads	all	500	0.03
ViT, one $8\times 8$ patch	digits	1 block, width 64, 4 heads	all	500	1.0
CNN	digits	3 conv. layers, channels 16/32/64	output	200	0.1

Appendix

Table 10: Settings of the architecture runs. Every run minimises the MSE through a frozen head with orthonormal rows, full batch. Stage 1 trains all non-head parameters with AdamW (learning rate

10^{-3}

, zero weight decay) for

T_{0}

steps and is omitted when

T_{0}=0

. Stage 2 trains only the first-layer (input-embedding) weight by

W\leftarrow W-\eta\operatorname{Polar}(G)

, where

G=\nabla_{W}L_{\mathrm{train}}

, for

1600

steps, scored on steps

200

–

1600

. The update coefficient is

\eta=\delta\|W_{\mathrm{start}}\|_{F}/\sqrt{M}

, where

W_{\mathrm{start}}

is the stage-2 initial weight and

M=\min(n_{\mathrm{rows}},n_{\mathrm{cols}})

after flattening any convolutional kernel to a matrix:

M=32,16,64,9

for the four rows below, respectively. Transformer data: sequences of

8

Gaussian tokens in

\mathbb{R}^{32}

; the teacher applies one attention layer and a ReLU layer of width

16

to the tokens’ projections onto a rank-

4

subspace, averages over tokens and outputs

4

standardised targets;

1024

training sequences, and ten fresh planted problems in the confirmation runs. Digit data:

500

training and

500

test images, one-hot targets.

part of Theorem 15	proved in	main ingredients
trajectory	Section G.2	Lemmas 16 , 17 , 19 ; Proposition 20
(P)	Theorem 41	Lemmas 30 , 35 , 38
(F) at $\alpha=0$	Theorem 51	Propositions 23 , 26 , 57 ; Lemmas 48 , 49 , 50 , 60 , 61 , 63 ; Proposition 64
(F) at small $\alpha>0$	Section G.7 , Step 3	Lemma 27 ; Proposition 28 (a)
(F) at $t=0$	Section G.7 , Step 4	Proposition 26 (4)
(R)	Corollary 44	Theorem 43 ; Lemmas 30 , 38

Appendix

Table 11: Where the parts of Theorem 15 are proved.

$r$	$[D_{\rm lo},D_{\rm hi}]$	$\tau_{F}$	$\kappa_{0}$	margin	$r$	$[D_{\rm lo},D_{\rm hi}]$	$\tau_{F}$	$\kappa_{0}$	margin
2	$[400,14930]$	0.013740	336	1.006223	41	$[2583,3363]$	0.078835	126	1.228737
2	$[14931,465350]$	0.014026	777	1.048716	42	$[2646,3408]$	0.079339	125	1.230167
3	$[400,119688]$	0.014191	286	1.012711	43	$[2709,3453]$	0.079824	124	1.231501
4	$[400,55355]$	0.017549	172	1.025338	44	$[2772,3497]$	0.080300	124	1.232778
5	$[400,32467]$	0.022016	136	1.029847	45	$[2835,3542]$	0.080758	123	1.233975
6	$[400,21679]$	0.025939	119	1.033781	46	$[2898,3587]$	0.081197	122	1.235096

Appendix

Table 12: The 80 region-N certificates at

\alpha=0

D

-interval,

\tau_{F}

, floor

\kappa_{0}

and certified margin

\min(\lambda_{1},\lambda_{c})/\lambda_{\perp}

t_{F}

row	$K$	$h$	$j$	$T$
2a	21600665	35	0	7920963464610295
2b	36059480	3	1	8085720989539134
3	29042409	35	0	8180717210882235
4	35333865	4	0	10116431691206438
5	34521219	4	0	12691884803174992
6	35458646	4	0	14952843575107892

Appendix

Table 13: Exact partition inputs for region N (rows 1–40).

row	$K$	$h$	$T$
41	90337968	3	45445623816640352
42	91790245	3	45735890048522056
43	93238329	3	46015601144698960
44	94680121	3	46290034673023472
45	96119816	3	46553913065643208
46	97556365	3	46807236322558152

Appendix

Table 14: Exact partition inputs for region N (rows 41–80).

row	$O$	$B$	$C$
2a	5292723980082718	32428288108055624	5608012396498452
2b	6759749330873526	47234666444723024	4567102900135213
3	6150004491384073	27557958287837056	5960833993768631
4	6813178771523468	24649586303920624	5919624845377477
5	6805068924342156	19562686959486248	5979297203968494
6	6961059657208525	17116075937116040	6087232431130972

Appendix

Table 15: Exact supersolution inputs for region N (rows 1–40).

row	$O$	$B$	$C$
41	12337113433222418	7633481250065179	5989283289725050
42	12446397115720632	7581022616188417	5984719452569983
43	12554428182284166	7530683027299399	5980389152653382
44	12661257956522734	7482341945772299	5976273220784853
45	12766998498356294	7435959038899650	5972366854639561
46	12871597276411164	7391309963979724	5968652097336570

Appendix

Table 16: Exact supersolution inputs for region N (rows 41–80).

Representation Learning

Representation learning develops data transformations that expose useful structure for prediction, transfer, generation, and control.

See all 813 papers

Awakening of the Buddha: Subspace Learning During Population-Loss Plateaus
Sep 30, 2026Akash KumarRepresentation Learning ReLU Neural Networks
Department of Computer Science & Engineering University of California-San Diego
Training-Loss Guarantees for Muon with Finite-Step Newton--Schulz Orthogonalization
Oct 2, 2026Amartya Roy, Souvik ChakrabortyMuon Optimizer Neural Network Training Dynamics
The School of Interdisciplinary Research Indian Institute of Technology Delhi, India · Robert Bosch GmbH, India · Department of Applied Mechanics Indian Institute of Technology Delhi, India +1
The Spectral Dynamics and Noise Geometry of Muon
Jun 7, 2026Pierfrancesco Beneventano, Mahmoud Abdelmoneum, Tomaso PoggioMuon Optimizer Matrix Optimization
Massachusetts Institute of Technology

Sep 30, 2026J/K move · Enter open · S save

Awakening of the Buddha: Subspace Learning During Population-Loss Plateaus

Akash Kumar

Department of Computer Science & Engineering University of California-San Diego

Representation Learning ReLU Neural Networks Neural Network Training Dynamics

Open paper AskPDF arXiv

Population loss can remain nearly constant while a neural network learns a substantially more predictive representation. We establish this separation for two-layer ReLU and leaky-ReLU networks trained on Gaussian inputs by simultaneous fixed-step population gradient descent on all parameters. For structured additive teachers whose links are positive mixtures of Gaussian-damped cubics in

H^1(γ)

, we give explicit conditions under which small IID Gaussian initialization yields a high-probability guarantee: at a checkpoint during a high-loss plateau, minimum alignment between the rank-

r

teacher subspace and the leading

r

-dimensional eigenspace of the predictor's average gradient outer product (AGOP) increases by at least

1/2

, and the minimum refit MSE under unchanged coefficient budgets decreases by more than

0.399

, both relative to initialization. The same trajectory subsequently attains a trained loss below every value in the plateau window. A complementary result treats unequal-weight cubic teachers and small additive Sobolev perturbations using projected-feature refits. For SwiGLU networks with an exactly fitted intercept, we prove leading-AGOP alignment during a loss plateau at fixed width and dimension as Gaussian initialization vanishes, for square-integrable teachers with nonzero Hermite content of degree one, two, or three. A rank-one cubic specialization also gives simultaneous unrestricted-refit gains at a prescribed width. An approximation lower bound further shows that certain interaction targets retain nonzero error when ridge neurons are restricted to shared orthogonal axes within the teacher subspace. Population-moment experiments with ReLU students across 21 teachers and 50 initializations per teacher complement the analysis.

Abstract

Representation Learning

Explore similar work

Awakening of the Buddha: Subspace Learning During Population-Loss Plateaus

Training-Loss Guarantees for Muon with Finite-Step Newton--Schulz Orthogonalization

The Spectral Dynamics and Noise Geometry of Muon