Hierarchische Sparse Attention richtig umgesetzt: Auf dem Weg zur unendlichen Kontextmodellierung

Zusammenfassung

Die Skalierung moderner großer Sprachmodelle (LLMs) auf lange Kontexte wird durch die quadratischen Rechenkosten und die schlechte Längenextrapolation dichter Attention begrenzt. Chunkweise spärliche Attention bietet eine vielversprechende Alternative, doch alle bestehenden Methoden bleiben aufgrund ihrer ungenauen Chunk-Auswahl hinter der vollen Attention zurück. Wir stellen die hierarchische Landmark-Sparse-Attention (HiLS-Attention) vor, einen auf Chunk-Ebene spärlichen Attention-Mechanismus, der die Chunk-Auswahl Ende-zu-Ende unter dem Sprachmodellierungsverlust (LM-Verlust) erlernt. HiLS faktorisiert Attention hierarchisch: Jede Query führt unabhängig Attention mit jedem abgerufenen Chunk durch, um Chunk-spezifische Informationen zu extrahieren, und die resultierenden Ausgaben werden gemäß den Chunk-Retrieval-Scores fusioniert. Durch die Einbeziehung der Retrieval-Scores in die Vorwärtsberechnung der Attention optimiert HiLS diese direkt mit dem LM-Verlust, was ein Ende-zu-Ende-Retrieval-Lernen und natives spärliches Training ermöglicht. Experimentelle Ergebnisse zeigen, dass HiLS-Attention eine mit voller Attention vergleichbare und in einigen Fällen sogar bessere Leistung bei domäneninternen Kontextlängen erzielt. Gleichzeitig extrapoliert HiLS-Attention mehr als das 64-fache der Trainingskontextlänge mit 90% Retrieval-Genauigkeit, weit über die volle Attention hinaus. Darüber hinaus können bestehende Modelle mit voller Attention durch leichtes fortgesetztes Vortraining in HiLS-Attention umgewandelt werden, wobei die domäneninterne Leistung erhalten bleibt und gleichzeitig eine ultra-lange Kontextextrapolation erworben wird. Zusammen mit seinem spärlichen KV-Zugriff und seiner spärlichen Berechnung durchbricht HiLS-Attention den üblichen Effizienz-Leistungs-Kompromiss und ermöglicht langkontextige LLMs, die sowohl effizienter als auch effektiver bei allgemeinen Langkontextaufgaben sind als ihre Pendants mit voller Attention.

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.