Subspace Optimization for Large Language Models with Convergence Guarantees

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: He, Yutong, Li, Pengrui, Hu, Yipeng, Chen, Chuyan, Yuan, Kun
Formato: Preprint
Publicado: 2024
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866910984808955904
author He, Yutong
Li, Pengrui
Hu, Yipeng
Chen, Chuyan
Yuan, Kun
author_facet He, Yutong
Li, Pengrui
Hu, Yipeng
Chen, Chuyan
Yuan, Kun
contents Subspace optimization algorithms, such as GaLore (Zhao et al., 2024), have gained attention for pre-training and fine-tuning large language models (LLMs) due to their memory efficiency. However, their convergence guarantees remain unclear, particularly in stochastic settings. In this paper, we reveal that GaLore does not always converge to the optimal solution and provide an explicit counterexample to support this finding. We further explore the conditions under which GaLore achieves convergence, showing that it does so when either (i) a sufficiently large mini-batch size is used or (ii) the gradient noise is isotropic. More significantly, we introduce GoLore (Gradient random Low-rank projection), a novel variant of GaLore that provably converges in typical stochastic settings, even with standard batch sizes. Our convergence analysis extends naturally to other subspace optimization algorithms. Finally, we empirically validate our theoretical results and thoroughly test the proposed mechanisms. Codes are available at https://github.com/pkumelon/Golore.
format Preprint
id arxiv_https___arxiv_org_abs_2410_11289
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Subspace Optimization for Large Language Models with Convergence Guarantees
He, Yutong
Li, Pengrui
Hu, Yipeng
Chen, Chuyan
Yuan, Kun
Machine Learning
Optimization and Control
Subspace optimization algorithms, such as GaLore (Zhao et al., 2024), have gained attention for pre-training and fine-tuning large language models (LLMs) due to their memory efficiency. However, their convergence guarantees remain unclear, particularly in stochastic settings. In this paper, we reveal that GaLore does not always converge to the optimal solution and provide an explicit counterexample to support this finding. We further explore the conditions under which GaLore achieves convergence, showing that it does so when either (i) a sufficiently large mini-batch size is used or (ii) the gradient noise is isotropic. More significantly, we introduce GoLore (Gradient random Low-rank projection), a novel variant of GaLore that provably converges in typical stochastic settings, even with standard batch sizes. Our convergence analysis extends naturally to other subspace optimization algorithms. Finally, we empirically validate our theoretical results and thoroughly test the proposed mechanisms. Codes are available at https://github.com/pkumelon/Golore.
title Subspace Optimization for Large Language Models with Convergence Guarantees
topic Machine Learning
Optimization and Control
url https://arxiv.org/abs/2410.11289