💬 NLP

Learning how to Forget: Fine-tuning for Long-Context Sparse Attention

Dit artikel introduceert een fine-tuning methode voor transformer-modellen die efficiënte long-context inferentie met sparse attention op gemiddelde hardware mogelijk maakt door modellen te laten co-adapteren met diverse KV cache-policies, wat vaak exact attention-benaderingen overtreft.

Oorspronkelijke auteurs: Matthias Seeger, Zeyu Zhang, Vihang Patil, Konstantinos Benidis, Sebastian Schelter

Gepubliceerd 2026-08-21

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Matthias Seeger, Zeyu Zhang, Vihang Patil, Konstantinos Benidis, Sebastian Schelter

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Moderne kunstmatige intelligentiesystemen die menselijke tekst genereren, vertrouwen op een mechanisme dat fungeert als een kortetermijngeheugen, waardoor ze kunnen onthouden wat er eerder in een gesprek of een lang document is gezegd. Dit geheugen wordt opgeslagen in een digitale buffer die groter wordt bij elk nieuw woord dat het systeem verwerkt. Om deze systemen over lange stukken tekst goed te laten functioneren, moet dit geheugen enorm zijn, maar de computerhardware die nodig is om het te huisvesten is duur en beperkt. Wanneer de geheugenbuffer vol raakt, moet het systeem beslissen welke oude stukken informatie het wegwerpt om ruimte te maken voor de nieuwe. Als het de verkeerde informatie weggooit, verliest het systeem zijn vermogen om te redeneren of vragen accuraat te beantwoorden. Dit creëert een moeilijke afweging: het klein houden van het geheugen bespaart geld en zorgt ervoor dat het systeem op standaardapparatuur kan draaien, maar het brengt het risico met zich mee dat de context die nodig is om slim te zijn, verloren gaat.

Onderzoekers proberen dit al lang op te lossen door het systeem te leren selectief te zijn over wat het behoudt, een proces dat bekende staat als 'sparse attention'. Echter, een nieuwe studie onthult een kritiek gebrek in hoe deze systemen tot nu toe zijn getraind. De meeste bestaande methoden trainen de AI met een perfect, onbeperkt geheugen en dwingen het vervolgens om later met een beperkt geheugen te werken. De onderzoekers ontdekten dat deze aanpak faalt omdat de AI nooit heeft geleerd hoe het moet functioneren onder de specifieke beperkingen waarmee het in de praktijk geconfronteerd zou worden. Door het model vanaf het begin te trainen om bewust te vergeten en zich aan te passen aan een vaste geheugengrootte, toonde het team aan dat het systeem aanzienlijk beter kan presteren dan diegenen die met onbeperkte middelen zijn getraind, zelfs wanneer het draait op één enkele, redelijk krachtige computerchip.

Het team, onder leiding van wetenschappers van Amazon Web Services en de Universiteit van Amsterdam, heeft een nieuwe manier ontwikkeld om deze grote taalmodellen te finetunen. In plaats van massale supercomputers te gebruiken om een perfect geheugen te simuleren, hebben ze de modellen geleerd om mee te adapteren met een specifief geheugenbeheerbeleid. Stel je een bibliothecaris voor die getraind is om boeken te organiseren in een bibliotheek met oneindige planken, om er later voor te worden verteld dat hij in een kleine kamer met slechts één plank moet werken. De bibliothecaris die getraind is in de grote bibliotheek zou waarschijnlijk moeite hebben met het prioriteren van wat hij moet bewaren in de kleine kamer. In tegenstelling hiertoe traint de methode die in dit artikel wordt voorgesteld de bibliothecaris direct in de kleine kamer, waardoor hij precies leert welke boeken hij moet houden en welke hij moet weggooien op basis van de regels van die specifieke ruimte. Dit stelt het model in staat om het ritme van zijn eigen beperkingen te leren kennen, in plaats van te proberen de gewoontes van het hebben van te veel ruimte te afleren.

De onderzoekers testten deze aanpak op een model met vier miljard parameters, een omvang die substantieel maar beheersbaar is. Ze voerden hun experimenten uit op een enkele grafische kaart met 40 gigabyte aan geheugen, een opstelling die betaalbaar is voor veel organisaties in vergelijking met de clusters van tientallen kaarten die vereist zijn door eerdere methoden. Ze vergeleken deze nieuwe trainingsmethode met de standaardmethode, die een techniek genaamd 'sequence parallelism' gebruikt om de geheugenbelasting over meerdere dure apparaten te verdelen. De resultaten lieten zien dat de modellen die met de nieuwe methode zijn getraind, vaak beter presteerden dan de standaardmodellen, vooral wanneer de taak vereiste dat het systeem specifieke, beknopte antwoorden genereerde in plaats van lange, uitweidende teksten. In verschillende tests met complexe vragen en data-extractie produceerde de standaardmethode outputs die veel te lang waren en vol stonden met willekeurige, zinloze getallen, terwijl de nieuwe methode leerde om op het juiste moment te stoppen en de correcte enkele waarde te verstrekken.

