🤖 machine learning

FOAM: Blocked State Folding for Memory-Efficient LLM Training

Ursprüngliche Autoren: Ziqing Wen, Jiahuan Wang, Ping Luo, Dongsheng Li, Tao Sun

Veröffentlicht 2026-05-14

📖 5 Min. Lesezeit🧠 Tiefgang

Ursprüngliche Autoren: Ziqing Wen, Jiahuan Wang, Ping Luo, Dongsheng Li, Tao Sun

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie versuchen, einem riesigen, superschlauen Roboter (ein Large Language Model) beizubringen, zu schreiben, zu chatten und zu reasoning. Um dies zu tun, zeigen Sie ihm eine massive Bibliothek voller Bücher (Trainingsdaten). Der Roboter lernt, indem er Fehler macht, seine Arbeit überprüft und seine internen „Gehirneinstellungen" (Parameter) anpasst, um besser zu werden.

Die Standardmethode für diese Anpassung ist wie ein sehr sorgfältiger Buchhalter namens Adam. Adam führt für jede einzelne Gehirneinstellung ein detailliertes Hauptbuch, das nicht nur den aktuellen Fehler, sondern auch die Geschichte vergangener Fehler verfolgt, um zu entscheiden, wie stark die Dinge geändert werden müssen.

Das Problem: Der Buchhalter ist zu schwer
Das Problem ist, dass dieser „Buchhalter" (der Optimierer) unglaublich schwer ist. Für einen Roboter mit Milliarden von Gehirneinstellungen nimmt das Hauptbuch des Buchhalters doppelt so viel Speicherplatz ein wie das Gehirn des Roboters selbst.

Analogie: Stellen Sie sich vor, Sie versuchen, ein Haus zu verlegen. Das Mobiliar (das Modell) ist groß, aber der Umzugswagen (die Optimierer-Zustände) ist noch größer. Wenn Sie nur einen kleinen Lkw haben (begrenzter Computerspeicher), können Sie das Haus gar nicht verlegen, egal wie gut das Haus ist. Dies ist die „Speicherengpass", der Forscher daran hindert, größere, schlauere Roboter zu trainieren.

Die alten Lösungen: Ecken schneiden
Bisherige Versuche, den Umzugswagen zu verkleinern, hatten Mängel:

Teile des Hauses einfrieren: Sie bewegen einige Möbel nicht mehr und stellen nur ein paar kleine Kartons ein. Dies spart Platz, macht das Haus aber weniger flexibel (geringere Leistung).
Unschärfe Fotos machen: Statt das eigentliche Mobiliar zu bewegen, machen Sie niedrig aufgelöste Fotos davon, um Platz zu sparen. Aber das Berechnen dieser Fotos kostet viel Zeit und Energie (Rechen-Overhead).
Hauptbücher teilen: Sie lassen verschiedene Zimmer dasselbe Notizbuch teilen. Dies spart Platz, macht die Anpassungen aber weniger präzise.

Die neue Lösung: FOAM (Die intelligente Faltmethode)
Die Arbeit stellt FOAM (Folded Optimizer with Approximate Moment) vor. Stellen Sie sich FOAM als einen Meister-Packer vor, der eine clevere „Falt"-Technik verwendet, um den riesigen Umzugswagen in einen winzigen Lieferwagen zu packen, ohne wichtige Details zu verlieren.

So funktioniert FOAM, unter Verwendung einfacher Metaphern:

1. Die „Block-Faltung" (Komprimierung)

Statt jede einzelne Gehirneinstellung individuell zu verfolgen, gruppiert FOAM sie in kleine Blöcke (wie ein Gitter aus 8 oder 16 Einstellungen).

Die Metapher: Stellen Sie sich vor, Sie haben eine lange Reihe von 100 Personen, und Sie müssen ihre durchschnittliche Körpergröße merken. Statt 100 separate Zahlen aufzuschreiben, gruppieren Sie sie in 10 Teams von je 10 Personen. Sie schreiben nur eine Zahl auf: die durchschnittliche Körpergröße jedes Teams.
Das Ergebnis: Dies verringert den für das „Hauptbuch" benötigten Speicherplatz enorm (bis zu 90 % weniger).

