PaperlayerView exact revision on arXiv

Source-linked AI summary

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, Nakyung Lee, Minsik Seol, Idris Nechnech, Jaehyeon Kim, Giho Lee, Jungwoo Lee

arXiv:2610.01172v1cs.LG

TL;DR

Practical hybrid Transformer-SSM models still lack evidence for LR transfer under finite-width, simplified-ZOH implementations, despite separate theoretical SSM scaling rules. This paper tests the original µP prescription with AdamW and finds near-zero transfer gaps across tested widths and depths, including on a production hybrid. It attributes the pattern to global update-to-weight invariance together with local per-component balance.

  • Problem

    SSM scaling rules were derived for pure SSMs in an infinite-width, growing-state, true-ZOH regime, leaving LR transfer for practical finite-width hybrids unaddressed.

    Source: §2.2 · p. 3

  • Method

    The paper empirically tests naïve µP with AdamW across hybrid architectures and analyzes LR transfer through global update-to-weight invariance and local per-component balance.

    Source: §1 · p. 2

  • Results

    Naïve µP with AdamW achieves near-zero LR transfer gaps across tested hybrid architectures, widths 256–2048, and depths 4–32; Nemotron-H transfers at a 0.0% gap.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The reported hybrid transfer behavior is explained by µP’s global update-to-weight invariance together with AdamW’s local per-component normalization.

    Source: §1 · p. 2

    Whether near-zero transfer persists when billion-scale training, the full training horizon, and multiple seeds are combined remains open; larger architectures, including Mixture-of-Experts, remain untested.

    Source: §6 · p. 10

Abstract

from arXiv · show

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 $\mu$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 $\mu$P correctness. We attribute this to a two-condition decomposition of LR transfer in hybrid architectures: a global update-to-weight invariance, enforced by $\mu$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.

1 Introduction

The paper addresses the costly per-scale LR tuning required for hybrid Transformer-SSM models, where existing SSM scaling theory does not directly match practical finite-width implementations. It reports near-zero LR transfer with naïve µP and AdamW, and proposes that global update-to-weight invariance and local per-component balance jointly explain the result.

1 Introduction

  • Hybrid-model training still requires expensive per-scale hyperparameter tuning, while prior SSM scaling rules address pure SSMs rather than practical finite-width hybrids.The SSM rules were derived for infinite width, growing state size, SGD, and true-ZOH discretization.Source: §1 · pp. 1–2
  • Naïve µP with AdamW achieves near-zero LR transfer gaps across tested hybrids, widths 256–2048, and depths 4–32, although all tested parameterizations fail the coordinate check.µP-SSM’s optimal LR drifts with width on SSM-bearing architectures, while naïve µP transfers cleanly.Source: §1 · p. 2
  • The proposed decomposition assigns global update-to-weight invariance to µP and local per-component balance to AdamW’s normalization, supported by an optimizer ablation.The per-component scaling test finds Vankadara’s prescription works at initialization for B and C but fails for Alog under simplified-ZOH Mamba.Source: §1 · p. 2
  • Nemotron-H, a production hybrid, transfers at a 0.0% gap without architecture-specific tuning.Source: §1 · p. 2

2 Background and Related Work

Prior work establishes µP width scaling for representable architectures and derives separate SSM rules, but those analyses do not settle LR transfer in practical hybrids. Related work also connects Adam normalization to robustness across heterogeneous components and distinguishes LR transfer from production-oriented per-component tuning.

2.1 Maximal Update Parameterization (µP)

  • µP scales hidden-layer initialization as N(0, 1/d) and learning rates as ηbase/d, keeping the update-to-weight ratio O(1) as width grows.Its Tensor Programs basis applies to architectures composed of representable operations; prior validation includes pure Transformers.Source: §2.1 · p. 3

2.2 µP-SSM: Scaling Rules for State-Space Models

  • SSM operations violate Tensor Programs representability conditions, motivating µP-SSM rules that keep each component’s output contribution constant with width.The cited analysis concerns infinite-width SSMs under SGD and Adam, with empirical validation on pure S4 and Mamba under SGD; its practical hybrid gap includes true-ZOH with growing state size versus simplified-ZOH at fixed Nx=16.Source: §2.2 · p. 3

