Hiërarchische Sparse Attention Correct Uitgevoerd: Naar Oneindige Contextmodellering

Samenvatting

Het opschalen van moderne grote taalmodellen (LLM's) naar lange contexten wordt beperkt door de kwadratische rekenkosten en de slechte lengte-extrapolatie van dichte aandacht. Stuksgewijze sparse aandacht biedt een veelbelovend alternatief, maar alle bestaande methoden schieten tekort ten opzichte van volledige aandacht vanwege hun onnauwkeurige chunkselectie. Wij stellen Hiërarchische Landmark Sparse (HiLS) Aandacht voor, een stuksgewijs spaarzaam aandachtsmechanisme dat chunkselectie end-to-end leert onder het taalmodelleringsverlies (LM-verlies). HiLS factoriseert aandacht hiërarchisch: elke query voert zelfstandig aandacht uit met elke opgehaalde chunk om chukspecifieke informatie te extraheren, en de resulterende uitvoer wordt samengevoegd op basis van de chunkretrievalscores. Door retrievalscores op te nemen in de voorwaartse aandachtsberekening, optimaliseert HiLS deze direct met het LM-verlies, wat end-to-end retrievalleren en native sparse training mogelijk maakt. Experimentele resultaten tonen aan dat HiLS-Aandacht prestaties levert die vergelijkbaar zijn met, en in sommige gevallen beter dan, volledige aandacht bij contextlengtes binnen het domein. Tegelijkertijd extrapoleert HiLS-Aandacht meer dan 64 keer de trainingscontextlengte met 90% retrievalnauwkeurigheid, ver voorbij volledige aandacht. Bovendien kunnen bestaande volledige-aandachtsmodellen worden omgezet naar HiLS-Aandacht met een lichte voortgezette pretraining, waarbij de prestaties binnen het domein behouden blijven terwijl ultra-lange-context-extrapolatie wordt verkregen. Samen met zijn sparse KV-toegang en -berekening doorbreekt HiLS-Aandacht de gebruikelijke afweging tussen efficiëntie en prestaties, waardoor lange-context-LLM's mogelijk worden die zowel efficiënter als effectiever zijn bij algemene lange-contexttaken dan hun tegenhangers met volledige aandacht.

English

Scaling modern large language models (LLMs) to long contexts is limited by the quadratic computation cost, and poor length extrapolation of dense attention. Chunk-wise sparse attention offers a promising alternative, but all existing methods fall short of full attention because of their inaccurate chunk selection. We propose Hierarchical Landmark Sparse (HiLS) Attention, a chunk-wise sparse attention mechanism that learns chunk selection end-to-end under the language-modeling (LM) loss. HiLS factorizes attention hierarchically: each query performs attention independently with each retrieved chunk to extract chunk-specific information, and the resulting outputs are fused according to chunk retrieval scores. By incorporating retrieval scores into the forward attention computation, HiLS optimizes them directly with the LM loss, enabling end-to-end retrieval learning and native sparse training. Experimental results show that HiLS-Attention achieves performance comparable to, and in some cases better than, full attention at in-domain context lengths. Meanwhile, HiLS-Attention extrapolates more than 64times the training context length with 90% retrieval accuracy, far beyond full attention. Moreover, existing full-attention models can be converted to HiLS-Attention with lightweight continued pretraining, preserving in-domain performance while acquiring ultra-long-context extrapolation. Together with its sparse KV access and computation, HiLS-Attention breaks the usual efficiency-performance trade-off, enabling long-context LLMs that are both more efficient and more effective on general long-context tasks than their full-attention counterparts.