Technische Samenvatting: Onzekerheidsgestuurde Selectie voor Blok-schaarse Aandacht

Probleemstelling

Block-sparse aandachtmechanismen (zoals SSA, Quest, H2O, SnapKV) schalen lang-context taalmodellen door de $O(N^2)$ softmax te vervangen door een per-query top-$k$ selectie over sleutelblokken. Hoewel efficiënt, is dit selectieproces myopisch (kortzichtig). Wanneer de scores van het $k$-de en $(k+1)$-de blok bijna gelijk zijn, legt de selector een beslissing vast zonder extra budget. Als het gedropte blok cruciale bewijslast bevat (bijv. in multi-hop redeneren of needle-in-a-haystack retrieval), gaat de informatie onherroepelijk verloren, aangezien downstream lagen dit niet kunnen herstellen. Deze fout is structureel, niet louter ruisachtig, en degradeert de prestaties met name in scenario's waar relevantie afhangt van latente context in plaats van oppervlakkige query-key matches.

Bestaande methoden behandelen de afkapwaarde als een vaste beslisregel. Ze houden geen rekening met de onzekerheid van die beslissing of de informatiewaarde (Value of Information, VoI) die verkregen kan worden door de aandachtset uit te breiden wanneer de afkapwaarde ambigu is.

Methodologie

