🤖 machine learning

Pro-KLShampoo: Projected KL-Shampoo with Whitening Recovered by Orthogonalization

Pro-KLShampoo ist ein neuartiger Optimierer, der KL-Shampoo verbessert, indem er die beobachtete „Spitzen-und-Flach"-Eigenwertstruktur seiner Prädikonditionierer ausnutzt, um eine vollständige Spektralverfolgung auf einem niedrigdimensionalen Unterraum mit einer Orthogonalisierung in den verbleibenden Richtungen zu kombinieren, wodurch in mehreren LLM-Skalen eine überlegene Leistung hinsichtlich Validierungsverlust, Speichereffizienz und Trainingsgeschwindigkeit erzielt wird.

Ursprüngliche Autoren: Ruotong Sun, Ermin Wei

Veröffentlicht 2026-05-08

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Ruotong Sun, Ermin Wei

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, komplexen Roboter (ein Large Language Model) beizubringen, menschliche Sprache zu sprechen. Um dies zu tun, zeigen Sie ihm Millionen von Beispielen und lassen ihn Fehler machen. Jedes Mal, wenn er einen Fehler macht, geben Sie ihm einen „Schubs" (einen Gradienten), um seinen Pfad zu korrigieren.

Das Problem ist, dass diese Schübe unordentlich sind. Manchmal braucht der Roboter einen winzigen, präzisen Schubs in eine Richtung und einen gewaltigen Stoß in eine andere. Wenn Sie ihn einfach zufällig schieben, lernt er langsam. Um schnell zu lernen, benötigen Sie einen intelligenten „Vorkonditionierer" – ein Werkzeug, das diese Schübe so umgestaltet, dass sie perfekt ausbalanciert und effizient sind.

Zwei beliebte Werkzeuge für diese Aufgabe sind KL-Shampoo und Muon.

KL-Shampoo ist wie ein Meisterbildhauer. Er versucht, die perfekte Form für jeden einzelnen Schubs zu schnitzen. Er ist sehr genau, erfordert aber viel schwere Arbeit (Rechenleistung und Speicher), um die Form jedes einzelnen Puzzleteils zu berechnen.
Muon ist wie ein Turner. Er versucht nicht, jedes einzelne Stück umzuformen; stattdessen richtet er einfach den Impuls der Schübe aus und sorgt dafür, dass sie sich auf eine saubere, orthogonale (rechtwinklige) Weise bewegen. Er ist schnell und leicht, könnte aber einige der subtilen Details verpassen, die der Bildhauer einfängt.

Die große Entdeckung: Das „Spike-and-Flat"-Muster
Die Autoren dieses Papiers betrachteten die Arbeit des „Bildhauers" (den KL-Shampoo-Vorkonditionierer) genau und bemerkten ein seltsames, schönes Muster. Sie stellten fest, dass die „Schub-Formen" nicht zufällig sind. Stattdessen sehen sie aus wie eine Spike-and-Flat-Landschaft:

Der Spike: Einige wenige Richtungen haben riesige, dominante Werte (wie ein paar hohe Berge).
Das Flat: Der Rest der Richtungen hat alle ungefähr die gleiche Höhe (wie eine weite, flache Ebene).

Dies geschieht über alle Schichten des Roboterhirns hinweg, von der ersten bis zur letzten. Es ist, als ob sich der Roboter nur wirklich um ein paar spezifische Richtungen kümmert und überall sonst nur „flaches" Rauschen herrscht.

Die Lösung: Pro-KLShampoo
Die Autoren entwickelten einen neuen Optimierer namens Pro-KLShampoo (Projected KL-Shampoo). Stellen Sie es sich als ein hybrides Werkzeug vor, das das Beste aus beiden Welten vereint, indem es die „Spike-and-Flat"-Entdeckung nutzt.

So funktioniert es, mit einer einfachen Analogie:

