Technische Samenvatting: FOAM – Geblokkeerde Toestand-Vouwing voor Geheugenefficiënte LLM-training

1. Probleemstelling

Grote Taalmodellen (LLMs) hebben opmerkelijke prestaties behaald door middel van een toename van het aantal parameters en trainingsdata, maar hun schaal introduceert aanzienlijke geheugenflessenhalsen tijdens de training. De de facto-optimizer, Adam (en AdamW), vereist het opslaan van schattingen van momenten van de eerste en tweede orde voor elke parameter, wat resulteert in een overhead van optimizer-toestandsgeheugen dat tweemaal de modelgrootte bedraagt. Dit maakt het vooraf trainen en fijnafstellen van grote modellen (bijv. met 7B+ parameters) geheugenbeperkt, waarbij vaak honderden gigabytes VRAM nodig zijn, zelfs bij kleine batchgroottes.

Bestaande geheugenefficiënte oplossingen staan voor afwegingen:

• Parameterbevriezing (bijv. LoRA): Beperkt updates tot subruimten met lage rang, wat potentieel de uitdrukkingskracht en prestaties van het model kan beperken.
• Projecties met lage rang (bijv. GaLore, APOLLO): Comprimeren optimizer-toestanden met behulp van Singuliere Waarde Decompositie (SVD) of willekeurige projecties. Hoewel effectief, introduceert SVD aanzienlijke rekenkundige overhead, en willekeurige projecties kunnen de prestaties verslechteren onder strikte geheugenbeperkingen.
• Blok-deling (bijv. Adam-Mini): Deelt schattingen van momenten van de tweede orde over blokken om opslag te verminderen, maar laat vaak de overhead van momenten van de eerste orde hoog of beperkt de geheugenreductie tot ongeveer 50%.

2. Methodologie: FOAM

De auteurs stellen FOAM (Folded Optimizer with Approximate Moment) voor, een gestructureerd raamwerk voor het vouwen van toestanden dat is ontworpen om optimizer-toestanden te comprimeren zonder de optimalisatie-nauwkeurigheid te offeren. FOAM werkt op het principe van geblokkeerde averaging gecombineerd met residucorrectie.

Kernmechanisme

Gegeven een gradiëntmatrix $G_t \in \mathbb{R}^{m \times n}$ op tijdstap $t$ en een vouwniveau $l$:

1. Geblokkeerde Averaging (Vouwen): De gradiënt wordt opgedeeld in blokken van $2^l$ opeenvolgende elementen. Een vouwoperator $A^{(l)}$ berekent het gemiddelde van elk blok, waardoor de gradiëntdimensie wordt gecomprimeerd van $n$ naar $n/2^l$. De momenten van de eerste en tweede orde ($\tilde{M}_t, \tilde{V}_t$) worden bijgewerkt op basis van deze gecomprimeerde gradiënten.
2. Ontvouwen: Om updates toe te passen op de volledige parameterruimte, kopieert een vouwoperator $E^{(l)}$ de gecomprimeerde invoer terug naar de oorspronkelijke dimensies.
3. Residucorrectie: Het direct kopiëren van updates zorgt ervoor dat alle elementen binnen een blok identieke waarden delen, waarbij uitsluitend op initialisatie wordt vertrouwd voor differentiatie. Om verloren informatie te herstellen, berekent FOAM een residu-term $R_t = G_t - \tilde{G}_t E^{(l)}$, die de fout vertegenwoordigt tussen de oorspronkelijke gradiënt en de gevouwen-ontvouwen versie.
   • Het moment van de eerste orde wordt gereconstrueerd als: $M_t = \tilde{M}_t E^{(l)} + R_t$.
   • Het moment van de tweede orde wordt gereconstrueerd als: $V_t = \tilde{V}_t E^{(l)} + R_t^2$.
   • Cruciaal wordt $R_t$ on-the-fly berekend en direct vrijgegeven; het wordt niet in het geheugen opgeslagen.

Hybride Toepassing

