Technische Samenvatting: Role-Decoupled Attention Residuals

Probleemstelling

In standaard Transformer-architecturen fungeert de residustraat als een vast geheugen waarbij elke laag de onmiddellijk voorafgaande staat overneemt via seriële optelling. Recente depth-routing architecturen, zoals Block Attention Residuals (Block AttnRes), verbeteren dit door sublagen toe te staan een inhoudsafhankelijke mix van eerdere representaties op te halen in plaats van alleen de onmiddellijke voorganger. Dit transformeert diepte van een vaste computationele volgorde naar een geleerde retrieval-as.

Block AttnRes introduceert echter een specifieke ontwerpbeperking: het gebruikt een enkele diepteverdeling om de gedeelde input voor alle drie de aandachtsprojecties (Queries, Keys en Values) te construeren. Dit koppelt twee functioneel verschillende rollen:

1. Matching Geometry: Bepaald door Queries en Keys, die beslissen waar te kijken (welke tokenposities overeenkomen).
2. Content Retrieval: Bepaald door Values, die beslissen wat voor inhoud wordt opgehaald.

Het artikel stelt dat het afdwingen van de representatie die het beste is voor het beslissen "waar te kijken", ook de beste moet zijn voor het beslissen "wat te lezen", een onnodige koppeling is. De centrale vraag is of het ontkoppelen van deze rollen de pretraining van taalmodellen verbetert wanneer alle andere factoren constant worden gehouden.

Methodologie: Role-Decoupled Attention Residuals (RD-AttnRes)

De auteur introduceert Role-Decoupled Attention Residuals (RD-AttnRes), een minimale architecturale relaxatie van Block AttnRes.

• Architectuur: RD-AttnRes behoudt de bronset, bloktopologie, normalisatie, MLP-routes en writeback-mechanismen van de ouderarchitectuur. De enige wijziging vindt plaats in de attention-zijde van de depth read.
• Ontkoppelingsmechanisme:
  • Queries en Keys (Q/K): Delen een enkele geleerde route query ($q^{QK}_\ell$) om een mix ($h^{QK}_\ell$) te berekenen uit het residuele dieptegeheugen. Dit behoudt een coherente matching-geometrie.
  • Values (V): Leren een onafhankelijke route query ($q^V_\ell$) om een aparte mix ($h^V_\ell$) te berekenen over exact dezelfde set residuele bronnen.
• Parameter-efficiëntie: De wijziging voegt slechts één vector van dimensie $d$ (modelbreedte) per laag toe. Voor de geteste schalen (120M en 343M parameters) vertegenwoordigt dit een verwaarloosbare toename in parameters (<0,01%).
• Exacte Recoverability: Als de value route query gekoppeld is aan de Q/K query ($q^V_\ell = q^{QK}_\ell$), reduceert RD-AttnRes exact tot de ouderlijke Block AttnRes architectuur. Dit garandeert dat het voorstel een strikte, controleerbare uitbreiding is in plaats van een systeem met meerdere gelijktijdige wijzigingen.
• Computationele Kosten: De methode voegt $O(BTK_\ell d)$ werk toe per attention-laag (waarbij $B$ de batchgrootte is, $T$ de sequentielengte en $K_\ell$ het aantal bronnen), maar introduceert geen nieuwe $O(T^2d)$ token-naar-token attention termen.

Experimenteel Ontwerp

De studie maakt gebruik van een rigoureus, gepaard experimenteel protocol om het effect van rolontkoppeling te isoleren:

• Datasets: FineWeb-Edu (10BT sample), met een deterministische split voor training, validatie en test.
• Modellen: Twee decoder-only Pre-Norm Transformer schalen:
  • S: 120M parameters (16 lagen, breedte 576).
  • M: 343M parameters (24 lagen, breedte 832).
