Same Function, Different Trajectory: No Standard Optimizer Is Equivariant to Transformer Reparameterization

Bum Jun Kim

cs.LGOct 13, 2026

Same Function, Different Trajectory: No Standard Optimizer Is Equivariant to Transformer Reparameterization

Authors: Bum Jun Kim

Organizations: The University of Tokyo

Abstract

We show that two Transformers representing the same function can follow different training trajectories. For example, doubling a query matrix and halving its key matrix preserves the model's function but changes the gradient scales seen by the optimizer. Our theoretical analysis establishes that standard block-local first-order optimizers fail to scale updates with parameters under a shared learning rate. For AdamW, transforming the first- and second-moment estimates alone does not correct the mismatch between parameter and update scales. To make runs restarted from equivalent checkpoints follow the same trajectory, we balance matrix pairs and transform their moment estimates together, bringing distinct but equivalent training states to the same starting point. To correct scale imbalance between paired Transformer matrices without changing predictions, we propose Symmetry-Aware Gradient Geometry Stabilization (SAGGS), which periodically equalizes the Frobenius norms of exact symmetry pairs, and its damped variant, Soft SAGGS. SAGGS reduces perplexity (PPL) by $0.374\pm0.026$ at 101M parameters on WikiText-103, and Soft SAGGS reduces PPL by $0.780\pm0.040$ in a 31M modern-architecture replication, both under scale imbalance. In a 123.8M checkpoint-resume experiment, reciprocally rescaling matrix pairs barely changes validation loss at restart but increases final PPL by $3.78\pm0.42$ ; hard and soft rebalancing recover to within $0.02$ PPL of the control.

Figures & tables

Figure 1: Same function, different trajectory. (1) Multiplying

\bm{A}

c

and dividing its partner

\bm{B}

c

leaves every prediction unchanged. The rectangle’s sides are the Frobenius norms of

\bm{A}

and

\bm{B}

, so the rectangle changes shape but keeps its area. (2) The same numerical update causes unequal relative changes in the two blocks. (3) Resuming training from the rescaled checkpoint, whose pair norms are imbalanced but whose function is unchanged, ends with much worse perplexity (PPL) than resuming from the original. (4) Hard and Soft SAGGS rebalance the pairs without changing the current function and recover from this degradation.

Sub-layer	$\bm{A}$	$\bm{B}$	Mechanism	Exact
Self-attention	$\bm{W}^{Q}_{\ell}$	$\bm{W}^{K}_{\ell}$	Bilinear in logits	✓ Yes
Self-attention	$\bm{W}^{V}_{\ell}$	$\bm{W}^{O}_{\ell}$	Linear composition	✓ Yes
MLP (ReLU)	$\bm{W}^{1}_{\ell}$	$\bm{W}^{2}_{\ell}$	Positive homogeneity	✓ Yes
MLP (GELU)	$\bm{W}^{1}_{\ell}$	$\bm{W}^{2}_{\ell}$	Non-homogeneous	✗ No
SwiGLU	$\bm{W}^{\mathrm{up}}_{\ell}$	$\bm{W}^{2}_{\ell}$	Linear in value path	✓ Yes

Table 1: Scalar pair rescalings

(\bm{A},\bm{B})\mapsto(c\bm{A},c^{-1}\bm{B})

c>0

. Rows marked ✓ Yes preserve the forward pass exactly and are used by SAGGS.

$\lambda$	Vanilla	SAGGS	$\bar{\Delta}$	$d$	$p_{\mathrm{sf}}^{+}$
$0.0$	$22.881$	$22.906$	$-0.025\pm 0.034$	$-0.75$	$0.9453$
$0.2$	$23.277$	$22.904$	$+0.374\pm 0.026$	$14.21$	$0.0078$

Table 2: 101M WikiText-103 comparison with

n=7

. Method columns report mean PPL;

\bar{\Delta}

reports the mean

\pm

sample SD of paired Vanilla-minus-SAGGS differences.

d

is paired Cohen’s

d

and

p_{\mathrm{sf}}^{+}

is the one-sided sign-flip

p

-value.

