Exclusive Self-Attention Beyond AdamW: Stability and Generalization under Muon Optimization

Technische Samenvatting: Exclusieve Self-Attention voorbij AdamW

Probleemstelling
Standaard self-attention mechanismen in Transformers maken het mogelijk dat de output voor een token op positie $i$ een gewogen mengeling is van value-vectoren, inclusief de eigen value-vector van het token ($v_i$). Hoewel expressief, kan dit ontwerp leiden tot een "representatieve shortcut" waarbij de attention-output sterk uitgelijnd blijft met $v_i$, zelfs wanneer de residual pathways al lokale tokeninformatie behouden. Dit fenomeen, getypeerd als "attention similarity bias", werd door Zhai [2] aangepakt via Exclusive Self-Attention (XSA), wat expliciet het component van de attention-output dat parallel loopt aan de eigen value-vector van het token verwijdert.

De oorspronkelijke XSA-studie evalueerde de methode echter uitsluitend binnen een NanoGPT-trainingsstack met de AdamW-optimizer. De auteurs van het oorspronkelijke werk identificeerden de compatibiliteit van XSA met alternatieve optimizers, zoals Muon, expliciet als een openstaande vraag. Dit is een kritieke kloof, aangezien modern onderzoek naar kleine taalmodellen (bijv. de nanochat-stack) steeds vaker hybride optimalisatiestrategieën gebruikt (Muon voor matrixparameters, AdamW voor embeddings) en architecturale kenmerken zoals value embeddings en rotary position encoding. Het blijft onduidelijk of de geometrische beperkingen van XSA stabiel of voordelig blijven wanneer ze worden overgedragen naar deze moderne, op Muon gebaseerde trainingsregimes.

Methodologie
De studie voert een gecontroleerde reproductie- en transferexperiment uit met behulp van de nanochat-trainingsstack. De methodologie omvat:

1. Implementatie: Een geparametriseerde vorm van XSA is geïntegreerd in nanochat. De operatie wordt gedefinieerd als $z_i = y_i - \alpha \cdot \text{proj}_{v_i}(y_i)$, waarbij $y_i$ de standaard attention-output is, $v_i$ de eigen value-vector van het token is, en $\alpha \in [0, 1]$ de projectiekracht regelt. Bij $\alpha=1$ wordt het zelf-uitgelijnde component volledig verwijderd.
2. Validatie van Correctheid: De implementatie wordt numeriek gevalideerd om te waarborgen dat bij $\alpha=0$ het XSA-pad equivalent is aan standaard attention. Korte trainingstrajecten verifiëren dat de baseline en $\alpha=0$ paden een overeenkomend vroeg stadium in het verliesgedrag vertonen.
3. Experimentele Opzet: Twee modeldieptes worden geëvalueerd onder identieke trainingscondities (dezelfde datavolgorde, optimizer-familie, learning rate schedule en contextlengte):
   • d12: 12 lagen, verborgen dimensie 768, 12 heads.
   • d24: 24 lagen, verborgen dimensie 1024, 16 heads.
   • Optimizer: Matrixparameters worden geoptimaliseerd met Muon; resterende parameters gebruiken AdamW.
   • Data: ClimbMix-400B.
4. Metrieken: De studie volgt de trainingsverlies, validatie bits-per-byte (bpb), downstream aggregate scores (Base CORE en ChatCORE), en attention-diagnostiek (cosine alignment tussen output en zelf-value, diagonaal gedrag).

Belangrijkste Bijdragen

• Muon-Stack Reproductie: De eerste integratie van XSA in een op Muon gebaseerde trainingsstack, wat direct antwoord geeft op de openstaande vraag over optimizer-compatibiliteit die door de oorspronkelijke studie werd achtergelaten.
• Geparametriseerde Implementatie: Een flexibele implementatie die standaard attention, volledige XSA en partiële projectie via een sterkteparameter $\alpha$ toestaat, naast controles voor laagselectie.
• Correctheid en Instrumentatie: Validatie van de $\alpha=0$ limiet en correctie van logging-artefacten (stap-index en smoothing conventies) die eerder de vroege vergelijkingen van het trainingsproces vertekenden.
• Evaluatie op Twee Dieptes: Een vergelijkende analyse bij d12 en d24 dieptes met een uitgebreid scala aan metrieken.
• Mechanistische en Systeem-analyse: Diagnostiek van attention self-alignment en een evaluatie van de runtime-overhead.

Resultaten
De empirische resultaten wijzen erop dat XSA stabiel is binnen de op Muon gebaseerde stack, maar wisselende downstream voordelen oplevert afhankelijk van de modeldiepte en de specifieke metriek:

• Optimalisatiestabiliteit: De trainingsverlies-trajecten voor XSA en baseline-modellen blijven nauw met elkaar verbonden bij beide dieptes. Er is geen bewijs van divergentie, NaNs of instabiliteit, wat bevestigt dat XSA de update-geometrie van Muon niet verstoort.
• Validatieprestaties (Bits-per-Byte): XSA verbetert consistent de helder gehouden validatie bpb:
  • d12: Verminderd van 0.8484 naar 0.8457 (0,32% relatieve verbetering).
  • d24: Verminderd van 0.7193 naar 0.7171 (0,31% relatieve verbetering).
• Downstream Metrieken (CORE):
  • d24: XSA verbetert zowel Base CORE (0,2593 naar 0,2606, +0,50%) als ChatCORE (0,3638 naar 0,3682, +1,21%).
  • d12: Hoewel bpb verbetert, daalt Base CORE van 0,1602 naar 0,1508 (-5,87%).
• Mechanistische Diagnostiek: Attention-diagnostiek laat verminderde self-alignment en verminderd diagonaal gedrag zien onder XSA, wat consistent is met het beoogde mechanisme om contextuele mixing te bevorderen.
• Runtime Overhead: De overhead wordt als bescheiden gekarakteriseerd. Hoewel een initiële schatting rond de 5% lag, plaatsten herhaalde throughput-metingen het verschil nabij de systeem-niveau variabiliteit van korte runs, wat suggereert dat er geen sprake is van een prohibitieve straf.

Betekenis en Claims
Het artikel concludeert dat Exclusive Self-Attention een compatibele architecturale inductieve bias is voor op Muon gebaseerde trainingsstacks, maar dat de bruikbaarheid ervan niet universeel is. De primaire betekenis ligt in het aantonen dat de representationele beperking van XSA zuiver overdraagt naar een modern optimalisatieregime zonder het trainen te destabiliseren.

De auteur benadrukt echter dat de downstream winsten schaal- en metriek-afhankelijk zijn. De verbetering in validatie bpb bij beide dieptes suggereert een representationeel voordeel, maar de degradatie van Base CORE bij d12 geeft aan dat de voordelen van het verwijderen van self-alignment kunnen afhangen van de modeldiepte, de capaciteit van residual/value pathways om lokale informatie te behouden, of de specifieke evaluatietaak. De studie claimt niet dat XSA een universeel superieure vervanging is voor standaard attention, maar eerder een stabiele, met weinig verstoring gepaard gaande beperking die voordelen biedt in specifieke contexten (met name bij d24 in deze studie).

Het werk dient als een robuustheidscontrole in plaats van een byte-voor-byte reproductie, waarbij wordt gewezen op de complexe interacties tussen architecturale wijzigingen, optimizer-geometrie en evaluatiemetrieken in modern taalmodelonderzoek.