• Protocol: Vijf gematchte seeds per schaal. Zowel Block AttnRes als RD-AttnRes zijn getraind met identieke initialisatie, datavolgorde, optimalisatieschema's en een budget van 2,0B tokens.
• Controles: Directionele controles bij 0,5B tokens zijn:
  • Tied: Forceert de value route om de Q/K route te matchen (herstelt de ouder).
  • Average: Voert twee routes uit maar middelt hun mixes (test of dubbele executie alleen helpt).
  • Frozen-V: Houdt de value route uniform (test of een vaste route helpt).
• Diagnostiek: Jensen–Shannon (JS) divergentie wordt gemeten tussen de Q/K en V diepteverdelingen om te verifiëren of de routes verschillende strategieën leren.

Belangrijkste Resultaten

1. Verbetering in Validatieverlies: RD-AttnRes verbetert de validatie negative log-likelihood (NLL) in alle 10 gepaarde vergelijkingen (5 seeds × 2 schalen).
   • 120M Schaal: Gemiddelde NLL-reductie van 0,0301, wat overeenkomt met een reductie in perplexiteit van 2,97%.
   • 343M Schaal: Gemiddelde NLL-reductie van 0,0247, wat overeenkomt met een reductie in perplexiteit van 2,43%.
2. Persistentie van Winst: De verbetering is zichtbaar bij alle evaluatiemijlpalen (0,25B, 0,5B, 1,0B en 2,0B tokens), hoewel de absolute grootte van het gat licht afneemt naarmate de training vordert.
3. Controle-validatie:
   • Tied en Average controles presteren vergelijkbaar met de ouderlijke Block AttnRes, wat aangeeft dat de winst niet komt door het aantal parameters of dubbele berekening.
   • Frozen-V presteert slecht, wat bevestigt dat de value route geleerd moet worden, niet vastgezet.
4. Route Divergentie: De geleerde Q/K en V verdelingen convergeren niet naar dezelfde oplossing. De gemiddelde JS divergentie over de lagen is ongeveer 0,218 (120M) en 0,215 (343M), wat bevestigt dat het model onderscheidende diepte-reads leert voor matching versus inhoudelijke extractie.

Betekenis en Claims

Het artikel maakt een nauwe, begrensde claim in plaats van een universele claim. De auteur beweert niet dat RD-AttnRes superieur is aan alle recente routing-architecturen of dat het universeel schaalgedrag vaststelt.

In plaats daarvan ondersteunen de resultaten de specifieke conclusie dat binnen de geteste trainingsregime (FineWeb-Edu, 2B tokens, specifieke modelmaten), de gedeelde QKV diepte-read in Block AttnRes een consequente beperking is. Het bewijs suggereert dat matching (Q/K) en inhoudelijke extractie (V) baat hebben bij verschillende diepte-reads.

Erkende Beperkingen:

• De studie is beperkt tot één Engelse webcorpus en modellen met minder dan een miljard parameters.
• De absolute prestatiekloof wordt kleiner over langere trainingshorizonten; toekomstig werk is nodig om te bepalen of de kloof standhoudt of groeit bij 4B–10B tokens.
• Er is een systeem-kanttekening: de huidige implementatie brengt een throughput-penalty op (~9–10%) en verhoogd geheugengebruik met zich mee vergeleken met de referentie Block AttnRes, hoewel gefuseerde kernels dit kunnen verzachten.
• Vergelijkingen met andere recente architecturen (bijv. MUDDFormer, Delta AttnRes) zijn conceptueel eerder dan gematchte empirische baselines.

Conclusie:
Het artikel stelt vast dat wanneer diepte een geleerde retrieval-as wordt, het delen ervan de downstream functie (rol) moet volgen in plaats van de implementatiegemakkelijkheid. Door de value route te ontkoppelen met minimale parameter-overhead, bereikt RD-AttnRes consistente verbeteringen in de pretraining loss, waarmee de hypothese wordt gevalideerd dat matching-geometrie en inhoudelijke extractie verschillende taken zijn die verschillende diepte-representaties vereisen.