Resume condition	Final PPL
Control resume	$20.84\pm 0.25$
Surgery resume, no SAGGS	$24.62\pm 0.58$
Surgery + SAGGS	$20.85\pm 0.25$
Surgery + Soft SAGGS	$20.85\pm 0.25$

Table 3: 123.8M checkpoint-surgery resume with

n=7

. Values are final PPL, reported as mean

\pm

sample SD across seeds. Paired contrasts appear in Table 11 .

Appendix figures & tables31 assets

Supplementary material from the paper’s appendix.

Appendix

Name	$d$	$d_{\mathrm{ff}}$	$L$	Heads	Params	LR
GPT-2-style
Small	256	1024	12	4	22M	$6\times 10^{-4}$
Medium	512	2048	24	8	101M	$6\times 10^{-4}$
Deep	256	1024	48	4	51M	$3\times 10^{-4}$
Xlarge	1024	4096	24	16	354M	$3\times 10^{-4}$
1B (SwiGLU)	2048	5504	22	16	1.22B	$3\times 10^{-4}$

Appendix

Table 4: Model configurations. LR is the peak learning rate; modern-architecture head counts are pairs

(h_{q},h_{kv})

of query heads and key-value (KV) heads. The first four parameter counts assume two-projection ReLU or GELU MLPs. Context lengths and SwiGLU variants are specified in the text.

Setting	$\lambda$	Contrast	$\bar{\Delta}$	$p_{\mathrm{sf}}^{+}$	Purpose
Primary
101M WT-103	$0.2$	V $-$ SAGGS	$+0.374\pm 0.026$	$0.0078$	Core language-model protection
31M modern	$0.2$	V $-$ Soft	$+0.780\pm 0.040$	$0.0078$	Architecture and tokenizer
101M Lion	$0.2$	V $-$ Soft	$+34.465\pm 25.170$	$0.0078$	Optimizer transfer
Workflow
123.8M surgery		No-SAGGS $-$ control	$+3.775\pm 0.416$	$0.0078$	State-reset resume

Appendix

Table 5: Primary PPL comparisons and checkpoint-resume diagnostics.

\bar{\Delta}

is the paired left-minus-right difference, reported as mean

\pm

sample SD. Rows marked † use three paired seeds; all others use seven. Soft denotes Soft SAGGS, WT-103 denotes WikiText-103, and V denotes Vanilla.

$\lambda$	Vanilla	SAGGS	$\bar{\Delta}$
$0.0$	$22.88\pm 0.060$	$22.92\pm 0.033$	$-0.044\pm 0.026$
$0.2$	$23.25\pm 0.024$	$22.87\pm 0.036$	$+0.380\pm 0.025$

Appendix

Table 6: Three-seed 101M WikiText-103 comparison. Table 2 gives the seven-seed extension. Method columns report PPL as mean

\pm

sample SD;

\bar{\Delta}

reports paired Vanilla-minus-SAGGS differences in the same format.

	Mean PPL		Paired difference
$\lambda$	Vanilla	Soft, $\alpha=0.5$	$\bar{\Delta}$	$\sigma_{\Delta}$	$t$	$d$	$p_{\mathrm{sf}}^{+}$
$0.0$	$20.83\pm 0.06$	$20.81\pm 0.06$	$+0.019$	$0.024$	$+2.11$	$+0.80$	$0.0391$
$0.2$	$21.58\pm 0.07$	$20.80\pm 0.06$	$+0.780$	$0.040$	$+52.18$	$+19.72$	$0.0078$

Appendix

Table 7: Seven-seed modern architecture and cross-tokenizer replication. The paired block is Vanilla minus Soft SAGGS with

\alpha=0.5

Seed	$\mathrm{PPL}_{\mathrm{V}}$	$\mathrm{PPL}_{\mathrm{S}}$	$\Delta$
1	$55.992$	$23.954$	$+32.038$
2	$113.248$	$24.065$	$+89.183$
3	$61.429$	$23.963$	$+37.466$

Appendix

Table 8: Three-seed hard-rebalancing Lion comparison at

\lambda=0.2

: 101M parameters,

L=24

, and 15K steps. Rows index the paired seeds. V denotes Vanilla, S denotes SAGGS, and

