Symmetry Breaking in Transformers for Efficient and Interpretable Training

Probleemstelling

De standaard implementatie van het attention-mechanisme in Transformers bevat overbodige rotatie-vrijheidsgraden (rotational degrees of freedom). Dit betekent dat de query- en key-matrices gezamenlijk kunnen worden geroteerd zonder dat dit invloed heeft op de attention-scores (die slechts afhankelijk zijn van inproducten) of de modelactivaties.

Hoewel deze richtingen geen gradiëntsignaal dragen, vormen ze een symmetrie die de optimalisatiedynamiek beïnvloedt. Het artikel identificeert twee specifieke problemen hiermee:

1. Beperking van energiebehoudende optimalisatie: Voor de Energy Conserving Descent (ECD) optimizer, een geheugenefficiënte methode gebaseerd op Hamiltoniaanse dynamica, leiden deze symmetrieën tot behoudswetten (behoud van hoekmomentum). Volgens de stelling van Noether beperken deze behouden grootheden de chaotische verkenning in de parameterruimte die nodig is voor ECD om effectief af te dalen in de verliesfunctie. Hierdoor presteert ECD in Transformers vaak slechter dan adaptieve methoden zoals AdamW of SOAP.
2. Gebrek aan interpretabiliteit: Deze rotatie-vrijheidsgraden worden doorgaans als "ruis" beschouwd, maar het artikel stelt dat ze potentieel kunnen worden gebruikt om semantisch betekenisvolle tokenklassen te versterken of te onderdrukken.

Methodologie

De auteurs introduceren een eenvoudige, theoretisch gemotiveerde architecturale wijziging: een symmetriebreking-protocol via niet-geleerde (unlearned) bias-termen.

• Mechanisme: Er worden vaste, willekeurige bias-termen ($b_Q$ en $b_V$) toegevoegd aan respectievelijk de query- en value-projecties binnen de attention-heads.
  • $q = W_Q x + b_Q(\text{batch})$
  • $v = W_V x + b_V(\text{batch})$
• Implementatie: De bias-componenten worden per trainingsbatch opnieuw bemonsterd uit een normale verdeling ($N(\mu, \sigma^2)$). Tijdens inferentie worden de gemiddelde waarden ($\mu$) gebruikt.
• Doel:
  1. Het doorbreken van de continue rotatiesymmetrie ($O(d)$) in de attention-heads, waardoor de behoudswetten voor hoekmomentum worden opgeheven.
  2. Het creëren van een "voorkeursrichting" ($\mathbb{E}[b_Q]$) waar het model zijn key-vectoren ($k = W_K x$) aan kan uitlijnen. Dit zorgt voor een exponentiële modulatie van de attention-weights ($e^{k \cdot b_Q}$), waardoor het model tokenklassen kan versterken of onderdrukken.

Het experimentele kader omvat het voor-trainen van GPT-2 (124M parameter) modellen op de FineWeb-Edu dataset (500M tokens) met vier verschillende optimalisatiealgoritmen: AdamW, SOAP, SGDM en ECD. De prestaties worden gemeten aan de hand van validatieverlies en downstream logisch redeneren (14 logische puzzeltaken).

Belangrijkste Bijdragen

1. Hamiltoniaanse verklaring voor ECD-falen: De auteurs tonen aan dat de rotatiesymmetrie in attention-heads leidt tot behoud van hoekmomentum. Dit belemmert de chaotische menging die essentieel is voor ECD, waardoor de optimizer vastloopt in suboptimale richtingen.
2. Symmetriebreking zonder geheugenoverhead: Ze stellen een methode voor die de symmetrie breekt zonder de geheugenefficiëntie van ECD te schaden (geen extra trainbare parameters nodig, alleen batch-wise sampling).
3. Interpreteerbaarheid via uitlijning: Ze tonen aan dat het model de vaste bias-richting leert te benutten om attention te versterken op semantisch waardevolle tokens (zoals zinsbeginwoorden, leestekens en functionele woorden) en ruis (zoals Unicode-artefacten) te onderdrukken.
4. Prestatieverbetering: Ze demonstreren dat deze minimale wijziging de kloof tussen geheugenefficiënte methoden (ECD) en zware adaptieve methoden (SOAP/AdamW) dichtmaakt.

Resultaten

• Validatieverlies:
  • Zonder symmetriebreking presteert ECD aanzienlijk slechter dan AdamW en SOAP.
  • Met symmetriebreking ($b_Q + b_V$) verbetert ECD drastisch en wordt het concurrerend met SOAP (bijv. validatieverlies daalt van ~3.93 naar ~3.35 bij GPT-2 124M).
  • SGDM profiteert ook, terwijl AdamW (dat geen rotatiesymmetrie respecteert) minder of geen winst boekt en soms zelfs achteruitgaat.
• Logisch Redeneren:
  • De verbetering in redeneertaken is heterogeen: sommige seeds verbeteren, andere blijven gelijk of verslechteren licht.
  • Cruciale bevinding: De prestatieverbetering correleert niet direct met het verlies, maar wel met de semantische uitlijning. Modellen die succesvol zijn, tonen een sterke versterking van structurele markers (leestekens, zinsstarters) en een sterke onderdrukking van ruis.
  • Bijvoorbeeld: Seed 123, die slechter presteerde op logica, toonde een zwakke versterking van leestekens en extreme onderdrukking van functionele woorden.
• Interpretatie: De analyse van de key-vectoren toont aan dat modellen actief leren om $W_K$ uit te lijnen met de gemiddelde bias $\mathbb{E}[b_Q]$. Dit leidt tot een selectieve versterking van tokens die logische structuur dragen.

Betekenis en Conclusie

Dit werk toont aan dat een zorgvuldige analyse van architecturale symmetrieën en leer-dynamica kan leiden tot eenvoudige, principieel onderbouwde aanpassingen die zowel efficiëntie als interpretabiliteit verbeteren.

• Efficiëntie: Het maakt geheugenefficiënte optimalisatoren (zoals ECD) bruikbaar voor grote taalmodellen, wat belangrijk is voor schaalbaarheid en kostenreductie.
• Interpretabiliteit: Het biedt een nieuw venster om te begrijpen hoe modellen attention toewijzen. De bias-richting fungeert als een "semantische filter" die het model kan leren manipuleren.
• Toekomst: De auteurs suggereren dat dit protocol breder getest moet worden op verschillende schalen en dat de relatie tussen symmetriebreking en redeneervermogen verder onderzocht moet worden.

Kortom, door een kleine, niet-geleerde bias toe te voegen, kunnen de "verborgen" rotatie-vrijheidsgraden van Transformers worden omgezet van een obstakel voor optimalisatie naar een krachtig mechanisme voor semantische controle en efficiënt leren.