Escaping the Subspace Trap: The Role of Optimizer Geometry in Model Width Expansion
Jiabei Chen, Haoyu Wang, Yang Yu, Yao Xu, Liangdong Wang, Guang Liu, Shizhu He, Jun Zhao, Kang Liu
Abstract
Pre-training large language models from scratch is prohibitively expensive as model scales increase. A practical alternative is Model Width Expansion (MWE), which grows a larger model from a well-pretrained ''seed'' model to inherit existing capabilities at initialization. However, we identify a phenomenon termed the Subspace Trap : during continual pre-training, parameter updates largely stagnate within a low-dimensional subspace aligned with the initialization, limiting the effective capacity of the expanded model. Our theoretical analysis investigates this issue by attributing it to the function-preserving properties of width expansion. In particular, element-wise adaptive optimizers remain confined to the trap, whereas optimizers that yield an isotropic geometry of parameter updates can escape. To demonstrate the impact of the subspace trap on model performance, we conduct empirical experiments across different model sizes and model families, which show that escaping the trap is principally effective in improving training efficiency and overall model performance. Detailed mechanistic analyses further confirm that escaping the trap indeed activates the new dimensions to encode general knowledge. Our code is available at https://github.com/A-PolarBear/Model-Width-Expansion.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 039e3c49-1f3c-4348-a4da-573cc05079f8Builds on13
- Language Models are Few-Shot LearnersTom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah et al.NeurIPS 2020 · 64,255 citations
- Measuring Massive Multitask Language UnderstandingDan Hendrycks, Collin Burns, Steven Basart, Andy Zou et al.ICLR 2021 · 7,905 citations
- WinoGrande: An Adversarial Winograd Schema Challenge at ScaleKeisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, Yejin ChoiAAAI 2020 · 3,037 citations
- PIQA: Reasoning about Physical Commonsense in Natural LanguageYonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao et al.AAAI 2020 · 2,916 citations
- Staged Training for Transformer Language ModelsSheng Shen, Pete Walsh, Kurt Keutzer, Jesse Dodge et al.ICML 2022 · 52 citations
Related papers
- Beyond Sunk Costs: Boosting LLM Pre-training Efficiency via Orthogonal Growth of Mixture-of-ExpertsRuizhe Wang, Yucheng Ding, Xiao Liu, Yaoxiang Wang et al.ICML 2026
- ADEPT: Continual Pretraining via Adaptive Expansion and Dynamic Decoupled TuningJinyang Zhang, Yue Fang, Hongxin Ding, Weibin Liao et al.ICLR 2026 · 5 citations
- Masked Structural Growth for 2x Faster Language Model Pre-trainingYiqun Yao, Zheng Zhang, Jing Li, Yequan WangICLR 2024 · 30 citations
- LDAdam: Adaptive Optimization from Low-Dimensional Gradient StatisticsThomas Robert, Mher Safaryan, Ionut-Vlad Modoranu, Dan AlistarhICLR 2025
- LESA: Learnable LLM Layer Scaling-UpYifei Yang, Zouying Cao, Xinbei Ma, Yao Yao et al.ACL 2025 · 6 citations