\Delta=\mathrm{PPL}_{\mathrm{V}}-\mathrm{PPL}_{\mathrm{S}}

	Mean PPL		Paired difference
$\lambda$	Vanilla	Soft SAGGS	$\bar{\Delta}$	$\sigma_{\Delta}$	$t$	$d$	$p_{\mathrm{sf}}^{+}$
$0.0$	$24.15\pm 0.39$	$24.20\pm 0.58$	$-0.055$	$0.232$	$-0.62$	$-0.24$	$0.6953$
$0.2$	$58.32\pm 25.20$	$23.86\pm 0.09$	$+34.465$	$25.170$	$+3.62$	$+1.37$	$0.0078$

Appendix

Table 9: Seven-seed Lion comparison on the 101M WikiText-103 model. PPL is reported as mean

\pm

sample SD; the paired block reports Vanilla-minus-Soft SAGGS differences.

$\lambda$	Vanilla	SAGGS	Soft	State
$0.0$	$23.92\pm 0.02$	$30.12\pm 10.56$	$23.97\pm 0.11$	$24.14\pm 0.15$
$0.2$	$76.89\pm 31.60$	$23.99\pm 0.06$	$23.86\pm 0.12$	$23.97\pm 0.08$

Appendix

Table 10: Lion state-covariant ablation on the 101M model with

n=3

. The upper block reports PPL as mean

\pm

sample SD, and the lower block reports paired Vanilla-minus-State differences. Soft denotes Soft SAGGS; State denotes hard SAGGS with covariantly transformed momentum buffers.

Contrast	$\bar{\Delta}$	$\sigma_{\Delta}$	$t$	$d$	$p_{\mathrm{sf}}^{+}$
Surgery no SAGGS $-$ control	$+3.775$	$0.416$	$+24.00$	$+9.07$	$0.0078$
Surgery no SAGGS $-$ SAGGS	$+3.765$	$0.413$	$+24.14$	$+9.12$	$0.0078$
Surgery no SAGGS $-$ Soft	$+3.766$	$0.414$	$+24.06$	$+9.09$	$0.0078$
Control $-$ SAGGS	$-0.010$	$0.008$	$-3.53$	$-1.33$	$0.9922$
Control $-$ Soft	$-0.009$	$0.007$	$-3.36$	$-1.27$	$0.9844$

Appendix

Table 11: Branched checkpoint-surgery full paired contrasts on the 123.8M modern architecture with

n=7

Resume condition	Final PPL
Control, optimizer reset	$23.27\pm 0.05$
Surgery, optimizer reset	$23.54\pm 0.06$
Surgery reset + Soft SAGGS	$23.27\pm 0.06$
Control, optimizer preserved	$23.92\pm 0.06$
Surgery, state-covariant optimizer	$23.97\pm 0.07$

Appendix

Table 12: State-preserved checkpoint-surgery diagnostic on the 31M modern-architecture WikiText-103 model with

n=7

. The lower block reports paired contrasts.

Resume condition	Final PPL
Control, preserved state	$23.91\pm 0.05$
Surgery + stale AdamW state	$23.96\pm 0.07$
Surgery + covariant AdamW state	$23.96\pm 0.07$
Surgery + stale state + Soft SAGGS	$23.90\pm 0.05$

Appendix

Table 13: Stale AdamW-state diagnostic on the 31M model with

n=7

: final PPL as mean

\pm

sample SD and paired contrasts after an exact

\lambda=0.2

edit. Stale branches retain the pre-edit moments.

Resume condition	Final PPL
Control state preserved	$23.90\pm 0.07$
Exact edit, covariant state	$23.96\pm 0.07$
Exact edit, all state stale	$23.96\pm 0.07$
Exact edit, deep-half stale	$23.96\pm 0.07$
Deep-half stale + Soft SAGGS	$23.90\pm 0.06$

Appendix

Table 14: Partial AdamW-state diagnostic on the 31M model with

n=7

: final PPL as mean

\pm

sample SD and paired contrasts. In deep-half-stale branches, only shallow-layer moments are co-transformed.

