Attention Residuals uitgelegd: Transformer-diepte heroverwegen

Moderne LLM's zijn dieper, breder en rekenintensiever geworden, maar meer Transformer-lagen stapelen levert niet altijd evenredige winst op. Een reden is dat standaard residu-verbindingen laagoutput aggregeren met vaste eenheidsgewichten, waardoor elke laag een uniforme som erft van alles wat eraan voorafging. Dit kan eerdere representaties verdunnen, de grootte van verborgen toestanden vergroten en het voor het netwerk lastiger maken om de meest bruikbare tussentijdse features selectief te hergebruiken.

In plaats van diepte te behandelen als een vaste additieve recursie, laat het Attention Residuals-paper van het Kimi-team elke laag zich richten op eerdere laagoutput met behulp van geleerde softmax-gewichten. In dit artikel leggen we uit waarom standaard residu-aggregatie een bottleneck wordt, hoe Attention Residuals werkt, waarom de blokvariant ertoe doet en wat de resultaten daadwerkelijk suggereren over het schalen van diepere taalmodellen.

Als je meer wilt leren over enkele ideeën achter moderne Transformer-architecturen, raad ik de DataCamp-cursus Transformer Models Tutorial in PyTorch aan.

Het probleem: het verdrinkende signaal

Residu-verbindingen zijn al jaren fundamenteel voor deep learning. In Transformers is de standaard residu-update:

hl=hl−1​+fl−1​(hl−1​)

Dit helpt gradiënten door zeer diepe netwerken stromen. Maar residu's zijn niet alleen een truc voor gradiënten. Ze bepalen ook hoe informatie over de diepte wordt geaggregeerd. Als we de recursie uitrollen, krijgen we:

hl​=h1​+∑i=1l−1​fi​(hi​)

Dit betekent dat de verborgen toestand op laag l gewoon de embedding is en een uniform gewogen som van alle eerdere laagoutputs. Elke bijdrage krijgt dus in feite hetzelfde gewicht.

Dat wordt op schaal een probleem. Het paper laat zien dat in PreNorm-architecturen ongewogen residu-opstapeling ervoor zorgt dat de grootte van verborgen toestanden met de diepte toeneemt, ruwweg als O(L). 

Naarmate de residu-stroom groeit, raken outputs van eerdere lagen steeds meer verdund in een grotere lopende som. Zodra een vroege-laagsignaal in deze opstapeling is gemengd, kunnen diepere lagen het niet selectief terughalen. Ze opereren dus alleen op de geaggregeerde toestand.

Dit leidt tot wat het paper het “verdrinkende signaal”-effect noemt. Een sterke empirische aanwijzing voor deze inefficiëntie komt uit laag-pruningstudies: een aanzienlijk deel van de lagen in getrainde modellen kan vaak worden verwijderd met minimale impact op de prestaties. Dit suggereert dat hoewel we de diepte blijven vergroten, modellen een effectief mechanisme missen om die volledig te benutten. In plaats van een hiërarchische redeneerketen te vormen, gedraagt het netwerk zich meer als een redundante estafette waarin vroege signalen geleidelijk worden verdund.

Er is ook een trade-off in de forward pass. Naarmate de residu-stroom groeit, moeten latere lagen mogelijk outputs met grotere amplitude produceren om de geaccumuleerde toestand merkbaar te beïnvloeden. Het paper verbindt dit aan PreNorm-gedrag, waar de grootte van verborgen toestanden monotoon toeneemt met de diepte onder standaard residu-aggregatie.

Residu-verbindingen helpen gradiënten stromen, en de gradiënt met betrekking tot een tussentijdse verborgen toestand kan worden geschreven als:

De identiteitsterm behoudt hier een direct gradiëntpad. Maar de residu's dwingen het forward-aggregatiepad nog steeds om elke eerdere laag te behandelen met een vast gewicht van 1,0. Dat is de structurele beperking die AttnRes probeert op te lossen.

Van additieve residu's naar dieptegewijze attention

De interessantste conceptuele stap in het paper is tijd–diepte-dualiteit. Residu-verbindingen comprimeren informatie over de diepte, zoals recurrente netwerken informatie over de tijd comprimeren. In sequentiemodellering verving attention recursie door elke positie selectief toegang te geven tot eerdere posities. AttnRes past dezelfde overgang toe op netwerkdiepte. In plaats van de volgende verborgen toestand te definiëren als een vaste som over eerdere lagen, laat AttnRes elke laag letten op eerdere laagoutputs:

waar de gewichten αi→l softmax-attentiongewichten over de diepte zijn en optellen tot 1. Die gewichten worden als volgt berekend:

waarbij het paper gebruikt:

Elke laag krijgt een geleerde pseudo-queryvector wl, en die query richt zich op sleutels en waarden die zijn opgebouwd uit eerdere laagoutputs. RMSNorm wordt toegepast op de sleutels zodat lagen met van nature grotere outputgroottes de softmax niet domineren puur vanwege hun schaal.

De kernverschuiving in het paper is dat de standaard residu's worden behandeld als een soort dieptegewijze lineaire attention, terwijl AttnRes dat opwaardeert naar dieptegewijze softmax-attention. In plaats van uniforme opstapeling krijgen we selectieve retrieval over de diepte.

Een klein maar belangrijk implementatiedetail is initialisatie. De auteurs stellen voor om alle pseudo-queryvectoren op nul te initialiseren, zodat de initiële attentiongewichten uniform zijn en het model de training begint als een gelijkgewogen gemiddelde in plaats van een willekeurig bevooroordeelde attentionmechaniek.

Figuur 1: Overzicht van Attention Residuals (Attention Residuals-paper)

Volledige Attention Residuals

Bij Volledige Attention Residuals richt elke laag zich op alle eerdere laagoutputs, zodat informatie selectief over de diepte kan worden geaggregeerd. Dit geeft het model maximale flexibiliteit, zodat een diepere laag de nadruk kan leggen op zijn directe voorganger, de oorspronkelijke embedding of elke eerdere laag als dáár het meest bruikbare signaal zit.

Bij standaardtraining is de geheugenoverhead van Volledige AttnRes kleiner dan het op het eerste gezicht lijkt, omdat veel laagoutputs toch al worden bewaard voor backpropagatie. Maar grootschalige training verandert het plaatje. Zodra je activeringsherberekening en pipeline-parallelisme introduceert, moeten die eerdere outputs expliciet worden bewaard en gecommuniceerd zodat latere lagen erop kunnen letten, en dat wordt kostbaar.

Volledige AttnRes is dus de beste manier om het kernidee te begrijpen, maar het is niet de versie die de meeste teams op schaal zouden willen inzetten.

Block Attention Residuals

Het paper introduceerde het idee van Block Attention Residuals om geheugen- en communicatie-overhead te verminderen. De lagen van het model worden opgedeeld in blokken, zodat binnen elk blok outputs worden gecombineerd met standaard additieve residu-opstapeling, maar over blokken heen voert het model attention uit over blokniveausamenvattingen in plaats van over elke individuele eerdere laag.

Als Bn de set lagen in blok n is, dan wordt de blokrepresentatie gegeven door:

Vervolgens let het model op de embedding b0=h1 eerdere bloksamenvattingen, en de partiële som van het huidige blok naarmate de berekening vordert. Dit reduceert geheugen en communicatie van O(Ld) naar O(Nd), waarbij N het aantal blokken is.

Het paper meldt dat met ongeveer acht blokken het grootste deel van het voordeel van de volledige versie wordt teruggewonnen, en dat de prestatiedelta tussen Volledige AttnRes en Block AttnRes kleiner wordt naarmate de schaal toeneemt. Blokgroottes zoals S=2,4,8 blijven allemaal dicht bij de volledige versie, terwijl veel grovere groeperingen terug richting baselinegedrag bewegen.

De nuttigste les uit het paper is dat je geen volledige dieptegewijze attention over elke laag nodig hebt om het grootste deel van de winst te behalen. Hoewel Block AttnRes dieptegewijze attention computationeel haalbaar maakt, vereist efficiënte inzet op schaal nog steeds zorgvuldige systeemontwerpkeuzes. Het paper introduceert ook verschillende optimalisaties die AttnRes praktisch maken voor training en inferentie in de echte wereld.

Figuur 2: Voorbeeld van cache-gebaseerde pipelinecommunicatie met 4 fysieke rangen en 2 virtuele stages per rang, waarbij gearceerde vakken het einde van AttnRes-blokken aangeven (Attention Residuals-paper)

Hier zijn drie belangrijke systeemsoptimalisaties die van AttnRes een drop-invervanging voor standaard residu's maken:

• Tweedelige berekening: Door de pseudo-query los te koppelen van de verborgen toestand kunnen we queries voor alle lagen binnen een blok batchen. Fase 1 berekent interblok-attention parallel, terwijl Fase 2 sequentiële intrablok-afhankelijkheden afhandelt met een Online Softmax Merge. Dit vermindert de I/O per laag aanzienlijk vergeleken met eerdere multistream-methoden zoals mHC.
• Cross-stage caching: In pipeline-parallelisme elimineren we redundante gegevensoverdrachten door eerder ontvangen blokken lokaal te cachen. Dit levert een V×-verbetering op (waarbij V het aantal virtuele stages is, zoals in Figuur 2) in de piekkosten per transitie, waardoor communicatie volledig kan overlappen met berekening.
• Geheugenefficiënt prefilling: Door blokrepresentaties te sharden over tensor-parallelle apparaten wordt de geheugendruk voor lang-context (128K+) sequenties met ordegroottes verkleind (bijv. van 15GB naar 0,3GB per apparaat).

Resultaten en prestatieanalyse

Het paper valideert AttnRes op meerdere niveaus, zoals schaalwetten, ablatiestudies, trainingsdynamiek en downstream-benchmarks.

Figuur 3: Schaalwetcurves voor Attention Residuals. (Attention Residuals-paper)

Het belangrijkste resultaat komt echter uit de schaalwet-experimenten. Over vijf modelgroottes behalen zowel Volledige AttnRes als Block AttnRes consequent een lagere validatieverlies dan de baseline. Op basis van de gefitte curves bereikt Block AttnRes hetzelfde verlies als een baseline die is getraind met ongeveer 1,25× meer compute. Dit suggereert dat AttnRes niet alleen theoretisch beter is, maar ook reken-efficiënter.

Op downstream-taken verbetert AttnRes ten opzichte van de baseline over alle geëvalueerde benchmarks. Enkele van de grootste winst zit op taken die veel redenering vereisen, zoals getoond in Tabel 1:

Tabel 1: Prestatievergelijking van AttnRes met de baseline, beide na hetzelfde pre-trainingsrecept (Attention Residuals-paper)

De winst is vooral sterk bij meerstapsredeneren, wiskunde en code, wat past bij de kernhypothese van het paper dat, als latere lagen eerdere representaties selectief kunnen terughalen in plaats van een vervaagde som te erven, compositiereasoning moet verbeteren.

Vergeleken met de baseline toont AttnRes gedurende de hele training een lager validatieverlies, meer begrensde outputgroottes over de diepte en meer uniforme gradiëntgroottes over lagen. De baseline lijdt aan het gebruikelijke PreNorm-verdundingspatroon, waarbij de grootte van verborgen toestanden monotoon toeneemt met de diepte, terwijl Block AttnRes een meer gecontroleerd, begrensd patroon oplevert dankzij selectieve aggregatie op blokgrenzen.

Er is ook een bredere bevinding op architectuurniveau dat, in een sweep met vaste compute en vaste parameters, de optimale configuratie voor AttnRes verschuift naar een dieper, smaller model vergeleken met de baseline. Dit suggereert dat AttnRes extra diepte nuttiger kan maken dan onder standaard residu-aggregatie.

Conclusie

Residu's worden vaak gezien als een noodzakelijke trainingstruc, een identity-snelkoppeling die gradiënten in leven houdt. AttnRes suggereert dat dit beeld te beperkt is. Residu-paden zijn ook het mechanisme waarmee informatie over de diepte wordt gerouteerd, en vaste additieve opstapeling is mogelijk te primitief voor zeer diepe modellen.

Attention Residuals gaat minder over het plakken van een bekend probleem en meer over het upgraden van een verwaarloosde ontwerpkeuze. Sequentiemodellering evolueerde van recursie naar attention omdat vaste recursie te beperkend was. AttnRes stelt dat diepte nu klaar kan zijn voor dezelfde overgang.

Het model vertoont nog steeds sterke localiteit, waarbij lagen vaak het meest letten op nabije voorgangers, maar het leert ook niet-triviale overslagschema's, behoudt een persistente weging op de embedding en handhaaft verschillend gedrag tussen pre-attention- en pre-MLP-lagen. 

De praktische les is niet dat alle modellen residu's moeten vervangen door volledige dieptegewijze attention. In plaats daarvan kadert het residu-aggregatie opnieuw als een kernontwerpruimte voor architectuur.