🤖 machine learning

COBS: Cumulant Order Block Sparse Attention

本論文は、圧縮された二次の統計量を用いる新規なセレクターによってアテンションの質量をより良く近似することで、長文脈リトリーバルの性能を向上させ、ハードウェア効率を維持しつつ高密度アテンションとの品質格差を大幅に縮小するブロック疎なアテンション手法であるCOBSを紹介するものである。

原著者： Alexander Tian, Aditya Ghai, Sanjit Neelam, Zaal Vasania, Akshay Mishra

公開日 2026-07-13

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Alexander Tian, Aditya Ghai, Sanjit Neelam, Zaal Vasania, Akshay Mishra

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

巨大な干し草の山の中から特定の針を見つけようとしている場面を想像してみてください。しかし、その干し草の山があまりにも巨大であるため、すべての藁（わら）を一つずつ調べていると、あなたの脳（あるいはコンピュータ）が処理しきれなくなってしまいます。これが、現代のAIモデルが非常に長い文書を読もうとする際に直面している問題です。モデルはこれまでに読んだすべての内容を記憶しておく必要がありますが、すべてのメモリトークンを一つずつチェックするのは時間がかかり、コストもかさみます。

この問題を解決するために、研究者たちは「ブロック疎行列アテンション（Block Sparse Attention）」と呼ばれるショートカットを試みました。すべての藁を一つずつ見る代わりに、藁を小さな束（「ブロック」と呼びます）としてまとめ、その中から最も興味深いものだけを選び出すことにしたのです。これは、偵察員を雇って、いくつかの束をスキャンさせ、「どの束に針が含まれていそうか」を報告させるようなものです。

問題点：偵察員が単純すぎた

この論文では、「NSA（Native Sparse Attention）」と呼ばれる一般的な手法について研究しています。このシステムでは、偵察員は藁の束を見て、それが重要かどうかを素早く推測します。しかし、この論文では、この偵察員が非常に単純なトリックを使っていることが判明しました。それは、束の中にある藁の「平均的な位置」だけを見ているという点です。

これを次のように考えてみてください。2つの藁の束があるとします。

束A は、藁が中央にぎゅっと密集しています。
束B は、藁がバラバラに散らばっており、左端から右端まで広く分布しています。

もし「平均的な位置」だけを見るなら、両方の束は全く同じに見えてしまいます。しかし実際には、より広い範囲をカバーしている束Bの方が、針が含まれている可能性が高いのです。従来の偵偵員（一次統計的手法）は、この「広がり」や「曲率」に対して盲目でした。彼らは、雲の形をその中心点を見るだけで推測しようとしているようなものでした。つまり、実際に重要な「ふわふわとした端の部分」を見逃していたのです。

解決策：COBS（より賢い偵察員）

著者らは、COBS（Cumulant Order Block Sparse Attention） と呼ばれる新しい手法を提案しています。COBSの偵察員は、単なる平均的な位置だけでなく、藁が「どのように広がっているか」を示す、極めて小さく圧縮されたマップを携行しています。

数学的な用語で言えば、論文ではこれを「二次統計量」または「共分散」と呼んでいます。私たちの例えに当てはめると、これは偵察員が「おや、この束は幅が広く乱雑だ。だから針を持っている確率が高いぞ！」と気づくようなものです。この追加の情報（ただし、容量を使いすぎないように圧縮されたもの）を保持することで、COBSはより優れた推測を行うことができます。

結果：劇的な進歩

チームは、32k RULERベンチマーク（11種類の異なる長文コンテキスト検索タスクによるテスト）を用いて検証を行いました。結果は以下の通りです。

旧来の方法（NSA MLP）： 単純な偵察員によるスコアは 0.2999 でした。針を見つけるのに苦戦していました。
完璧な方法（OSA）： もし魔法を使ってショートカットなしで正解を知ることができれば（これを「オラクル」と呼びます）、スコアは 0.9040 になります。
新しい方法（COBS）： 広がりマップを持つ賢い偵察員によるスコアは 0.8195 でした。

