Technische Samenvatting: Neural Attention Search Linear (NAtS-L)

Probleemstelling

De kwadratische computationele complexiteit ($O(L^2)$) van standaard softmax-gebaseerde zelf-attentie in Transformers is een significante bottleneck geworden voor scenario's met lange context. Hoewel lineaire aandachtmodellen een veelbelovend alternatief bieden door verleden Key-Value (KV) waarden te comprimeren in een vaste verborgen staat om een lineaire complexiteit ($O(L)$) te bereiken, lijden ze vaak aan beperkte expressiviteit vanwege de beperkingen van deze vaste staat-grootte. Daarentegen bewaren standaard Transformers alle KV-contextinformatie, wat zorgt voor betere langetermijn-retrieval, maar dit gaat ten koste van hoge geheugen- en computationele kosten.

Bestaande hybride benaderingen proberen dit te verzachten door softmax- en lineaire aandachtlagen af te wisselen of door vaste schema's te gebruiken om tussen hen te schakelen. Deze methoden behouden echter vaak vaste structuren die mogelijk niet flexibel genoeg zijn om alle vereiste informatie over variërende contexten heen te vangen, en hun efficiëntie wordt nog steeds beperkt door de aanwezigheid van volledige softmax-lagen. Bovendien vertrouwen veel sparse attention-varianten op vooraf gedefinieerde menselijke heuristieken, waardoor correlaties die niet door hun sparse maps worden gedekt, mogelijk worden genegeerd.

Methodologie

Het artikel stelt Neural Attention Search Linear (NAtS-L) voor, een framework dat zowel lineaire als softmax-aandachtoperaties binnen dezelfde laag toepast, maar op verschillende tokens. In plaats van een vaste architectuur, bepaalt NAtS-L automatisch de optimale aandachtoperatie voor elk tokenblok op basis van de inputcontext.

Kernmechanismen

1. Token-niveau Hybridisatie: De inputsequentie wordt opgedeeld in chunks. Een Attention Score Layer (bestaande uit mean-pooling gevolgd door een lineaire laag) brengt elke chunk in kaart naar een reeks scores die verschillende operatietypes vertegenwoordigen (Softmax Attention versus Gated DeltaNet, een lineaire aandachtvariant). De operatie met de hoogste score wordt aan die chunk toegewezen.
2. Verenigde Zoekruimte: Het framework verenigt niet-lineaire softmax en lineaire aandachtoperaties, waardoor ze binnen een enkele laag kunnen worden verwerkt. Het leert het belang van elk tokenblok te begrijpen, waarbij softmax-aandacht wordt toegepast op tokens met een grote langetermijnimpact (die langetermijn-retrieval vereisen) en lineaire aandacht op tokens met een kortetermijnimpact.
3. Chunk-wijze Berekening: Om hardware-efficiëntie te behouden, maakt NAtS-L gebruik van een chunk-wijze parallel benadering. Binnen een chunk worden zowel lineaire als softmax-correlaties berekend. Het model houdt verborgen staten bij voor beide operatietypes. Wanneer een nieuwe chunk arriveert, voert het model een pseudo-update uit voor beide staten; echter, voor chunks die niet zijn geselecteerd als de doeloperatie, worden de verborgen staten gereset of teruggedraaid naar het begin van de chunk, zodat alleen de relevante KV-waarden worden behouden voor het geselecteerde operatietype.
4. Output Merging: De outputs van de geselecteerde softmax- en lineaire operaties worden individueel genormaliseerd (met behulp van RMS-normalisatie) en vervolgens gecombineerd via een gewogen som. De gewichten voor deze combinatie worden dynamisch gegenereerd uit de query ($Q$) projecties, waardoor het model leert de bijdragen van beide aandachtstypes te balanceren.
5. Gradiënt-gebaseerde Optimalisatie: De selectie van aandachtstypes wordt end-to-end geleerd. Gradiënten worden berekend voor de attention score layer, waardoor het model de operatie-selectie en de modelgewichten gezamenlijk kan optimaliseren zonder dat er hulp-losses of vooraf getrainde modellen nodig zijn.