2.3 Adam and Hyperparameter Transfer

  • Adam’s coordinatewise normalization can make updates depend on gradient signs, and the paper reports that g/√v is essential in its hybrid ablation while removing weight decay changes loss by at most 0.4%.The weight-decay ablation lasts 2k steps, shorter than the weight-decay timescale at target widths, so it does not cover longer horizons examined by Kosson et al. [2026].Source: §2.3 · p. 3

2.4 Hybrid Architectures and Production Training

  • Production hybrid training generally uses per-scale tuning; Falcon-H1 tunes 35 per-component multipliers for stability and performance rather than µP-style LR transfer.Nemotron-H is another production hybrid reliant on per-scale tuning, with Mamba-2 SSD blocks, about 8% attention layers, and squared-ReLU FFN.Source: §2.4 · p. 4

3 Experimental Setup

The experiments test whether µP-based LR transfer persists across hybrid block interfaces, model scales, and training conditions. They compare parameterization schemes using controlled corpora and define transfer by the relative loss difference between proxy- and target-optimal learning rates.

3 Experimental Setup

  • The study probes transfer first at heterogeneous block interfaces, then across depth and width, and finally in larger training runs with downstream evaluation.Source: §3 · p. 4

3.1 Hybrid Transformer-SSM Architectures

  • The architecture set includes Transformer, Mamba, and parallel-fusion hybrid blocks, with serial, parallel, and hybrid-Transformer patterns; varying the Transformer-to-SSM ratio showed no effect on LR transfer.Reversed block orders are included for transfer coverage, not to compare orderings.Source: §3.1 · p. 4

3.2 Parameterization Schemes

  • Four schemes are compared: SP, naïve µP, block-wise µP, and an approximated block-wise variant using µP-style scaling for selected SSM learning rates.The block-wise schemes apply Vankadara et al. [2024]’s SSM prescription, while the approximated variant retains per-component initialization and ηC scaling but changes ηA and ηB scaling.Source: §3.2 · pp. 4–5

3.3 Training Details

  • Training uses AdamW, fixed seeds, WIKITEXT-2 and WIKITEXT-103 for short-horizon and convergence runs, and FINEWEB-EDU for larger-corpus validation.Unless varied, sweeps use sequence length 128 and batch size 64; transferred-LR convergence runs also use sequence length 2048, and no cross-dataset LR-optimality claim is made.Source: §3.3 · p. 5

3.4 Evaluation Metrics

  • The transfer gap is the relative loss difference between training a target at the proxy’s optimal LR and at its own optimal LR.A 0% gap means the proxy-optimal LR is also optimal at the target within the tested LR resolution.Source: §3.4 · p. 5

4 Results

The results section validates LR transfer through interface and depth tests, with Table 1 reporting optimal learning rates by scheme, architecture, and width. The supplied findings contrast scheme stability and transfer behavior across interfaces.

4.1 Unit Interface and Depth Validation

  • For WIKITEXT-103 and FINEWEB-EDU sweeps recording both losses, training- and validation-loss optima match in 209 of 218 cells; on WIKITEXT-2, validation prefers LR 0.1 versus training-loss LR 1.0 for naïve µP.The WIKITEXT-2 run traverses the data about seven times in 2k steps.Source: §4.1 · p. 5
  • Interface tests report transfer statuses ranging from no shifts to drift on all six SSM-bearing cases, while the table organizes optimal LRs by scheme, architecture, and width.The table uses 2k-step WIKITEXT-2 training; its caption says uniform shading within a row indicates transfer across the corresponding widths.Source: Table 1 · p. 6
  • Naïve µP remains stable beyond the highest tested LR, η=1.0, across widths and interfaces, whereas SP’s stability window shrinks from 0.30–1.0 at width 128 to 0.03–0.10 at width 2048.The approximated block-wise scheme collapses to η≤0.01 at width ≥256 for serial SSM interfaces MM and MT.Source: §4.1 · p. 6

Appendix B.

Naïve µP transfers across width and depth in the tested hybrids, including joint scaling to hundreds of millions of parameters and validation on Nemotron-H. The appendix also compares alternative parameterizations and tests transfer on a curated web-scale corpus.