Een cruciaal onderdeel van dit succes was het verbeteren van de "heavy-hitter oracle", een populaire strategie om te beslissen welke informatie behouden moet worden. Deze strategie werkt door bij te houden aan welke stukken informatie het model in de loop van de tijd de meeste aandacht besteedt. De onderzoekers ontdekten dat de originele versie van deze strategie traag en inefficiënt was. Ze hebben de onderliggende code herschreven om veel sneller te werken, waardoor het systeem deze belangscores kan berekenen zonder het hele proces te vertragen. Deze optimalisatie betekende dat het systeem slimme beslissingen kon nemen over wat het moest vergeten in realtime, zonder de enorme rekenkracht te vereisen die gewoonlijk bij dergelijke taken hoort. Het team heeft ook een nieuwe open-source softwarebibliotheek uitgebracht om deze technieken beschikbaar te maken voor anderen, met als doel de drempel te verlagen voor iedereen die lang-contextuele AI-systemen wil bouwen zonder een fortuin aan hardware nodig te hebben.

De studie benadrukt dat de manier waarop een model wordt getraind net zo belangrijk is als de hardware waarop het draait. Wanneer de onderzoekers het model dwongen om met exact dezelfde geheugenbeperkingen te trainen als waar het tijdens het gebruik mee te maken krijgt, leerde het om effectief met die beperkingen om te gaan. In één specifieke test met JSON-data faalde de standaardmethode volledig, waarbij het niet in staat was de juiste datapunten te vinden, terwijl de nieuwe methode erin slaagde deze ongeveer de helft van de tijd te identificeren. Dit suggereert dat het vermogen om lange contexten aan te kunnen niet alleen een kwestie is van meer geheugen hebben, maar van het systeem leren hoe het het geheugen dat het heeft, moet beheren. De bevindingen wijzen erop dat voor veel toepassingen de meest effectieve weg vooruit niet is om grotere computers te bouwen, maar om de software efficiënter te leren omgaan met de middelen die het al bezit.

Technisch Overzicht: Leren hoe te vergeten: Fine-tuning voor Long-Context Sparse Attention

1. Probleemstelling

