Learning Rate Transfer for Hybrid Transformer-SSM Architectures

Jimin Seo; Gyubok Lee; Yeonsik Jo; Kiwoong Yoo; Yeongoon Kim; Minhae Oh; Jin Woo Koo; Suhwan Kim

cs.LGOct 1, 2026

Learning Rate Transfer for Hybrid Transformer-SSM Architectures

Authors: Jimin Seo, Gyubok Lee, Yeonsik Jo, Kiwoong Yoo, Yeongoon Kim, Minhae Oh, Jin Woo Koo, Suhwan Kim, +6 more

Organizations: Seoul National University · SB Intuitions · LG AI Research · Hodoo AI

Abstract

We study learning rate (LR) scaling for hybrid architectures combining Transformer and State-Space Model (SSM) blocks, a class adopted by several recent production language models. In particular, we focus on the gap between the theoretical scaling rules derived for SSMs under zero-order-hold (ZOH) discretization at infinite width with growing state size, and the field-standard practical implementations using simplified-ZOH Mamba at fixed state size. Surprisingly, in this practical regime hybrid architectures achieve a near-zero LR transfer gap across widths 256-2048 and depths 4-32 up to billion-parameter scale using only the original $μ$ P prescription, even though SSM operations fall outside its Tensor Programs representability conditions and every parameterization we test fails the standard coordinate-check diagnostic of $μ$ P correctness. We attribute this to a two-condition decomposition of LR transfer in hybrid architectures: a global update-to-weight invariance, enforced by $μ$ P's initialization and LR scaling; and a local per-component balance, provided by AdamW's per-parameter normalization. Our observations show that the optimal LR is invariant to width up to 8 $\times$ , that this width invariance holds across depth, sequence length, batch size, and Transformer-to-SSM ratio, and that it transfers to Nemotron-H, a production hybrid outside our custom architecture set. We hope these findings fill the gap between theoretical scaling rules and practical hybrid implementations, and stimulate further research toward bridging it.

Figures & tables

Figure 1 : The coordinate check and LR transfer disagree on hybrids ( MHMH , $d{=}16$ ). Top: mean absolute activation per block after 10 steps on WikiText-103 , relative to $w{=}128$ ; dotted line is width invariance, bold line the median over blocks. Bottom: training loss against learning rate on WikiText-2 (2k steps), each width’s optimum dashed and diverged runs open. The diagnostic separates SP from both $\mu\text{P}$ variants, while transfer separates naïve $\mu\text{P}$ , which holds $\eta^{*}{=}1.0$ at every width, from the other two.

	Naïve $\mu\text{P}$			Block-wise $\mu\text{P}$ (approximated)			Block-wise $\mu\text{P}$			SP
Arch	$w{=}256$	$w{=}512$	$w{=}1024$	$w{=}256$	$w{=}512$	$w{=}1024$	$w{=}256$	$w{=}512$	$w{=}1024$	$w{=}256$	$w{=}512$	$w{=}1024$
MM	1.0	1.0	1.0	1.0	3.0	3.0	3.0	3.0	1.0	0.010	0.003	0.003
TT	1.0	1.0	1.0	1.0	1.0	1.0	1.0	1.0	1.0	0.003	0.001	0.0003
MTMT	1.0	1.0	1.0	1.0	1.0	1.0	1.0	3.0	1.0	0.010	0.003	0.003
MHMH	1.0	1.0	1.0	1.0	3.0	3.0	3.0	3.0	1.0	0.003	0.003	0.001
HMHM	1.0	1.0	1.0	1.0	1.0	3.0	3.0	3.0	1.0	0.003	0.003	0.001

Table 1: Optimal LR at each (scheme, architecture, width); uniform shading within each row means LR transfer to the corresponding widths. Values from 2k-step training on WikiText-2 . The SP grid for MM and MTMT starts at

\eta{=}0.003

Published baselines (The Pile, 300B tokens)
Model	Params	HellaSwag	PIQA	WinoGr	ARC-E	ARC-C	LAMBADA	Avg
		acc_n	acc	acc	acc	acc_n	acc
Mamba-370M	370M	46.5	69.5	55.3	55.1	28.0	55.6	51.7
Pythia-410M	410M	40.6	66.9	53.8	52.1	24.6	51.4	48.2
Retrained on FineWeb-Edu (5B tokens)
Mamba-370M	370M	34.6	61.6	53.8	55.0	26.2	19.9	41.9

Table 2: Zero-shot downstream evaluation. Top: published baselines (300B tokens, The Pile). Bottom: all models retrained on FineWeb-Edu 5B tokens. MHMH variants:

\dagger

MHMH -312M uses the GPT-NeoX tokenizer (50304 vocab, tokenizer-matched to Pythia);

\ddagger

MHMH -415M uses tiktoken cl100k_base (100277 vocab).

Depth	Scheme	Best LR	$w{=}256$	$w{=}512$	$w{=}1024$	Gap
16	Naïve $\mu\text{P}$	1.0	3.760	3.344	3.369	0.0%
16	SP	0.001–0.003	3.775	3.382	3.356	1.9%
32	Naïve $\mu\text{P}$	1.0	3.527	3.313	3.138	0.0%
32	SP	0.001–0.003	3.542	3.350	3.152	1.9%

Table 3: Nemotron-H width transfer ( WikiText-103 , 20k steps, warmup=1000;

\eta\in\{1,3,5\}

for naïve

\mu\text{P}

\{0.001,0.003,0.01\}

for SP). Naïve

\mu\text{P}

