Das KI-Labor mit einem Schätzwert von 84 Milliarden Yuan setzt erneut ein neues Zeichen: Sie wollen den Large Language Models einen "Halsband" aufsetzen.

Gerade jetzt hat das von Mira Murati, der ehemaligen CTO von OpenAI, gegründete Thinking Machines Lab erneut ein Ergebnis veröffentlicht!

Dies ist ihr zweiter Forschungsartikel nach "Defeating Nondeterminism in LLM Inference" - "Modular Manifolds".

Blog-Adresse: https://thinkingmachines.ai/blog/modular-manifolds/

Das Training von großen neuronalen Netzen ist wie das Gehen auf einer Seilbahn. Man muss vorsichtig sein, um die interne "Gesundheit" des Netzes aufrechtzuerhalten und zu verhindern, dass wichtige Tensoren wie Gewichte, Aktivierungswerte oder Gradienten zu groß oder zu klein werden, um eine Reihe von Problemen wie numerische Überläufe zu vermeiden.

Einer der wichtigen Ansätze besteht darin, den großen Modellen eine einheitliche Größenordnungsmessung zu bieten.

Zunächst muss man die Basis stabilisieren.

Es ist eine übliche Methode, die Layer Norm-Technik zu verwenden, um die Ausgabe jeder Schicht in einen geeigneten Bereich zurückzuziehen und die Aktivierungsvektoren zu normalisieren.

Es ist auch üblich, die Gradientenaktualisierung zu normalisieren. Beispielsweise führt der Muon-Optimierer eine Spektralnormalisierung der Aktualisierung durch, um die Amplitude jeder Aktualisierungsschritt zu kontrollieren.

Im nächsten Schritt muss man direkt die Gewichte selbst "kontrollieren".

Die Normalisierung der Gewichtsmatrix ist ein vielversprechender Ansatz.

Der Artikel bietet einen neuen Blickwinkel auf die Optimierungsalgorithmen: Die Gewichtstensoren werden auf einer bestimmten Untermannigfaltigkeit (submanifold) eingeschränkt, um die Optimierungsalgorithmen in Zusammenarbeit mit diesen Mannigfaltigkeitsbeschränkungen zu gestalten.

Das ist wie das Wechseln von "Feuerlöschen" zu "Feuerverhütung":

Man platziert die Parameter von Anfang an in einem gesunden Bereich, um das Training stabiler und interpretierbarer zu machen, so dass der große Modell effizienter und stabiler trainiert werden kann.

Die Form des Mannigfaltigkeitsoptimierers

Wir wissen, dass eine Mannigfaltigkeit einfach eine Fläche ist, die lokal flach aussieht.

Wenn man sie stark genug vergrößert, sieht sie wie eine gewöhnliche Ebene aus.

Der lokal flache Raum in der Nähe eines Punktes auf der Mannigfaltigkeit wird "Tangentialraum" (tangent space) genannt.

Wie in Abbildung 1 gezeigt, ist eine dreidimensionale Kugel oder eine höherdimensionale Hyperkugel eine Mannigfaltigkeit. Der rote Teil in der Abbildung zeigt die Tangentialebene an einem bestimmten Punkt.

Um sicherzustellen, dass die Gewichte in der angegebenen Mannigfaltigkeit bleiben, ist eine einfache Methode, einen gewöhnlichen Optimierer zu verwenden und die Gewichte nach jeder Aktualisierungsschritt auf die Mannigfaltigkeit zu projizieren.

Das Problem ist jedoch, dass wenn der Optimierungsschritt zu weit von der Mannigfaltigkeit abweicht und dann erzwungenermaßen zurückprojiziert wird, die nominale Lernrate nicht mehr der tatsächlichen Verschiebung der Parameter auf der Mannigfaltigkeit entspricht, was unsere Intuition über die Beziehung zwischen "Schrittgröße - Effekt" schwächt.

Um einen Trainingsalgorithmus auf der Mannigfaltigkeit zu entwerfen, muss man zunächst klären, wie man die "Entfernung" im Tangentialraum misst.

Ein Lösungsansatz besteht darin, direkt im Tangentialraum zu optimieren. Auf diese Weise bewegt man sich bei jedem Schritt entlang der "Oberfläche" der Mannigfaltigkeit, und die Lernrate kann besser der "tatsächlichen Verschiebung" entsprechen.

Eine übliche Wahl ist die euklidische Distanz, aber man kann auch andere Distanzmaße wählen, wie in Abbildung 2 gezeigt.

Es ist wichtig zu beachten, dass die Wahl des Distanzmaßes direkt die Richtung des optimalen Optimierungsschrittes beeinflusst.

In Abbildung 3 zeigt der rosa Pfeil den ursprünglichen Gradienten - d.h. die partielle Ableitung der Verlustfunktion nach den Gewichten.

Das bedeutet, dass wir nicht unbedingt genau in Richtung des Gradienten bewegen müssen.

Um diesen Prozess mathematisch auszudrücken, können wir die "optimale Aktualisierungsrichtung unter Mannigfaltigkeitsbeschränkungen und einem bestimmten Distanzmaß" als ein beschränktes Optimierungsproblem betrachten. Wir können eine Hyperkugel mit euklidischer Norm als Beispiel nehmen.