Die „VIP-Lounge" (Der Unterraum):
Der Algorithmus identifiziert die „Spike"-Richtungen – die wenigen wichtigen Berge. Er bringt diese in eine spezielle „VIP-Lounge" (einen verfolgten Unterraum). In dieser Lounge verwendet er das schwere, präzise Bildhauerwerkzeug (KL-Shampoo), um die perfekte Form für diese wenigen kritischen Richtungen zu erhalten. Er verschwendet keine Zeit mit dem Rest.
Das „Offene Feld" (Das Komplement):
Für die „Flat"-Richtungen (den Rest der Landschaft) hört er auf, jeden einzelnen Kieselstein zu schnitzen. Stattdessen verwendet er einen cleveren Trick namens Orthogonalisierung (übernommen vom Muon-Werkzeug).
- Der magische Trick: Die Autoren bewiesen mathematisch, dass wenn man die Schübe in diesem flachen Bereich einfach „aufrichtet" (orthogonalisiert), dies magisch exakt das gleiche algebraische Ergebnis liefert, als hätte man dort das schwere Schnitzen durchgeführt. Es ist, als würde man erkennen, dass man für ein flaches Feld keine Karte braucht; man muss nur geradeaus laufen, und man landet genau an derselben Stelle, als hätte man jede Koordinate berechnet.

Warum ist das besser?

Geschwindigkeit: Indem das schwere Schnitzen für die „flachen" Teile ignoriert und sie einfach ausgerichtet werden, läuft der Algorithmus viel schneller. Er spart viel Zeit auf der Computer-Uhr (echte Weltzeit).
Speicher: Er muss die riesigen, komplexen Formen für die flachen Teile nicht speichern. Er speichert nur die kleinen „VIP"-Formen und eine einzige Zahl für den Rest. Dies spart viel Computerspeicher.
Leistung: In Tests mit verschiedenen Robotergrößen (GPT-2 und LLaMA) lernte Pro-KLShampoo schneller und erreichte eine niedrigere Fehlerrate als der ursprüngliche KL-Shampoo, bei gleichzeitig geringerem Speicherverbrauch.

Zusammenfassung
Das Papier sagt: „Wir haben festgestellt, dass die komplexe Mathematik hinter dem modernen KI-Training ein einfaches Muster hat: ein paar große Spitzen und einen flachen Schwanz. Wir haben ein neues Werkzeug gebaut, das die Spitzen mit äußerster Sorgfalt behandelt und den flachen Schwanz mit einem einfachen, schnellen Trick. Dieser Trick funktioniert genauso gut wie die harte Mathematik, ist aber viel schneller und günstiger auszuführen."

Das Ergebnis ist ein intelligenterer, schnellerer Weg, KI-Modelle zu trainieren, der Zeit und Computerressourcen spart, ohne die Qualität zu beeinträchtigen.

Technische Zusammenfassung: Pro-KLShampoo

Problemstellung

Das Vorab-Training moderner Large Language Models (LLM) stützt sich stark auf Optimierer, die die Matrixstruktur von Gradienten ausnutzen. Zwei unterschiedliche Fronten haben sich herausgebildet:

Explizite Kronecker-faktorierte Vorkonditionierung: Vertreten durch KL-Shampoo, das den Vorkonditionierer schätzt, indem die Kullback-Leibler-Divergenz (KL) zwischen der wahren Gradientenkovarianz und einer Kronecker-strukturierten Approximation minimiert wird.
Orthogonalisierung des Gradientenmoments: Veranschaulicht durch Muon, das das Moment orthogonalisiert, um den steilsten Abstieg unter der Spektralnorm zu realisieren.

Diese Ansätze werden typischerweise isoliert entwickelt. Während KL-Shampoo durch KL-Minimierung eine starke theoretische Fundierung bietet, entstehen hohe Rechen- und Speicherkosten durch die Aufrechterhaltung vollständiger Kronecker-Faktoren. Umgekehrt sind Orthogonalisierungsmethoden wie Muon effizient, operieren jedoch unter einem anderen geometrischen Rahmen (steilster Abstieg der Spektralnorm) ohne die Perspektive der KL-Schätzung. Die Arbeit adressiert die Notwendigkeit, diese Ansätze zu vereinen, um die statistische Effizienz von KL-Shampoo mit der Recheneffizienz von Orthogonalisierungsmethoden zu erreichen.