Appendix B.

  • Across depths 4–32, naïve µP keeps η∗=1.0 in all 32 architecture-depth cells, whereas full block-wise µP shifts to η∗=3.0 in 25 of 28 SSM-bearing cells.Full block-wise µP destabilizes at η≥10 in all 28 SSM-bearing cells; the approximated variant’s 20k-step loss stays within 0.5% of naïve µP at width 1024.Source: Appendix B · p. 6

4.2 Width Transfer and Scale Validation

  • η∗=1.0 at every tested width for all seven architectures under naïve µP, the only scheme among those tested that transfers cleanly.The other three schemes drift with width; the joint depth-width grid was therefore tested only under naïve µP.Source: §4.2 · p. 6
  • 0.0% transfer gap holds from proxy to target across an approximately 11×–14× parameter increase, and across three tested architectures on FINEWEB-EDU.For MHMH and HHHH, the joint depth-width transfer uses η=1.0; on FINEWEB-EDU, SP shifts 3× and proxy-LR transfer incurs a 1.1–8.7% penalty.Source: §4.2 · p. 6

4.3 Generalization to Nemotron-H

  • Nemotron-H tests transfer on a production hybrid composed of Mamba-2, squared-ReLU FFN, and attention-only blocks with grouped-query attention and no rotary embeddings.Its block types differ from those in the custom architecture set.Source: §4.3 · p. 7
  • On Nemotron-H, naïve µP has η∗=1.0 and a 0.0% transfer gap at every tested depth-width combination, while SP’s optimum shifts 3× across width.The tested learning rates η∈{1, 3, 5} for naïve µP lie in a broad basin, all within approximately 3% of optimum.Source: §4.3 · p. 7

5 Analysis

The analysis finds that hybrid LR transfer can hold despite failed coordinate checks, and proposes global update-to-weight invariance plus local per-component update balance as explanatory conditions. In the tested simplified-ZOH, fixed-state-size regime, AdamW supplies the local balance, while µP-SSM rules do not; broader theoretical questions remain open.

5.1 Naïve µP Violates Coordinate Checks but Learning Rate Transfers

  • Every tested parameterization fails the coordinate check on hybrids, yet naïve µP achieves a 0.0% transfer gap across the architectures in Table 1.Block-wise µP’s optimum drifts with width on every SSM-bearing architecture, so coordinate checks do not predict transfer in this setting.Source: §5.1 · p. 7

5.2 Two Conditions for Scalable Transfer in Hybrid Architectures

  • The proposed decomposition requires width-invariant update-to-weight ratios within parameter groups and preserved per-component update ratios between proxy and target widths.Under naïve µP, ρ=O(η) is width-independent; a width shift in β indicates rebalancing across components and can break proxy optimality.Source: §5.2 · p. 8

5.3 Optimizer Choice Under the Two Conditions

  • Under AdamW, per-parameter normalization supplies local update balance without component-specific LR scaling, while adding Vankadara’s scaling over-corrects the Alog/FFN pair.For that pair, the gap is 20–21% with AdamW and naïve µP, versus 166–436% when block-wise scaling is added; under SGD with naïve µP it is 235–342%.Source: §5.3 · p. 8
  • Under SGD, Vankadara’s rule yields width-invariant step-1 updates for B and C but not Alog, and the B/C invariance does not persist through training.The Alog update has α=+1.00 at step 1; drift from step 1 to the stable window reaches +2.4 for Alog, +1.34 for B, and +0.68 for C.Source: §5.3 · pp. 8–9

5.4 Summary

  • In the tested simplified-ZOH Mamba regime at fixed Nx=16, AdamW’s normalization supplies Condition 2, but whether the decomposition suffices at infinite width remains open.A faithful µP-SSM test under true-ZOH SSMs with growing Nx also remains open.Source: §5.4 · p. 9

6 Conclusion

  • Naïve µP with AdamW transfers across widths up to billion-parameter scale in the experimental setting and generalizes to Nemotron-H, with robustness across several individually varied settings.The 1.24B sweep, full training horizon, and multiple seeds were validated individually rather than jointly; whether the near-zero gap persists when combined remains open.Source: §6 · pp. 9–10

A.1 Compute Budget and Tuning-Cost Savings

Using naïve µP transfer can substantially reduce per-architecture LR-tuning costs compared with sweeping at the target width. The reported cost estimates use single-A6000 configurations and vary by training horizon.