2. Die „Residual-Korrektur" (Das Sicherheitsnetz)

Wenn Sie nur die Durchschnitte verwenden würden, würden Sie die einzigartigen Details jeder Person verlieren. Vielleicht ist eine Person in einem Team sehr groß und eine andere sehr klein. Der Durchschnitt verbirgt das.

Die Metapher: FOAM ist schlau. Nachdem es den „Team-Durchschnitt" aufgeschrieben hat, berechnet es schnell die Differenz (den Fehler) zwischen den echten Personen und dem Durchschnitt. Es nennt dies den „Residual".
Der Trick: Es behält diese Differenz nicht für immer. Es berechnet sie, nutzt sie, um das Update für diesen spezifischen Moment zu korrigieren, und wirft sie dann weg. Es ist wie ein Koch, der eine Suppe probiert, das Salz anpasst und dann den genauen Geschmack des Löffels vergisst, anstatt eine Aufzeichnung über jeden je probierten Löffel zu führen.
Warum es wichtig ist: Dies stellt sicher, dass der Roboter immer noch die einzigartigen Details jeder einzelnen Gehirneinstellung lernt, obwohl der Speicher komprimiert ist.

3. Das „Entfalten" (Wiederherstellung)

Wenn es Zeit ist, das Gehirn des Roboters tatsächlich zu aktualisieren, nimmt FOAM diese komprimierten „Team-Durchschnitte" und erweitert sie wieder auf die volle Größe, wobei es die „Residual"-Korrekturen oben drauflegt.

Das Ergebnis: Der Roboter erhält eine präzise Aktualisierung, genau wie mit dem schweren Buchhalter, aber der Computer musste nur die winzige gefaltete Version tragen.

Was die Arbeit behauptet (Die Ergebnisse)
Die Autoren testeten FOAM beim Training verschiedener Roboter-Modelle (von kleinen 60-Millionen-Parameter-Modellen bis hin zu großen 7-Milliarden-Parameter-Modellen). Hier ist, was sie herausfanden:

Massive Speichereinsparungen: FOAM reduziert den für den Optimierer benötigten Speicher um bis zu 90 %. Insgesamt verringert es den Speicherbedarf des Trainings um etwa 50 %. Das bedeutet, Sie können größere Modelle auf derselben Hardware trainieren oder dieselben Modelle viel schneller trainieren.
Kein Leistungsverlust: Trotz der starken Komprimierung performen die mit FOAM trainierten Modelle genauso gut (oder manchmal sogar besser) als Modelle, die mit dem Standard-schweren Buchhalter trainiert wurden. Sie lernen schneller und erreichen eine höhere Genauigkeit.
Geschwindigkeit: Da es keine schwere Last tragen muss, ist der Trainingsprozess schneller.
Vielseitigkeit: Es funktioniert bei verschiedenen Arten von Roboter-Architekturen (wie LLaMA, Qwen und anderen) und kann sogar mit anderen speichersparenden Tricks kombiniert werden.

Auf den Punkt gebracht
FOAM ist wie ein magischer Koffer. Es ermöglicht Ihnen, eine massive Menge an Informationen (den Speicher des Optimierers) in einen winzigen Raum zu packen, indem Sie ihn ordentlich falten und eine schnelle „Korrektur-Notiz" hinzufügen, um sicherzustellen, dass nichts verloren geht. Dies ermöglicht es Forschern, intelligentere, größere KI-Modelle zu bauen, ohne super-teure, massive Computer zu benötigen.

Technische Zusammenfassung: FOAM – Blockiertes Zustandsfalten für speichereffizientes LLM-Training

1. Problemstellung

