HiLS Attentionとは？LM損失で学習するスパース注意で訓練長64倍を外挿

階層的ランドマークトークンによるチャンク選択をLM損失でエンドツーエンド学習し、ヒューリスティック手法が抱える精度不足を根本から解消
訓練長8Kのモデルが512K（64倍超）まで外挿しても90%以上の検索精度を維持し、全注意比で最大15.7倍の推論高速化を達成
既存の全注意モデルへの軽量な継続事前学習で変換でき、1Mトークン超の長文脈LLM応用への実用的な選択肢を提供

研究の背景

大規模言語モデル（LLM）を長い文書や会話に対応させるには、コンテキスト長の大幅な拡張が不可欠です。しかし標準的なTransformerが採用する全注意（Full Attention）機構は、入力長の2乗に比例して計算量が増大するため、数十万トークンを超える処理では現実的なコストに収まりません。

この問題へのアプローチとして、関連度の高いチャンク（数十トークン単位のブロック）だけを選んで注意計算する「スパース注意」が研究されてきました。ところが、NSA（Native Sparse Attention）などの従来手法のほとんどは、どのチャンクを選ぶかをヒューリスティック（経験則）に頼って判断しており、チャンク内のトークン間で注意の分布が不均一な場合に選択精度が大きく低下するという課題を抱えていました。

HiLS Attentionの仕組み

Tencent HY Teamや上海科技大学などの研究チームが提案するHiLS-Attention（Hierarchical Landmark Sparse Attention）は、この問題をエンドツーエンド学習で解決します。核心となるアイデアは、各チャンクの末尾にランドマークトークンと呼ばれる特殊なトークンを配置し、そこから圧縮チャンクキー（compressed chunk key）を導出することです。

図1: HiLS-Attentionの全体構造。チャンク間ソフトマックスとチャンク内ソフトマックスの2段階に注意を因数分解し、ランドマークトークン由来の圧縮キーでチャンク質量を効率よく推定する

注意の計算は2段階に分解されます。まず「チャンク間ソフトマックス」でクエリが各チャンクに割り当てる注意質量の合計を決定し、次に「チャンク内ソフトマックス」でその質量をチャンク内のトークン間に分配します。この因数分解により、全QK計算を経ずにチャンク選択が実現します。

最も重要な特徴は、圧縮チャンクキーが次トークン予測損失（LM損失）から直接勾配を受け取れる構造になっている点です。チャンク選択スコアが順伝播の注意重みに組み込まれているため、言語モデル損失によるエンドツーエンド学習が可能になります。モデルは何を選ぶべきかをタスクから自律的に学習でき、別途のヒューリスティックな選択モジュールは不要です。

位置エンコーディングにはHoPE（Hybrid Positional Encoding）を組み合わせます。RoPEの一部の次元を訓練長以内の回転周期に限定し、残りをNoPE（位置エンコーディングなし）に置き換えることで、訓練長を超えたコンテキストへの外挿を支えます。超長文脈の扱い方として、AgenticSTSのようにメモリ量を意図的に抑える方向もありますが、HiLS-Attentionはモデル自体の処理窓を大幅に広げるアプローチで補完的な関係にあります。

実験結果と性能

345MパラメータのモデルをHiLS-Attention + HoPEで訓練長8Kとして学習したところ、同一設定の全注意モデルと同等のパープレキシティ（4.94）を達成しました。そのモデルをコンテキスト長512K（訓練長の64倍超）で評価した結果、長文脈検索ベンチマーク「RULER」で90%以上の精度を維持しました。

図2: 1.4Bモデルの訓練ステップごとのパープレキシティとRULER精度。HiLS-Attention + HoPEの組み合わせが訓練長をはるかに超えたコンテキストでも高い性能を保つことを示す

推論効率の面では、H800 GPU上での同一Tritonカーネル実装による比較で、コンテキスト長512KにおいてプリフィルはHiLS-Attentionが13.5倍、デコードは15.7倍高速でした。コンテキスト長が約16K以上になった時点で全注意より速くなり、長文になるほど差が拡大します。

図3: 345MモデルにおけるHiLS-Attentionと全注意の推論レイテンシ比較（NVIDIA H800、bf16）。16K超でHiLS-Attentionが上回り始め、512Kではプリフィル13.5倍・デコード15.7倍の高速化を達成

7Bモデルでの継続事前学習（50Bトークン）実験では、LongBenchの複数タスクでHiLS-Attention + HoPEがYaRN拡張の全注意ベースラインに匹敵する性能を示しました。また、隣接するクエリトークン間で読み込むチャンクIDの重複率が高いことも確認されており、グループ内でのキャッシュ再利用が効果的に機能する設計になっています。

まとめ

HiLS-Attentionは、スパース注意の弱点だったチャンク選択精度の低さを、ランドマークトークンとLM損失によるエンドツーエンド学習によって解決した手法です。訓練長の64倍超まで高精度を維持しながら推論を大幅に高速化し、既存モデルへの軽量な適用も可能という実用性を兼ね備えています。

コードはGitHub（Tencent-Hunyuan/HiLS-Attention）で公開されており、1Mトークンを超える文書処理や超長期の会話履歴を扱うLLMアプリケーションへの応用が期待されます。チャンク選択のエンドツーエンド学習という設計思想は、今後のスパース注意研究にも広く影響を与えるでしょう。