Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation

Die Arbeit stellt LoRA-Pre vor, einen neuartigen Low-Rank-Optimierer, der durch die Umformulierung von Momentum als Online-Linearregression den Speicherverbrauch erheblich reduziert und dabei beim Vor- und Feinabstimmen von Llama-Modellen die Leistung bestehender Methoden übertrifft.

Das Wesentliche

Das große Problem: Der überladene Rucksack

Stell dir vor, du trainierst einen riesigen Roboter (eine sogenannte "Large Language Model" oder KI), damit er Sprache versteht und schreibt. Damit dieser Roboter lernt, braucht er einen sehr klugen Lehrer, den wir Optimierer nennen (bekannte Namen sind Adam oder Muon).

Der Lehrer hat eine besondere Aufgabe: Er muss sich merken, welche Wege der Roboter in der Vergangenheit erfolgreich gegangen sind und welche nicht. Diese Erinnerungen nennt man Impuls (Momentum).

Das Problem: Bei modernen, riesigen KIs ist dieser "Erinnerungs-Rucksack" des Lehrers so schwer, dass er fast so viel Platz einnimmt wie der Roboter selbst!

• Der Roboter (die KI) ist schon riesig.
• Der Rucksack (die Optimierer-Zustände) ist doppelt so schwer wie der Roboter.
• Ergebnis: Um große KIs zu trainieren, braucht man extrem teure und große Computer. Das ist teuer und langsam.

Die geniale Idee: Vom Archiv zum Regler

Die Autoren des Papiers haben etwas Entdeckendes gemacht. Sie haben sich gefragt: "Was macht dieser Lehrer eigentlich genau, wenn er sich die Impulse merkt?"

Statt eines riesigen Archivs, in dem jede einzelne Bewegung gespeichert wird, haben sie erkannt: Der Lehrer führt eigentlich eine Art "Online-Rechnung" durch.

Stell dir vor, der Lehrer versucht nicht, eine Bibliothek zu füllen, sondern er versucht, einen einfachen Regler (einen Linearen Regler) zu bauen. Dieser Regler soll die vergangenen Bewegungen so gut wie möglich vorhersagen.

• Die alte Methode: "Ich speichere jeden einzelnen Schritt in einem riesigen Buch." (Sehr viel Platz nötig).
• Die neue Erkenntnis: "Ich baue einen kleinen Regler, der die Tendenz der Schritte berechnet." (Wenig Platz nötig).

Die Lösung: LoRA-Pre (Der Low-Rank-Trick)

Basierend auf dieser Erkenntnis haben die Forscher LoRA-Pre entwickelt. Das ist wie ein Zaubertrick, um den Rucksack des Lehrers zu verkleinern, ohne dass er seine Intelligenz verliert.

Die Metapher: Das Orchester
Stell dir den Impuls des Lehrers als ein riesiges Orchester vor, das 10.000 Instrumente spielt.

• Normalerweise: Man muss für jedes Instrument einen eigenen Musiker und ein eigenes Notenblatt bezahlen. Das ist teuer (viel Speicher).
• Mit LoRA-Pre: Die Forscher sagen: "Warte mal! Die Musik, die dieses Orchester spielt, ist eigentlich gar nicht so komplex. Sie lässt sich durch nur zwei einfache Melodien beschreiben, die zusammen gespielt werden."

Sie zerlegen das riesige Orchester in zwei kleine Gruppen (zwei Matrizen), die zusammen fast dasselbe machen wie das große Orchester, aber nur einen Bruchteil des Platzes brauchen.

• Statt 10.000 Musiker braucht man nur noch 200.
• Das Ergebnis (die Musik/der Lernfortschritt) klingt fast genauso gut, aber es ist viel billiger und schneller.

Warum ist das so besonders?

1. Es funktioniert von Anfang an: Bisherige Tricks (wie LoRA) waren gut, um fertige KIs nur noch ein bisschen zu verbessern (Feinabstimmung). Aber sie scheiterten oft, wenn man eine KI von Null an lernte (Pre-Training). LoRA-Pre funktioniert in beiden Fällen perfekt.
2. Es ist extrem effizient: Die Forscher haben gezeigt, dass sie mit nur 1/8 der Größe (des "Rangs") der alten Methoden genauso gute Ergebnisse erzielen. Das ist, als würde man ein Auto bauen, das mit einem Achtel des Treibstoffs genauso schnell fährt.
3. Es passt sich ständig an: Andere Methoden aktualisieren ihre "kleinen Gruppen" nur alle paar Tage. LoRA-Pre passt die Gruppen bei jedem einzelnen Schritt an. Es ist wie ein Navigator, der den Weg sofort korrigiert, statt erst morgen zu schauen, ob man noch auf der richtigen Straße ist.