A.1 Compute Budget and Tuning-Cost Savings

  • The reported per-configuration wall times are estimated on one A6000 for HHHH at depth 16; other architectures and GPUs are within ±20%.The proxy-to-target wall-time ratio is 5.5× at 2k steps and about 12× at 20k.Source: §A.1 · p. 14, Table 8 · p. 15
  • At w=1024, proxy sweeping plus one target verification is estimated at 0.44 GPU-hours for 2k steps and 4.0 GPU-hours for 20k steps.The SP baseline sweeping seven LRs at w=1024 is estimated at 2.0 and 22.3 GPU-hours, respectively.Source: §A.1 · p. 15
  • 78% lower tuning cost at 2k steps and 82% at 20k steps are estimated for naïve µP transfer; a width-2048 extrapolation reaches about 84% savings.The corresponding speedups are 4.5×, 5.6×, and an extrapolated 6.4×, respectively.Source: §A.1 · p. 15

A.2 Dataset Sizes and Training Coverage

The appendix describes the dataset coverage used to interpret the transfer results and the approximated block-wise µP scaling choices. The approximation retains selected per-component rules while adjusting the A and B learning-rate scaling.

A.2 Dataset Sizes and Training Coverage

  • The canonical dataset-coverage accounting uses seq_len=128 and batch_size=64, or 8,192 tokens per optimizer step.Table 9 reports corpus token counts, tokenizer, and epoch-equivalent step coverage; the supplied passages do not give its individual corpus values.Source: Table 9 · p. 15, §A.2 · p. 15
  • The 0.0% WIKITEXT-2 transfer gap reflects about seven passes over the small corpus, while the 20k-step WIKITEXT-103 runs cover about one epoch.FINEWEB-EDU sweeps use about one-third of the 500M-token subset; the 5B-token run uses a larger effective batch than the canonical accounting.Source: §A.2 · p. 15

A.3 Approximated Block-wise µP: Formula

  • The approximated block-wise variant keeps per-component initialization and ηC scaling but applies naïve-µP exponents to ηA and ηB.It drops ηA's Θ(Nu) amplification and replaces ηB's Θ(1/√Nu) scaling with Θ(1/Nu), avoiding the redundant AdamW normalization described in the passage.Source: §A.3 · pp. 15–16
  • At w=1024 after 20k WIKITEXT-103 steps, the approximation stays within 0.5% of naïve µP, while full block-wise µP has 5.6–7.7% higher loss across cells.The cited comparison is under AdamW on hybrid architectures.Source: §A.3 · p. 16

A.4 Applying µP-SSM Beyond the Pure-SSM SGD Setting

The appendix tests µP-SSM rules beyond their original pure-SSM, SGD, growing-state derivation regime and reports additional interface, depth, and width-transfer evidence. In the hybrid AdamW setting, the full rule has worse loss and width-dependent LR behavior than naïve µP, while the approximation stays close in the reported convergence comparison.

A.4 Applying µP-SSM Beyond the Pure-SSM SGD Setting

  • The µP-SSM rule was derived for pure SSMs with SGD, growing Nx, and true-ZOH discretization, unlike the hybrid AdamW experiments with fixed Nx and simplified ZOH.Whether the same exponents are theoretically optimal under simplified ZOH is not established.Source: §A.4 · pp. 16, 18
  • With proportional Nu/Nx scaling under SGD, the Table 1 rule has lower loss than naïve µP at η=0.3–30 across widths, while naïve µP catches up at η=100.In fixed-Nx pure Mamba under SGD, the rule is lower-loss at low-to-mid LRs but diverges at η=100 for w=128; the reported setup differs from the original growing-Nx regime.Source: §A.4 · pp. 16–17
  • At 20k steps on WIKITEXT-103, full block-wise µP has 5.6–7.7% higher loss than naïve µP in all nine hybrid cells; its optimum also drifts with width on SSM-bearing architectures.The approximated variant is within 0.5% of naïve µP at w=1024; under AdamW, the explicit component LR rule is redundant with g/√v and slows training.Source: §A.4 · p. 17, Table 11 · p. 17

