Learning Rate Transfer for Hybrid Transformer-SSM Architectures
Training hybrid Transformer-SSM models requires expensive per-scale learning rate (LR) tuning. This paper empirically shows that "naïve" Maximal Update Parameterization (μP) — as originally prescribed for pure Transformers, with no SSM-specific modifications — combined with AdamW achieves near-zero LR transfer gap across hybrid architectures from width 256 to 2048 (8×) and depth 4 to 32 (8×), up to billion parameters and on a production Nemotron-H model. The authors propose a two-condition decomposition (global update-to-weight invariance from μP, local per-component balance from AdamW's g/√v normalization) to explain why coordinate checks fail while transfer still succeeds.
This is a solid empirical paper that fills a practical gap: it establishes that naïve μP transfers cleanly to hybrid Transformer-SSM architectures under AdamW in the practical simplified-ZOH Mamba regime, contradicting the intuition that SSM-specific extensions are necessary. The extensive grid sweeps across widths, depths, architectures, batch sizes, and sequence lengths are compelling, and generalization to Nemotron-H is a meaningful stress test. The two-condition decomposition is useful as an organizing framework, but the paper carefully acknowledges it is not a formal proof and does not cover their setting. The main limitations are that most conclusions rest on short-horizon (2k step / 16M–164M token) sweeps rather than full-scale training, and the comparison to μP-SSM is made across several regime mismatches (SGD vs AdamW, true-ZOH vs simplified-ZOH, growing vs fixed N_x=16) that confound diagnosing whether μP-SSM is truly unsuited or simply misapplied.
The core empirical finding appears robust: naïve μP achieves η*=1.0 across all tested widths and architectures in Table 1, including MHMH, HHHH, and Nemotron-H, while SP and block-wise μP-SSM drift. The broad sensitivity analysis (batch size, sequence length, T:M ratio, depth, multi-seed) and the partial validation at billion parameters with downstream evaluation on FineWeb-Edu strengthen the practical relevance. The critical observation that "every parameterization we test fails the coordinate check on hybrids, yet naïve μP transfers cleanly" (Section 5.1) is theoretically interesting and well-supported by the data in Table 1 and Figure 8 (Appendix E.1). The explicit regime-gap diagnosis of μP-SSM (simplified-ZOH, fixed N_x=16, AdamW vs SGD) is honest and carefully documented in Appendix A.4.
First, the headline transfer conclusions are drawn overwhelmingly from 2k-step and 20k-step LR sweeps (16M–164M tokens) on small corpora. The limitation section acknowledges that billion scale, full horizon, and multiple seeds are validated individually rather than jointly. A single 1.24B model trained to 24.9B tokens at the transferred LR is promising but thin evidence for production-scale claims. Second, the comparison to Chennuru Vankadara et al.'s μP-SSM is confounded by simultaneous regime mismatches: their derivation assumed pure SSMs, plain SGD, true-ZOH discretization, and growing N_x→∞, whereas the tested regime uses hybrid architectures, AdamW, simplified-ZOH Mamba, and fixed N_x=16. The paper diagnoses why μP-SSM fails but cannot cleanly separate regime mismatch from theoretical inadequacy. Third, the optimizer ablation supporting Condition 2 is incomplete: the SGD grid does not establish transfer because the optimum lies above η=100 (Appendix E.3.2), and the 2k-step weight-decay ablation explicitly "does not cover the longer training horizons examined by Kosson et al. (2026)." This weakens the claim that "g/√v is essential" relative to weight decay.
The evidence mostly supports the main claim. The width-transfer grid across 7 architectures (Table 1), the joint depth-width transfer (Appendix C.1), the FineWeb-Edu web-scale validation (Appendix C.2), and the Nemotron-H generalization (Table 3) together form strong evidence that naïve μP transfers under AdamW. The comparison to related work is mostly fair: the paper correctly distinguishes itself from Falcon-H1's 35 per-component multipliers and acknowledges concurrent work. However, the downstream evaluation (Table 2) has a notable caveat: retrained Mamba-370M and Pythia-410M on 5B tokens of FineWeb-Edu achieve much lower average accuracy (41.9% and 42.4%) than published versions trained on The Pile (300B tokens; 51.7% and 48.2%), confirming that the 5B-token budget is far from converged, though the within-setup comparison is still valid. The paper could more explicitly quantify how much performance is left on the table by not tuning LR at scale beyond naïve μP.
The paper provides extensive appendices with per-experiment details, hyperparameters, and an experiment inventory (Table 7). Parameterization formulas are given in Appendix A.3, and the wall-time/cost analysis in Appendix A.1 is helpful. However, the code is not released, and several architectural and training hyperparameters are spread across appendices rather than consolidated (e.g., exact Mamba block configuration, full optimizer state). The key hyperparameters — AdamW (β1=0.9, β2=0.95), cosine warmup, gradient clipping at 1.0, fixed seed — are stated, but reproducibility would be significantly improved by code release. The use of a fixed seed (single seed for most experiments, three seeds for some) means variance estimates are limited. Multiple widths and depths are tested, but the final billion-parameter 24.9B-token run uses a single seed, making it hard to assess reproducibility at that scale.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.