Large Language Models (LLMs) haben durch erhöhte Parameterzahlen und Trainingsdaten bemerkenswerte Leistungen erzielt, doch ihre Skalierung führt während des Trainings zu erheblichen Speicherengpässen. Der De-facto-Optimierer Adam (und AdamW) erfordert die Speicherung von Schätzungen erster und zweiter Ordnung für jeden Parameter, was zu einem Speicher-Overhead des Optimiererzustands führt, der doppelt so groß ist wie die Modellgröße. Dies macht das Pre-Training und Fine-Tuning großer Modelle (z. B. mit 7B+ Parametern) speicherlimitiert und erfordert oft Hunderte von Gigabyte VRAM, selbst bei kleinen Batch-Größen.

Bestehende speichereffiziente Lösungen gehen Kompromisse ein:

Parametereinfrieren (z. B. LoRA): Beschränkt Updates auf Unterräume mit niedrigem Rang, was möglicherweise die Ausdruckskraft und Leistung des Modells einschränkt.
Projektion mit niedrigem Rang (z. B. GaLore, APOLLO): Komprimiert Optimiererzustände mittels Singulärwertzerlegung (SVD) oder zufälliger Projektionen. Obwohl effektiv, führt SVD zu erheblichem Rechenaufwand, und zufällige Projektionen können unter strengen Speicherbeschränkungen die Leistung verschlechtern.
Block-Sharing (z. B. Adam-Mini): Teilt Schätzungen zweiter Ordnung über Blöcke hinweg, um den Speicherbedarf zu reduzieren, lässt jedoch oft den Overhead für Momente erster Ordnung hoch oder begrenzt die Speicherreduktion auf etwa 50 %.

2. Methodik: FOAM

Die Autoren schlagen FOAM (Folded Optimizer with Approximate Moment) vor, ein strukturiertes Rahmenwerk zum Falten von Zuständen, das entwickelt wurde, um Optimiererzustände zu komprimieren, ohne die Optimierungsgenauigkeit zu beeinträchtigen. FOAM basiert auf dem Prinzip des blockierten Mittelwerts in Kombination mit einer Residualkorrektur.

Kernmechanismus

Gegeben eine Gradientenmatrix $G_t \in \mathbb{R}^{m \times n}$ zum Zeitpunkt $t$ und ein Faltlevel $l$ :

Blockierter Mittelwert (Falten): Der Gradient wird in Blöcke aus $2^l$ aufeinanderfolgenden Elementen partitioniert. Ein Faltoperator $A^{(l)}$ berechnet den Mittelwert jedes Blocks und komprimiert die Gradientendimension von $n$ auf $n/2^l$ . Die Momente erster und zweiter Ordnung ( $\tilde{M}_t, \tilde{V}_t$ ) werden basierend auf diesen komprimierten Gradienten aktualisiert.
Entfalten: Um Updates auf den vollständigen Parameterraum anzuwenden, repliziert ein Faltoperator $E^{(l)}$ die komprimierten Einträge zurück auf die ursprünglichen Dimensionen.
Residualkorrektur: Das direkte Replizieren von Updates führt dazu, dass alle Elemente innerhalb eines Blocks identische Werte teilen und sich ausschließlich auf die Initialisierung zur Differenzierung verlassen. Um verlorene Informationen wiederzugewinnen, berechnet FOAM einen Residualterm $R_t = G_t - \tilde{G}_t E^{(l)}$ $R_{t} = G_{t} - \tilde{G}_{t} E^{(l)}$ , der den Fehler zwischen dem ursprünglichen Gradienten und der gefalteten-entfalteten Version darstellt.
- Das Moment erster Ordnung wird rekonstruiert als: $M_t = \tilde{M}_t E^{(l)} + R_t$ .
- Das Moment zweiter Ordnung wird rekonstruiert als: $V_t = \tilde{V}_t E^{(l)} + R_t^2$ .
- Entscheidend ist, dass $R_t$ on-the-fly berechnet und sofort freigegeben wird; er wird nicht im Speicher gespeichert.

Hybride Anwendung

FOAM wird spezifisch auf die MLP- und Attention-Module angewendet, während für Embedding-Schichten standardmäßiges Adam verwendet wird. Ein Skalierungskoeffizient $\alpha$ passt das Lernratenverhältnis zwischen FOAM- und Adam-Modulen an, eine gängige Praxis bei speichereffizienten Optimierern.