Belangrijkste Bijdragen

1. Verenigd Framework: De auteurs bieden een verenigde beschrijving voor niet-lineaire softmax en lineaire aandacht, waardoor ze binnen een enkele laag kunnen worden verwerkt en geselecteerd.
2. Adaptieve Architectuur: NAtS-L wordt voorgesteld als een framework dat automatisch de optimale aandachtoperatie bepaalt voor inputcontextinformatie op een token-niveau (chunk-niveau), in plaats van te vertround op vaste laag-wijze ratio's.
3. Efficiëntie en Expressiviteit: Het artikel toont aan dat NAtS-L een sterke balans bereikt tussen prestaties en efficiëntie, waarbij het een betere langetermijncontext-modellering biedt dan pure lineaire modellen, terwijl het lagere computationele latentie behoudt dan volledige softmax Transformers.

Experimentele Resultaten

De auteurs evalueerden NAtS-L op academische schaal pre-training taken met de Fineweb-Edu dataset met modelgroottes van 380M en 800M parameters.

• Taalmodellering & Redeneren: Op zero-shot commonsense reasoning benchmarks (bijv. LAMBADA, PIQA, Hellaswag) presteerden NAtS-L en de hybride variant vergelijkbaar met of iets beter dan pure Gated DeltaNet (GDN), Mamba2 en Transformer baselines.
• Retrieval Taken: Op real-world retrieval taken (bijv. SWDE, SQD, DROP) behaalde NAtS-L Hybrid de beste gemiddelde scores, waarbij het pure lineaire modellen (die moeite hadden met specifieke taken zoals FDA) en pure Transformers (die moeite hadden met SQD) overtrof. Dit onderstreept de robuustheid van het model in het behouden van noodzakelijke contextinformatie.
• Lengte Extrapolatie: Geëvalueerd op contexten tot 65.536 tokens (PG19, NarrativeQA, CodeParrot), behield Nats-L Hybrid de perplexiteit-niveaus, zelfs wanneer de inputcontext de trainingslengte (4096 tokens) overschreed, terwijl pure lineaire hybriden (GDN Hybrid) niet effectief konden extrapoleren.
• Long-Context Benchmarks: Op RULER en LongBench behaalde N-AtS-L Hybrid state-of-the-art prestaties onder de vergeleken modellen, met name bij 8k en 16k contextlengtes.
• Inference Latentie: N-AtS-L Hybrid vertoonde een 5.4x versnelling in pre-filling vergeleken met Transformers en een 2.3x versnelling in decoding bij een 128k contextlengte, terwijl het slechts 1.66x trager was dan pure GDN-modellen.

Betekenis en Claims

Het artikel claimt dat N-AtS-L een flexibele hybride architectuur biedt die tokens adaptief toewijst aan ofwel lineaire of softmax-aandacht op basis van hun specifieke informatiebehoeften. Door deze distributie end-to-end te leren, vermijdt het model de rigiditeit van vaste hybride structuren en de inefficiëntie van volledige softmax-aandacht.

De auteurs benadrukken dat deze aanpak het model in staat stelt om:

• Noodzakelijke langetermijncontext-gerelateerde informatie te behouden via softmax-aandacht waar nodig.
• De algehele computationele kosten te verlagen via lineaire aandacht voor tokens met een kortetermijnimpact.
• Een betere nauwkeurigheid-efficiëntie trade-off te bereiken zonder dat er vooraf getrainde modellen of aanvullende beperkingen op de ratio van de aandachtstypes nodig zijn.

Het werk suggereert dat een token-niveau mengeling van lineaire en softmax-aandacht een levensvatbare weg is naar schaalbare, lang-context taalmodellen die geen concessies doen aan retrieval-capaciteiten of inferentie-snelheid.