Pair	$L$	$\lambda=0.01$	$\lambda=0.05$	$\lambda=0.1$
Q–K (ReLU)	12	$4.7\times 10^{-7}$	$2.7\times 10^{-7}$	$2.4\times 10^{-7}$
	24	$7.4\times 10^{-5}$	$4.9\times 10^{-5}$	$1.2\times 10^{-4}$
	48	$1.8\times 10^{-4}$	$2.2\times 10^{-4}$	$1.3\times 10^{-4}$
V–O (ReLU)	12	$1.2\times 10^{-7}$	$9.8\times 10^{-8}$	$1.0\times 10^{-7}$
	24	$8.6\times 10^{-5}$	$7.6\times 10^{-5}$	$8.6\times 10^{-5}$
	48	$1.1\times 10^{-5}$	$4.7\times 10^{-5}$	$5.2\times 10^{-5}$

Appendix

Table 15: Maximum relative gradient-norm residual

\varepsilon_{\mathrm{rel}}

(Eq. 26 ) in FP32, by depth, tilt, and pair type. GELU MLPs are included as a non-exact comparison.

Model	Params	$L$	Q–K $G$	V–O $G$	MLP $G$
Small	22M	12	$0.200\pm 0.026$	$0.068\pm 0.004$	$0.099\pm 0.002$
Medium	101M	24	$0.136\pm 0.003$	$0.047\pm 0.002$	$0.037\pm 0.002$
Deep	51M	48	$0.163\pm 0.021$	$0.120\pm 0.026$	$0.093\pm 0.003$

Appendix

Table 16: Increase in maximum absolute symmetry-pair log-ratio during standard ReLU training at

\lambda=0

. The reported

G=\Delta_{\max}^{(29.5\mathrm{K})}-\Delta_{\max}^{(0)}

uses

\delta=10^{-12}

as in Appendix G.4 . Values are mean

\pm

population SD over two seeds.

Figure 2: Growth

G

in maximum absolute pair log-ratio from step 0 to 29.5K during standard ReLU training. Q–K pairs show the largest mean growth. The ordering of MLP and V–O differs across model configurations. Error bars show

\pm 1

population SD over two seeds.

$\lambda$	$c_{47}$	Vanilla	SAGGS+P	$\Delta$
$0.0$	$1$	$30.34\pm 0.02$	$30.27\pm 0.01$	$+0.07$
$0.2$	$\approx 1.2\times 10^{4}$	$61.46\pm 2.99$	$30.31\pm 0.04$	$+31.15$
$0.3$	$\approx 1.3\times 10^{6}$	$150.58\pm 0.96$	$30.35\pm 0.06$	$+120.23$

Appendix

Table 17: Deep ReLU model with

L=48

under Q–K rescaling over 30K steps. Values are best validation PPL, reported as mean

\pm

population SD over two seeds;

\Delta

denotes Vanilla minus SAGGS+P.

$\lambda$	$c_{23}$	Vanilla	SAGGS+P	$\Delta$	Vanilla	SAGGS+P
		PPL			Gradient slope
$0.0$	$1$	$23.53$	$23.59$	$-0.06$	$0.018$	$0.020$
$0.1$	$\approx 10$	$23.89$	$23.56$	$+0.33$	$0.097$	$0.018$
$0.2^{\dagger}$	$\approx 100$	$25.84\pm 0.04$	$23.58\pm 0.06$	$+2.26$	$0.187$	$0.017$

Appendix

Table 18: 354M GELU model with Q–K rescaling over 15K steps. Values are best validation PPL, with

\Delta

denoting Vanilla minus SAGGS+P; gradient slopes are seed means at the last logged state, step 14.5K. The row marked † uses three seeds, with PPL as mean

\pm

sample SD; the other rows use a single seed.

Deep, $L=48$
$\lambda$	Method	PPL	MLP imbalance	Gradient slope
$0.0$	Vanilla (baseline)	$28.48$	$0.052$	$0.001$
	SAGGS+P	$28.42$	$<10^{-4}$	$0.001$
$0.2$	Vanilla	$63.24$	$8.70$	$0.212$
	SAGGS+P	$28.46$	$<10^{-4}$	$0.0002$