Het artikel stelt een onzekerheidsgestuurde router voor die bovenop bestaande block-scoring backbones (zoals SSA-stijl mean-pooling of Quest's K-min/K-max upper bounds) werkt zonder hertraining of extra parameters te vereisen. De methode werkt in vijf stappen:

1. Block Scoring (Backbone Agnostisch): Sleutels worden gegroepeerd in blokken. Scores worden berekend met een bestaande backbone (bijv. mean-pooled inwendig product of coördinaat-wijze upper bounds).
2. Per-Tile Selectie: Om de computationele overhead te verminderen, wordt selectie uitgevoerd per Query Tile (een groep opeenvolgende query-rijen) in plaats van per rij. De tile-score is de maximale score over de rijen in de tile. Een top-$k$ selectie bepaalt de initiële set blokken ($kv\_idx$).
3. Onzekerheidsmeting (VoI Signaal): De methode berekent een genormaliseerde afkapmarge ($\sigma$) voor elke tile en head:
   $$\sigma = \frac{s_{(k-1)} - s_{(k)}}{s_{(0)} - s_{(k)}}$$
   waarbij $s_{(i)}$ de gesorteerde block scores zijn.
   • $\sigma \to 1$: De afkapwaarde is beslissend (de behouden set ligt ver boven de staart).
   • $\sigma \to 0$: De afkapwaarde is ambigu (het $k$-de en $(k+1)$-de blok zijn gelijk).
     Deze $\sigma$ dient als een proxy voor het verwachte verlies van de top-$k$ beslissing; een lage $\sigma$ duidt op een hoge waarde in het verwerven van meer informatie.
4. Aggregatie en Triggering: Per-head $\sigma$-waarden worden gemiddeld tot een per-tile signaal ($\bar{\sigma}$). Een trigger wordt geactiveerd als $\bar{\sigma}$ onder een laag-specifieke empirische kwantiel-drempelwaarde ($q$) valt. Dit identificeert de "risicovolle" tiles waar de beslissing ambigu was.
5. Gecontroleerde Expansie: Getriggerde tiles breiden hun aandachtset uit met een factor $\rho$ (bijv. $\rho=2$), waarbij ze aandacht besteden aan de top-$\rho k$ blokken in plaats van top-$k$. Niet-getriggerde tiles behouden het basisbudget. De implementatie gebruikt een gefuseerde kernel die niet-getriggerde tiles opvult met sentinels om uniforme dispatch-vormen te behouden, wat ervoor zorgt dat de kosten sublineair schalen met de expansie.

Belangrijkste Bijdragen

1. VoI Formulering: Een nieuwe, backbone-agnostische formulering van de selector afkapwaarde als een beslissing onder onzekerheid. Het voegt slechts één scalar ($\sigma$) per tile toe en één kwantiel-drempelwaarde per laag, wat geen hertraining of extra parameters vereist.
2. Composability: Empirische demonstratie dat de router onderscheidende scoring backbones (zoals SSA-stijl K-mean en Quest's K-max) kan combineren. De router verbetert consequent de prestaties van welke backbone dan ook superieur is voor een specifieke model, wat bewijst dat "betere scoring" en "beter budgetallocatie" orthogonale verbeteringen zijn.
3. Gefuseerde Implementatie: Een aangepaste Triton kernel die selectie en expansie in één enkel codepad afhandelt, waarbij vorm-uniforme dispatch wordt behouden. Dit zorgt ervoor dat het systeem bij 128K context draait op $0.62\times$ tot $0.80\times$ de wall time van dense aandacht.
4. Diagnostische Benchmarking: Introductie van de Pointer-Chase Haystack (PCH) benchmark om de kwaliteit van de selector te isoleren van de modelcapaciteit, en het gebruik van LongBench-v1 als een negatieve controle om aan te tonen dat de voordelen van de router afhankelijk zijn van het feit dat het selector-budget krap is ten opzichte van de dichtheid van bewijs.

Resultaten

De methode werd geëvalueerd op vier modellen (Qwen2.5-14B, Qwen2.5-7B-1M, Mistral-Nemo, Qwen3.6-35B) over contexten van 32K tot 128K.

• Prestatiewinsten:
  • Op LongBench-v2 medium (32K context) bereikte de router-op-Quest configuratie een gepaarde recall van 0.75 vergeleken met 0.47 voor de top-$k$ baseline (+28 procentpunten), met een McNemar $p < 0.01$.
  • Op RULER NIAH multi-key bereikte de router-op-Quest 0.98 nauwkeurigheid (vs. 0.93 voor top-$k$) op Qwen2.5-14B en 1.00 op Qwen2.5-7B-1M.
  • De router verbeterde consequent de winnende backbone voor elk model. Opmerkelijk genoeg versprong de optimale backbone op basis van architectuur: Quest's K-max won op Qwen2.5/Mistral, terwijl SSA-stijl K-mean won op Qwen3.6 (vanwege de QK-Norm regularisatie van magnitude variantie). De router liftte beide.
• Efficiëntie:
  • Bij 128K context behield de router 0.81 van de dense nauwkeurigheid op Qwen2.5-7B-1M en 0.89 op Qwen3.6, terwijl hij respectievelijk draaide op 0.62$\times$ en 0.80$\times$ de wall time van dense aandacht.
  • De "Pareto win" (betere snelheid bij vergelijkbare kwaliteit) ontstaat bij langere contexten (64K–128K) naarmate het aandel van de aandacht in de totale berekening toeneemt, consistent met de wet van Amdahl.
• Budget vs. Selectiviteit: Ablatie-studies laten zien dat de lift van de router komt van zowel een verhoogd budget als een selectieve allocatie. Bij modellen nabij het prestatieplafond (bijv. Qwen-14B) is de lift voornamelijk budgetgedreven; bij modellen met meer ruimte (bijv. Mistral-Nemo) levert de selectiviteit (het budget alleen uitgeven op ambigue tiles) het grootste deel van de winst op.

Betekenis en Claims

Het artikel claimt dat de onzekerheidsgestuurde router een generieke, plug-and-play laag biedt voor block-sparse aandacht die de structurele foutmodus van myopische top-$k$ selectie mitigeert. De betekenis ligt in:

• Robuustheid: Het past zich aan de specifieke scoring-karakteristieken van verschillende modelarchitecturen aan (bijv. het afhandelen van het QK-Norm effect in Qwen3.6) zonder handmatige afstemming.
• Efficiëntie: Het bereikt bijna dense nauwkeurigheid bij 128K context met aanzienlijk verminderde wall time, wat lange-context inferentie praktischer maakt.
• Theoretische Onderbouwing: Het frameert het selectieprobleem vanuit een Value-of-Information lens, waarbij een finite-sample regret bound wordt geboden onder Gaussische ruis aannames (hoewel het artikel opmerkt dat deze bound momenteel zwak is in de door ruis gedomineerde regimes waarin de methode opereert).

De auteurs stellen expliciet dat de lift van de router conditioneel is: het activeert alleen wanneer het per-query budget klein is ten opzichte van de locatie van de antwoord-relevante bewijslast. In scenario's waar het budget ruim is (bijv. LongBench-v1 met natuurlijke korte prompts), biedt de router verwaarloosbaar voordeel, wat dient als een negatieve controle die de specificiteit van de methode voor tight-budget, lang-context scenario's valideert.