achieves

\eta^{*}{=}1.0

at all widths (matching our other architectures); SP optimum shifts from

\eta{=}0.003

w{=}256

\eta{=}0.001

w{=}1024

		$A_{\log}/$ FFN		in_proj $/$ FFN
Optimizer	Scheme	$128{\to}256$	$128{\to}512$	$128{\to}256$	$128{\to}512$
AdamW	Naïve $\mu\text{P}$	20%	21%	4%	5%
AdamW	Block-wise $\mu\text{P}$	166%	436%	15%	20%
SGD	Naïve $\mu\text{P}$	235%	342%	0%	0%
SGD	Block-wise $\mu\text{P}$	233%	835%	23%	31%

Table 4: Condition 2 gap

\Delta\beta_{M/T}(c_{\text{proxy}}\!\to\!c_{\text{target}})

in %, with proxy width

128

. The

A_{\log}/

FFN pair is the most width-divergent under Chennuru Vankadara et al. (2024) ’s

\eta_{A}\!\in\!\Theta(N_{u})

rule; in_proj

/

FFN is a near-isotropic Mamba projection. MHMH

d{=}16

, ssm_expand

=

N_{x}{=}16

\eta_{\text{base}}{=}1.0

; mean over steps

50

–

200

and

3

seeds.

Optimizer	Adaptive LR	Momentum	Best loss $w{=}256$	$w{=}1024$	$w{=}1024$ lower
AdamW	Yes	Yes ( $\beta_{1}{=}0.9$ )	4.392	4.099	7/7
Adam no-momentum	Yes	No ( $\beta_{1}{=}0$ )	4.470	4.465	7/7
SGD + momentum	No	Yes (0.9)	7.732	7.955	1/7
SGD no-momentum	No	No	9.398	10.803	2/7

Table 5: Optimizer ablation at a fixed parameterization ( HHHH at

d{=}16

, naïve-

\mu\text{P}

parameter groups for every optimizer, 2k steps on WikiText-103 ,

\eta\in\{0.001,\dots,1.0\}

; per-LR sweeps in Appendix E.3 ). Every best loss falls at

\eta{=}1.0

, the top of this grid. The last column counts the LRs at which

w{=}1024

reaches a lower loss than

w{=}256

Component	Scheme	$\alpha(N_{u})$ step 1	$\alpha(N_{u})$ stable	$\Delta\alpha$
$A_{\log}$	naïve	$+0.86$	$+1.55$	$+0.68$
$A_{\log}$	faithful	$\mathbf{+1.00}$	$+3.42$	$+2.42$
$B$	naïve	$+0.99$	$+1.33$	$+0.34$
$B$	faithful	$\mathbf{+0.006}$	$+1.35$	$+1.34$
$C$	naïve	$+1.00$	$+1.34$	$+0.33$
$C$	faithful	$\mathbf{+0.029}$	$+0.71$	$+0.68$

Table 6: Effective-update

N_{u}

scaling at step 1 vs the stable window (steps 100–199) under SGD, base LR

\eta\!=\!1

. Vankadara’s framework predicts

\alpha(N_{u})\!\approx\!0

at initialization for the faithful prescription (width-invariant update). The faithful values for

B

and

C

at step 1 are bolded;

A_{\log}

is bolded to flag the predicted-but-not-observed cancellation.

Appendix figures & tables47 assets

Supplementary material from the paper’s appendix.

Appendix

§	Experiment	#	Key finding
A.4	Faithful $\mu\text{P}$ -SSM: pure-Mamba SGD + hybrid AdamW	33	Lower loss in Vankadara regime; $5$ – $8\%$ higher loss on hybrids
B.1	Max stable LR across 9 unit interfaces $\times$ widths $\times$ schemes	195	Naïve $\mu\text{P}$ stable to $\eta{=}1$ on all; SP shrinks with width
B.2	Depth $4$ – $32$ across 8 architectures ( $w{=}512$ )	128	Naïve $\mu\text{P}$ keeps $\eta^{*}{=}1$ at all depths
C.1	Joint depth $\times$ width grid (proxy $\to$ target)	54	$\eta{=}1$ transfers $d{=}8\,w{=}256\to d{=}32\,w{=}1024$
C.2	FineWeb-Edu width transfer (500M tokens)	132	$0.0\%$ gap; SP $3\times$ shift
C.3	Width 2048 ( $8\times$ scaling; MHMH 1.24B, HHHH 1.66B params)	20	$0.0\%$ gap at $8\times$

Appendix

Table 7: Experiment inventory across the appendix. “#” is the number of configurations in the design; “naïve / approx. / faithful” denote naïve

\mu\text{P}

, approximated block-wise

\mu\text{P}

, and the literal Vankadara prescription, respectively.

Width	Params (HHHH)	2k steps / 20k steps
$w{=}256$	71M	3.1 min / ${\sim}16$ min
$w{=}512$	181M	7.1 min / ${\sim}55$ min
$w{=}1024$	518M	17.1 min / ${\sim}191$ min

Appendix

Table 8: Estimated mean wall time per config on a single A6000 (HHHH,

d{=}16

). Times rise super-linearly because the SSM scan and attention costs are non-quadratic in width; the proxy-to-target ratio is

5.5\times

2

k steps and

{\sim}12\times

20