Small, $L=12$

Appendix

Table 19: Single-seed SwiGLU MLP rescaling: best validation PPL over 30K steps, with MLP pair imbalance and gradient-ratio depth slope measured at the last logged state, step 29.5K.

Figure 3: Single-seed SwiGLU MLP rescaling at (a)

L=12

and (b)

L=48

: best validation PPL for Vanilla and SAGGS+P (Table 19 ). The deeper model uses only

\lambda\in\{0,0.2\}

. Vanilla markers shift left and SAGGS+P markers shift right by

0.007

for visibility; the bracket in (b) gives Vanilla minus SAGGS+P at

\lambda=0.2

Pair type	Maximum relative error (%)
	$\lambda=0.05$	$\lambda=0.1$	$\lambda=0.2$	$\lambda=0.3$
Q–K	$0.152$	$0.118$	$0.129$	$0.146$
V–O	$0.110$	$0.107$	$0.110$	$0.156$
MLP (SwiGLU)	$0.030$	$0.042$	$0.038$	$0.054$

Appendix

Table 20: 1.2B-parameter gradient law verification in BF16 mixed precision. Values are the maximum relative error across all 22 layers for each symmetry pair type and

\lambda

. All errors are below

0.16\%

$\lambda$	Vanilla PPL	SAGGS+P PPL	$\Delta$
$0.1$	$22.91$	$22.90$	$-0.01$
$0.2$	$23.24$	$22.89$	$-0.35$
$0.3$	$25.91$	$22.92$	$-3.00$

Appendix

Table 21: Single-seed rescaling-strength sweep at 101M parameters and

L=24

over 30K steps. Here

\Delta

is SAGGS+P minus Vanilla. The

\lambda=0.3

difference is

-2.9954

before rounding.

Scale	$\lambda$	Vanilla PPL	SAGGS+P PPL	$\Delta$
Small, 22M	$0$	$34.39\pm 0.08$	$34.38\pm 0.08$	$-0.00\pm 0.06$
Small, 22M	$0.2$	$34.20\pm 0.04$	$34.30\pm 0.06$	$+0.11\pm 0.05$
Medium, 101M	$0^{\dagger}$	$22.84\pm 0.05$	$22.88\pm 0.04$	$+0.042\pm 0.003$
Medium, 101M	$0.2$	$23.25\pm 0.03$	$22.87\pm 0.03$	$-0.38\pm 0.05$
354M	$0^{\dagger}$	$20.66\pm 0.03$	$20.67\pm 0.04$	$+0.012\pm 0.006$
354M	$0.2^{\dagger}$	$21.82\pm 0.02$	$20.65\pm 0.03$	$-1.17\pm 0.01$

Appendix

Table 22: Multi-seed validation across scales. Rows marked † use two seeds; the others use three.

\Delta

is SAGGS+P minus Vanilla; paired SD uses within-seed differences. Values are mean

\pm

SD: Medium and 354M entries use sample SD over the complete specified seed sets, while small-model entries retain their original recorded summaries.

	$\lambda=0$	$\lambda=0.2$
Metric	Vanilla	Vanilla	SAGGS
Final validation PPL	$23.81$	$221.24$	$\mathbf{23.70}$
Best validation PPL	$23.81$	$221.24$	$\mathbf{23.70}$
Mean gradient norm	$0.68$	$27.79$	$0.68$
Q–K max-abs log-ratio	$0.032$	$4.222$	$0.55\times 10^{-4}$
V–O max-abs log-ratio	$0.033$	$4.117$	$1.23\times 10^{-4}$

Appendix

Table 23: Single-seed 1.2B-parameter SwiGLU training over 10K steps with DDP on 4

\times

H100 GPUs. Gradient norms are measured before clipping at

1.0

Figure 4: Maximum absolute pair log-ratio across layers during 1.2B training: (a) Q–K and (b) MLP. The shared vertical scale is linear from

0

10^{-4}

and logarithmic above; the break marks the transition. Post-rebalance SAGGS samples are at most

4.8\times 10^{-7}

and appear at zero. Samples align with rebalancing events and end at 9.9K; Table 23 reports the final 10K-step measurements after the intervening updates.

	Mean PPL		Paired difference