Moderne Large Language Models (LLM's) vereisen het vermogen om extreem lange contexten te verwerken voor taken zoals multi-turn conversaties, toolgebruik en chain-of-thought redeneren. Terwijl standaard transformer-implementaties kwadratisch schalen in rekenkracht en lineair in geheugen met de contextbreedte, zijn sparse attention-methoden voorgesteld om lineaire tijd en constante geheugenschaalbaarheid te bereiken. Deze methoden maken gebruik van een vastgelegde Key-Value (KV) cache waarbij informatie wordt verwijderd op basis van specifieke beleidsregels (bijv. Heavy-Hitter Oracle, H2O) zodra de cache vol is.

Er bestaat echter een aanzienlijke kloof tussen inferentie en training op dit gebied:

Trainingsbeperkingen: Het fine-tunen van modellen met sparse attention-beleidsregels is bijzonder moeilijk op gemiddelde hardwarebudgetten. Standaard automatische differentiatie (autograd) vereist het opslaan van de volledige computationele graaf, inclusocief KV-cachebuffers voor elk deel (chunk) van de sequentie. Voor lange contexten overschrijdt deze geheugeneis de beschikbare GPU-geheugencapaciteit, wat vaak dure multi-GPU-opstellingen met sequentie- of contextparallelisme (SP/CP) noodzakelijk maakt.
Co-adaptatiekloof: Modellen die getraind zijn met exacte attention (met behulp van SP/CP) presteren vaak aanzienlijk slechter wanneer ze worden ingezet met sparse attention-inferentie. Dit komt omdat het model nooit is blootgesteld aan de specifieke geheugenbeperkingen en de verwijderingslogica die het tijdens de inferentie zal tegenkomen.
Implementatie-inefficiëntie: Vooraanstaande sparse attention-beleidsregels zoals H2O worden onderbenut in state-of-the-art (SotA) inferentiebibliotheken (bijv. vLLM) omdat de huidige implementaties traag zijn en geen geoptimaliseerde kernels missen voor de specifieke vereisten van sparse attention (bijv. het teruggeven van gesommeerde attention-gewichten).

2. Methodologie

Het paper stelt een nieuwe fine-tuning methode voor waarmee transformer-modellen kunnen co-adapteren met willekeurige KV-cache beleidsregels op gemiddelde hardware (gedemonstreerd op een enkele Nvidia A100 40GB GPU).

2.1 Geneste Activation Checkpointing met CPU Offloading

Om de geheugenbottleneck bij het trainen op lange sequenties ( $N \gg N_C$ , waarbij $N_C$ de cachegrootte is) aan te pakken, gebruikt de auteur een geneste activation checkpointing strategie:

Outer Loop (Layers): De backward pass iterereert over de modellagen.
Inner Loop (Cells): De sequentie wordt verdeeld in "cells" (groepen chunks). Voor elke cell wordt een forward pass uitgevoerd om de noodzakelijke KV-cachebuffers naar het CPU-geheugen op te slaan.
Replay Mechanisme: In plaats van te differentiëren door de KV-cache beleidsregel (die vaak niet-differentieerbaar of duur is), logt de forward pass alle verwijderingsbeslissingen. Tijdens de backward pass worden deze beslissingen geherhaald ("replayed") met "replay caches" die de forward verwijderingslogica nabootsen zonder dat er gradiëntberekening door de beleidsregel zelf nodig is.

2.2 Delta Encoding voor KV Cache Buffers

De meest cruciale geheugenoptimalisatie betreft de manier waarop KV-cachebuffers binnen de autograd-graaf worden opgeslagen.

Lineaire Recurrentie: De auteurs observeren dat de KV-cachebuffer voor een nieuwe chunk ($keys'$) een lineaire transformatie is van de vorige buffer ($keys$), waarbij alleen de waarden van de nieuwe tokens en de specifieke slots die door de verwijderingsbeleidsregel worden overschreven, verschillen.
Delta Opslag: In plaats van de volledige $keys$ en $values$ tensoren voor elke chunk in de computationele graaf op te slaan, slaat de methode alleen de delta (het verschil) en de indices van de overschreven slots op.
Autograd Saved Tensors Hooks: Gebruikmakend van PyTorch's autograd_saved_tensors_hooks, implementeren de auteurs een "packing" mechanisme. Tijdens de forward pass worden volledige tensoren vervangen door referenties naar (index, delta) paren. Tijdens de backward pass worden de volledige tensoren gereconstrueerd ("unpacked") vanuit de delta's en opeenvolgende buffers. Dit vermindert de geheugenvoetafdruk van de autograd-graaf met een factor $k$ (het aantal chunks per cell), waardoor deze op het niveau van de inferentie komt te liggen.

2.3 H2O Beleidsverbeteringen

Het paper introduceert specifieke verbeteringen aan de Heavy-Hitter Oracle (H2O) beleidsregel:

Batch Onafhankelijkheid: De oorspronkelijke H2O-implementatie aggregeerde scores over de batch-dimensies. De nieuwe methode staat toe dat verwijderingsbeslissingen onafhankelijk worden genomen voor elke batch-positie.
Normalisatie: Een genormaliseerde H2O-score is geïntroduceerd om bias naar tokens die langer in de cache aanwezig zijn te voorkomen.
Kernel Optimalisatie: Om efficiënte training en inferentie mogelijk te maken, leveren de auteurs Triton-code die gesommeerde attention-gewichten teruggeeft naast een FlashInfer Scaled Dot Product Attention (SDPA) kernel. Dit stelt H2O in staat om efficiënt te draaien zonder terug te vallen op trage, naïeve SDPA-implementaties.

3. Belangrijkste Bijdragen

Nieuwe Fine-Tuning Methode: Een methode om transformer-modellen te fine-tunen met sparse attention en willekeurige KV-cache beleidsregels op middelen die vergelijkbaar zijn met sparse attention inferentie (bijv. een enkele A100). Het combineert geneste activation checkpointing, CPU offloading en delta encoding via autograd hooks.
H2O Implementatie & Optimalisatie: Methodologische verbeteringen aan de H2O-beleidsregel (batch onafhankelijkheid, normalisatie) en een hoogwaardige implementatie met speciale SDPA-kernels die gesommeerde attention-gewichten teruggeven, wat een belangrijke stap is richting latenties die competitief zijn met SotA inferentiebibliotheken.
Uitgebreide Evaluatie: Experimenten op de Helmet benchmark suite (64k en 128k context windows) tonen aan dat hoewel de prestaties variëren per dataset, de voorgestelde methode modellen in staat stelt om te co-adapteren met KV-cache beleidsregels, waarbij ze vaak beter presteren dan modellen die getraind zijn met sequentieparallelisme wanneer ze worden geëvalueerd onder sparse attention inferentie.
KeysAndValues Library: De release van een open-source bibliotheek (awslabs/keys_values) die performante, gemakkelijk te gebruiken code biedt voor long-context inferentie en fine-tuning met sparse attention.

4. Experimentele Resultaten

De auteurs evalueerden een Qwen3-4B-Instruct model op de Helmet benchmark over 64k en 128k contextlengtes.

Vergelijking: Modellen werden getraind met:
- SP: Sequence Parallelism met exacte attention (SotA baseline).
- US: De voorgestelde methode met de specifieke KV-cache beleidsregel in werking.
- NO: De basis checkpoint zonder fine-tuning.
Bevindingen:
- Gemengde resultaten op Table 1 Datasets: Voor de datasets in Tabel 1 waren de resultaten gemengd en inconclusief. De SP-baseline (exacte attention) behaalde de beste resultaten voor nq en hotpot qa. Voor trivia qa presteerde de basis checkpoint zonder fine-tuning (no) het best, wat aangeeft dat fine-tuning in dit specifieke geval niet hielp. De voorgestelde methode (us) presteerde beter dan SP op pop qa.
- Sterke prestaties op Table 2 Datasets: Voor de datasets in Tabel 2 presteerde us sterk beter dan zowel sp als no.
- Foutmodi van SP: SP-getrainde modellen vertoonden een foutmodus waarbij de outputs excessief lang waren en willekeurige onzin bevatten wanneer ze werden onderworpen aan sparse attention verwijdering, waarschijnlijk omdat ze nooit getraind waren om de "vergetelheid"-mechanisme te hanteren. In contrast leerde us wel correct te stoppen en gaf meestal een enkel getal of het correcte doelwit uit.
- Beleid Variantie: Resultaten varieerden aanzienlijk per dataset. Hoewel H2O-varianten goed presteerden op verschillende taken, presteerde de lastrec (lr) beleidsregel beter dan H2O op datasets zoals nq en hotpot qa. Daarnaast behaalde de SP-baseline (exacte attention) voor nq en hotpot qa vaak de beste resultaten, terwijl voor pop qa en andere datasets in Tabel 2 de US-methode met sparse attention-beleidsregels superieur was.
- Efficiëntie: De H2O-implementatie met de nieuwe kernel-ondersteuning vertegenwoordigt een grote verbetering in efficiëntie en is een grote stap richting latenties die competitief zijn met SotA-bibliotheken, hoewel de auteurs erkennen dat de huidige sparse attention latentie nog steeds aanzienlijk hoger is dan methoden die gebruikmaken van sequentie- of contextparallelisme.

5. Betekenis en Claims

Het paper claimt dat sparse attention in de praktijk onderbenut is, voornamelijk vanwege latentieproblemen en de moeilijkheid om modellen te fine-tunen om te adapteren aan specifieke verwijderingsbeleidsregels.

Haalbaarheid: De auteurs demonstreren dat hoogwaardige fine-tuning voor long-context sparse attention haalbaar is op gemiddelde hardware, wat de noodzaak wegneemt voor enorme gedistribueerde trainingsclusters uitsluitend voor het adapteren aan een cache-beleid.
Noodzaak van Co-adaptatie: De resultaten suggereren sterk dat modellen moeten worden getraind onder dezelfde beperkingen (verwijderingsbeleidsregels) die ze tijdens de inferentie zullen ervaren om prestatievermindering te voorkomen.
Toekomstige Richting: Het paper betoogt dat door de KV-cache structuur te vereenvoudigen en de SDPA kernel-vereisten te verduidelijken (bijv. het teruggeven van gesommeerde gewichten), de community sparse attention beter kan integreren in SotA inferentiebibliotheken.
Beperkingen: De auteurs erkennen dat hoewel hun methode de latentie verbetert, de sequentiële aard van sparse attention-beslissingen nog steeds tot hogere latenties kan leiden vergeleken met RingAttention of Sequence Parallelism voor zeer grote hardwarebudgetten. Ze zijn van plan dit in toekomstig werk aan te pakken via kernel fusie en asynchrone CPU offloading.

Samenvattend biedt het paper een praktisch pad om "te leren hoe te vergeten", waardoor modellen kunnen specialiseren voor long-context inferentie op toegankelijke hardware. Hoewel de prestaties ten opzichte van sequentieparallelisme variëren per dataset, presteert de methode vaak beter dan exact-attention training in scenario's die aanpassing aan specifieke sparse attention-beleidsregels vereisen.