SinkTrack: Attention Sink based Context Anchoring for Large Language Models

Probleemstelling

Grote Taalmodellen (LLMs) en Multimodale LLMs (MLLMs) kampen met twee fundamentele en hardnekkige beperkingen die hun inzetbaarheid in kritieke toepassingen belemmeren:

1. Hallucinaties: Het genereren van inhoud die feitelijk onjuist is of niet overeenkomt met de input (bijvoorbeeld het beschrijven van objecten die niet in een afbeelding aanwezig zijn).
2. Contextvergetelheid: Het verlies van de oorspronkelijke instructies of context tijdens het genereren van lange teksten of complexe redeneringen.

De auteurs identificeren attentie-drift als de primaire oorzaak van deze problemen. Tijdens de autoregressieve generatie verschuift de aandacht van het model progressief naar de nieuw gegenereerde tokens, terwijl de aandacht voor de initiële tokens (die de cruciale context, instructies of beeldinformatie bevatten) drastisch afneemt.

Methodologie: SINKTRACK

De kern van de oplossing is gebaseerd op een intrinsiek kenmerk van LLMs: de Attention Sink. Hoewel de aandacht voor de meeste initiële tokens afneemt, behoudt het zeer eerste token (meestal het <BOS>-token, "Beginning of Sequence") consistent een zeer hoge aandachtswaarde gedurende de hele generatie, ondanks dat het semantisch vaak leeg is.

SINKTRACK transformeert deze passieve "attention sink" in een actief mechanisme voor context-ankering. De methode omvat de volgende stappen en evolutie:

1. Concept: Het <BOS>-token wordt gebruikt als een informatieanker. Sleutelcontextuele kenmerken (zoals visuele features van een afbeelding of de gebruikersinstructie) worden in de representatie van dit token geïnjecteerd.
2. Ontwikkelingsfasen:
   • Hard Injection: Direct vervangen van de Value-vector van <BOS> in de KV-cache door een samengevoegde informatieve vector. Dit leidde tot een modelcrash, wat aantoont dat de interne rekenstroom van het model intact moet blijven.
   • Soft Injection: Een gewogen som van de originele <BOS>-state en de informatieve vector. Hoewel dit werkte, vereiste het handmatig afgestemde hyperparameters (injectiestrength) en introduceerde het ruis door alle context gelijkmatig te aggregeren.
   • SINKTRACK (Definitieve Oplossing): Een dual-track cross-attention mechanisme dat de beperkingen van soft injection overwint.
     • Track 1 (Adaptieve Injectie): Voor het <BOS>-token wordt een cross-attention mechanisme gebruikt. Het <BOS>-token fungeert als Query (Q), terwijl de Key (K) en Value (V) worden afgeleid van de externe, contextuele informatie (bijv. beeldfeatures). Dit stelt het model in staat om adaptief relevante informatie uit de context op te halen zonder handmatige tuning.
     • Track 2 (Behoud van Native Stroom): Alle andere tokens (de reguliere tokens) ondergaan de standaard causale self-attention. Dit zorgt ervoor dat de oorspronkelijke rekenstroom van het LLM ongestoord blijft.
     • Samenstelling: De uiteindelijke output is een concatenatie van de bijgewerkte <BOS>-representatie en de reguliere token-representaties.
3. Efficiëntie: De injectie is een eenmalige, lichte operatie die plaatsvindt voordat de generatie begint. De injectie wordt opgeslagen in de KV-cache, wat betekent dat er geen extra rekenkosten zijn tijdens de autoregressieve stappen. Het is dus training-vrij en plug-and-play.

Belangrijkste Bijdragen

• Nieuw Mechanisme: SINKTRACK is de eerste methode die de synergie tussen "attention drift" en "attention sink" benut om contextverlies tegen te gaan door het <BOS>-token te gebruiken als een dynamisch informatieanker.
• Training-vrij en Efficiënt: Het vereist geen fine-tuning (wat het risico op catastrofale vergeetvermogen of modelinstorting elimineert) en voegt verwaarloosbare inferentie-overhead toe.
• Generaliseerbaarheid: De methode werkt over verschillende architecturen (decoder-only), schalen (van 3B tot 12B parameters) en modaliteiten (tekst en multimodaal).
• Adaptiviteit: In tegenstelling tot eerdere benaderingen met statische wegingen, gebruikt SINKTRACK cross-attention om context adaptief op te halen, wat de noodzaak voor handmatige hyperparameter-tuning elimineert.

Resultaten

Experimenten op diverse benchmarks tonen consistente en aanzienlijke verbeteringen ten opzichte van directe prompting en Chain-of-Thought (CoT) redenering:

• Multimodale Taken:
  • Op de M3CoT benchmark (multimodaal, multi-hop redenering) steeg de nauwkeurigheid van Qwen2.5-VL-7B met +22,8% ten opzichte van CoT.
  • Op POPE (hallucinatiedetectie) verbeterde het zowel de nauwkeurigheid als de Macro-F1 score aanzienlijk, wat aantoont dat het minder niet-bestaande objecten genereert.
• Tekstgebaseerde Taken:
  • Op SQuAD2.0 (leesbegrip) verbeterde Llama3.1-8B-Instruct met +21,6%.
  • Op QuAC (lange context conversatie) behield SINKTRACK een sterke prestatie terwijl de contextlengte toenam, terwijl baseline-methoden degradeerden door attention drift. Bij de langste context (7.354 tokens) behaalde Llama3.1-8B een nauwkeurigheid van 56,2% (+7,1% t.o.v. CoT).
• Kwalitatieve Analyse: SINKTRACK leidt tot logischere redeneerketens en voorkomt tegenstrijdigheden die vaak voorkomen bij CoT (waarbij de conclusie in strijd is met eerdere stappen).

Betekenis en Impact

SINKTRACK biedt een elegante, efficiënte oplossing voor de fundamentele beperkingen van huidige LLMs zonder de noodzaak van kostbare training of complexe externe modules. Door de bestaande "attention sink" te herprogrammeren als een actief informatieanker, lost het het probleem van contextverlies en hallucinaties op op een manier die inherent robuust is voor verschillende modelgroottes en -types. Dit maakt het een veelbelovende techniek voor het verbeteren van de betrouwbaarheid van AI-systemen in real-world toepassingen, zoals medische diagnose, juridische analyse en complexe vraag-antwoord systemen.