B.1 Max Stable Learning Rate

  • Naïve µP remains stable at η=1.0 across every unit-interface and width cell reported in Table 12.The table covers nine ordered T, M, and H pairs; SP's maximum stable LR shrinks with width.Source: §B.1 · p. 18, Table 12 · p. 19

B.2 Depth Scaling

  • Loss decreases with depth for all four schemes on MHMH and HHHH, while the block-wise variants have higher 2k-step WIKITEXT-2 loss.The approximated variant is within 0.5% of naïve µP after 20k steps on WIKITEXT-103 at w=1024; the faithful variant retains a 5–8% gap at convergence.Source: §B.2 · p. 19, Figure 2 · p. 20

C.1 Joint Depth-Width Transfer

  • At w=1024 and d=32, η=1.0 is optimal for HHHH and MHMH, with a 0.0% gap from the proxy model at w=256 and d=8.The proxy is about 11–14× smaller in parameters; the supplied table caption cautions that its three-point grid misses the basin between η=1 and η=3.Source: §C.1 · p. 20, Table 13 · p. 21

C.2 FineWeb-Edu Width Transfer

  • The FINEWEB-EDU width-transfer evaluation includes HHHH, MHMH, and MTMT under SP and naïve µP, using 164M training tokens from a 500M-token subset.The associated validation-loss trajectory is for the MHMH-415M, 5B-token run; the caption reports convergence from 5.45 at step 500 to 2.96 without retuning η=1.0 from the proxy.Source: §C.2 · p. 21, Figure 4 · p. 22

C.3 Width 2048 and Billion Scale

The section extends width and depth tests to billion-parameter models, refines the LR sweep, and checks other mixers and a Nemotron-H reproduction. Naïve µP transfer remains strong in the reported settings, with explicit scope limits on the longest-horizon and reproduction claims.

C.3 Width 2048 and Billion Scale

  • At w=2048, naïve µP's 2k-step optimum is η=1.5 for MHMH and HHHH; at 20k steps, MHMH's optimum returns to η=1.0.The models contain 1.24B and 1.66B parameters, respectively; the associated validation, seed, and horizon checks are performed separately rather than jointly, and the full-horizon run uses one LR and seed.Source: §C.3 · p. 22

C.4 Fine-Grained LR Sweep

  • The fine-grained sweep finds η=1.5 optimal at every width after 2k steps, with η=1.0 costing at most 0.6%; at 20k steps, η=1.0 is optimal across widths.The 20k-step basin spans η=0.5–1.5, with per-cell spread ≤0.7%.Source: §C.4 · p. 23, Table 19 · p. 23

C.5.1 Mamba-2 Compatibility

  • For Mamba-2 blocks in MHMH, naïve µP's basin minimum is η=2.0 at both tested widths, while SP's optimal LR shifts 10× from w=256 to w=1024.SP's optima are η=0.01 and η=0.001, respectively; η=0.03 nearly diverges at w=1024.Source: §C.5.1 · p. 24

C.5.2 SP Fair Baseline

  • The SP fair-baseline comparison is reported for WIKITEXT-103 at 20k steps.The supplied passage identifies the comparison table but gives no result values.Source: Table 22 · p. 24

C.5.3 Gated Linear Attention

  • With GLA interleaved with Transformer blocks, η=1.0 is optimal at every tested width on training and validation loss, yielding a 0.0% transfer gap.This GTGT experiment uses WIKITEXT-103 for 20k steps and one seed.Source: §C.5.3 · p. 24, Table 23 · p. 24

C.6 Nemotron-H Block Composition

  • Both Nemotron-H attention reproductions show approximately 0% transfer gap, supporting transfer across hybrid designs outside the paper's T/M/H block family.The result is for dense attention and Nx=16, whereas production Nemotron-H uses sparse attention and Nx=128; MHA and GQA runs use different horizons, so their absolute losses are not comparable.Source: §C.6 · p. 25, Table 24 · p. 26

D.1 Convergence at 20k Steps

  • At 20k steps on WIKITEXT-103, naïve µP has a 0.0% transfer gap; SP divergence rates rise with training duration, reaching 100% in the measured 20k-step grid.A mid-training interval briefly favors η=0.3 for width 256, but the passage describes this as self-correcting.Source: §D.1 · p. 26, Table 25 · p. 26, Figure 5 · p. 26

