Technische Samenvatting: Verschillende Lagen, Verschillende Manifolds

Probleemstelling

De geometrie van de gewichtsruimte is een cruciale factor in de optimalisatie van neurale netwerken, maar huidige methoden voor training met manifold-restricties passen doorgaans een enkele, uniforme geometrische restrictie toe (bijv. orthogonaliteit of Stiefel-restricties) op alle gewichtsmatrices in een Transformer. Dit werk onderzoekt of verschillende Transformer-modules—specifiek attention-lagen en Multi-Layer Perceptron (MLP) blokken—voorkeur geven aan verschillende gewichtsruimte-geometrieën. De auteurs stellen dat de computationele rollen en niet-lineaire mechanismen van deze modules significant verschillen, wat suggereert dat een uniforme manifold-toewijzing suboptimaal of zelfs destabiliserend kan zijn.

Methodologie

De studie maakt gebruik van Manifold Muon, een optimizer die matrix-genormaliseerde updates toepast op de gewichten van verborgen lagen, terwijl deze worden beperkt tot specifieke matrix-manifolds. De auteurs vergelijken twee primaire restrictiefamilies:

1. Stiefel-restrictie: Dwingt kolom-orthogonaliteit af ($W^\top W = I$), wat de singuliere waarden strikt begrenst en sterke spectrale controle biedt.
2. DGram-restrictie: Een zwakkere Gram-ruimte restrictie die vereist dat de off-diagonaal elementen van de Gram-matrix nul zijn ($\text{Off}(W^\top W) = 0$). Dit behoudt kolom-orthogonaliteit maar staat toe dat kolomnormen (en dus singuliere waarden) variëren, wat extra vrijheidsgraden in schaal biedt.

De auteurs evalueren vijf verschillende toewijzingsstrategieën op een GPT-2 small model (124M parameters) die is voorgetraind op OpenWebText:

• UNCONSTRAINED: Geen manifold-restricties.
• ALL-STIEFEL: Stiefel-restricties op zowel de Attention- als de MLP-gewichten.
• ALL-DGRAM: DGram-restricties op zowel de Attention- als de MLP-gewichten.
• HETERO: Stiefel-restricties op Attention, DGram op MLP.
• HETERO-INV: DGram op Attention, Stiefel op MLP.

Alle configuraties delen identieke modelarchitecturen, trainingsschema's en optimizer-hyperparameters om het effect van de manifold-toewijzing te isoleren.

Belangrijkste Resultaten

De experimenten onthullen een duidelijke asymmetrie in hoe verschillende modules reageren op manifold-restricties:

• Optimale Configuratie: De HETERO-toewijzing (Stiefel voor Attention, DGram voor MLP) bereikte de laagste stabiele validatieverlies (3.3544), waarmee het zowel de onbeperkte baseline als de uniforme ALL-STIEFEL configuratie overtrof.
• Instabiliteit in DGram-Attention: Configuraties die DGram-restricties toewijzen aan attention-lagen (ALL-DGRAM en HETERO-INV) werden instabiel onder de gedeelde hyperparameter-instelling. Deze runs vertoonden een explosieve groei in het validatieverlies en slaagden er niet in te herstellen.
• Stabiliteit in DGram-MLP: Omgekeerd bleven configuraties die DGram-restricties toewijzen aan MLP-lagen (wanneer gekoppeld aan Stiefel-attention) stabiel en verbeterden ze de prestaties ten opzichte van de uniforme Stiefel-restricties.

Mechanistische Analyse

Het artikel biedt een mechanistische verklaring voor de instabiliteit die wordt waargenomen bij DGram-geconstraineerde attention:

1. Groei van Singuliere Waarden: In tegen tegenover Stiefel-restricties staat DGram toe dat de singuliere waarden van gewichtsmatrices willekeurig groot worden. In attention-lagen vermenigvuldigen de query ($W_Q$) en key ($W_K$) projecties zich om attention-logits te vormen ($Z = X W_Q W_K^\top X^\top / \sqrt{d}$).
2. Logit-amplificatie: Ongecontroleerde groei in de spectrale normen van $W_Q$ en $W_K$ versterkt direct de magnitude van de attention-logits.
3. Softmax-verzadiging: Grote logit-magnitudes duwen de softmax-functie in een verzadigd regime, waarbij de output nadert tot een one-hot vector. Dit veroorzaakt dat de Jacobiaan van de softmax exponentieel naar nul gaat, wat leidt tot gradiënt-degeneratie en trainingsinstabiliteit.
4. MLP-tolerantie: MLP-lagen tolereren DGram-restricties omdat hun activatiefuncties (bijv. GELU) element-gewijs opereren en geen token-scores koppelen via een genormaliseerd waarschijnlijkheidssimplex. Bijgevolg dwingt de schaalvergroting in MLP-gewichten de gehele module niet in een verzadigd routing-regime.

Bijdragen

De auteurs identificeren drie primaire bijdragen:

1. Empirische Evaluatie: Ze evalueren module-specifieke manifold-toewijzingen voor Manifold Muon in GPT-2 pretraining, waarmee zij aantonen dat uniforme restricties onvoldoende zijn.
2. Ontdekking van Asymmetrie: Ze tonen aan dat attention- en MLP-lagen verschillende voorkeuren hebben voor gewichtsruimte-geometrie, waarbij attention sterke spectrale controle (Stiefel) vereist en MLP profiteert van de vrijheid van zwakkere schaalbehoudende (DGram) geometrie.
3. Instabiliteitsmechanisme: Ze identificeren de groei van singuliere waarden en de daaropvolgende softmax-verzadiging als een plausibel mechanisme voor de instabiliteit van DGram-geconstraineerde attention.

Betekenis en Beperkingen

Het artikel betoogt dat symmetrie-bewuste en geometrie-bewuste optimalisatie voor Transformers module-specifiek moet zijn in plaats van uniform. Dit suggereert dat gewichtsruimte-geometrie een praktisch ontwerpaspect is voor het verbeteren van de trainingsstabiliteit en dynamiek in moderne netwerken.

De auteurs houden een bescheiden standpunt aan met betrekking tot hun claims:

• Zij beweren niet dat DGram-attention inherent instabiel is onder alle omstandigheden, maar dat het een instabiel pad blootlegt onder de geteste gedeelde hyperparameter-instelling zonder aanvullende schaalcontrole.
• De experimenten zijn beperkt tot GPT-2 small en een specifieke set seeds; het verschil in validatieverlies tussen de beste heterogene en uniforme configuraties is bescheiden.
• De mechanistische analyse dient als een verklaring voor de waargenomen instabiliteit in plaats van een volledig bewijs dat DGram-attention niet gestabiliseerd kan worden (bijv. via aangepaste learning rates of expliciete schaal-regularisatie).

Toekomstig werk wordt gesuggereerd om automatische methoden te ontwikkelen voor het toewijzen van manifold-restricties over verschillende modules, architecturen en schalen.