Research map / Architectures and efficiency / Attention and transformers

Learning Rate Transfer for Hybrid Transformer-SSM Architectures

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Sparse efficient transformers for long context · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Shows that naive muP with AdamW gives near-zero learning-rate transfer gap for Transformer-Mamba hybrids up to 2048 width and 1.24B params (and Nemotron-H), and proposes a two-condition decomposition (global update-to-weight invariance plus AdamW's per-parameter normalization) to explain why coordinate checks fail yet transfer holds.

Problem

Learning rate must be retuned per scale for hybrid Transformer-SSM architectures; theoretical muP-SSM scaling rules derived for pure SSMs under true-ZOH with growing state size do not match practical simplified-ZOH Mamba at fixed state size, and it is unclear whether muP transfers to hybrids at prac

Approach

Original muP (naive, without SSM-specific rules) under AdamW already transfers the optimal LR across widths and depths for hybrid architectures, even though SSM ops violate Tensor Programs representability and coordinate checks fail. The authors attribute this to AdamW's per-parameter normalization providing local per-component balance, while muP's init and LR scaling provide global update-to-weig

Key results

• Naive muP with AdamW achieves 0.0% LR transfer gap (optimal eta=1.0 at every width) across 7 hybrid architectures at widths 256-1024 and 8x width scaling to 2048 with eta* invariant.
• Block-wise muP-SSM drifts optimum with width on 6/6 SSM-bearing architectures and yields 5.6-7.7% higher loss than naive muP at 20k steps on WikiText-103.
• On Nemotron-H, naive muP achieves 0.0% gap across depths 16 and 32 and widths 256-1024, while SP shifts optimum 3x and incurs 1.9% penalty.

Surprising findings

• Every parameterization tested fails the standard coordinate check on hybrids, yet naive muP transfers LR with near-zero gap.
• Under SGD, the muP-SSM rule's intended per-component cancellation does not appear at step 1 for A_log, amplifying width dependence instead.

Limitations

• Billion-scale, full training horizon, and multiple seeds are validated individually rather than jointly; Mixture-of-Experts and larger models untested.
• Muon and Mamba-3 not transferred under the same recipe; complex-valued states lie outside the real-valued muP rules.
• Experiments use simplified-ZOH Mamba at fixed state size N_x=16; the true-ZOH growing-N_x regime of muP-SSM was not tested, and the two-condition decomposition is an explanatory diagnostic rather than a proven theorem.

Experimental setup

Datasets: WikiText-2 WikiText-103 FineWeb-Edu

Benchmarks: WikiText-2 WikiText-103 FineWeb-Edu HellaSwag PIQA WinoGrande ARC-E ARC-C LAMBADA

Metrics: learning rate transfer gap training loss validation loss zero-shot accuracy update-to-weight ratio balance (Delta beta)

Baselines: Standard Parameterization Block-wise muP Block-wise muP (approximated) Pythia-410M Mamba-370M Falcon-H1 (related)

Builds on: Maximal Update Parameterization muP-SSM (Vankadara et al.) Mamba AdamW Nemotron-H

Closest papers by meaning

• Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models (NeurIPS 2024 · similarity 0.88)
• Peri-LN: Revisiting Normalization Layer in the Transformer Architecture (ICML 2025 · similarity 0.87)
• CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs (NeurIPS 2026 · similarity 0.86)
• Dense Structural Compression of Transformers via Gauge-Correct Channel Removal (NeurIPS 2026 · similarity 0.86)
• Unraveling the Gradient Descent Dynamics of Transformers (NeurIPS 2024 · similarity 0.86)
• On Feature Learning in Structured State Space Models (NeurIPS 2024 · similarity 0.86)