D.2 Multi-Seed Validation

  • Across three seeds, η=1.0 is optimal in all 12 per-seed conditions, with a 0.0% transfer gap and cross-seed CV ≤1.7% at the optimum.The experiments cover HHHH and MHMH at two widths; Figure 6 also reports anomalously high variance for MHMH at w=256 with η=3.0.Source: §D.2 · p. 26, Figure 6 · p. 27

D.3 Batch Size Robustness

Naïve µP preserves width-transfer performance across batch sizes, sequence lengths, and Transformer-to-SSM ratios in the tested settings. The best-loss heterogeneous ratio also favors more SSM capacity in later layers.

D.3 Batch Size Robustness

  • All tested batch sizes preserve width transfer, with transfer gaps no greater than 1.1%.Smaller batches shift the optimal LR downward, but the shift is width-independent.Source: §D.3 · p. 27, Table 26 · p. 27

D.4 Sequence Length Scaling

  • η=1.0 is optimal across sequence lengths 128–2048, all three architectures, and both tested widths, with a 0.0% width-transfer gap.No divergence occurred across 90 runs.Source: §D.4 · p. 28, Table 27 · p. 28

D.5 T:M Ratio Transfer

  • Naïve µP achieves a 0.0% transfer gap across all six tested uniform and heterogeneous Transformer-to-SSM ratio patterns.The sweep covers three widths and 252 configurations.Source: §D.5 · p. 28, Figure 7 · p. 29
  • For SP, the expand=1 layout has a 1144% transfer gap, while naïve µP remains at η=1.0.SP transfers on five of six tested ratios because proxy and target optima coincide on those particular layouts.Source: §D.5 · p. 29, Figure 7 · p. 29
  • het_taper achieves the best loss (0.114), beating uniform expand=2 (0.130) by approximately 12%.The taper shifts from Transformer-dominant early layers to SSM-heavy late layers.Source: §D.5 · p. 29

D.6 Serial T:M Ratio Sweep

Naïve µP transfers across the tested serial Transformer-to-SSM ratios, while alternative parameterizations show width-dependent drift, especially on SSM-heavy layouts. The section also tests transfer boundaries, clipping robustness, and activation-coordinate behavior.

D.6 Serial T:M Ratio Sweep

  • Naïve µP keeps η*=1.0 across the full serial ratio range, whereas SP and the faithful block-wise prescription drift on SSM-bearing layouts.The full prescription shifts on every SSM-bearing layout, with η=3.0 diverging at width 1024.Source: Table 29 · p. 29, §D.6 · pp. 29–30
  • Approximated block-wise µP incurs residual loss penalties of 4.2% on M1T3 and 24–36% on heavily SSM-dominated layouts when transferring η=1.0 to width 1024.TT and MTMT remain at η*=1.0 throughout the tested widths.Source: §D.6 · p. 30, Table 30 · p. 30

D.7 Gradient Clipping Threshold

  • Changing grad_clip from off to 5.0 leaves η=1.0 optimal and final loss within 0.020 nats across clip values at both tested widths.Disabling clipping does not destabilize training in the tested configurations.Source: §D.7 · p. 30

D.8 Where Naïve µP Transfer Does Not Extend

  • Standard parameterization shifts its optimal LR 3–10× with width; on Jamba, reusing the proxy LR yields catastrophic loss at width 1024.The Jamba optimum moves from η=0.01 at width 256 to η=0.003 at width 1024.Source: §D.8 · p. 31
  • Approximated block-wise µP drifts on SSM-bearing serial layouts, with 4–36% residual loss penalties on SSM-dominated layouts at width 1024.It transfers with a 0% gap on balanced MTMT.Source: §D.8 · p. 31
  • Faithful block-wise µP under AdamW produces 5.6–7.7% higher loss than naïve µP in the tested hybrids and drifts on every SSM-bearing layout.Only pure-Transformer TT transfers cleanly in the 2k-step sweep.Source: §D.8 · p. 31
  • Plain SGD does not show the same transfer behavior, localizing the tested µP transfer property to an optimizer that supplies per-parameter normalization.The stated decomposition assigns the global condition to µP and the local condition to the optimizer.Source: §D.8 · p. 31