これは、COBSが「苦戦していた旧来の手法」と「完璧な手法」の間のギャップの約 86% を埋めたことを意味します。これは驚異的な改善です！

コスト：それだけの価値はあるのか？

通常、賢くなるということは、より多くの作業を必要とすることを意味します。しかし、COBSは効率的です。

旧来の方法は、ある一定量のデータを読み込みました。
完璧な方法（すべてを読み込む方法）は、COBSよりも 15.15倍 多いデータを読み込みます。
COBSは、旧来の苦戦していた方法よりもわずか 1.21倍 のデータしか読み込みません。

つまり、COBSは旧来の単純な方法よりもほんの少しの労力を増やすだけで、完璧に近い精度を手に入れているのです。

論文が「ノー」と言っていること

著者らは、一見すると良さそうなショートカットに見えても、実際には機能しないいくつかのアイデアを慎重に排除しました。

平均値に複雑さを加えること： 「平均」の推測をもっと賢くするために、高度なニューラルネットワーク（MLB）を使う試みも行われましたが、あまり効果はありませんでした。問題は平均値の複雑さではなく、平均値そのものが「間違った道具」であったことです。必要なのは平均ではなく、「広がり」の情報なのです。
「広がり」を単純な箱で捉えること： 別の手法では、最小値と最大値を見て広がりを推測する（「箱」を作る）方法が試されました。これは多少の効果はありましたが、COBSのような精密な広がりマップには及びませんでした。
さらに複雑な数学（三次統計量）を加えること： 著者らは、「歪度（スキュー：広がりがどれくらい偏っているか）」を加えるテストも行いました。驚くべきことに、これは低複雑度のレベルでは状況をむしろ悪化させ、モデルを混乱させました。これは根本的な解決策ではなく、単なる「絆創膏」のようなものでした。著者らは、シンプルに「広がり（二次統計量）」にとどめることが最適であると判断しました。

信頼性はどの程度か？

著者らは、管理された実験を行ったため、これらの数値に非常に自信を持っています。単に推測したのではなく、32k RULER テストで性能を測定し、COBSが一貫して旧来の手法を上回ることを確認しました。また、これがモデルの短い文章の理解能力を損なっていないことも確認しており（損なっていません）、長文における次の単語の予測能力も旧来の手法より向上していることを確認しています。

ただし、著者らは限界についても正直に述べています。

この検証は、約 12億パラメータ のモデルを用いて行われました。これが、大手テック企業が使用するような超巨大なモデルでも全く同じように機能するかどうかは断定できませんが、数学的には可能であると示唆されています。
モデルは、長文コンテキスト能力をテストするための特定の合成データ（RULIZERスタイル）で学習されています。これは標準的なテスト方法ですが、現実世界のデータは少し異なる挙動を示す可能性があります。

まとめ

この論文は、干し草の山から効率的に針を見つけるためには、単に束の中心を見るだけでは不十分であることを示しています。その束がどのように広がっているかを知る必要があります。その広がりの情報を、極めて小さく圧縮されたマップとして加えることで、COBS はAIモデルが速度を落とすことなく、より正確に長い文書を読めるようにします。それは、「そこそこの精度」と「完璧」の間の溝を、最小限の努力で埋める画期的な手法なのです。

技術要約: COBS (Cumulant Order Block Sparse Attention)

1. 問題提起

長文脈（ロングコンテキスト）における大規模言語モデル（LLM）の推論は、計算量よりもメモリ帯域幅によってボトルネックとなります。具体的には、デコーディング速度を制限し、ハードウェアの利用率を低下させる、キー・バリュー（KV）キャッシュのデンス・アテンション（密なアテンション）による読み出し操作が問題となります。ブロック疎（ブロック・スパース）アテンションは、連続したトークンのブロックのみを読み出すことでハードウェアに優しい解決策を提供しますが、主要なオープンウェイトモデルでは依然として十分に活用されていません。既存の手法は、デンス・アテンション、きめ細かなトークン選択、またはKVエビクション（除去）に依存することが多いです。