$\lambda$	Vanilla	Soft, $\alpha=0.5$	$\bar{\Delta}$	$\sigma_{\Delta}$	$t$	$d$	$p_{\mathrm{sf}}$
$0.0$	$55.92\pm 0.07$	$55.95\pm 0.03$	$-0.032$	$0.081$	$-0.69$	$-0.40$	$0.500$
$0.2$	$74.28\pm 0.41$	$56.00\pm 0.14$	$+18.278$	$0.437$	$+72.37$	$+41.79$	$0.250$

Appendix

Table 24: OpenWebText paired comparison with

n=3

seeds. The paired block is Vanilla minus Soft SAGGS with

\alpha=0.5

and reports the two-sided sign-flip

p

-value

p_{\mathrm{sf}}

	Mean PPL		Paired difference
$\lambda$	Vanilla	Soft, $\alpha=0.5$	$\bar{\Delta}$	$\sigma_{\Delta}$	$t$	$d$	$p_{\mathrm{sf}}^{+}$
$0.0$	$19.86\pm 0.06$	$19.81\pm 0.06$	$+0.055$	$0.028$	$+3.40$	$+1.96$	$0.1250$
$0.2$	$41.12\pm 0.65$	$19.96\pm 0.09$	$+21.156$	$0.667$	$+54.90$	$+31.70$	$0.1250$

Appendix

Table 25: 123.8M modern-architecture scale replication, 15K AdamW steps,

n=3

. The paired block is Vanilla minus Soft SAGGS with

\alpha=0.5

	Mean PPL		Paired difference
$\lambda$	Vanilla	Soft, $\alpha=0.5$	$\bar{\Delta}$	$\sigma_{\Delta}$	$t$	$d$	$p_{\mathrm{sf}}^{+}$
$0.0$	$25.51\pm 0.04$	$25.42\pm 0.10$	$+0.088$	$0.052$	$+2.96$	$+1.71$	$0.1250$
$0.2$	$111.45\pm 4.83$	$25.64\pm 0.10$	$+85.813$	$4.730$	$+31.42$	$+18.14$	$0.1250$

Appendix

Table 26: 296.7M modern-architecture scale replication, 12K AdamW steps,

n=3

. The paired block is Vanilla minus Soft SAGGS with

\alpha=0.5

Figure 5: Pair-scale variation in released checkpoints. (a) Eleven checkpoints, using the shape-adjusted log-ratio $\widetilde{r}_{\ell}=\log(\left\lVert A_{\ell}\right\rVert_{\mathrm{F}}/\left\lVert B_{\ell}\right\rVert_{\mathrm{F}})-(1/2)\log(n_{A}/n_{B})$ , where $n_{A}$ and $n_{B}$ count matrix entries. Red denotes larger root mean square (RMS) entry magnitudes in $A_{\ell}$ and blue denotes larger RMS entry magnitudes in $B_{\ell}$ ; zero means equal RMS, which differs from Frobenius balance for unequal-size GQA pairs. Blank SwiGLU rows denote GELU MLPs. (b) Four Pythia pre-training runs and layerwise Q–K log-ratios for 410M. Equal matrix sizes make RMS and Frobenius log-ratios coincide. Appendix I.6 gives the protocols.

Figure 6: Fine-tuning from forward-identical GPT-2 774M checkpoints: 200 AdamW updates with fresh moments and 16,384 held-out targets. (a) Top-1 disagreement and (b) mean absolute token-NLL difference, relative to the canonicalized branch. Canonicalizing either start gives the same trajectory, shown in green. The dashed reference uses different sampled training windows; the reference is not an uncertainty band. (c) Net Q–K movement relative to each branch’s starting norm, a cumulative rather than one-step measurement. (d) Q–K log-ratios: layer mean and min–max band. Soft SAGGS uses

\alpha=0.5

at handoff and every ten steps. Appendix I.6 gives the protocols.

