正しく実装された階層的スパースアテンション：無限コンテキストモデリングに向けて

要旨

近年の大規模言語モデル（LLM）の長文脈への拡張は、二次計算コストや密な注意機構における長さ外挿性の低さによって制限されている。チャンク単位のスパース注意機構は有望な代替手段を提供するが、既存の手法はいずれもチャンク選択の不正確さから完全注意機構には及ばない。本稿では、階層的ランドマーク・スパース注意機構（HiLS-Attention）を提案する。これはチャンク単位のスパース注意機構であり、言語モデリング（LM）損失の下でエンドツーエンドにチャンク選択を学習する。HiLS-Attentionは注意機構を階層的に分解する。各クエリは、個々に取得したチャンクに対して独立に注意を実行することでチャンク固有の情報を抽出し、その出力はチャンク取得スコアに従って統合される。取得スコアを順方向の注意計算に組み込むことで、HiLS-AttentionはそれらをLM損失によって直接最適化し、エンドツーエンドの取得学習とネイティブなスパース学習を可能にする。実験結果は、HiLS-Attentionがドメイン内の文脈長において完全注意機構と同等、場合によってはそれを上回る性能を達成することを示している。さらに、HiLS-Attentionは90%の取得精度で訓練時の文脈長の64倍以上の外挿を実現し、完全注意機構を大きく上回る。また、既存の完全注意機構モデルは、軽量な継続事前学習によりHiLS-Attentionに変換可能であり、ドメイン内性能を維持しつつ超長文脈の外挿能力を獲得する。スパースなKVアクセスと計算と相まって、HiLS-Attentionは通常の効率と性能のトレードオフを打破し、完全注意機構を用いたモデルと比較して、一般的な長文脈タスクにおいてより効率的かつ効果的な長文脈LLMを実現する。

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.