Subspace Optimization for Large Language Models with Convergence Guarantees (2410.11289v2)

Abstract: Subspace optimization algorithms, such as GaLore (Zhao et al., 2024), have gained attention for pre-training and fine-tuning 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.

• The paper introduces GoLore, a new subspace optimization method that addresses convergence limitations in stochastic LLM training environments.
• It employs a dynamic mix of SVD-based and random low-rank projections to mitigate gradient noise while optimizing memory usage.
• Empirical and theoretical analyses confirm GoLore’s effectiveness with an O(1/√T) convergence rate, enabling scalable LLM training.

Subspace Optimization for LLMs with Convergence Guarantees

The paper "Subspace Optimization for LLMs with Convergence Guarantees" addresses the application of subspace optimization algorithms to LLMs with a focus on improving memory efficiency during pre-training or fine-tuning. While existing methods like GaLore show promise in reducing memory usage, their convergence properties are overshadowed by inherent stochastic uncertainties. The authors not only identify scenarios where these methods fall short but also propose alternatives that ensure convergence.

Subspace optimization has been hailed for its ability to reduce memory consumption, a critical factor given the ever-increasing scale of LLMs. Algorithms like GaLore are front and center due to their innovative approach of projecting full-parameter gradients into low-memory subspaces. However, this study highlights a pivotal flaw: GaLore does not always converge to optimal solutions in stochastic environments, which are typical in LLM training. Through a constructed counterexample, the authors demonstrate how GaLore fails to converge under standard non-convex optimization conditions like lower boundedness and L-smoothness.

To tackle the convergence issue, the authors put forth GoLore (Gradient random Low-rank projection), a novel variant designed to perform more reliably in stochastic settings. By integrating random projections, GoLore effectively mitigates the influence of gradient noise—an insightful advancement over GaLore's deterministic subspace setup. The paper provides rigorous convergence guarantees for GoLore, establishing its ability to reach stationary solutions at a $\mathcal{O}(1/\sqrt{T})$ rate—a hallmark for non-convex stochastic problems. This enhancement in convergence without increasing batch sizes is essential since larger batch sizes are typically infeasible due to memory constraints.

Through numerical experiments and comprehensive theoretical analysis, the paper validates that GoLore, by employing a dynamic mix of GaLore and GoLore strategies at different training stages, notably increases empirical performance in LLM tasks. This hybrid approach allows initial reliance on SVD-derived subspaces for capturing meaningful gradients, subsequently switching to random projections as the model approaches a local optimum.

From a practical standpoint, the implications of this research are significant. It enables the further scaling of LLMs by optimizing memory usage without compromising training fidelity. This means that more expansive and intricate models can be trained or fine-tuned on existing hardware capabilities, offering computational accessibility without sacrificing speed or accuracy.

Theoretically, this work proposes a new direction for convergence analysis in memory-efficient algorithms. It challenges the field to reconsider assumptions in stochastic optimization, especially in the context of modern deep learning paradigms where variance reduction through large-batch methods is not viable. By showcasing the applicability of random projections, the authors lay the groundwork for future exploration of similar techniques across various domains of machine learning.

In conclusion, the paper makes a substantial contribution by addressing a previously less-understood aspect of memory-efficient LLM training—convergence. The insights and solutions presented, notably GoLore, provide a critical step forward, ensuring that as models and datasets grow, the training methodologies remain robust and resource-conscious. Future research may focus on refining these approaches, exploring reduced computational costs, and expanding the principles to other subspace learning tasks.