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
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:17249-17263, 2026.

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-chen26eq, title = {Escaping the Subspace Trap: The Role of Optimizer Geometry in Model Width Expansion}, author = {Chen, Jiabei and Wang, Haoyu and Yu, Yang and Xu, Yao and Wang, Liangdong and Liu, Guang and He, Shizhu and Zhao, Jun and Liu, Kang}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {17249--17263}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/chen26eq/chen26eq.pdf}, url = {https://proceedings.mlr.press/v306/chen26eq.html}, 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.} }
Endnote
%0 Conference Paper %T Escaping the Subspace Trap: The Role of Optimizer Geometry in Model Width Expansion %A Jiabei Chen %A Haoyu Wang %A Yang Yu %A Yao Xu %A Liangdong Wang %A Guang Liu %A Shizhu He %A Jun Zhao %A Kang Liu %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-chen26eq %I PMLR %P 17249--17263 %U https://proceedings.mlr.press/v306/chen26eq.html %V 306 %X 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.
APA
Chen, J., Wang, H., Yu, Y., Xu, Y., Wang, L., Liu, G., He, S., Zhao, J. & Liu, K.. (2026). Escaping the Subspace Trap: The Role of Optimizer Geometry in Model Width Expansion. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:17249-17263 Available from https://proceedings.mlr.press/v306/chen26eq.html.

Related Material