Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers
Diese Arbeit zeigt auf, dass mit Muon trainierte Transformer einen „Post-Grokking-Kollaps“ aufweisen, bei dem die Generalisierung nach dem Erlernen modularer Arithmetik aufgrund von Instabilitäten an der Schnittstelle zwischen Repräsentation und Readout versagt, ein Phänomen, das durch divergente Optimierer-Dynamiken und Maskierungseffekte charakterisiert ist und durch das Einfrieren spezifischer Modellkomponenten oder das Reskalieren des aufgabenorientierten Schaltkreises behoben werden kann.
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 bringen einem Roboter beibringen, ein sehr spezifisches mathematisches Rätsel zu lösen: zwei Zahlen addieren und den Rest bei der Division durch eine große Primzahl finden. Zuerst scheint der Roboter die Antworten einfach nur auswendig zu lernen, wie ein Schüler, der Karteikarten paukt. Aber dann geschieht etwas Magisches, das man „Grokking“ nennt. Plötzlich hört der Roboter auf, nur auswendig zu lernen, und beginnt, das Muster tatsächlich zu verstehen. Er kann das Rätsel perfekt lösen, selbst bei Zahlen, die er noch nie zuvor gesehen hat. Das ist eine große Sache für Wissenschaftler, denn es hilft ihnen zu verstehen, wie künstliche Intelligenz lernt zu denken, anstatt nur auswendig zu lernen.
Es gibt jedoch einen Haken. Manchmal, nachdem der Roboter die Lösung gefunden hat, beginnt er sie wieder zu vergessen, obwohl er immer noch am selben Rätsel trainiert wird. Das ist so, als würde ein Schüler, der endlich Algebra verstanden hat, dann aber bei denselben Aufgaben wieder dumme Fehler macht. Wissenschaftler nutzen spezielle Werkzeuge, um in das „Gehirn“ des Roboters zu schauen, um zu sehen, was dort passiert. Sie schauen sich die spezifischen Muster an, die der Roboter zum Lösen des mathematischen Problems verwendet (wie Musiknoten in einem Lied), um zu sehen, ob der Roboter immer noch die richtigen Noten spielt. Die große Frage ist: Warum vergisst der Roboter? Liegt es daran, dass er aufgehört hat, die Noten zu lernen, oder ist etwas anderes im Spiel?
Das Mysterium des verschwindenden Genies
In dieser Studie untersuchten Forscher einen speziellen Typ von KI-Roboter namens Transformer, der darauf trainiert wurde, modulare Additionsrätsel zu lösen. Sie nutzten zwei verschiedene „Lehrer“ (Optimizer), um den Roboter zu unterrichten. Ein Lehrer namens AdamW ist der Standard, der zuverlässige Instruktor. Der andere, Muon, ist ein neuerer, schnellerer Lehrer, der dem Roboter anscheinend hilft, das Rätsel viel schneller zu verstehen.
Die Forscher fanden heraus, dass Muon tatsächlich ein Geschwindigkeitsdieb ist. Er hilft dem Roboter, den „Aha!“-Moment des Grokkings etwa halb so schnell zu erreichen wie AdamW. Aber hier kommt die Wendung: Die Lösungen, die Muon findet, sind unglaublich zerbrechlich. Sobald der Roboter das Rätsel gelöst hat, beginnt er oft wieder zu scheitern, wobei die Genauigkeit von 100 % auf nahezu Null abfällt, nur um sich später wieder zu erholen. Dies geschieht immer wieder, obwohl der Roboter immer noch am selben Problem trainiert wird. AdamW ist auch nicht völlig sicher, aber er scheitert viel seltener und weniger schwerwiegend.
Der gebrochene Handschlag
Was geht hier also schief? Die Forscher entdeckten, dass das Problem nicht darin besteht, dass der Roboter die Mathematik vergisst. Der Roboter beherrscht die Lösung in seinem „Gehirn“ immer noch perfekt. Das Problem ist ein gebrochener Handschlag zwischen zwei Teilen des Roboters: dem Teil, der denkt (die Repräsentation), und dem Teil, der antwortet (der Readout).
Stellen Sie sich vor, das Gehirn des Roboters ist eine Bibliothek voller Bücher (die Gedanken) und der Mund des Roboters ist der Bibliothekar, der die Bücher laut vorliest (die Antwort).
- Die Bibliothek (die verborgenen Schichten) schreibt jede Sekunde neue Notizen in die Bücher.
- Der Bibliothekar (der Output-Head) versucht, diese Bücher zu lesen, um die Antwort zu geben.
Mit dem Muon-Lehrer schreibt die Bibliothek die Bücher jedes Mal in einer leicht anderen Sprache um, wenn sie aktualisiert wird, aber der Bibliothekar versucht weiterhin, die Bücher in der alten Sprache zu lesen. Zuerst verstehen sie sich noch. Aber weil die Bibliothek die Sprache schneller ändert, als der Bibliothekar sich anpassen kann, verstehen sie sich schließlich nicht mehr. Die Bibliothek schreibt immer noch die richtige Mathematik, aber der Bibliothekar liest die falschen Wörter.
Die Forscher bewiesen dies, indem sie eine Seite einfroren. Als sie aufhörten, die Sprache der Bibliothek ändern zu lassen, blieb der Roboter perfekt. Als sie aufhörten, den Lesestil des Bibliothekars zu ändern, blieb der Roboter ebenfalls perfekt. Aber als beiden erlaubt wurde, sich frei zu bewegen, driften sie auseinander und der Roboter begann zu scheitern.
Der Geist in der Maschine
Hier ist der verwirrendste Teil: Selbst wenn der Roboter scheitert, enthält die „Bibliothek“ immer noch die perfekte Lösung. Die Forscher verwendeten einen speziellen Filter, um nur den Teil des Gehirns zu betrachten, der die Mathematik beherrscht. Als sie dies taten, war die Lösung immer noch zu 100 % korrekt!
Es stellt sich heraus, dass der Roboter unter Circuit Masking leidet. Stellen Sie sich vor, die Bibliothek schreit die richtige Antwort, aber der Rest des Roboters schreit eine falsche Antwort mit genau demselben Volumen. Der Bibliothekar hört ein Durcheinander und gibt die falsische Antwort. Die richtige Antwort ist immer noch da, aber sie wird durch Lärm übertönt.
Die Forscher zeigten, dass man, wenn man einfach die Lautstärke des korrekten Teils (der Mathematik-Familie) erhöht, ohne sonst etwas zu ändern, der Roboter sofort wieder zu 100 % genau ist. Die Mathematik ist nicht verschwunden; sie wurde nur im Lärm verloren.
Warum die Standardwerkzeuge es nicht bemerkten
Normalerweise prüfen Wissenschaftler, ob ein Roboter funktioniert, indem sie nach den „Noten“ schauen, die er zur Lösung des Problems verwendet. Sie prüfen, ob die richtigen Noten vorhanden sind und ob sie laut genug sind. In dieser Studie sagten diese Werkzeuge, der Roboter sei in Ordnung. Die Noten waren da und sie klangen fast exakt so wie zuvor.
Aber der Roboter scheiterte trotzdem. Das liegt daran, dass die Werkzeuge auf die Noten schauten, aber nicht auf das Volumen oder die Sprache, die der Bibliothekar verwendete. Die Noten waren perfekt, aber der Bibliothekar hörte ihnen nicht mehr korrekt zu. Es ist, als würde man ein Radio überprüfen und feststellen, dass der Sender klar sendet, aber nicht merken, dass das Radio auf die falsche Frequenz eingestellt ist.
Der Tiefgang
Die Forscher testeten auch, was passiert, wenn sie den Roboter „tiefer“ machen (mehr Denk-Schichten hinzufügen). Sie fanden heraus, dass Muon noch schneller lernt, aber das Problem des gebrochenen Handschlags schlimmer wird. In tieferen Modellen schreibt das Gehirn des Roboters die Mathematik in einem Abschnitt, aber der Teil, der die Antwort liest, befindet sich in einem anderen Abschnitt, und sie driften schneller auseinander.
Sie versuchten auch, den Muon-Lehrer zu korrigieren, indem sie seinen speziellen „Normalisierungs“-Trick (eine Methode, um die Updates stabil zu halten) entfernten. Ohne diesen Trick lernte der Roboter die Mathematik auf eine sehr enge, konzentrierte Weise, aber er stürzte schließlich komplett ab und konnte sich nicht mehr erholen. Dies deutet darauf hin, dass der spezielle Trick dem Roboter zwar hilft, schnell zu lernen und sein Wissen zu verbreiten, ihn aber auch anfälliger für diese spezifische Art des driftenden Scheiterns macht.
Das Fazit
Die wichtigste Lehre aus dieser Arbeit ist, dass schnelles Lernen gefährlich sein kann, wenn die Teile des Roboters nicht synchronisiert sind. Der Roboter hat die Mathematik nicht vergessen; er hat nur die Verbindung zwischen dem Teil, der die Mathematik kennt, und dem Teil, der sie ausspricht, verloren.
Die Forscher fanden heraus, dass der Roboter stabil bleibt und niemals vergisst, wenn man den „sprechenden“ Teil (die Outputs und Embeddings) einfriert, nachdem der Roboter die Lösung gelernt hat. Dies deutet darauf hin, dass die Instabilität nicht darin liegt, dass der Roboter schlecht in Mathematik ist, sondern dass die zwei Hälften des Roboters sich mit unterschiedlichen Geschwindigkeiten bewegen und ihre gemeinsame Sprache verlieren.
Kurz gesagt: Der Roboter ist ein Genie, das plötzlich anfängt zu stottern, weil sein Gehirn und sein Mund verschiedene Dialekte sprechen. Die Mathematik ist im Inneren immer noch perfekt, aber ohne eine Möglichkeit, die Dialekte aufeinander abzustimmen, kann das Genie seine Fähigkeiten nicht zeigen.
Ertrinken Sie in Arbeiten in Ihrem Fachgebiet?
Erhalten Sie tägliche Digests der neuesten Arbeiten passend zu Ihren Forschungsbegriffen — mit technischen Zusammenfassungen, in Ihrer Sprache.