Saved in:
| Main Authors: | , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2502.05790 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866909957259001856 |
|---|---|
| author | Zhang, Haochen Yin, Junze Wang, Guanchu Liu, Zirui Yang, Lin F. Zhang, Tianyi Shrivastava, Anshumali Braverman, Vladimir |
| author_facet | Zhang, Haochen Yin, Junze Wang, Guanchu Liu, Zirui Yang, Lin F. Zhang, Tianyi Shrivastava, Anshumali Braverman, Vladimir |
| contents | Low-rank optimization has emerged as a promising approach to enabling memory-efficient training of large language models (LLMs). Existing low-rank optimization methods typically project gradients onto a low-rank subspace, reducing the memory cost of storing optimizer states. A key challenge in these methods is selecting suitable subspaces to ensure an effective optimization trajectory. Most existing approaches select the dominant subspace to preserve gradient information, as this intuitively provides the best approximation. However, we find that in practice, the dominant subspace stops changing during pretraining, thereby constraining weight updates to similar subspaces. In this paper, we propose importance sampling for low-rank optimization in LLM pretraining with a provable convergence guarantee, which the dominant subspace approach does not have. Empirically, we demonstrate that our method significantly outperforms previous methods in LLM pretraining tasks. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2502_05790 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Breaking the Frozen Subspace: Importance Sampling for Low-Rank Optimization in LLM Pretraining Zhang, Haochen Yin, Junze Wang, Guanchu Liu, Zirui Yang, Lin F. Zhang, Tianyi Shrivastava, Anshumali Braverman, Vladimir Machine Learning Low-rank optimization has emerged as a promising approach to enabling memory-efficient training of large language models (LLMs). Existing low-rank optimization methods typically project gradients onto a low-rank subspace, reducing the memory cost of storing optimizer states. A key challenge in these methods is selecting suitable subspaces to ensure an effective optimization trajectory. Most existing approaches select the dominant subspace to preserve gradient information, as this intuitively provides the best approximation. However, we find that in practice, the dominant subspace stops changing during pretraining, thereby constraining weight updates to similar subspaces. In this paper, we propose importance sampling for low-rank optimization in LLM pretraining with a provable convergence guarantee, which the dominant subspace approach does not have. Empirically, we demonstrate that our method significantly outperforms previous methods in LLM pretraining tasks. |
| title | Breaking the Frozen Subspace: Importance Sampling for Low-Rank Optimization in LLM Pretraining |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2502.05790 |