E.1 Coordinate Check

  • Hidden activations vary with width across the shown hybrid µP variants, while their logit magnitudes remain comparable; activation magnitude also increases with depth.The coordinate-check metric is the median across blocks of each block’s maximum-to-minimum activation ratio across widths; 1 indicates width invariance.Source: Figure 8 · p. 32, Figure 9 · p. 33, §E.1 · p. 33

E.2 Gradient Measurements

Gradient measurements show that SSM components have disparate scales, and the faithful component-wise scaling rules do not maintain their predicted invariance during training. Optimizer ablations distinguish AdamW’s transfer behavior from SGD’s in the tested hybrids.

E.2 Gradient Measurements

  • Adam’s per-parameter 1/√v normalization compensates for widely differing gradient scales across components.The measurements compare per-component gradient RMS across widths.Source: Figure 11 · p. 34, §E.2 · p. 34
  • The measurements support a two-condition account: naïve µP supplies global update-to-weight invariance, while Adam normalization restores per-component balance despite gradient imbalance.The aggregate gradient norm at initialization is much lower under µP than SP at width 2048.Source: §E.2 · pp. 34–35
  • At initialization, the faithful prescription yields near-zero width dependence for B (α=+0.006) and C (α=+0.029), but Alog grows with Nu (α=+1.00).Under naïve µP, Alog already has α=+0.86, so the added ηA scaling amplifies rather than cancels the observed dependence.Source: §E.2 · p. 35
  • The predicted B and C invariance does not persist during training: their faithful exponents drift to +1.35 and +0.71, respectively, across steps 100–199.The passage characterizes the initialization-level replication as a one-step phenomenon.Source: §E.2 · pp. 35–36

E.3.1 SGD vs AdamW at 2k Steps; E.3.2 SGD at 20k Steps with Extended LR Range

  • At 2k steps SGD trains inconclusively for transfer comparison, but at 20k it requires at least 100× Adam’s LR and shows reversed width scaling.At the best tested LR, MHMH width 256 has loss 5.121 under SGD versus 3.518 under AdamW; AdamW retains a 0.0% gap.Source: §E.3.1 · p. 37, §E.3.2 · p. 38, Table 33 · p. 38

E.3.3 Approximated Block-wise µP with SGD

  • Approximated block-wise µP does not fix SGD’s reversed width scaling; for MHMH, its width gap is +5.6% versus +1.1% under naïve µP.Both variants have higher loss at width 1024 than at width 256.Source: §E.3.3 · p. 38, Table 35 · p. 38, Figure 13 · p. 39

E.3.4 Adam without Weight Decay

  • Removing weight decay shifts loss by at most 0.4% and leaves η=1.0 optimal across the three tested architectures.This supports that weight decay is not load-bearing for transfer in this 2k-step setting.Source: §E.3.4 · p. 39, Table 36 · p. 40

E.3.5 Momentum Ablation

Momentum removal, epsilon variation, and alternative optimizers probe which optimizer features accompany transfer in the tested settings. The results point to adaptive per-parameter normalization rather than momentum as the transfer-relevant feature.

E.3.5 Momentum Ablation

  • Removing Adam’s momentum raises loss by 2–9%, while three of four architecture-width cells retain η*=1.0.The one borderline shift is to η=0.3 for Adam with β1=0 on MHMH at width 256.Source: §E.3.5 · p. 40, Table 37 · p. 40

E.3.6 Adam Epsilon Sensitivity

  • At ϵ=10^-10, 10^-8, or 10^-6, η=1.0 stays within 0.5% of the best LR; ϵ=10^-4 degrades loss by 0.10–0.30 nats.The passage attributes the degradation at large ϵ to suppressed per-parameter normalization.Source: §E.3.6 · p. 41, Table 38 · p. 41

E.3.7 Sophia and Muon

  • Sophia keeps its optimum across widths, whereas Muon’s optimum moves at width 1024; their comparison does not isolate the normalizer.The tested prescriptions differ as well as their normalizers.Source: §E.3.7 · p. 41, Table 39 · p. 41

E.3.8 Confound Resolution Summary

  • Table 40 summarizes which Adam component contributes a required feature versus an incidental one.The supplied caption does not specify the table’s row or column encoding.Source: Table 40 · p. 42

Related papers

Loading 2610.01172v1…