LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models

Titel

LoSA: Locality Aware Sparse Attention voor Block-Wijze Diffusie Taalmodellen

1. Het Probleem: De "KV Inflation" in Diffusie Taalmodellen

Block-wise Diffusion Language Models (DLMs) genereren tekst in blokken van meerdere tokens tegelijk, in tegenstelling tot autoregressieve modellen die tokens sequentieel genereren. Hoewel dit voordelen biedt voor complexe redeneertaken, lopen DLM's vast in een prestatieprobleem bij lange contexten: de KV-cache (Key-Value cache) wordt een bottleneck.

• Het probleem van naïeve sparse attention: Bestaande methoden voor sparse attention (ontwikkeld voor autoregressieve modellen) proberen de hoeveelheid te laden KV-gegevens te verminderen door voor elke query slechts een subset van de cache te selecteren.
• KV Inflation: In een DLM-blok selecteren verschillende queries (tokens binnen hetzelfde blok) vaak verschillende subsets van de KV-cache. Omdat de attention-mechanisme geheugengebonden is, wordt de totale kosten bepaald door de unie van alle geselecteerde posities. Zelfs als elke query slechts $k$ posities kiest, kan de unie van alle $B$ queries in het blok veel groter zijn (tot $B \times k$). Dit leidt tot een "KV Inflation", waarbij het geheugentransport niet daalt zoals verwacht, waardoor de versnelling door sparsiteit teniet wordt gedaan.

2. De Observatie: Localiteit van Representatieveranderingen

De auteurs identificeren een unieke structuur in block-wise diffusie die ontbreekt in autoregressieve decoding: Localiteit van Representatieveranderingen over de denoising-stappen heen.

• Tussen twee opeenvolgende denoising-stappen ($t-1$ en $t$) ondergaat slechts een klein deel van de tokens in het blok significante veranderingen in hun verborgen staat (hidden states). Deze worden Active tokens genoemd.
• De overgrote meerderheid van de tokens, de Stable tokens, vertoont nauwelijks veranderingen in hun query-vector. Hun attention-patronen en output blijven dus nagenoeg constant.
• Deze observatie wordt visueel onderbouwd met Mean Squared Error (MSE) metingen tussen stappen, waarbij blijkt dat de veranderingen sterk gelokaliseerd zijn.

3. Methodologie: LoSA (Locality-Aware Sparse Attention)

Op basis van deze observatie stellen de auteurs LoSA voor, een methode die de geheugenoperaties reduceert door slimme caching en selectieve berekening.

Kernprincipes:

1. Locality Pruning: Aan het begin van elke denoising-stap wordt voor elk token de verandering in de query-vector gemeten (via MSE). Tokens met een hoge verandering worden geclassificeerd als Active, de rest als Stable.
2. Hergebruik voor Stable Tokens: Voor Stable tokens wordt de attention-output van de vorige stap direct hergebruikt. Er wordt geen nieuwe sparse attention berekening uitgevoerd voor deze tokens. Hierdoor behouden ze de volledige informatie van de dense attention (in tegenstelling tot naïeve sparse attention die informatie mist).
3. Sparse Attention voor Active Tokens: Alleen voor de Active tokens wordt sparse attention (bijv. met de QUEST-algoritme) toegepast om de relevante prefix-KV-posities te selecteren.
4. Verkleining van de Unie: Omdat sparse attention alleen wordt toegepast op de subset van Active tokens (in plaats van het hele blok), wordt de unie van de te laden KV-posities drastisch verkleind. Dit verlaagt direct het geheugentransport.

Werkstroom:

• Initialisatie: Bij de eerste iteratie van een blok wordt dense attention gebruikt om de cache te vullen.
• Iteratie:
  1. Bereken localiteitsscores en selecteer Active tokens.
  2. Gebruik de sparse selector alleen voor deze actieve tokens om een unie van KV-indices te vormen.
  3. Laad alleen deze unieke KV-indices.
  4. Bereken prefix-attention voor actieve tokens; hergebruik de output voor stabiele tokens.
  5. Voeg prefix- en suffix-bijdragen samen via online-softmax.

4. Belangrijkste Bijdragen

• Identificatie van KV Inflation: De auteurs tonen aan waarom bestaande sparse attention-methoden falen in DLMs door de explosieve groei van de unie van KV-indices.
• Ontdekking van Localiteit: Ze kwantificeren dat representatieveranderingen in diffusiestappen sterk gelokaliseerd zijn, wat hergebruik mogelijk maakt.
• LoSA Algoritme: Een nieuwe architectuur die cache-hergebruik combineert met dynamische sparse attention, wat leidt tot minder geheugenverkeer en hogere nauwkeurigheid.
• Uitgebreide Evaluatie: De methode is getest op meerdere modellen (Trado-8B, Trado-4B, SDAR-8B) en benchmarks.

5. Resultaten

De evaluatie toont aan dat LoSA zowel efficiëntie als nauwkeurigheid verbetert ten opzichte van baselines zoals QUEST, SparseD en Sparse-dLLM.

• Nauwkeurigheid: Op de LongBench-benchmark behoudt LoSA bijna dense nauwkeurigheid, zelfs bij agressieve sparsiteit.
  • Bij een budget van 128 (zeer agressief) behaalde LoSA een gemiddelde verbetering van +10,43 punten (Trado-8B) en +8,38 punten (SDAR-8B) ten opzichte van QUEST.
  • Op korte contexttaken (commonsense reasoning) presteert LoSA vergelijkbaar met de baselines, wat aantoont dat de methode geen schade toebrengt waar het probleem minder acuut is.
• Efficiëntie (Geheugendichtheid): LoSA verlaagt de attention-dichtheid (percentage geladen KV-cache) met gemiddeld 1,54x ten opzichte van QUEST.
• Snelheid (Latentie):
  • Op NVIDIA RTX A6000 GPU's bereikte LoSA een 4,14x versnelling in de attention-berekening ten opzichte van dense attention.
  • QUEST (naïeve aanpassing) bereikte slechts 3,26x.
  • De trend bleek schaalbaar naar nieuwere hardware (RTX 5090), met een snelheidswinst van 3,67x.

6. Significatie

Dit paper is significant omdat het een fundamentele beperking van block-wise Diffusion Language Models oplost. Door in te zien dat niet alle tokens in een blok evenveel "werk" nodig hebben bij elke stap, introduceert LoSA een mechanisme dat:

1. Het geheugenprobleem (KV Inflation) effectief omzeilt.
2. De nauwkeurigheid behoudt door dense informatie voor stabiele tokens te bewaren (in plaats van te verwaarlozen zoals bij traditionele sparse attention).
3. DLM's praktischer maakt voor lange contexten, wat essentieel is voor hun toepassing in complexe redeneertaken en documentanalyse.

De methode is orthogonal tot de keuze van de specifieke sparse selector (zoals QUEST) en kan dus worden geïntegreerd in bestaande systemen met minimale overhead.