Technische Samenvatting: Korte Venster-Attention Maakt Langetermijngeheugen Mogelijk

Probleemstelling

Moderne Large Language Models (LLM's) die gebaseerd zijn op softmax-attention, vertrouwen op een Key-Value (KV) cache die lineair groeit met de sequentielengte, wat leidt tot verbodsfijne reken- en geheugenkosten voor taken met lange contexten. Omgekeerd behouden lineaire Recurrente Neurale Netwerken (RNN's) en State Space Models (SSM's) constante reken- en geheugenkosten, maar lijden ze vaak onder inferieure herinneringsvermogens in vergelijking met Transformers.

Recente hybride architecturen proberen lokale sliding window attention (SWA) te combineren met globale attention of lineaire RNN-lagen om efficiëntie en prestaties in evenwicht te brengen. De impact van de sliding window-lengte op het samenspel tussen lokaal en langetermijngeheugen blijft echter onderbestudeerd. Specifiek is het onduidelijk hoe de venstergrootte de mogelijkheid van het model beïnvloedt om langetermijnafhankelijkheden te leren wanneer lineaire RNN's aanwezig zijn, en of grotere vensters inherent de prestaties verbeteren of de mogelijkheid van het model om te extrapoleren naar langere contexten belemmeren.

Methodologie

De auteurs onderzoeken hybride architecturen die Sliding Window Attention (SWA)-lagen afwisselen met lineaire RNN-lagen, waarbij specifiek xLSTM (een variant van mLSTM met matrixgeheugen) wordt gebruikt als langetermijngeheugencomponent. De voorgestelde architectuur, genaamd SWAX, wisselt tussen SWA-lagen en xLSTM-lagen in een verhouding van 1:1.

De studie richt zich op twee primaire methodologische bijdragen:

1. Venstergrootte-analyse: De auteurs trainen SWAX-modellen met vaste sliding window-groottes variërend van 128 tot 2048 tokens. Ze evalueren deze modellen op redeneerbenchmarks met korte context (bijv. ARC, Hellaswag, MBPP) en ophaaltaken met lange context (bijv. RULER Needle-In-A-Haystack).
2. Stochastische venstertraining: Om de afweging tussen korte-contextprecisie (bevoordeeld door grotere vensters) en lange-contextextrapolatie (bevoordeeld door kleinere vensters) aan te pakken, stellen de auteurs een trainingsstrategie voor waarbij de sliding window-grootte stochastisch wordt bemonsterd voor elke batch. In hun experimenten wisselt de venstergrootte tussen een korte grootte (128) en een lange grootte (2048) met een specifieke waarschijnlijkheid (0,5 voor 1,4B-modellen, 0,75 voor 7B-modellen). Daarnaast hanteren ze een afkoelingsschema waarbij de stochastische bemonstering wordt verwijderd voor de laatste 10% van de training om de korte-contextprestaties te stabiliseren.

De experimenten worden uitgevoerd op modellen met 1,4 miljard en 7 miljard parameters, getraind op 150 miljard tokens met een vaste sequentielengte van 16k tijdens de pre-training.

Belangrijkste Bijdragen

1. Tegen-intuïtieve bevinding over venstergrootte: Het artikel toont aan dat voor hybride architecturen die SWA en lineaire RNN's combineren, grotere sliding windows de prestaties met lange context kunnen verslechteren. In tegenstelling tot het geloof dat grotere vensters meer context bieden, vinden de auteurs dat grote vensters (bijv. 2048) het model toestaan om te sterk te vertrouwen op de precieze lokale softmax-attention, waardoor de lineaire RNN-lagen onvoldoende worden getraind op langetermijnafhankelijkheden.
2. Mechanisme van het voordeel van korte vensters: Korte vensters (bijv. 128) dwingen het model om voor informatie die buiten het lokale venster valt, te vertrouwen op de lineaire RNN-lagen. Dit "toezichtsignaal" dwingt de lineaire lagen zich te specialiseren in het modelleren van langetermijnafhankelijkheden, waardoor het vermogen van het model om te extrapoleren naar sequentielengtes die ver buiten de trainingsvenster vallen, verbetert.
3. Stochastische trainingsstrategie: De auteurs introduceren een trainingsprocedure waarbij de venstergrootte stochastisch varieert. Deze aanpak dwingt het model om zowel de korte-termijnprecisie van het venster als het langetermijngeheugen van de RNN te benutten, wat effectief werkt als een vorm van dropout op het attention-mechanisme.

Resultaten

• Prestaties met lange context: Modellen getraind met een vast groot venster (2048) vertonen een catastrofale daling in prestaties op ophaaltaken met lange context (bijv. RULER NIAH) wanneer ze worden getest op sequenties langer dan het venster. Daarentegen behouden modellen getraind met kortere vensters (128) een hoge nauwkeurigheid, zelfs bij sequentielengtes van 131k tokens. De stochastische trainingsmethode bereikt prestaties die vergelijkbaar zijn met of beter zijn dan die van modellen met alleen korte vensters op deze taken met lange context.
• Prestaties met korte context: Hoewel korte vensters (128) alleen soms onderpresteren op redeneerbenchmarks met korte context in vergelijking met grotere vensters, herstelt de stochastische trainingsmethode deze prestaties. Op zowel de 1,4B- als de 7B-schaal levert stochastische training korte-contextprestaties op die vergelijkbaar zijn met, of in sommige gevallen beter zijn dan, modellen getraind met een vast 2048-venster.
• Generalisatie: De bevindingen gelden voor verschillende lineaire attention-ruggengraten (xLSTM en Gated DeltaNet) en verschillende hybride configuraties (het afwisselen van SWA met volledige attention-lagen). De stochastische trainingsstrategie presteert consistent beter dan baselines met vaste vensters op zowel korte- als lange-contextbenchmarks.

Betekenis en Claims

Het artikel beweert dat de keuze van de sliding window-grootte in hybride architecturen niet zomaar een afweging is tussen rekenkosten en prestaties, maar een kritieke factor is in hoe het model leert geheugenverantwoordelijkheden te verdelen.

De auteurs stellen dat:

• Korte vensters essentieel zijn voor het trainen van langetermijngeheugen: Door het lokale attention-veld te beperken, wordt het model gedwongen zijn lineaire RNN-lagen te gebruiken voor langetermijnafhankelijkheden, wat leidt tot betere lengte-extrapolatie.
• Stochastische training de afweging oplost: Door de venstergrootte dynamisch te veranderen tijdens de training, leert het model om zowel de hoge-fideliteit lokale attention voor korte contexten als het globale lineaire geheugen voor lange contexten te benutten.
• Te grote afhankelijkheid van lokale attention schadelijk is: Als een model wordt getraind met een groot venster, faalt het om de noodzakelijke mechanismen te leren om informatie te halen buiten dat venster, wat leidt tot slechte generalisatie op langere sequenties.

Het werk suggereert dat hybride architecturen om state-of-the-art prestaties met lange context te bereiken zonder onbegrensde rekenkosten, de trainingsprocedure actief moet voorkomen dat het model te sterk vertrouwt op lokale attention-mechanismen, zelfs als die mechanismen preciezer zijn voor korte-range taken.