Addestramento Efficiente in Memoria di LLM con Discesa Online su Sottospazi

Memory-Efficient LLM Training with Online Subspace Descent

August 23, 2024

Autori: Kaizhao Liang, Bo Liu, Lizhang Chen, Qiang Liu

cs.AI

Abstract

Recentemente, una vasta gamma di algoritmi di addestramento efficienti in termini di memoria per LLM ha guadagnato una notevole popolarità. Questi metodi sfruttano la struttura a basso rango dei gradienti per proiettare gli stati dell'ottimizzatore in un sottospazio utilizzando una matrice di proiezione ottenuta tramite decomposizione ai valori singolari (SVD). Tuttavia, la convergenza di questi algoritmi dipende fortemente dalle regole di aggiornamento della loro matrice di proiezione. In questo lavoro, forniamo la prima garanzia di convergenza per regole di aggiornamento arbitrarie della matrice di proiezione. Questa garanzia è generalmente applicabile agli ottimizzatori che possono essere analizzati con Hamiltonian Descent, inclusi i più comuni, come LION e Adam. Ispirati dalla nostra comprensione teorica, proponiamo Online Subspace Descent, una nuova famiglia di ottimizzatori di discesa in sottospazio senza SVD. Invece di aggiornare la matrice di proiezione con autovettori, Online Subspace Descent aggiorna la matrice di proiezione con PCA online. Online Subspace Descent è flessibile e introduce solo un sovraccarico minimo durante l'addestramento. Dimostriamo che, per il compito di pre-addestramento di modelli LLaMA che vanno da 60M a 7B di parametri sul dataset C4, Online Subspace Descent raggiunge una perplexity inferiore e una migliore performance nei task downstream rispetto ai metodi di addestramento a basso rango all'avanguardia in diverse configurazioni, riducendo il divario con i baseline a rango pieno.

English

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 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.