ブロック疎アテンションにおける主要な課題は、ブロック選択です。すなわち、すべてのブロックのキーとバリューをすべて読み取ることなく、与えられたクエリに対してどのKVキャッシュのブロックが最も関連性の高い情報を含んでいるかを特定することです。DeepSeekのNative Sparse Attention (NSA) などの現在の最先端手法は、「一次（first-order）」の選択戦略を利用しています。これらの手法は、ブロックをクエリに依存しないベクトル（MLPや平均プーリングによる）に要約し、それらをクエリに対して線形にスコアリングします。本論文は、これらの一次近似がアテンション・マスの非線形な曲率を捉えることに失敗しており、その結果、デンス・アテンションと比較して大きな性能差が生じていると主張しています。

2. 手法

2.1 理論的基礎：選択オラクル（Selection Oracle）

著者らは、ブロック選択を「デンス・アテンションの出力を一部のブロックから再構成する問題」として定式化しました。彼らは、最適な選択基準は、ブロック内のアテンション・スコアの総和であるアテンション・マス（ $m_b$ ）によってブロックをランク付けすることであると導き出しました：
$m_b = \sum_{r \in b} e^{q^T k_r}$
すべてのキーを読み取ることで正確なアテンション・マスを算出する「オラクル疎アテンション（Oracle Sparse Attention, OSA）」は、デンス・アテンションの性能にほぼ一致することを示しており、選択基準自体は妥当であることを証明しています。ボトルネックは、コンパクトでクエリに依存しない要約から、このマスをいかに推定するかという点にあります。

2.2 累積モーメント（Cumulant Expansion）

本論文は、累積モーメント展開（対数モーメント生成関数のテイラー展開と同等）を用いたアテンション・マスの推定について分析しています。対数マスは以下のように展開されます：
$\ln m_b = \ln L + q^T \kappa_1 + \frac{1}{2} q^T \kappa_2 q + \dots$
ここで：

$\kappa_1$ はブロック内の平均キー（ $\bar{k}_b$ ）です。
$\kappa_2$ はブロック内のキーの共分散行列（ $\Sigma_b$ ）です。

既存の一次セレクター（NSAのMLPや平均プーリングなど）は、線形項（ $q^T \bar{k}_b$ ）のみを捉え、二次項の曲率（ $\frac{1}{2} q^T \Sigma_b q$ ）を破棄してしまいます。この欠落により、たとえ平均が類似していても、ブロック内のキーの分布がクエリの方向に対して大きな広がりを持つ場合、既存の手法は失敗します。

2.3 COBSの手法

COBS (Cumulant Order Block Sparse Attention) は、選択プロセスにこの二次項を組み込むことでこれに対処します。この手法は、圧縮、選択、およびスライディングウィンドウのブランチを含むNSAアーキテクチャに基づき、以下のイノベーションを取り入れています。

二次スコアリング（Second-Order Scoring）： 単なる平均ベクトルではなく、ブロックの平均キー（ $\bar{k}_b$ ）と、ブロック内の共分散行列（ $\Sigma_b$ ）の圧縮表現を保持します。選択スコアは以下のようになります：
$\text{score}_b \approx q^T \bar{k}_b + \frac{1}{2} q^T \Sigma_b q$
共分散圧縮（Covariance Compression）： フルサイズの $D \times D$ $D \times D$ 共分散行列を保持することはメモリ的に困難であるため、COBSは3つの圧縮技術を採用しています：
- 低ランク分解（Low-Rank Factorization）： $\Sigma_b$ の上位 $r$ 個の固有ベクトルと固有値のみを保持します。
- 部分空間投影（Subspace Projection）： 共分散を、より低次元のクエリ部分空間（クエリの二次モーメントの上位固有ベクトルによって張られる空間）に投影し、ストレージを $rD $から$ rs$ へ削減します。
- 量子化（Quantization）： 固有ベクトルをFP4 (E2M1) フォーマットでスケールして保存し、ブロック平均はBF16で保持します。
NoPE (No Position Encoding)： 圧縮および選択ブランチから回転位置エンコーディング（RoPE）を除去することで、ブロックの要約が内容のみに依存するようにし、位置情報のノイズが要約を損なうのを防ぎます。

