CRISP: Cliff-awaRe Input-adaptive Sparse Prefilling con instradamento basato sulla struttura della massa

Abstract degli autori

Nell’inferenza dei LLM a contesto lungo, il costo della fase di prefilling dell’attenzione cresce in modo quadratico, rendendo la self-attention un grave collo di bottiglia computazionale. I metodi tradizionali di attenzione sparsa attenuano il problema ricorrendo a schemi fissi o a profilazione offline, ma non hanno la flessibilità necessaria per adattarsi alla struttura dell’attenzione, che dipende dall’input. I metodi dinamici più recenti affrontano il problema instradando le teste verso schemi sparsi in tempo reale, ma si affidano a indicatori indiretti per l’instradamento, che comportano costi aggiuntivi, e a meccanismi di allocazione del budget che trascurano la gerarchia delle masse dopo la softmax. Presentiamo CRISP (Cliff-awaRe Input-adaptive Sparse Prefilling), che individua e affronta due sfide strutturali di questo paradigma di instradamento dinamico. In primo luogo, mostriamo che la decisione di instradamento può essere ricavata direttamente dalla struttura della mappa di attenzione proxy. Sostituiamo l’instradamento basato sulla divergenza di Jensen-Shannon (JSD) con C_struct, un indicatore strutturale che misura la massa nelle posizioni compatibili con Vertical-Slash e riproduce le decisioni di instradamento della JSD, eliminando sia il prodotto matriciale con aggregazione sia il costo aggiuntivo della successiva divergenza KL. In secondo luogo, formalizziamo il brusco calo della massa dopo la softmax e dimostriamo teoricamente che, nei contesti lunghi, le soglie di copertura strettamente cumulative accumulano rumore di fondo O(n). CRISP affronta il problema con una soglia che tiene conto dei sink e si basa sul livello del rumore di fondo. Nei test empirici su InfiniteBench, RULER e LongBench, condotti su due famiglie di modelli, CRISP è nel complesso il metodo di attenzione sparsa più efficace ed eguaglia o supera l’attenzione densa esatta nei benchmark incentrati sul recupero delle informazioni. Nei compiti di recupero ottiene fino a +28,0 punti percentuali rispetto alle baseline e, con 512k token, accelera l’attenzione fino a 5,30 volte, soprattutto grazie alla nostra eliminazione del rumore O(n) durante la selezione, che preserva l’integrità strutturale.