3. Hauptbeiträge

Innovatives Optimierer-Design: FOAM erreicht eine Reduktion des Speicher-Overheads für Optimiererzustände von bis zu 90 % und eine Reduktion des gesamten End-zu-End-Speicherbedarfs von etwa 50 % im Vergleich zu Full-Adam. Im Gegensatz zu projektionsbasierten Methoden benötigt es keine Projektionsmatrizen und unterstützt die Optimierung aller Parameter.
Theoretische Garantien: Unter Standardannahmen nicht-konvexer Optimierung (L-Glattheit, beschränkte Varianz, beschränkte Gradienten) beweisen die Autoren, dass FOAM die gleiche Konvergenzrate von $O(\frac{\log T}{\sqrt{T}})$ wie Vanilla-Adam erreicht. Das Residualfehlerverhältnis $\delta_l$ ist durch 1 beschränkt und beeinträchtigt die primäre Konvergenzrate nicht negativ.
Umfassende Evaluierung: Die Methode wird über Modellgrößen (60M bis 7B), Architekturen (LLaMA, Qwen, GPT, DeBERTa) und Aufgaben (Pre-Training auf C4, Fine-Tuning auf GLUE und MMLU) hinweg validiert.

4. Experimentelle Ergebnisse

Die Studie berichtet folgende empirische Befunde:

Pre-Training-Leistung: FOAM übertrifft speichereffiziente Baselines (GaLore, APOLLO, Adam-Mini, Muon) konsistent in Bezug auf die Validierungsperplexität (PPL), während gleichzeitig ein deutlich geringerer Speicherbedarf aufrechterhalten wird. Beispielsweise reduziert FOAM-3 auf LLaMA-1.3B den gesamten Speicherverbrauch im Vergleich zu Adam um ~50 %, bei gleichzeitig niedrigerer PPL.
Konvergenzgeschwindigkeit: FOAM zeigt eine schnellere Konvergenz und erfordert oft weniger Trainingiterationen, um eine Ziel-PPL im Vergleich zu anderen speichereffizienten Methoden zu erreichen.
Skalierbarkeit: FOAM behält auf größeren Modellen (LLaMA-3B, 7B) und unter Trainingsregimen mit langen Sequenzen (bis zu 2048 Tokens) eine überlegene Leistung bei.
Fine-Tuning: Auf GLUE- und MMLU-Benchmarks erreicht FOAM eine Leistung, die mit Full-Adam und anderen Baselines vergleichbar ist oder diese übertrifft, was seine Nützlichkeit über das Pre-Training hinaus demonstriert.
Robustheit: Ablationsstudien bestätigen, dass der Residualterm für die Leistung entscheidend ist; seine Entfernung führt zu langsamerer Konvergenz und höherer finaler PPL. FOAM bleibt zudem robust, wenn es mit 8-Bit-Quantisierung kombiniert wird und wenn es mit anderen Optimierern wie Adam-Mini und Muon integriert wird.

5. Bedeutung und Behauptungen

Die Studie behauptet, dass FOAM eine praktische, hocheffiziente und optimiereragnostische Lösung für die Speicherengpässe beim Training großer Sprachmodelle bietet. Durch die Eliminierung der Notwendigkeit zum Einfrieren von Gewichten oder teurer Projektionsberechnungen ermöglicht FOAM die Optimierung aller Parameter bei drastisch reduziertem Speicher-Overhead. Die Autoren positionieren FOAM als komplementäre Alternative zu Projektionstechniken mit niedrigem Rang und bieten eine Methode, die theoretisch die Konvergenz von Adam erreicht und empirisch überlegene Speichereffizienz sowie oft bessere Trainingsleistung liefert.

Die Arbeit legt nahe, dass FOAM das Training größerer Modelle auf bestehender Hardware ermöglichen oder die Kosten für das Training aktueller State-of-the-Art-Modelle senken kann, wodurch das Training leistungsstarker LLMs zugänglicher wird.