FOAM wordt specifiek toegepast op de MLP- en Attention-modules, terwijl standaard Adam wordt gebruikt voor Embedding-lagen. Een schalingscoëfficiënt $\alpha$ past de leerstapverhouding aan tussen FOAM- en Adam-modules, een veelvoorkomende praktijk bij geheugenefficiënte optimizers.

3. Belangrijkste Bijdragen

• Innovatief Optimizer-ontwerp: FOAM bereikt tot 90% reductie in de overhead van optimizer-toestandsgeheugen en ongeveer 50% reductie in de totale end-to-end geheugenvoetafdruk in vergelijking met Full-Adam. In tegenstelling tot op projectie gebaseerde methoden vereist het geen projectiematrices en ondersteunt het optimalisatie van volledige parameters.
• Theoretische Garantiën: Onder standaard aannames voor niet-convexe optimalisatie (L-gladheid, begrenste variantie, begrenste gradiënten) bewijzen de auteurs dat FOAM dezelfde convergentiesnelheid van $O(\frac{\log T}{\sqrt{T}})$ bereikt als vanilla Adam. De residufoutverhouding $\delta_l$ is begrensd door 1 en heeft geen negatief effect op de primaire convergentiesnelheid.
• Uitgebreide Evaluatie: De methode wordt gevalideerd over modelgroottes (60M tot 7B), architecturen (LLaMA, Qwen, GPT, DeBERTa) en taken (vooraf trainen op C4, fijnafstellen op GLUE en MMLU).

4. Experimentele Resultaten

Het artikel rapporteert de volgende empirische bevindingen:

• Pre-training Prestaties: FOAM presteert consistent beter dan geheugenefficiënte baselines (GaLore, APOLLO, Adam-Mini, Muon) wat betreft validatie-perplexiteit (PPL), terwijl het aanzienlijk lager geheugengebruik behoudt. Bijvoorbeeld, op LLaMA-1.3B reduceert FOAM-3 het totale geheugengebruik met ongeveer 50% in vergelijking met Adam, terwijl het een lagere PPL bereikt.
• Convergentiesnelheid: FOAM toont een snellere convergentie, waarbij vaak minder trainingsiteraties nodig zijn om een doel-PPL te bereiken in vergelijking met andere geheugenefficiënte methoden.
• Schaalbaarheid: FOAM behoudt superieure prestaties op grotere modellen (LLaMA-3B, 7B) en onder trainingsregimes met lange sequenties (tot 2048 tokens).
• Fijnafstelling: Op GLUE- en MMLU-benchmarks bereikt FOAM prestaties die vergelijkbaar zijn met of beter zijn dan Full-Adam en andere baselines, wat zijn nut beyond pre-training aantoont.
• Robuustheid: Ablatiestudies bevestigen dat de residu-term cruciaal is voor de prestaties; het verwijderen ervan leidt tot langzamere convergentie en een hogere finale PPL. FOAM blijft ook robuust wanneer gecombineerd met 8-bit kwantisatie en wanneer geïntegreerd met andere optimizers zoals Adam-Mini en Muon.

5. Betekenis en Claims

Het artikel beweert dat FOAM een praktische, uiterst efficiënte en optimizer-onafhankelijke oplossing biedt voor de geheugenflessenhalsen bij de training van grote taalmodellen. Door de noodzaak van gewichtsbewaring of dure projectieberekeningen te elimineren, stelt FOAM volledige parameter-optimalisatie mogelijk met drastisch verminderde geheugenoverhead. De auteurs positioneren FOAM als een complementair alternatief voor technieken met projecties met lage rang, en bieden een methode die theoretisch overeenkomt met de convergentie van Adam, terwijl het empirisch superieure geheugenefficiëntie en vaak betere trainingsprestaties levert.

Het werk suggereert dat FOAM het trainen van grotere modellen op bestaande hardware kan mogelijk maken of de rekenkosten van het trainen van huidige state-of-the-art modellen kan verlagen, waardoor high-performance LLM-training toegankelijker wordt.