Das Ergebnis

Die Forscher haben ihre Methode an verschiedenen KI-Modellen getestet (von kleinen bis zu sehr großen).

• Ergebnis: Die KIs lernten schneller und wurden schlauer als mit den alten Methoden, obwohl sie viel weniger Speicherplatz benötigten.
• Vergleich: Wenn man LoRA-Pre mit dem Standard-Adam-Optimierer vergleicht, gewinnt LoRA-Pre fast immer, besonders bei den großen Modellen.

Zusammenfassung in einem Satz

Die Forscher haben entdeckt, dass das Gedächtnis eines KI-Lehrers eigentlich nur eine einfache Rechnung ist, und haben einen Weg gefunden, dieses Gedächtnis auf ein winziges Stück Papier zu komprimieren, ohne dass die KI dabei dumm wird.

Das macht das Trainieren von super-intelligenten KIs für viel mehr Menschen und Unternehmen erschwinglich!

Technische Zusammenfassung

Titel: Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation

Veröffentlicht bei: ICLR 2026
Autoren: Zhengbo Wang, Jian Liang, Ran He, Zilei Wang, Tieniu Tan

---

1. Problemstellung

Moderne Optimierer wie Adam und Muon sind fundamental für das Training großer Sprachmodelle (LLMs). Diese Optimierer speichern jedoch nicht nur die Modellgewichte, sondern auch Schätzungen der ersten und zweiten Momente der Gradienten (Momentum und Varianz).

• Herausforderung: Die Speicherung dieser Zustände (States) vervielfacht den Speicherbedarf (oft um das Dreifache) und stellt einen erheblichen Engpass für die Skalierbarkeit und Recheneffizienz dar, insbesondere beim Pre-Training von Modellen mit Milliarden von Parametern.
• Bestehende Lösungen: Ansätze wie GaLore oder LoRA versuchen, den Speicherbedarf durch Projektion in niedrigrangige Unterräume zu reduzieren. Diese Methoden basieren oft auf periodischen Aktualisierungen der Projektionsmatrizen (z. B. via SVD).
• Nachteil bestehender Methoden: Durch die verzögerte Anpassung der Unterräume (Lag) kommt es zu einer Akkumulation von Fehlern in der Zustandsschätzung, was zu suboptimaler Leistung führt. Es fehlt an einer dynamischen Methode, die sich sofort an verändernde Gradientenräume anpasst.

2. Methodik: LoRA-Pre

Die Autoren schlagen LoRA-Pre vor, einen neuartigen, speichereffizienten Optimierer für das Pre-Training, der auf einer fundamentalen mathematischen Neuinterpretation des Momentum-Updates basiert.

A. Theoretische Grundlage: Momentum als Online-Linearregressor

Der Kern der Arbeit ist die Erkenntnis, dass das exponentielle gleitende Mittel (EMA) für Momentum mathematisch äquivalent zum Training eines Online-Linearregressors mittels Gradientenfluss ist.

• Das Standard-EMA-Update $m_{t+1} = \beta \cdot m_t + (1-\beta) \cdot g_t$ lässt sich umformen als ein Gradientenabstiegsschritt auf dem Parameter $m$ mit der Verlustfunktion $L(m; g) = \frac{1}{2} \|m - g\|_F^2$.
• Dies bedeutet: Momentum ist im Wesentlichen ein lineares Modell, das versucht, die Historie der Gradienten zu approximieren.

B. Low-Rank Faktorisierung

Anstatt den vollen Momentum-Matrix $m \in \mathbb{R}^{p \times q}$ zu speichern, wird dieser in zwei niedrigrangige Matrizen zerlegt:
$$m = m_B \cdot m_A$$
wobei $m_B \in \mathbb{R}^{p \times r}$ und $m_A \in \mathbb{R}^{r \times q}$ mit $r \ll \min(p, q)$.

• Speichereffizienz: Die Komplexität sinkt von $O(p \times q)$ auf $O((p+q) \times r)$.
• Update-Regeln: Anstatt Backpropagation durchzuführen, leiten die Autoren geschlossene Update-Regeln für $m_B$ und $m_A$ her, indem sie die Newton-Methode auf die Faktorisierungs-Loss-Funktion anwenden. Diese Regeln behalten die EMA-Struktur bei, arbeiten jedoch im komprimierten Raum.
• Anpassung an Adam und Muon: Die Methode wird für beide Optimierer adaptiert. Für die zweite Moment-Schätzung (Varianz) wird eine spezielle Re-Parametrisierung ($v = (v_B v_A)^{\circ 2}$) verwendet, um die notwendige Element-weise Positivität zu gewährleisten.

