Attention hiérarchique éparse bien faite : vers la modélisation de contexte infini

Résumé

La mise à l'échelle des grands modèles de langage modernes (LLMs) pour des contextes longs est limitée par le coût de calcul quadratique et la faible extrapolation de longueur de l'attention dense. L'attention parcimonieuse par blocs offre une alternative prometteuse, mais toutes les méthodes existantes sont inférieures à l'attention complète en raison d'une sélection de blocs imprécise. Nous proposons l'attention parcimonieuse hiérarchique par points de repère (HiLS), un mécanisme d'attention parcimonieuse par blocs qui apprend la sélection de blocs de bout en bout sous la perte de modélisation du langage (LM). HiLS factorise l'attention de manière hiérarchique : chaque requête effectue une attention indépendamment avec chaque bloc récupéré pour extraire des informations spécifiques au bloc, et les sorties résultantes sont fusionnées selon les scores de récupération des blocs. En intégrant les scores de récupération dans le calcul direct de l'attention, HiLS les optimise directement avec la perte LM, permettant un apprentissage de récupération de bout en bout et un entraînement parcimonieux natif. Les résultats expérimentaux montrent que l'attention HiLS atteint des performances comparables, voire supérieures dans certains cas, à l'attention complète pour des longueurs de contexte dans le domaine d'entraînement. Par ailleurs, l'attention HiLS extrapole plus de 64 fois la longueur de contexte d'entraînement avec une précision de récupération de 90 %, bien au-delà de l'attention complète. De plus, les modèles à attention complète existants peuvent être convertis en attention HiLS grâce à un pré-entraînement continu léger, préservant les performances dans le domaine tout en acquérant une extrapolation de contexte ultra-long. Combinée à son accès et à son calcul parcimonieux pour les clés-valeurs, l'attention HiLS brise le compromis habituel entre efficacité et performance, permettant des LLMs à long contexte à la fois plus efficaces et plus performants sur les tâches générales de long contexte que leurs homologues à attention complète.

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.