Dataset	Train tokens	Tokenizer (vocab)	Steps/epoch	2k steps	20k steps
WikiText-2	$2.39$ M	GPT-2 BPE ( $50{,}257$ )	$292$	$6.9$ epochs	$68$ epochs
WikiText-103	${\sim}120$ M	tiktoken cl100k_base ( $100{,}277$ )	$14{,}648$	$14\%$ epoch	$1.37$ epochs
FineWeb-Edu 500M	$500$ M	tiktoken cl100k_base ( $100{,}277$ )	$61{,}035$	$3.3\%$ epoch	$33\%$ epoch

Appendix

Table 9: Dataset sizes and training coverage at the canonical

\text{seq\_len}{=}128

\text{batch\_size}{=}64

setup (

8{,}192

tokens/step). The right two columns express the step budget as a fraction (or multiple) of one full epoch.

LR	naïve $w{=}128$	block-wise $w{=}128$	naïve $w{=}1024$	block-wise $w{=}1024$
0.3	10.39	8.23	8.83	8.34
1.0	9.52	7.01	7.92	7.43
3.0	7.88	6.64	7.38	6.79
10.0	6.58	6.30	6.88	6.34
30.0	5.96	6.05	6.39	5.95
100.0	5.55	DIV (step 942)	6.00	5.96

Appendix

Table 10: Faithful block-wise

\mu\text{P}

vs. naïve

\mu\text{P}

under plain SGD on pure Mamba (20k steps, WikiText-103). Final training loss.

Arch	Width	Naïve	Block-wise	Gap
HHHH	256	3.514	3.712	$+5.6\%$
HHHH	512	3.270	3.500	$+7.0\%$
HHHH	1024	3.192	3.414	$+7.0\%$
MHMH	256	3.741	3.957	$+5.8\%$
MHMH	512	3.317	3.520	$+6.1\%$
MHMH	1024	3.171	3.395	$+7.1\%$

Appendix

Table 11: Faithful block-wise

\mu\text{P}

vs. naïve

\mu\text{P}

on hybrids under AdamW (20k steps, WikiText-103 ,

\eta_{\text{base}}{=}1.0

, warmup

200

steps). Final training loss; relative gap is

