Memory-Efficient LLM Training with Online Subspace Descent

TL;DR

This paper introduces Online Subspace Descent, a memory-efficient optimizer with convergence guarantees, that improves large language model training by replacing SVD with online PCA, leading to better performance and lower perplexity.

Contribution

The paper provides the first convergence guarantee for arbitrary update rules in low-rank gradient methods and proposes a novel optimizer using online PCA for efficient LLM training.

Findings

Online Subspace Descent outperforms existing low-rank methods in LLaMA pretraining.

Achieves lower perplexity and better downstream task performance.

Narrower gap with full-rank training methods.

Abstract

Recently, a wide range of memory-efficient LLM training algorithms have gained substantial popularity. These methods leverage the low-rank structure of gradients to project optimizer states into a subspace using projection matrix found by singular value decomposition (SVD). However, convergence of these algorithms is highly dependent on the update rules of their projection matrix. In this work, we provide the \emph{first} convergence guarantee for arbitrary update rules of projection matrix. This guarantee is generally applicable to optimizers that can be analyzed with Hamiltonian Descent, including most common ones, such as LION, Adam. Inspired by our theoretical understanding, we propose Online Subspace Descent, a new family of subspace descent optimizer without SVD. Instead of updating the projection matrix with eigenvectors, Online Subspace Descent updates the projection matrix with online PCA. Online Subspace Descent is flexible and introduces only minimum overhead to training. We show that for the task of pretraining LLaMA models ranging from 60M to 7B parameters on the C4 dataset, Online Subspace Descent achieves lower perplexity and better downstream tasks performance than state-of-the-art low-rank training methods across different settings and narrows the gap with full-rank baselines.