Methodik: Pro-KLShampoo

Die Autoren schlagen Pro-KLShampoo (Projected KL-Shampoo mit durch Orthogonalisierung wiederhergestellter Whitening) vor. Die Kerninnovation ergibt sich aus einer strukturellen Beobachtung der Vorkonditionierer von KL-Shampoo: Ihre Eigenwertspektren weisen eine „Spike-and-Flat"-Form auf. Dies besteht aus wenigen dominanten Eigenwerten (dem „Spike"), gefolgt von einem annähernd einheitlichen Schwanz (dem „Flat"). Diese Struktur wird empirisch über Schichten und Trainingsstadien hinweg beobachtet und ist unter einem Signal-plus-Rauschen-Gradientenmodell mit Rang $\rho$ theoretisch exakt.

1. Strukturelle Einschränkung

Pro-KLShampoo beschränkt einen der Kronecker-Faktoren von KL-Shampoo (speziell den rechten Faktor $R$ für rechteckige Schichten, wo $m \le n$ ) auf eine parametrische Familie, die an die Spike-and-Flat-Struktur angepasst ist:

Projizierter Unterraum ( $r$ -dimensional): Der Vorkonditionierer behält die volle Spektralstruktur auf einem verfolgten Unterraum bei, der von $U \in \text{St}(n, r)$ aufgespannt wird.
Komplementärer Unterraum ( $n-r$ Dimensionen): Die verbleibenden Richtungen werden in einen einzigen gemeinsamen Eigenwert $\mu_\perp$ kollabiert.

Diese Einschränkung reduziert den Speicherbedarf des rechten Faktors von $O(n^2)$ auf $O(nr)$ und die Kosten der QR-Zerlegung von $O(n^3)$ auf $O(r^3)$ .

2. Wiederherstellung durch Orthogonalisierung

Eine naive Einschränkung auf einen skalaren Schwanz führt zu Informationsverlust, wenn der Schwanz nicht perfekt einheitlich ist. Um die algebraische Form des vollständigen KL-Shampoo-Vorkonditionierers im komplementären Unterraum wiederherzustellen, wendet Pro-KLShampoo Orthogonalisierung (polare Zerlegung) auf die komplementäre Gradientenkomponente an.

Die Update-Regel lautet:
$\Delta W = -\alpha_{kl} L^{-1/2} G U S^{-1/2} U^\top - c_a \text{polar}\left(\mu_\perp^{-1/2} L^{-1/2} G P_\perp\right)$
Wobei:

Der erste Term das vollständige KL-Shampoo-Update innerhalb des verfolgten Unterraums darstellt.
Der zweite Term die Orthogonalisierung auf das Komplement anwendet.
Die Autoren eine Identität (Gleichung 12) beweisen, die zeigt, dass diese Orthogonalisierung die algebraische Form der per-Richtung-Whitening des vollständigen KL-Shampoo-Vorkonditionierers im Komplement wiederherstellt und damit Muon-artige Orthogonalisierung effektiv mit dem KL-Schätzrahmen verbindet.

3. Unterraum-Verfolgung

Der optimale Unterraum $U$ wird als Eigenraum der Matrix $\Phi_L = E[G^\top L^{-1} G]$ identifiziert. Unter der Spike-and-Flat-Vermutung ist der Top- $r$ -Eigenraum von $\Phi_L$ die optimale Wahl. Der Algorithmus verfolgt diesen Unterraum mittels eines einzelnen Schritts der Potenziteration (via QR-Zerlegung) auf einer laufenden Schätzung von $\Phi_L$ .

Hauptbeiträge

Strukturelle Beobachtung & Parametrische Einschränkung: Die Arbeit identifiziert die „Spike-and-Flat"-Eigenwertstruktur in den Vorkonditionierern von KL-Shampoo (empirisch robust und unter Low-Rank-Gradientenmodellen theoretisch exakt) und nutzt sie, indem sie das KL-Ziel auf eine parametrische Familie einschränkt.
Algebraische Wiederherstellung durch Orthogonalisierung: Sie zeigt, dass die Anwendung von Orthogonalisierung auf den komplementären Unterraum die algebraische Form des vollständigen KL-Shampoo-Vorkonditionierers wiederherstellt. Dies bietet eine theoretische Brücke zwischen Muon-artiger Orthogonalisierung und KL-Schätzung.
Konvergenzgarantie: Die Autoren liefern eine nicht-konvexe Konvergenzgarantie mit einer Rate von $O(1/\sqrt{T})$ für den idealisierten Algorithmus unter Operator-Norm-Glattheit.
Empirische Leistung: Eine umfassende Evaluierung an GPT-2 (124M/350M) und LLaMA (134M/450M) Modellen zeigt, dass Pro-KLShampoo KL-Shampoo bei der Validierungsverlust, dem Spitzen-Speicherbedarf pro GPU und der Wandzeit zur Erreichung angepasster Verlustniveaus über alle getesteten Unterrang-Ränge ( $r \in \{32, 64, 128\}$ ) hinweg konsistent übertrifft.

Experimentelle Ergebnisse

Die Autoren bewerteten Pro-KLShampoo gegen AdamW, Muon und KL-Shampoo in vier Vorab-Trainingskonfigurationen.

Validierungsverlust: Pro-KLShampoo erzielte bei jedem getesteten Rang über alle Modellskalen hinweg den niedrigsten Validierungsverlust.
Speichereffizienz: Der Spitzen-Speicherbedarf pro GPU war niedriger als bei KL-Shampoo und vergleichbar mit Muon.
Wandzeit: Pro-KLShampoo erreichte spezifische Validierungsverlustniveaus schneller als KL-Shampoo.
- Bei GPT-2 wurde die Beschleunigung primär durch reduzierte Zeit pro Schritt getrieben (Ersetzung von $O(n^3)$ QR durch $O(r^3)$ ).
- Bei LLaMA wurde die Beschleunigung durch schnellere Konvergenz (niedrigere Trainingsverlustkurven) zusätzlich zur Effizienz pro Schritt getrieben.
Ablationsstudien:
- Das Entfernen der Orthogonalisierung (Smok-Hop-Variante) führte zu Leistungseinbußen, insbesondere bei größeren Skalen (LLaMA 450M), was die Notwendigkeit der Orthogonalisierung zur Wiederherstellung der per-Richtung-Whitening bestätigt.
- Weder die reine Unterraum- noch die reine Komplement-Variante erreichten die Leistung des vollständigen Pro-KLShampoo, was darauf hindeutet, dass beide Komponenten essenziell sind.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass Pro-KLShampoo zwei zuvor isolierte Optimierungsfelder erfolgreich vereint. Durch die Ausnutzung der intrinsischen „Spike-and-Flat"-Struktur von Gradientenkovarianzen erreicht es die statistische Effizienz von KL-Shampoo, während es Orthogonalisierung nutzt, um die Recheneffizienz aufrechtzuerhalten.

Die Autoren betonen, dass der Ersatz der skalaren Schwanz-Update durch Orthogonalisierung keine Heuristik ist, sondern eine algebraische Wiederherstellung der vollständigen KL-Shampoo-Form. Dies ermöglicht es der Methode, die Vorteile der KL-basierten Vorkonditionierung (wie Robustheit der Schrittweite) beizubehalten, während der mit der Aufrechterhaltung vollständiger Kronecker-Faktoren verbundene Rechenaufwand erheblich reduziert wird. Die Arbeit legt nahe, dass niedrigdimensionale Spektralstrukturen in der Optimierung tiefer Netzwerke allgegenwärtig sind und effektiv genutzt werden können, um effizientere Optimierer für das Vorab-Training von LLM zu entwerfen.

Die Arbeit bleibt bezüglich der Einschränkungen bescheiden und stellt fest, dass der Rang $r$ empirisch festgelegt wird, und künftige Arbeiten sind für eine schichtweise adaptive Kalibrierung des Mischgewichts $\alpha_{kl}$ sowie eine Konvergenzanalyse der praktischen Variante erforderlich, die Newton-Schulz-Iterationen und Nesterov-Momentum beinhaltet.