3. 主な貢献

選択の定式化： ブロック選択をアテンション・マスによるランキングへと還元し、正確なマスを使用するオラクルがデンス・アテンションの性能をほぼ完全に回復することを示しました。
累積モーメント次数の分析： 既存の一次セレクターの限界を理論的に説明し、破棄されたブロック内の共分散が性能差の原因であることを特定しました。
COBSアルゴリズム： ブロックごとに圧縮された二次統計量（共分散）を保持し、フルなキーを読み取ることなくアテンション・マスをより正確に近似する実用的な手法を提案しました。
実証的検証： 二次的な選択が、長文脈ベンチマークにおいて、疎なアテンションとデンス・アテンションの間のギャップを大幅に埋めることを示しました。

4. 実験結果

4.1 長文脈リトリーバル (32k RULER)

11のタスクからなる32k RULERベンチマークで評価：

NSA Baseline (MLP): 平均スコア 0.2999.
NSA Mean-Pool (NoPE): 平均スコア 0.5554 (NoPEの利点を示しています).
COBS (Subspace, $r=4$ , FP4): 平均スコア 0.8195.
Dense Attention: 平均スコア 0.9040.
解消されたギャップ： COBSは、NSAベースラインとデンス・アテンションの間のギャップの約 86% を埋めました。
トラフィック効率： COBSは、NSAベースラインの 1.21倍 のKVキャッシュ読み取りトラフィックのみを使用し、デンス・アテンションより 15.15倍 少ないトラフィックで動作します。

4.2 短文脈および言語モデリング

短文脈推論： COBSは短文脈の常識推論性能を維持し、標準的なベンチマーク（OpenBookQA, PIQAなど）においてデンス・アテンションと同等の性能を達成しています。
位置ごとのNLL (Position-wise NLL)： COBSは、平均次トークン負の対数尤度（NLL）が 1.633 と最も低く（デンスの場合は1.727）、リトリーバルの改善が長文脈の言語モデリング能力を犠牲にしていないことを示しています。

4.3 アブレーション研究

ランクのスイープ（Rank Sweep）： パフォーマンスはランク $r=8$ でピークに達し、 $r=16$ ではディストラクターター（妨害要素）からの符号なし分散質量の蓄積により低下します。著者らは $r \le 8$ を推奨しています。
部分空間 vs グローバル： 適応的なレイヤーごとの部分空間次元（ $s \approx 85$ ）は、グローバルな次元数128と同等の性能を示しながら、ストレージコストを削減しています。
量子化： 固有ベクトルのFP4量子化は、実質的にロスレスであることが示されています。

5. 意義と主張

本論文は、**累積モーメントの次数（cumulant order）**が、既存の疎アテンション手法に対する診断的なレンズであると同時に、改善のための実用的な設計軸でもあると主張しています。一次（平均）から二次（共分散）の要約へと移行することで、COBSはブロック疎アテンションが、メモリ帯域幅の使用量を大幅に削減しながらデンス・アテンションの品質に近づけることを実証しました。

著者らは、これは特定のトレーニングプロトコル（RULER形式のSFT）を用いた $\approx 1.2$ B パラメータ規模でのメカニズム研究であることを強調しています。これは、あらゆるLLMに即座に導入可能なシステムであると主張しているのではなく、ブロック選択における理論的なギャップを埋めることができるという証明です。この研究は、ブロック疎性が、選択メカニズムが十分に表現力を持っていれば（すなわち、より高い累積モーメント次数を持てば）、主要なモデルにおいても実現可能であることを示唆しています。

著者が述べている制限事項：

本研究は、特定のバックボーン規模および事前学習シーケンス長に限定されています。
比較には、オリジナルのNSAの報告とは異なるハイパーパラメータ（非重複ブロック、大きなMLPなど）を用いた制御されたNSAベースラインを使用しています。
疎なブランチでのNoPEの使用と、デンスなブランチでのRoPEの使用は、長文脈の比較において交絡変数となります。
結果はRULER形式のファインチューニングに依存しており、これがすべての長文脈シナリオに完全に汎用できるとは限りません。