C. Dynamische Subraum-Anpassung

Im Gegensatz zu Methoden mit periodischen Updates (wie GaLore) aktualisiert LoRA-Pre die niedrigrangigen Faktoren in jedem Trainingsschritt online. Dies eliminiert die Fehlerakkumulation und ermöglicht eine kontinuierliche Anpassung an den sich wandelnden Gradientenraum.

3. Wichtige Beiträge

1. Theoretische Äquivalenz: Etablierung der mathematischen Äquivalenz zwischen EMA-Momentum-Updates und dem Training eines Online-Linearregressors.
2. LoRA-Pre Algorithmus: Entwicklung eines neuen Optimierers, der Momentum durch Low-Rank-Faktorisierung komprimiert, ohne die Optimierungsdynamik zu beeinträchtigen.
3. Versatile Implementierung: Bereitstellung von Varianten für Adam und Muon, inklusive mathematischer Herleitung der Update-Regeln.
4. Umfassende Validierung: Experimente sowohl im Pre-Training als auch im Fine-Tuning über verschiedene Modellgrößen hinweg.

4. Experimentelle Ergebnisse

Pre-Training (Llama-Architekturen)

• Setup: Pre-Training von Llama-Modellen (60M bis 1B Parameter) auf dem C4-Datensatz.
• Vergleich: LoRA-Pre wurde gegen Full-Optimierer (Adam, Muon) und andere Low-Rank-Methoden (GaLore, LoRA, ReLoRA, Fira, etc.) getestet.
• Ergebnis: LoRA-Pre erreicht die beste Leistung über alle Modellgrößen hinweg.
  • LoRA-Pre (Adam) übertrifft die besten bisherigen effizienten Baselines (z. B. Fira) bei 130M, 350M und 1B Modellen deutlich (Verbesserung von 0.81 bis 2.45 Perplexity-Punkten).
  • Die Integration mit dem Muon-Optimierer (LoRA-Pre Muon) zeigt weitere Verbesserungen, insbesondere bei kleineren Modellen.

Fine-Tuning (Mathematische Aufgaben)

• Setup: Fine-Tuning von Llama-3.1-8B und Llama-2-7B auf MetaMathQA, Evaluation auf GSM8K und MATH-500.
• Ergebnis: LoRA-Pre übertrifft konsistent alle effizienten Fine-Tuning-Baselines (LoRA, rsLoRA, DoRA, GaLore).
  • Gegenüber Standard-LoRA erzielt LoRA-Pre signifikante Verbesserungen: +3.14 Punkte auf Llama-3.1-8B und +6.17 Punkte auf Llama-2-7B.

Rank-Effizienz (Ablationsstudie)

• LoRA-Pre zeigt eine außergewöhnliche Rank-Effizienz.
• Ein Modell mit LoRA-Pre bei Rang $r=16$ erreicht eine vergleichbare Leistung wie GaLore bei Rang $r=128$ (8-fache Reduktion des Ranges).
• Dies wird auf die kontinuierliche Subraum-Anpassung zurückgeführt, die es dem Optimierer erlaubt, auch mit sehr kleinen Rängen effektiv zu lernen, während projektionsbasierte Methoden größere Ränge benötigen, um durch Fehlerakkumulation bedingte Verluste auszugleichen.

5. Bedeutung und Fazit

Das Paper „Taming Momentum" bietet einen Paradigmenwechsel in der Optimierung von LLMs.

• Speichereffizienz: Durch die Reduktion des Speicherbedarfs für Optimierer-Zustände ermöglicht LoRA-Pre das Training größerer Modelle auf begrenzter Hardware oder die Nutzung größerer Batch-Größen.
• Leistung: Im Gegensatz zu vielen Low-Rank-Methoden, die oft nur für Fine-Tuning geeignet sind, funktioniert LoRA-Pre hervorragend auch beim Pre-Training von Grund auf (From-Scratch), wo volle Rang-Updates traditionell als notwendig galten.
• Robustheit: Die Methode ist robust gegenüber der Wahl des Ranges und funktioniert kompatibel mit verschiedenen Optimierern (Adam, Muon).

Zusammenfassend demonstriert LoRA-Pre, dass eine tiefgreifende theoretische Neuinterpretation von Optimierer-Zuständen (als lineare Regression) zu praktischen Durchbrüchen in der Effizienz und Leistungsfähigkeit von Large Language Models führen kann. Der Code ist öffentlich verfügbar.