Grokking: Transformer lernen in niedriger Dimension und Kurven steuern

In einer neuen Studie auf arXiv wird das rätselhafte Phänomen des „Grokking“ – dem verzögerten Übergang von bloßem Auswendiglernen zu echter Generalisierung bei kleinen algorithmischen Aufgaben – aus geometrischer Sicht beleuchtet. Forscher untersuchten Transformer, die auf modularer Arithmetik trainiert wurden, und analysierten die Dynamik ihrer Aufmerksamkeitsgewichte.

Durch Hauptkomponentenanalyse (PCA) der Gewichtstrajektorien zeigte sich, dass der Lernprozess überwiegend in einem stark reduzierten, niedrigen Dimensionsraum stattfindet. Ein einzelner Hauptkomponent erklärte zwischen 68 % und 83 % der Varianz der Trajektorien, was auf eine starke Konfokussierung des Trainings hinweist.

Um die Geometrie des Verlustlandschapes zu erfassen, wurden sogenannte Kommutator‑Defekte gemessen – ein Maß für die Nicht‑Kommutativität aufeinanderfolgender Gradienten­schritte. Diese Defekte wiesen eine rasche Zunahme der Krümmung in Richtungen orthogonal zum niedrigen Dimensionsraum auf, während die Trajektorie selbst weitgehend innerhalb dieses Raumes blieb. Interessanterweise kündigte die Krümmungszunahme die spätere Generalisierung an und folgte einer potenziellen Gesetzmäßigkeit, die sich über verschiedene Lernraten und Hyperparameter hinweg zeigte.

Durch gezielte Interventionen konnte gezeigt werden, dass Bewegung entlang des erkannten Subraums für das Grokking unerlässlich ist, während eine künstliche Erhöhung der Krümmung allein nicht ausreicht. Die Ergebnisse deuten darauf hin, dass Grokking ein geometrischer Escape aus einem metastabilen Zustand ist, in dem das Modell zunächst in einem niedrigen Dimensionsraum gefangen ist und durch akkumulierte transversale Krümmung schließlich die Generalisierung erreicht. Die Befunde wurden konsistent über mehrere Lernraten, ein langsames Regime mit anderen Hyperparametern und drei unabhängige Zufallsinitialisierungen hinweg reproduziert.