Wir bezeichnen den Gradienten mit g, den aktuellen Punkt auf der Hyperkugel mit w, die Aktualisierungsrichtung mit a und die Lernrate mit η. Das Problem, das wir lösen müssen, lautet:

Wenn wir uns wieder den Visualisierungen in Abbildung 1, 2 und 3 ansehen, bedeutet diese Formel: Der grüne Pfeil (d.h. die optimale Lösung für a) muss zwei Bedingungen erfüllen:

Er muss auf der roten Tangentialebene liegen und auf dem gelben Kreis mit Radius η.

Wir können die Methode der Lagrange-Multiplikatoren anwenden, um dieses Problem zu lösen.

Hierbei sind λ und μ die Lagrange-Multiplikatoren.

Wir leiten die Lagrange-Funktion nach a ab und setzen sie gleich Null. Dann lösen wir die Gleichungen in Kombination mit den beiden Beschränkungsbedingungen, um λ und μ zu bestimmen und so die optimale Aktualisierungsrichtung zu erhalten.

Einfach ausgedrückt, besteht die optimale Aktualisierungsstrategie darin, zunächst die radiale Komponente in Richtung von w aus dem Gradienten zu subtrahieren, d.h. den Gradienten auf den Tangentialraum zu projizieren, dann das Ergebnis zu normalisieren und schließlich mit der Lernrate zu multiplizieren.

Die so erhaltene Aktualisierungsrichtung liegt im Tangentialraum.

In Abbildung 4 wird dieser kleine Korrekturschritt als "Retraktionsabbildung" (retraction map) bezeichnet.

Der vollständige Mannigfaltigkeitsoptimierungsalgorithmus lautet wie folgt:

Zusammengefasst besteht der erste Ordnung Mannigfaltigkeitsoptimierer aus drei Schritten:

Man findet einen Tangentialvektor der Länge 1, der sich möglichst weit in Richtung des Gradienten bewegt;

Man multipliziert diese Richtung mit der Lernrate und subtrahiert das Ergebnis von den aktuellen Gewichten;

Man zieht die aktualisierten Gewichte über die Retraktionsabbildung auf die Mannigfaltigkeit zurück.

Beim Ausführen dieses Prozesses müssen wir uns entscheiden, welche Mannigfaltigkeit wir als Beschränkung wählen und wie wir das "Längenmaß" definieren.

Je nach diesen beiden Entscheidungen erhalten wir verschiedene Optimierungsalgorithmen, wie in der folgenden Tabelle gezeigt.

Manifold Muon

Die typische Gewichtsmatrix W in einem Transformer ist ein "Vektorwandler", d.h. sie transformiert den Eingangsvektor x in den Ausgangsvektor y = Wx.

Wir möchten eine Mannigfaltigkeitsbeschränkung und eine Distanzfunktion entwerfen, so dass die Matrix auf den Eingangsvektor auf eine sinnvolle Weise wirkt: Sie soll weder zu große noch zu kleine Ausgangswerte erzeugen, noch soll sie beim Aktualisieren der Gewichte starke oder kaum messbare Änderungen des Ausgangsvektors verursachen.

Eine gute Methode, um zu verstehen, wie eine Matrix auf einen Vektor wirkt, ist die Singulärwertzerlegung (SVD), wie in Abbildung 5 gezeigt.

Die SVD zeigt, wie die Matrix den Eingangsvektor entlang verschiedener Achsen streckt, indem sie die Matrix in ihre Bestandteile zerlegt.

Wir möchten, dass die "Streckwirkung" der Matrix nahe bei 1 liegt, daher haben wir uns für eine Matrixmannigfaltigkeit entschieden, bei der alle Singulärwerte gleich 1 sind.

Diese Matrixmannigfaltigkeit wird in der Mathematik als Stiefel-Mannigfaltigkeit bezeichnet. Unter der Annahme einer hohen Matrix (m ≥ n) kann sie äquivalent als folgende Menge definiert werden:

Um einen Optimierer für die Stiefel-Mannigfaltigkeit zu entwerfen, müssen wir auch eine geeignete Distanzfunktion wählen.

Um die maximale Streckwirkung der Gewichtsaktualisierung auf den Eingangsvektor zu begrenzen, ist die Spektralnorm (spectral norm), d.h. die Messung des größten Singulärwerts der Matrix, eine geeignete Wahl.

Obwohl sie nur die maximale Wirkung beschränkt, kann sie indirekt auch verhindern, dass die minimale Wirkung zu klein wird, da der Optimierer diese Obergrenze sättigt.

Genau diese Idee hat zur Entwicklung des Muon-Optimierers geführt.

Wenn diese Idee mit der Stiefel-Mannigfaltigkeitsbeschränkung kombiniert wird, entsteht das "manifold Muon" Problem.

Ein Schlüsselergebnis des Artikels ist ein konvexes Optimierungsproblem, das mit einer Standardmethode - der Dualen Steigungsverfahren (dual ascent) - gelöst werden kann.