Technische Samenvatting: Multi-Resolutie Residuele Routering voor Diepe Decoder-Only Transformers

Probleemstelling

Standaard decoder-only Transformers vertrouwen op PreNorm residuele verbindingen waarbij sublaag-updates worden geaggregeerd met vaste eenheidsgewichten. Hoewel deze structuur stabiel trainen mogelijk maakt, legt het een rigide accumulatieregel op die de bijdrage van individuele lagen kan verwateren naarmate de netwerkdiepte toeneemt. Recente ontwikkelingen, zoals Attention Residuals en Block Attention Residuals (Block AttnRes), pakken dit aan door vaste accumulatie te vervangen door geleerde, inhoudsafhankelijke routering over voorgaande outputs. Block AttnRes verbetert de efficiëntie specifiek door meerdere lagen te comprimeren naar blok-niveau residuele samenvattingen ($C_b$).

De auteurs identificeren echter een kritieke beperking in Block AttnRes: een enkele bloksamenvatting vertegenwoordigt slechts de laagfrequente (DC) component van de residuele trajectorie binnen een blok. Door alleen de som van de updates ($\sum u_{b,i}$) op te slaan, verwerpt het mechanisme waardevolle directionele structuur, zoals:

1. Onbalans tussen sublagen: Of het blok wordt gedomineerd door Attention- of MLP-updates.
2. Temporele dynamiek: Of vroege updates binnen het blok in een andere richting wijzen dan late updates.

De auteurs hypothetiseren dat naarmate de lengte van de residuele trajectorie groeit, deze verloren gegane directionele informatie steeds crucialer wordt voor effectieve diepte-wijze routering.

Methodologie: WAV v1

De auteurs stellen WAV v1 (Waveform Attention v1) voor, een lichtgewicht, multi-resolutie uitbreiding van Block AttnRes. In plaats van een blok uitsluitend te representeren door de totale som, vergroot WAV v1 elke blokrepresentatie uit met twee zero-sum directionele detailbases, terwijl de originele Attention- en MLP-modules behouden blijven.

1. Multi-Resolutie Representatie

Voor een blok $b$ dat $m$ sublaag-updates $\{u_{b,i}\}_{i=1}^m$ bevat, berekent WAV v1:

• State Basis ($C_b$): De standaard bloksom, die de laagfrequente staat vertegenwoordigt.
  $$C_b = \sum_{i=1}^m u_{b,i}$$
• Phase Basis ($D^{phase}_b$): Een detailbron die Attention- versus MLP-updates contrasteert. Het wijst een teken van $+1$ toe aan Attention-updates en $-1$ aan MLP-updates.
  $$D^{phase}_b = \sum_{i=1}^m a_i u_{b,i}, \quad \text{waarbij } a_i \in \{+1, -1\}$$
• Split Basis ($D^{split}_b$): Een detailbron die de eerste helft contrasteert met de tweede helft van het blok. Het wijst $+1$ toe aan de eerste helft en $-1$ aan de tweede helft.
  $$D^{split}_b = \sum_{i=1}^m r_i u_{b,i}, \quad \text{waarbij } r_i \in \{+1, -1\}$$

Deze bases worden online geaccumuleerd terwijl de updates worden geproduceerd, wat een verwaarloosbare computationele overhead met zich meebrengt.

2. Routeringsmechanisme

De depth-wise softmax mixer ontvangt een uitgebreide bronpool. Voor Block AttnRes is de pool $\{e, C_0, \dots, C_{b-1}, PC\}$. Voor WAV v1 wordt deze uitgebreid met de detailbases voor elk voltooide blok en het huidige gedeeltelijke blok:
$$S_{WAV} = \{e, C_0, \tilde{D}^{phase}_0, \tilde{D}^{split}_0, \dots, C_{b-1}, \tilde{D}^{phase}_{b-1}, \tilde{D}^{split}_{b-1}, PC, \tilde{P}^{phase}, \tilde{P}^{split}\}$$

3. Stabiliteitsmechanismen

Om destabilisatie tijdens de vroege fase van training te voorkomen, hanteert WAV v1 twee conservatieve strategieën:

• Negatieve Detail Bias: Detailbronnen worden geïnitialiseerd met een sterke negatieve bias ($\beta_D = -2.0$), wat ervoor zorgt dat het model zich bij initialisatie gedraagt als standaard Block AttnRes en details pas gebruikt wanneer ze nuttig zijn.
• Detached RMS Matching: Detail-tensoren worden geschaald ten opzichte van hun corresponderende bloksamenvatting ($C$) met behulp van een gedetacheerde RMS-ratio, afgekapt op een maximale schaalfactor $\rho$. Dit voorkomt dat details domineren simpelweg door grotere ruwe magnitudes.
• Conservatieve Readout: Standaard leest de finale predictie-mixer alleen de embedding en $C$-bronnen, en niet de detailbronnen, wat de stabiliteit verder waarborgt.

Belangrijkste Bijdragen

1. Formulering: Het artikel kadert blok-residuele routering als een multi-resolutie representatieprobleem, waarbij onderscheid wordt gemaakt tussen de laagfrequente staat-informatie ($C_b$) en de intra-blok trajectorie-informatie (directionele details).
2. Algoritme: Introductie van WAV v1, een drop-in uitbreiding die bloksamenvattingen aanvult met phase en split detailbases zonder de kern van de Attention- of MLP-functies te wijzigen.
3. Empirisch Bewijs: Voorlopige schalingsexperimenten tonen aan dat de bruikbaarheid van multi-resolutie routering diepte-afhankelijk is, en pas significant voordelig wordt in diepere architecturen.

Experimentele Resultaten

De auteurs evalueerden 12-, 24- en 48-laagse GPT decoder-only modellen op karakter-niveau datasets TinyStories en Text8.

• Diepte-afhankelijke Trend:
  • 12 Lagen: WAV v1 presteert onder Block AttnRes, wat suggereert dat directionele details redundant of ruisachtig zijn in ondiepe netwerken.
  • 24 Lagen: WAV v1 wordt competitief, met lichte verbeteringen op TinyStories en vergelijkbare prestaties op Text8.
  • 48 Lagen: WAV v1 bereikt de beste validatieverlies op beide datasets en presteert aanzienlijk beter dan Block AttnRes, ReZero en LayerScale.
• Kwantitatieve Winsten (48 Lagen):
  • TinyStories: Verminderde de validatieverlies van 0.4960 (Block AttnRes) naar 0.4738 (WAV v1), een reductie van 2,19% in PPL.
  • Text8: Verminderde de validatieverlies van 0.9363 naar 0.9305, een reductie van 0,57% in PPL.
• Trainingsdynamiek: Bij 48 lagen scheidt WAV v1 zich vroeg in de training van Block AttnRes en behoudt een persistent voordeel, wat aangeeft dat de verbetering geen laat-stadium artefact is maar een fundamentele verschuiving in de optimalisatietrajectorie.

Betekenis en Claims

Het artikel claimt dat diepe residuele routering niet alleen profiteert van het selecteren welk blok te lezen, maar ook van het routeren van de interne directionele structuur van elk blok. De resultaten suggereren dat naarmate de residuele trajectorie langer wordt, de "vorm" van de update-trajectorie (bijv. dominantie van aandacht of vroege- versus late dynamiek) een informatieve signaal wordt voor de diepte-wijze mixer.

De auteurs positioneren WAV v1 als een complementaire benadering van bestaande methoden:

• In tegenstelling tot ReZero of LayerScale, die de magnitude van updates moduleren, verandert WAV v1 de representatie die beschikbaar is voor de router.
• In tegenstelling tot Delta Attention Residuals, die routeren over individuele delta's (hoge kosten), biedt WAV v1 een blok-niveau, multi-resolutie benadering die efficiëntie balanceert met structurele rijkdom.

De auteurs concluderen dat toekomstige evaluaties van residuele routeringsmethoden rekening moeten houden met diepte, aangezien de voordelen van multi-resolutie informatie gekoppeld kunnen zijn aan de lengte en structuur van de residuele trajectorie. Ze erkennen huidige beperkingen, waaronder het gebrek aan standaarddeviatie over verschillende seeds, de restrictie tot kleine character-level modellen, en het gebruik van hand-ontworpen detailbases, wat suggereert dat toekomstig werk de exploratie van leerbare bases en grotere schaal token-level modellen moet verkennen.