Warum Muon bei QK-Gewichten scheitert und wie man es repariert
From Muon to Gradient Clipping: Some Thoughts on QK Stability
Der Muon-Optimierer, der auf Funktionsebene statt im Parameterraum arbeitet, zeigt vielversprechende Ergebnisse, versagt jedoch bei den Query- und Key-Matrizen von Transformern und führt zu instabilem Training. Dieser Artikel analysiert die Ursache: Die Kopplung von Q und K über das Aufmerksamkeits-Score-Produkt erfordert eine Beschränkung des Produkts, nicht der einzelnen Matrizen. Der Autor leitet eine modifizierte Optimierungsformel her, die die Änderung des Score-Produkts direkt begrenzt, und diskutiert die Herausforderungen bei der Lösung, einschließlich Entkopplung und Pseudoinverse. Der Beitrag verbindet theoretische Herleitung mit praktischen Erkenntnissen aus Arbeiten wie Kimi K2.
Da Muons Update vollen Rang hat, ist seine Energie gleichmäßig über alle Singulärrichtungen verteilt, sodass es viel wahrscheinlicher mit den Singulärvektoren von W kollidiert als ein Adam-Update mit niedrigem Rang.