Figure 7: Depth-resolved response to an exact attention-pair tilt $c_{\ell}=e^{0.2\ell}$ in GPT-2 355M with $L=24$ during 400 Lion steps. (a) First-step relative Q–K movement and the sign-term prediction $\eta\sqrt{n}/\left\lVert W\right\rVert_{\mathrm{F}}$ ; full updates also include decay. (b) Direction change $1-\cos(W_{t},W_{0})$ for tilted $W_{K}$ and $W_{O}$ , and $W_{K}$ with tilt + Soft every ten steps. Columns are 15 evaluation steps on a nonlinear time axis. (c) Attention-entropy change from handoff after 400 steps, averaged over 16 fixed sequences, heads, and query positions. (d) Mean held-out NLL. Rebalancing includes handoff and leaves buffers unchanged; Soft uses $\alpha=0.5$ . The figure shows a single-seed stress test; protocols are in Appendix I.6 .

Figure 8: Token-level restart differences in GPT-2 124M. (a), (b), and (c) Outcome-independent excerpts from three WikiText-103 validation articles. Columns are ground-truth BPE tokens conditioned on preceding text; a leading dot denotes a space and @-@ is the corpus hyphen. Rows show branch-minus-control token NLL at handoff or after 30 shared AdamW updates. Blue denotes increased and red denotes decreased target probability on a shared signed scale. The edit co-transforms moments; the canonical branch additionally balances weights and moments, while Soft SAGGS with $\alpha=0.5$ leaves buffers unchanged at rebalancing. Appendix I.6 gives the protocols.

Architecture and loss
Symbol	Meaning
$\bm{\theta}$	model parameters
$\mathcal{L}(\bm{\theta})$	data loss
$L$ , $\ell$	depth and layer index
$d$ , $d_{k}$ , $d_{\mathrm{ff}}$	hidden, head, and MLP widths
$\bm{W}^{Q}_{\ell}$ , $\bm{W}^{K}_{\ell}$ , $\bm{W}^{V}_{\ell}$ , $\bm{W}^{O}_{\ell}$	attention projections

Appendix

Table 27: Main mathematical notation used in the paper.

Explore similar work

Optimizer-Induced Low-Dimensional Drift and Transverse Dynamics in Transformer Training
Feb 27, 2026Yongzhong XuNeural Network Training Dynamics Transformer
Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers
May 18, 2026Tim Tsz-Kit Lau, Weijie SuNeural Network Optimization
University of Pennsylvania, Philadelphia, PA 19104, USA.
First-Order Predictable but Pairwise Fragile: Local Task Adaptation in Trained Transformers
Jul 18, 2026Irina Piontkovskaia, Sergey NikolenkoFine-Tuning Neural Network Optimization
1DAIMLD, Moscow, Russia · 2St. Petersburg Department of the Steklov Institute of Mathematics, St. Petersburg, Russia · 3St. Petersburg State University

Feb 27, 2026J/K move · Enter open · S save

Optimizer-Induced Low-Dimensional Drift and Transverse Dynamics in Transformer Training

Yongzhong Xu

Neural Network Training Dynamics Transformer Adaptive Gradient Methods Intrinsic Dimensionality

Open paper AskPDF arXiv

We analyze cumulative parameter trajectories of transformer training under AdamW and identify a dominant low-dimensional drift direction ("backbone") that captures 60--80% of long-horizon displacement from initialization. This direction is highly stable over rolling training windows yet reorients gradually across phases, particularly following objective reweighting. Per-batch gradients exhibit near-noise-floor alignment with the backbone, whereas optimizer-integrated updates align strongly with it, indicating that the structure emerges from accumulated optimizer dynamics rather than instantaneous gradient geometry. Replacing AdamW with SGD-family optimizers eliminates this structure, and reducing

β_2

smoothly degrades backbone dominance and reheating recoverability. Reheating experiments show that transverse probe modes can be transiently re-excited without substantially altering accumulated backbone drift. These results provide a trajectory-level characterization of optimizer-induced geometric structure in transformer training and shift attention from instantaneous gradient properties to cumulative update dynamics.

Abstract

Explore similar work

Optimizer-Induced Low-Dimensional Drift and Transverse Dynamics in Transformer Training

Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers

First-Order Predictable but Pairwise Fragile: Local Task Adaptation in Trained Transformers