今回の論文
今回取り上げるのは、Jingyang Yuan、Huazuo Gao、Damai Dai、Junyu Luo、Liang Zhao、Zhengyan Zhang、Zhenda Xie、Y. X. Wei、Lean Wang、Zhiping Xiao、Yuqing Wang、Chong Ruan、Ming Zhang、Wenfeng Liang、Wangding Zeng による論文「Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention」です。2025年2月16日に arXiv で公開され、2月27日に改訂版が公開されています。公開元は arXiv、研究分野は長文LLM、attention最適化、推論高速化、長文学習です。URL は https://arxiv.org/abs/2502.11089 です。
この論文を選んだ理由は、単に「attentionを間引いて速くする」話ではなく、学習段階からスパースattentionを前提に設計し、しかもGPUで実際に速く動く形まで詰めているからです。長文RAG、長い会話履歴を抱えるAIエージェント、コード理解のような実務的な課題にそのままつながりやすく、開発のヒントが多い論文です。
どんな技術か
Native Sparse Attention は、長いコンテキストを扱うLLMで、すべての過去トークンに毎回フルattentionをかけるのではなく、重要そうな情報だけを3つの経路で選んで見る 技術です。
その3つとは、全体を粗く眺める「圧縮トークン」、重要ブロックだけを細かく見る「選択トークン」、直近の文脈を確実に拾う「スライディングウィンドウ」です。これらを別々にattentionし、最後にゲートで混ぜ合わせます。
ポイントは、後付けでKVキャッシュを削る推論専用テクニックではなく、最初からこの疎な見方で学習できる ことです。つまり「学習時はフルattention、推論時だけ近似」というギャップを減らしつつ、長文処理を軽くしようという設計です。
課題
この技術が解決しようとしている課題は、長文LLMでフルattentionを使い続けると、計算量もメモリアクセス量も重くなりすぎることです。
何が難しいのかというと、attentionは長さが伸びるほど見る相手が増えます。特に64Kのような長文になると、推論時には過去のKVキャッシュを大量に読み出す必要があり、計算そのものよりメモリ帯域がボトルネックになりやすいです。学習時もprefill時もコストは急増します。
既存の方法にも、KVキャッシュを捨てる手法や、一部ブロックだけを参照する手法があります。ただし、それらにはいくつか限界があります。まず、prefillだけ速い、あるいはdecodeだけ速いなど、片側しか効かない手法が多いです。さらに、既存モデルに後付けで疎化を入れる方法では、もともとdense attentionで学習されたモデルの振る舞いとズレやすく、精度低下が起きやすいです。加えて、ランダムアクセスの多い疎化はGPU実装で思ったほど速くならないことがあります。
なぜこの課題を解く必要があるのかというと、実際のAIシステムでは長文を扱う場面が増えているからです。RAGで大量の文書片を入れる、エージェントが長い履歴とツール出力を保持する、コードアシスタントがリポジトリ全体を読む、といった場面では、長い文脈を扱えても遅すぎれば使いにくくなります。精度だけでなく、長文を現実的な速度で回せるか が重要です。
用語解説
- Sparse Attention
- 全トークン同士を毎回比較するのではなく、一部の重要なトークンやブロックにだけattentionを計算する考え方です。NSAはこの疎化を後付けではなく、モデル設計の中心に置いています。
- KVキャッシュ
- 自己回帰生成で、過去トークンのkeyとvalueを保存して再利用する仕組みです。長文推論ではこの読み出し量が大きなボトルネックになるため、NSAの高速化はここに強く関係します。
- Grouped-Query Attention (GQA)
- 複数のquery headでkey/valueを共有するattention構造です。近年の実運用LLMでよく使われます。NSAはGQA前提で、headごとにバラバラの選択をさせず、グループ単位で同じKVブロックを読むよう設計しています。
- Arithmetic Intensity
- 計算量とメモリアクセス量の比率です。GPU最適化では重要な概念で、学習・prefillでは計算寄り、decodeではメモリアクセス寄りの最適化が必要になります。NSAはこの違いを意識して設計されています。
- Sliding Window Attention
- 直近の一定範囲だけは常に参照するattentionです。NSAでは、局所パターンを確実に拾うための専用分岐として使われ、圧縮や選択の枝が近傍依存だけに引っ張られないようにしています。
技術の仕組み
NSAの中核は、過去のkey/value全体をそのまま見るのではなく、より小さくて情報密度の高い表現に写してからattentionする ことです。論文では、クエリごとに再構成した key/value 集合に対してattentionを計算します。
基本アイデア
通常のattentionは、各クエリが過去の全トークンを見る設計です。NSAではここを、圧縮、選択、局所窓の3経路に分けます。
- 圧縮経路: 長い履歴をブロック単位でまとめて粗く見る
- 選択経路: 圧縮経路の結果を手掛かりに、重要ブロックだけ細かく見る
- 局所窓経路: 直近トークンは必ず見る
最後に各経路の出力をゲートで重み付けして足し合わせます。要するに「全体把握」と「精密確認」と「近傍優先」を1回のattentionブロックの中に同居させた形です。
圧縮トークンで全体を粗く見る
まず過去のkey/value列を連続ブロックに区切り、それぞれを1個の圧縮表現にまとめます。論文では、ブロック内位置情報を含むMLPで圧縮キーを作っています。これにより、長い履歴の全体像を少数の圧縮トークンで眺められます。
ここで重要なのは、最初から全トークンを細かく比較しないことです。粗い地図を先に作ることで、どのあたりに重要情報がありそうかを安く見積もれます。
選択トークンで重要ブロックだけ深掘りする
圧縮だけだと細かな情報が消えるので、NSAは次に「どのブロックを精密に見るか」を決めます。おもしろいのは、追加の重い選別器を作るのではなく、圧縮attentionで得た中間スコアを流用して重要ブロックを選ぶ点です。
ブロック選択はトークン単位ではなく、連続ブロック単位で行います。これは意味的な近接性だけでなく、GPUで連続メモリ読み出しをしやすくするためです。選ばれた上位 n 個のブロックだけ、本来の細かいkey/valueを取り出してattentionします。
この設計により、ランダムなトークン散布ではなく、意味的にもハードウェア的にもまとまりのある疎化 ができます。
GQA前提でグループ単位に選ぶ
最近のLLMではGQAがよく使われます。GQAでは複数headが同じKVを共有するので、headごとに別々のブロックを選ぶと、結局読むべきKVの合併集合が大きくなってしまいます。
NSAはここを意識して、同じGQAグループ内のheadが共有の重要度スコアを作り、グループ単位で同じブロックを読む ようにしています。これは論文の大事な工夫で、理論上の疎化だけでなく、decode時の実メモリアクセス削減に直結します。
スライディングウィンドウで近傍情報を確実に残す
局所的な依存関係は多くのタスクで非常に重要です。そのためNSAは、最近の w トークンを別枝として必ず参照します。これにより、近傍情報が圧縮や選択の枝に埋もれません。
論文では、近傍パターンだけで学習が完結してしまうショートカットを防ぐ意図も説明しています。局所処理を専用枝に切り出すことで、他の枝はより長距離の情報を学びやすくなります。
3枝をゲートで統合する
圧縮、選択、局所窓の各attention出力は、そのまま単純平均せず、MLPとsigmoidで得るゲート値で重み付けして統合されます。これにより、クエリごとに「今回は局所を見るべきか、遠距離を探すべきか」を調整できます。
この仕組みがあるため、NSAは固定パターンの疎attentionではなく、入力に応じて粗視化と精視化の比率を変えられる動的疎attention と言えます。
学習可能であることの意味
多くの疎attention手法は推論時の近似として使われますが、NSAは学習時からその疎構造で逆伝播できます。離散的すぎる選択やクラスタリングに頼りすぎず、ブロック設計と演算子設計を工夫しているためです。
これは実務的にかなり大きいです。学習時はdense、推論時だけsparseだと、推論時の近似誤差をモデルが前提にしていません。NSAはそこを減らし、疎な見方そのものをモデル能力に組み込む 方向を取っています。
実験と結果
この論文は、精度だけ、あるいは速度だけを見ていません。一般ベンチマーク、長文評価、推論付き数学評価、さらに学習速度とdecode速度まで確認しています。
何を検証したのか
主に次の点を検証しています。
- フルattentionより精度を落とさず学習できるか
- 長文ベンチマークで有効か
- 針探し型の長文検索で取りこぼしが減るか
- 推論付きタスクでも長い思考連鎖を支えられるか
- 学習・prefill・decodeの各段階でどれだけ速くなるか
どんなモデルと評価を使ったのか
論文では、GQAとMoEを組み合わせた27B総パラメータ、3Bアクティブパラメータのバックボーンを260Bトークンで事前学習しています。一般評価には MMLU、MMLU-PRO、CMMLU、BBH、GSM8K、MATH、DROP、MBPP、HumanEval を使用しています。長文評価では LongBench と 64K の needle-in-a-haystack テスト、推論評価では AIME 2024 を使っています。
つまり、単なるperplexity比較ではなく、知識、推論、コード、長文検索、長文推論 まで見ている構成です。
一般ベンチマークで平均性能が上がった
一般ベンチマークでは、Full Attention の平均が 0.443、NSA の平均が 0.456 でした。9指標中7指標でNSAが上回っています。特に reasoning 系で改善が大きく、GSM8K は 0.486 から 0.520、DROP は 0.503 から 0.545 に伸びています。
ここが重要なのは、NSAが「長文のときだけ速いが、普段は弱い近似法」ではなかった点です。短めの一般ベンチでも平均性能が落ちていないどころか、むしろ上がっています。論文は、疎な見方を前提に学習したことで、不要なattentionノイズが減り、重要情報へ集中しやすくなった可能性を示唆しています。
LongBenchでFull Attentionを上回った
LongBenchでは、NSAの平均が 0.469、Full Attention が 0.437、Exact-Top が 0.423 でした。単に他の疎attention法より良いだけでなく、フルattentionよりも平均で上回っています。
特に多段推論が必要なHPQでは 0.350 から 0.437、2Wikiでは 0.305 から 0.356、コード理解のLCCでは 0.163 から 0.232 に改善しています。長文になるほど、全トークンを見ること自体より、どこを見るかをうまく絞ること が効いていると読めます。
64Kのneedle-in-a-haystackで全位置100%検索
needle-in-a-haystackでは、64Kコンテキストのどの位置に答えを埋めてもNSAは完全な検索精度を示したと論文は報告しています。長文RAGや長会話の履歴から根拠を拾う用途を考えると、かなり示唆的です。
もちろん、これは合成タスクなので実業務そのものではありません。ただし「粗い圧縮で候補領域を探し、必要な場所だけ細かく見る」という設計が、長い文脈から一点を拾う問題に強い ことは読み取れます。
AIMEでも長い推論に強かった
AIME 2024の supervised fine-tuning 後評価では、8K上限で Full Attention-R が 0.046、NSA-R が 0.121、16K上限では Full Attention-R が 0.092、NSA-R が 0.146 でした。
ここから言えるのは、NSAが単なる検索用attentionではなく、長い推論過程の中で重要な中間状態を保持する方向にも効いている可能性があることです。数学推論にそのまま一般化しすぎるべきではありませんが、長い思考連鎖を崩しにくい疎化 としては有望です。
速度面でもかなり大きい
速度面では、64K長で最大 9.0倍のforward高速化、6.0倍のbackward高速化が報告されています。decodeではメモリアクセス量ベースで最大 11.6倍相当の高速化が期待されると示されています。
decodeで重要なのは、1トークン生成のたびに読むKV量です。論文設定では、64Kのフルattention相当の読み出しを、NSAでは圧縮トークン、選択ブロック、局所窓に絞ることで大きく削減しています。長文会話や長文推論で速度差が広がるのはこのためです。
何に使える？
NSAは研究向けのアーキテクチャに見えますが、実際には長文AIアプリのボトルネックにかなり近いところを触っています。以下は論文の結果を踏まえた応用可能性です。一部は推測を含みますが、その場合は明示します。
長文RAGの基盤モデル設計
RAGでは、関連文書を多く詰めるほど回答が良くなるとは限りません。不要なチャンクが増えると、モデルが重要箇所を見失いやすくなります。NSAのように、まず粗く全体を見て、重要ブロックだけ細かく見る設計は、長文RAGの基盤モデルに向いています。
特に、法務文書、社内規程、障害対応ログ、研究資料のように、長い文書の中から一部だけ拾えばよいケースで効きそうです。needle-in-a-haystackやLongBenchの結果を見る限り、根拠探索と長文理解の両立 に相性があります。
長い履歴を抱えるAIエージェント
エージェントは会話履歴、ツール出力、途中メモ、計画、コード断片などでコンテキストがすぐ膨らみます。NSAの設計は、全部を均等に見るのではなく、遠い過去は圧縮して眺め、必要そうな箇所だけ細かく読む方向です。
これは、長時間タスクを進めるコーディングエージェントや業務支援エージェントで特に有効そうです。ここは論文の直接評価ではありませんが、GQA前提のdecode効率や長文推論結果を踏まえると、履歴肥大化に強いモデル設計 として参考になります。
コード理解とリポジトリ規模の支援
LongBenchのLCC改善は、コード文脈の長距離依存に強い可能性を示しています。大規模リポジトリで、定義箇所と利用箇所が遠い、設定ファイルと実装が散らばっている、といったケースでは「全部見る」より「重要ファイル群を当てる」ほうが実用的です。
コードアシスタントの基盤モデルや、IDE補助の長文推論モデルを考えるうえで、NSAの発想はかなり参考になります。
長文推論やレポート生成
長い資料を読みながら要約やレポートを書くタスクでも、直近文脈と遠距離参照の両方が必要です。NSAの3枝構成は、近傍の文脈を捨てずに、遠くの重要断片も拾う設計なので、調査支援やレポート自動生成にも向いている可能性があります。
開発や事業へのヒント
この論文から得られる一番大きなヒントは、長文性能を上げる方法が「コンテキストを増やす」だけではないということです。どの粒度で情報を見るかを設計する ことが、精度と速度の両方に効きます。
アプリ層でも3段階の情報設計を真似できる
NSAの発想は、必ずしもモデルを自前学習しなくても応用できます。たとえばRAGでも、全文をいきなり投げるのではなく、まず粗い要約や章単位索引で候補を絞り、その後に候補チャンクだけ詳細投入し、最後に直近会話を必ず残す設計にできます。
これはまさに、圧縮、選択、局所窓の3段階に対応しています。モデル内部のattentionをそのまま再現するわけではありませんが、情報の見せ方を階層化する という考え方は小規模プロダクトでも使えます。
ハードウェア都合を無視しない
理論上の計算削減が、そのまま速度改善になるとは限りません。NSAが強いのは、連続ブロックで読ませる、GQAグループで同じKVを使う、Tensor Coreに乗りやすい形にする、といった実装都合まで含めて設計しているからです。
これはAIプロダクト開発でも同じです。推論コスト改善を考えるときは、アルゴリズムだけでなく、実行基盤でどんなメモリアクセスになるか を見る必要があります。
後付け最適化には限界がある
既存の大規模モデルに後から圧縮や枝刈りを入れるだけでは、精度が崩れやすいことがあります。NSAは、最初から疎なattentionで学習することで、そのギャップを減らしています。
事業やプロダクトの観点では、あとから無理に最適化するより、最初から用途に合う制約を前提に設計する ほうが長期的には強い、という教訓として読めます。長文業務に特化した内製モデルや蒸留モデルを作るときに特に効く考え方です。
今後注目すべき方向性
今後は、dense attentionを前提にした巨大モデル一本より、用途に応じて「どこをどう疎にするか」を設計したモデルが増える可能性があります。これは推測ですが、長文RAG、エージェント、コード支援はいずれも長い履歴と速度制約を同時に持つためです。
その意味でNSAは、単なる1本の疎attention論文というより、長文LLMを実用速度で成立させる設計指針 として見ておく価値があります。
限界
NSAにも限界や注意点があります。
まず、既存のフルattention系エコシステムに比べると、実装の複雑さは上がります。圧縮、選択、局所窓、ゲート、さらに専用カーネル設計まで含むので、単純なattention置換では済みません。
次に、論文は27B/3B active規模で有望な結果を示していますが、すべてのモデル規模、すべてのデータ分布、すべてのマルチモーダル設定で同じ効果が出るとはまだ言えません。特に超大規模学習や他アーキテクチャへの一般化は今後の検証が必要です。
また、学習可能な疎attentionとはいえ、選択ブロック数、圧縮ブロック長、局所窓サイズなど、設計ハイパーパラメータは多めです。用途に応じた調整が必要で、実装や学習レシピの難易度は高い部類です。
さらに、needle-in-a-haystackのような合成評価で強いことが、そのまま複雑な現実タスクのすべてに直結するわけではありません。実運用ではretrievalの質、ツール呼び出し、プロンプト設計など他要因も大きく効きます。
よくある質問
Q. Native Sparse Attention は、既存のKVキャッシュ圧縮手法と何が違うのですか？
A. 大きな違いは、推論時だけ後付けでKVを減らすのではなく、学習時から疎attentionを前提に設計している点です。さらにNSAは、圧縮トークンで候補領域を探し、その情報を使って選択ブロックを決め、近傍は別枝で必ず残します。単純なキャッシュ削減より、長文理解の精度維持まで狙った設計です。
Q. なぜトークン単位ではなくブロック単位で選ぶのですか？
A. 理由は2つあります。1つは、attention重要度が近い位置でまとまりやすいことです。もう1つは、GPUで連続メモリ読み出しをしやすく、実際の速度改善につながることです。ランダムなトークン選択は理論上疎でも、実装上は遅くなりやすいです。
Q. RAGにそのまま導入できますか？
A. 既存の一般公開モデルへ直接差し込むのは簡単ではありません。NSAはモデル内部のattention設計そのものを変えるため、基本的には事前学習や再学習が関わります。ただし、アプリ層では「粗く絞る→詳細確認→直近文脈保持」という設計思想をそのまま真似できます。
Q. 長文なら常にフルattentionより良いのですか？
A. 論文では平均的に強い結果が出ていますが、常にすべてで勝つとまでは言えません。タスクによってはフルattentionの単純さが有利なこともありますし、ブロック長や窓サイズの設定次第で結果は変わります。特に短文中心の用途では、複雑化に見合うかを見極める必要があります。
Q. 小規模チームがこの論文からすぐ使える学びは何ですか？
A. もっとも実践しやすいのは、情報を一段で詰め込まず、粗い索引、重要候補、直近文脈の3段階で扱うことです。モデルを自前で再学習しなくても、RAGやエージェントの入力設計、再ランキング、要約パイプラインにこの考え方を移せます。
今日の学び
この論文は、長文LLMでフルattentionが重すぎるという課題を扱っています。解決策として、圧縮、選択、局所窓の3経路を組み合わせた学習可能な疎attentionを提案し、精度と速度の両立を狙いました。
そこから得られるヒントは明快です。長文性能は、単にコンテキスト長を増やすだけではなく、どの粒度で情報を見るか、どの情報を残すか の設計で大きく変わります。RAGでもエージェントでも、情報を階層的に扱う発想はすぐ応用できます。