(\mathcal{L}_{\text{block-wise}}-\mathcal{L}_{\text{{na\"{i}ve} }})/\mathcal{L}_{\text{{na\"{i}ve} }}

		Width ( $d_{\text{model}}$ )
Interface	Scheme	128	256	512	1024	2048
TT	SP	1.0	0.30	0.30	0.30	0.03
	Naïve $\mu\text{P}$	1.0	1.0	1.0	1.0	1.0
	Block-wise (approx.)	1.0	1.0	1.0	1.0	1.0
	Block-wise (faithful)	1.0	1.0	1.0	1.0	1.0
TM	SP	0.30	0.30	0.30	0.10	0.03

Appendix

Table 12 : Max stable LR at the unit interface across all 9 (T, M, H) ordered pairs (darker = less stable), 1k steps (warmup 50) on WikiText-2 . LR grid:

\{0.001,0.003,0.01,0.03,0.1,0.3,1.0\}

; an entry of

1.0

indicates stability at the grid maximum (the true max stable LR may be higher). Naïve

\mu\text{P}

: stable at

\eta{=}1.0

across every (interface, width) cell. SP: max stable LR shrinks with width. Block-wise

\mu\text{P}

(approximated): stable across every cell, but with

\eta_{A}

instead set to

\eta_{\text{base}}N_{u}

it collapses for serial-SSM interfaces ( TM , MT , MM ) where the per-component LR correction stacks with Adam’s

g/\sqrt{v}

. Block-wise

\mu\text{P}

(faithful, the literal Vankadara prescription): stable across every cell, in contrast to the approximated variant with

\eta_{A}{=}\eta_{\text{base}}N_{u}

Figure 2 : Depth scaling results comparing all four schemes (SP, naïve $\mu\text{P}$ , block-wise $\mu\text{P}$ approximated, block-wise $\mu\text{P}$ faithful) at $w{=}512$ , 2k steps on WikiText-2 . (a) Loss vs. depth at each scheme’s optimal LR for MHMH and HHHH . (b) Loss vs. depth for MHMH with the optimal LR labeled at each point. (c) Loss at maximum depth ( $d{=}32$ ) for all 8 architectures. All schemes show decreasing loss with depth; the two block-wise variants have higher loss at the 2k-step horizon on WikiText-2 ; the approximated variant’s 20k-step loss on WikiText-103 stays within $0.5\%$ of naïve $\mu\text{P}$ at $w{=}1024$ (Appendix A.4 ). The faithful variant retains a $5$ – $8\%$ loss gap relative to naïve $\mu\text{P}$ even at convergence (Table 11 ).

Figure 3 : Loss at $\eta{=}1.0$ across the width $\times$ depth grid (training loss, 2k steps; values from Table 13 ). Best: MHMH $(w{=}1024,d{=}32)$ at 4.096.

Arch	Depth $\times$ Width	$\eta$ =0.3	$\eta$ =1.0	$\eta$ =3.0	Gap @ $\eta{=}1$
MHMH	$d$ =8, $w$ =256	4.508	4.304	4.318	0.00%
	$d$ =8, $w$ =512	4.488	4.253	4.230	0.55%
	$d$ =8, $w$ =1024	4.403	4.140	4.204	0.00%
	$d$ =16, $w$ =256	4.411	4.219	4.241	0.00%
	$d$ =16, $w$ =512	4.495	4.262	4.843	0.00%
	$d$ =16, $w$ =1024	4.453	4.164	4.246	0.00%

Appendix

Table 13: Depth-width joint transfer for MHMH and HHHH ( WikiText-103 , 2k steps, naïve

\mu\text{P}

). Gap = relative loss penalty of using the proxy-transferred LR (

\eta{=}1

) instead of each row’s best-of-grid value,

(\mathcal{L}(\eta{=}1)-\mathcal{L}^{*})/\mathcal{L}^{*}

. The 3-point grid does not sample the basin between

\eta{=}1

and

\eta{=}3

(Tab 19 shows the true basin lies near

\eta\!\approx\!1.5

–

2

); cells where

\eta{=}3

wins by less than 1% should be read as basin-edge sampling artifacts, not as preferring high LR.

		$w{=}256$	$w{=}1024$
Arch	Scheme	Loss (opt. LR)	Opt. LR	Transferred	Retuned
HHHH	Naïve $\mu\text{P}$	4.467 ( $\eta{=}1.0$ )	1.0	4.115	—
HHHH	SP	4.452 ( $\eta{=}0.003$ )	0.001	4.194	4.149
MHMH	Naïve $\mu\text{P}$	4.313 ( $\eta{=}1.0$ )	1.0	3.993	—
MHMH	SP	4.315 ( $\eta{=}0.003$ )	0.001	4.339	3.998
MTMT	Naïve $\mu\text{P}$	4.451 ( $\eta{=}1.0$ )	1.0	4.033	—

Appendix

Table 14: FineWeb-Edu width transfer (20k steps, 164M training tokens from a 500M-token subset). SP’s optimal LR shifts

3\times

(

0.003\!\to\!0.001

); transferring the proxy optimal LR incurs

1.1

–

8.7\%

penalty (largest on MHMH and MTMT , smallest on HHHH ).

Figure 4 : Validation loss curve for MHMH 415M on FineWeb-Edu (5.0B tokens, 38k steps). Smooth convergence from 5.45 (step 500) to 2.96; $\eta{=}1.0$ transferred from proxy ( $w{=}256$ , 65M params) without retuning.

Arch	$\eta{=}0.3$	$\eta{=}0.5$	$\eta{=}0.7$	$\eta{=}1.0$	$\eta{=}1.5$	$\eta{=}2.0$	$\eta{=}3.0$	Gap @ $\eta{=}1$
MHMH	4.234	4.070	3.989	3.934	3.907	3.984	4.044	0.70%
HHHH	4.480	4.296	4.209	4.161	4.121	4.173	4.270	0.97%

Appendix

Table 15: Width 2048, naïve

\mu\text{P}

sweep.

Arch	$\eta{=}0.001$	$\eta{=}0.003$	$\eta{=}0.01$
MHMH	3.935	4.102	11.3
HHHH	4.211	4.341	10.8

Appendix

Table 16: Width 2048, SP sweep.

	20k steps (164M tokens)				60k steps (492M tokens)
$\eta$	seed 42	seed 123	seed 456	mean	seed 42	seed 123	seed 456	mean
0.5	–	–	–	–	3.290	3.284	3.275	3.283
0.7	3.234	3.237	3.238	3.236	3.260	3.257	3.256	3.258
1.0	3.223	3.231	3.229	3.228	3.215	3.227	3.220	3.220
1.5	3.636	3.415	3.500	3.517	3.315	3.301	3.254	3.290
2.0	3.661	3.563	3.481	3.568	–	–	–	–

Appendix

Table 17: MHMH at

w{=}2048

(1.24B parameters) on WikiText-103 , three seeds: validation loss at 20k and 60k steps. Bold marks the argmin of each column; – marks an LR not run at that horizon.

Depth	Parameters	$\eta^{*}$	$\Delta$ (nats)	$\Delta$ (%)
48	834M	1.0	0.000	0.00%
64	1.04B	1.0	0.000	0.00%
96	1.46B	0.7	0.043	1.06%
128	1.88B	1.5	0.004	0.09%

Appendix

Table 18: MHMH at

w{=}1024

and depth 48–128 ( WikiText-103 , 2k steps, effective batch 64, one seed, validation loss).

\eta^{*}

is the best of

\{0.3,0.5,0.7,1.0,1.5,2.0,3.0\}

;

\Delta

is the loss of

\eta{=}1.0

above it.

Arch	Width	$\eta$ =0.3	$\eta$ =0.5	$\eta$ =0.7	$\eta$ =1.0	$\eta$ =1.5	$\eta$ =2.0	$\eta$ =3.0	Gap @ $\eta{=}1$
MHMH	256	4.411	4.290	4.262	4.218	4.202	4.209	5.523	0.38%
	512	4.495	4.354	4.301	4.265	4.240	4.249	4.255	0.59%
	1024	4.453	4.285	4.221	4.170	4.145	4.145	4.264	0.60%
HHHH	256	4.562	4.468	4.421	4.403	4.387	4.399	4.443	0.36%
	512	4.504	4.368	4.294	4.268	4.255	4.286	4.349	0.31%
	1024	4.386	4.217	4.142	4.107	4.083	4.126	4.174	0.59%

Appendix

Table 19: 7-point loss vs. LR for MHMH and HHHH ( WikiText-103 , 2k steps, naïve

\mu\text{P}

). Gap =

(\mathcal{L}(\eta{=}1)-\mathcal{L}^{*})/\mathcal{L}^{*}

in %, the relative penalty of using the proxy-transferred LR (

\eta{=}1

) instead of the row’s best-of-grid value.

W	$\eta{=}0.5$	$\eta{=}0.7$	$\eta{=}1.0$	$\eta{=}1.5$	$\eta{=}2.0$	$\eta{=}3.0$	$\eta{=}5.0$	Gap @ $\eta{=}1$
256	4.190	4.142	4.113	4.117	4.093	4.096	4.318	0.48%
1024	4.360	4.299	4.253	4.235	4.232	4.332	4.480	0.51%

Appendix

Table 20: Mamba-2 MHMH , naïve

\mu\text{P}

sweep.

W	$\eta{=}0.0003$	$\eta{=}0.001$	$\eta{=}0.003$	$\eta{=}0.01$	$\eta{=}0.03$
256	5.163	4.541	4.247	4.198	4.825
1024	4.624	4.342	4.384	4.580	11.3

Appendix

Table 21: Mamba-2 MHMH , SP sweep.

	$\eta$ =0.0003	$\eta$ =0.001	$\eta$ =0.003	$\eta$ =0.01	$\eta$ =0.03
MHMH $w$ =256	4.181	3.826	3.774	3.858	4.505
MHMH $w$ =1024	3.267	3.210	3.635	3.853	11.5
MTMT $w$ =256	4.004	3.685	3.615	3.671	4.462
MTMT $w$ =1024	3.226	3.186	3.559	11.5	11.5

Appendix

Table 22: SP fair baseline ( WikiText-103 , 20k steps).

Width	Parameters	$\eta{=}0.3$	$\eta{=}1.0$	$\eta{=}3.0$	$\eta^{*}$
256	60M	3.732	3.708	3.785	1.0
512	137M	3.575	3.556	3.587	1.0
1024	342M	3.488	3.450	3.501	1.0

Appendix

Table 23: GTGT (GLA interleaved with Transformer blocks, depth 16) under naïve

\mu\text{P}

on WikiText-103 (20k steps, one seed): validation loss. Bold marks the argmin of each row.

Variant	Depth	Best LR	$w{=}256$	$w{=}512$	$w{=}1024$	Gap
MHA, no RoPE	16	$3.0$	$4.219$	$4.120$	$\mathbf{4.151}$	$\mathbf{0.07\%}$
MHA, no RoPE	32	$3.0$	$4.453$	$4.113$	$\mathbf{4.067}$	$\mathbf{0.00\%}$
GQA, no RoPE	16	$1.0$	$3.760$	$3.344$	$\mathbf{3.369}$	$\mathbf{0.00\%}$
GQA, no RoPE	32	$1.0$	$3.527$	$3.313$	$\mathbf{3.138}$	$\mathbf{0.00\%}$

Appendix

Table 24: Nemotron-H attention-implementation comparison (naïve

\mu\text{P}

, WikiText-103 ; MHA rows 2k steps with warmup=100, GQA rows 20k steps with warmup=1000). Both variants achieve clean LR transfer; absolute losses and the optimum LR are not comparable across the two horizons.

	Loss at $\eta{=}1.0$			SP divergence rate
Arch	$w{=}256$	$w{=}512$	$w{=}1024$	@2k	@20k
HHHH	3.514	3.270	3.191	67%	100%
MHMH	3.741	3.317	3.171	67%	100%
MTMT	3.607	3.495	3.158	78%	100%

Appendix

Table 25: Convergence validation (20k steps, WikiText-103 , naïve

\mu\text{P}

). Optimal LR is

\eta{=}1.0

at all 9 (arch, width) conditions. SP divergence rate is on the 3-LR

\times

3-width grid (

\eta\in\{0.3,1.0,3.0\}

w\in\{256,512,1024\}

9

runs per arch); rates rise with training duration.

Figure 5 : Optimal LR at each training checkpoint (20k steps, WikiText-103). At $w{=}256$ , $\eta{=}0.3$ briefly leads during mid-training (steps 3k–14k) before $\eta{=}1.0$ recovers. By convergence, $\eta{=}1.0$ is universally optimal.

Figure 6 : Final training loss (mean $\pm$ std) across 3 seeds (2k steps, WikiText-103 ). $\eta{=}1.0$ (blue) is optimal in all 12 per-seed conditions. Cross-seed CV $\leq$ 1.7% at the optimum, confirming high reproducibility. MHMH $\eta{=}3.0$ at $w{=}256$ shows anomalously high variance (CV $\approx$ 5.8%).

Batch size	MHMH gap	HHHH gap	MHMH optimal LR	HHHH optimal LR
8	0.5%	0.0%	0.7–1.0	0.5
16	0.0%	1.1%	1.0	0.7–1.0
32	0.3%	0.0%	1.5–2.0	1.0
64	0.0%	0.0%	1.5–2.0	1.5

Appendix

Table 26: Width transfer gap at each batch size (naïve

\mu\text{P}

, 2k steps, WikiText-103 , training loss). Gap: loss at w=256 optimal LR applied to w=1024 vs. best at w=1024.

Arch	Width	seq=128	seq=256	seq=512	seq=1024
HHHH	$w{=}256$	4.399	4.100	4.074	4.477
HHHH	$w{=}1024$	4.102	4.262	4.215	4.509
MHMH	$w{=}256$	4.213	4.388	4.580	4.084
MHMH	$w{=}1024$	4.161	4.323	4.300	4.338
TT	$w{=}256$	4.258	4.235	4.454	4.267
TT	$w{=}1024$	4.255	4.459	4.207	4.547

Appendix

Table 27: Sequence length scaling at

\eta{=}1.0

(naïve

\mu\text{P}

, 2k steps, WikiText-103 ). Final loss at each (arch,

w

, seq_len).

\eta{=}1.0

is optimal at every cell; LR transfer is sequence-length invariant for tested lengths 128–1024.

expand	$d_{\text{inner}}$	Attention	SwiGLU FFN	T total	M total	T:M Ratio	H block
1	512	1.05M	2.16M	3.21M	0.85M	3.8:1	4.06M
2	1024	1.05M	2.16M	3.21M	1.69M	1.9:1	4.91M
4	2048	1.05M	2.16M	3.21M	3.39M	0.9:1	6.60M

Appendix

Table 28 : H-block parameter breakdown by ssm_expand (

d_{\text{model}}=512

Figure 7 : Transfer gap across T:M ratios (naïve $\mu\text{P}$ and SP, WikiText-2 , 2k steps, warmup 100; training loss). Naïve $\mu\text{P}$ : 0.0% gap for all 6 configurations. SP: 0.0% for 5/6 but catastrophic (1144%) for expand1 where the optimal LR shifts with width. LR=1.0 is invariant to the T:M ratio; no re-tuning needed when changing ssm_expand .

		Naïve $\mu\text{P}$			Block-wise $\mu\text{P}$ (approx.)			Block-wise $\mu\text{P}$			SP
Pattern	%M	$w{=}256$	$w{=}512$	$w{=}1024$	$w{=}256$	$w{=}512$	$w{=}1024$	$w{=}256$	$w{=}512$	$w{=}1024$	$w{=}256$	$w{=}512$	$w{=}1024$
TT	0%	1.0	1.0	1.0	1.0	1.0	1.0	1.0	1.0	1.0	0.003	0.003	0.003
M1T3	25%	1.0	1.0	1.0	1.0	3.0	3.0	1.0	3.0	1.0	0.003	0.003	0.003
MTMT	50%	1.0	1.0	1.0	1.0	1.0	1.0	1.0	3.0	1.0	0.010	0.003	0.003
M3T1	75%	1.0	1.0	1.0	1.0	1.0	3.0	1.0	3.0	1.0	0.010	0.003	0.003
JAMBA (1:7)	87.5%	1.0	1.0	1.0	1.0	3.0	3.0	3.0	3.0	1.0	0.010	0.003	0.003

Appendix

Table 29: Serial T:M ratio transfer at depth

16

. Optimal

\eta^{*}

at each (scheme, pattern, width); shading mirrors Table 1 , where uniform color within a row indicates clean LR transfer. Naïve

\mu\text{P}

holds

\eta^{*}{=}1.0

across the full ratio range; approximated block-wise

\mu\text{P}

drifts on the SSM-bearing layouts; SP drifts on every layout where the SSM density is non-trivial. The full block-wise prescription shifts on every SSM-bearing layout, with

\eta{=}3.0

diverging at

w{=}1024

on each.

Pattern	%M	Loss at $\eta{=}1.0$	Best loss (best $\eta$ )	Penalty
TT	0%	1.048	1.048 ( $\eta{=}1.0$ )	$0.0\%$
M1T3	25%	0.755	0.725 ( $\eta{=}3.0$ )	$4.2\%$
MTMT	50%	0.734	0.734 ( $\eta{=}1.0$ )	$0.0\%$
M3T1	75%	0.818	0.600 ( $\eta{=}3.0$ )	$36.4\%$
JAMBA	87.5%	0.852	0.630 ( $\eta{=}3.0$ )	$35.2\%$
MM	100%	1.144	0.916 ( $\eta{=}3.0$ )	$24.8\%$

Appendix

Table 30: Approximated block-wise

\mu\text{P}

residual loss penalty at

w{=}1024

( WikiText-2 , 2k steps). Penalty is the fractional loss increase incurred by transferring the proxy optimum

\eta{=}1.0

instead of using each row’s per-cell optimum at

w{=}1024

. Naïve

\mu\text{P}

has

0\%

penalty on every row.

Figure 8 : Coordinate check for HHHH ( $d{=}16$ , widths 128–2048). Row 1: naïve $\mu\text{P}$ , hidden activations vary with width (median block width-spread $2.78\times$ , Table 31 ; violates textbook $\mu\text{P}$ ). Row 2: block-wise $\mu\text{P}$ (approximated), similar pattern. Row 3: block-wise $\mu\text{P}$ (literal Vankadara Table 1), hidden activations predominantly $10^{3}$ – $10^{4}\times$ larger in absolute magnitude than naïve at deeper blocks, with the largest width-spread. Row 4: SP, activations explode. Despite the hidden-state divergence, logit-level magnitudes (rightmost column) are comparable across all $\mu\text{P}$ variants ( $2.0$ – $5.4$ at step 10, with width-spread $\leq$ $2.5\times$ within each variant), suggesting layer-wise normalization absorbs much of the width-dependence; naïve $\mu\text{P}$ achieves $0.0\%$ LR transfer under Adam. Each row has its own y-axis.

Figure 9 : Coordinate check for MHMH ( $d{=}16$ , widths 128–2048). Same pattern as HHHH. Each row has its own y-axis.

Arch	Scheme	Median width-spread
TT	Naïve $\mu\text{P}$	1.52
	SP	819
GTGT	Naïve $\mu\text{P}$	1.23
	SP	$7.3{\times}10^{3}$
HHHH	Naïve $\mu\text{P}$	2.78
	Block-wise (approx.)	2.45

Appendix

Table 31: Median width-spread of the coordinate check (step 10, widths 128–2048, five seeds,

d{=}16

): for each hidden block, the largest over the smallest seed-averaged mean absolute activation across widths, and the median over the 16 blocks. A value of 1 means width-invariant.

Figure 10 : Coordinate check for TT (pure Transformer, $d{=}16$ , widths 128–2048) with the same instrument. Row 1: naïve $\mu\text{P}$ , width lines overlap (median block width-spread $1.52\times$ , Table 31 ). Row 2: SP, activations diverge with width.

Figure 11 : Gradient RMS per component vs. width for MHMH at step 1, under SP (widths 128–512) and naïve $\mu\text{P}$ (128–1024). Components differ widely in gradient scale. Under Adam, the $1/\sqrt{v}$ denominator normalizes these disparate scales.

	Width
Scheme	128	256	512	1024	2048
SP	0.6	0.9	1.2	1.7	2.5
Naïve $\mu\text{P}$	0.1	0.2	0.3	0.4	0.5
Block-wise $\mu\text{P}$	0.1	0.2	0.3	0.4	0.5

Appendix

Table 32: Gradient norm at initialization across widths and schemes (MT interface).

Figure 12 : Stable-window scatter diagnostic for MHMH naïve $\mu\text{P}$ , $A_{\log}$ . (A) Per-run log-gradient trajectories over steps 100–199 show isolated single-step collapses (downward spikes) rather than smooth noise. (B) Pooled log-residual histogram against a Gaussian with matching standard deviation. The residuals are strongly heavy-tailed (excess kurtosis $+4.8$ ). (C) Top-1% $|$ residual $|$ values scattered throughout the stable window; no concentration at early or late steps. (D) Outlier concentration by $(d,\texttt{expand})$ cell: non-uniform, with $\texttt{expand}=2,4$ at $d=128$ and $d=512$ concentrating outliers 2–3%; $\texttt{expand}=1$ uniformly quiet. The outliers are transient update-instability events structured by configuration, not uniform measurement noise.

Arch	Optimizer $\times$ Width	$\eta$ =0.3	$\eta$ =1.0	$\eta$ =3.0	$\eta$ =10.0	$\eta$ =30.0	$\eta$ =100.0
MHMH	AdamW $w$ =256	3.566	3.518	3.616	4.018	9.909	5.445
	AdamW $w$ =1024	3.226	3.129	3.567	3.667	3.838	10.5
	SGD $w$ =256	6.927	6.503	6.176	5.668	5.238	5.121
	SGD $w$ =1024	7.296	6.736	6.329	6.147	5.687	5.176
HHHH	AdamW $w$ =256	3.505	3.487	3.581	div	4.542	10.9
	AdamW $w$ =1024	3.160	3.082	3.476	3.691	3.829	10.9

Appendix

Table 33: SGD vs AdamW for MHMH and HHHH (naïve

\mu\text{P}

, WikiText-103 , 20k steps, warmup 1000).

Optimizer $\times$ Width	$\eta$ =100	$\eta$ =300	$\eta$ =1000	$\eta$ =3000	$\eta$ =7000
SGD $w$ =256	5.122	5.077	5.067	div	–
SGD $w$ =1024	5.174	5.033	5.010	5.004	div

Appendix

Table 34: HHHH under SGD above

\eta{=}100

, in the configuration of Table 33 (naïve-

\mu\text{P}

parameter groups, WikiText-103 , 20k steps, one seed): final training loss. The

\eta{=}100

column repeats Table 33 . Bold marks the argmin of each row; div : diverged; –: not run.

Parameterization	Optimizer	Best LR @256	Loss @256	Best LR @1024	Loss @1024
Naïve $\mu\text{P}$	AdamW	1.0	3.518	1.0	3.129
Naïve $\mu\text{P}$	SGD	100.0	5.121	100.0	5.176
Approximated block-wise $\mu\text{P}$	SGD	100.0	5.122	100.0	5.408

Appendix

Table 35:

2{\times}2

matrix: parameterization

\times

optimizer (MHMH, best LR at each width).

Figure 13 : Training curves for three parameterization–optimizer combinations. AdamW (blue) separates by width (wider = better); SGD variants (pink, orange) overlap or show reversed scaling. Approximated block-wise $\mu\text{P}$ does not improve SGD over naïve $\mu\text{P}$ .

Figure 14 : Width gap (loss at $w{=}1024$ minus loss at $w{=}256$ ) over training. Positive values indicate reversed scaling (wider = worse). AdamW (blue) achieves healthy scaling; both SGD variants (pink, orange) show persistent reversed scaling.

	$\eta$ =0.3	$\eta$ =1.0	$\eta$ =3.0
HHHH $w$ =256 WD=0	4.573	4.411	4.517
HHHH $w$ =256 WD=0.1	4.561	4.404	5.013
HHHH $w$ =1024 WD=0	4.392	4.090	4.471
HHHH $w$ =1024 WD=0.1	4.381	4.102	4.161
MHMH $w$ =256 WD=0	4.411	4.234	4.296
MHMH $w$ =256 WD=0.1	4.411	4.217	4.234

Appendix

Table 36: Adam no-WD ablation ( WikiText-103 , 2k steps, naïve

\mu\text{P}

Arch	Optim $\times$ Width	$\eta$ =0.001	$\eta$ =0.003	$\eta$ =0.01	$\eta$ =0.03	$\eta$ =0.1	$\eta$ =0.3	$\eta$ =1.0
MHMH	AdamW $w$ =256	9.968	8.215	6.638	5.625	4.898	4.412	4.219
	AdamW $w$ =1024	9.303	7.950	6.665	5.706	5.020	4.450	4.160
	Adam ( $\beta_{1}{=}0$ ) $w$ =256	9.951	8.194	6.726	5.653	4.954	4.448	4.511
	Adam ( $\beta_{1}{=}0$ ) $w$ =1024	9.289	7.935	6.718	5.737	5.059	4.608	4.418
	SGD $w$ =256	11.5	11.5	11.5	11.4	9.999	8.462	7.626
	SGD $w$ =1024	11.5	11.5	11.5	11.4	10.9	8.821	7.911

Appendix

Table 37 : Optimizer ablation for MHMH and HHHH (2k steps, WikiText-103 , naïve

\mu\text{P}

). Final loss at each LR; bold marks the optimum within each row.

$\eta$	$w$	$\epsilon{=}10^{-10}$	$\epsilon{=}10^{-8}$	$\epsilon{=}10^{-6}$	$\epsilon{=}10^{-4}$
$0.5$	$256$	4.306	4.293	4.260	4.557
$0.5$	$1024$	4.287	4.286	4.281	4.852
$1.0$	$256$	4.236	4.209	4.177	4.316
$1.0$	$1024$	4.176	4.167	4.150	4.447
$1.5$	$256$	4.219	4.211	4.163	4.243
$1.5$	$1024$	4.156	4.148	4.209	4.302

Appendix

Table 38: AdamW

\epsilon

sensitivity for MHMH ( WikiText-103 , 2k steps, naïve

\mu\text{P}

). Each

\eta

block has paired rows for

w{=}256

and

w{=}1024

so width-invariance can be read directly across the row pair; bold marks the optimum across the four

\epsilon

values within each row.

Optimizer	Arch	$\eta^{*}$ $w{=}256$	$w{=}512$	$w{=}1024$	Gap (mean)	Gap (per seed)
Sophia	HHHH	0.1	0.1	0.1	0.00%	0.00% (3/3)
Sophia	MHMH	0.1	0.1	0.1	0.00%	0.00% (3/3)
Muon	HHHH	0.01	0.01	0.005	3.58%	2.29–5.74%
Muon	MHMH	0.01	0.01	0.005	7.83%	5.62–10.38%

Appendix

Table 39: Sophia and Muon on WikiText-103 (20k steps, three seeds, validation loss).

\eta^{*}

is the argmin of the seed-mean loss; the gap is the transfer gap from

w{=}256

w{=}1024

, for the seed mean and per seed. LR grids: Sophia

\{0.03,0.1,0.3\}

, Muon

\{0.002,0.005,0.01,0.02,0.05,0.1\}

Adam component	Role for LR transfer	Evidence
Adaptive LR ( $g/\sqrt{v}$ )	Load-bearing	Only Adam variants transfer
Momentum ( $\beta_{1}$ )	Not load-bearing	2–9% degradation, transfer preserved
Weight decay ( $\lambda$ )	Not load-bearing	$\leq$ 0.4% effect

Appendix

Table 40: Adam-component contribution to LR transfer: which feature is required vs. incidental.

Explore similar work

Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate
May 20, 2026Dayal Singh Kalra, Maissam BarkeshliHyperparameter Large Language Model Training
Department of Computer Science, University of Maryland, College Park · Department of Physics, University of Maryland, College Park · Joint Quantum Institute, University of Maryland, College Park +1
Priming: Hybrid State Space Models From Pre-trained Transformers
May 8, 2026Aditya Chattopadhyay, Elvis Nunez, Prannay Kaul +6Generative Pretrained Transformers Priming
AWS Agentic AI
Long-Context Modeling via GSS-Transformer Hybrid Architecture with Learnable Mixing
Jun 15, 2026Kuzey Torlak, Hüseyin Arda Arslan, Anıl Dervişoğlu +2Efficient Long-Context Inference Hybrid Transformer
Kadıköy Anadolu High School · Politecnico di Torino · Istanbul Technical University +2

May 20, 2026, cs.LGJ/K move · Enter open · S save

Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate

Dayal Singh Kalra, Maissam Barkeshli

Department of Computer Science, University of Maryland, College Park · Department of Physics, University of Maryland, College Park · Joint Quantum Institute, University of Maryland, College Park+1

Hyperparameter Large Language Model Training Transfer Learning Extrapolation

Open paper AskPDF arXiv

Hyperparameter transfer allows extrapolating optimal optimization hyperparameters from small to large scales, making it critical for training large language models (LLMs). This is done either by fitting a scaling law to the hyperparameters or by a judicious choice of parameterization, such as Maximal Update (

μ

P), that renders optimal hyperparameters approximately scale invariant. In this paper, we first develop a framework to quantify hyperparameter transfer through three metrics: (1) the quality of the scaling law fit, (2) the robustness to extrapolation errors, and (3) the asymptotic loss penalty due to choice of parameterization. Next, we investigate through a comprehensive series of ablations why

μ

P appears to offer high-quality learning rate transfer relative to standard parameterization (SP), as existing theory is inadequate. We find that the overwhelming benefit of

μ

P relative to SP when training with AdamW arises simply from maximizing the learning rate of the embedding layer. In SP, the embedding layer learning rate acts as a bottleneck that induces training instabilities; increasing it by a factor of width to match

μ

P dramatically smooths out training while improving hyperparameter transfer. We also find that weight decay improves the scaling law fits, while, in the fixed token-per-parameter setting, it hurts the robustness of the extrapolation.

Abstract

Explore similar work

Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate

Priming: Hybrid State Space Models From Pre-trained Transformers

Long-Context Modeling via GSS-Transformer Hybrid Architecture with Learnable Mixing