AI Navigate

IndexCache: クロスレイヤーのインデックス再利用によるスパースアテンションの高速化

arXiv cs.CL / 2026/3/13

💬 オピニオンDeveloper Stack & InfrastructureIdeas & Deep AnalysisModels & Research

要点

  • IndexCacheは、クロスレイヤーの上位k個のインデックスを再利用することにより、注意計算を $O(L^2)$ から $O(Lk)$ に削減することで、スパースアテンションの非効率を低減します。
  • 層を、インデクサーを備えた少数のFull層と、最も近いFull層の上位kインデックスを再利用する多数のShared層に分割します。
  • 訓練不要のIndexCacheは、キャリブレーションセット上で言語モデリング損失を直接最小化する貪欲探索アルゴリズムを適用し、重みの更新を必要としません。
  • 訓練対応のIndexCacheは、保持されたインデクサーを、それが担当するすべての層の平均アテンション分布に対して訓練する多層蒸留損失を導入し、単純な交互配置パターンでも全インデクサーの精度に匹敵させることを可能にします。
  • 30B の Deep Sparse Attention モデルを対象とした実験は、IndexCacheがインデクサー計算の75%を削減し、品質の劣化はほとんどなく、プレフィルで最大1.82×、デコードで最大1.48×のスピードアップを達成しました。これらの肯定的な結果は、図1に示す本番規模のGLM-5モデルでの予備的証拠によってさらに裏付けられています。

Abstract

長い文脈を持つエージェント的ワークフローは、大規模言語モデルの定義的なユースケースとして浮上しており、推論速度とサービングコストの両方にとってアテンションの効率性が極めて重要になっています。スパースアテンションはこの課題に効果的に対処します。DeepSeek Sparse Attention(DSA)は代表的な本番運用グレードの解決策であり、軽量のインデクサーがクエリごとに上位k個の最も関連するトークンを選択し、コアのアテンションを O(L^2) から O(Lk) に削減します。しかし、インデクサー自体は O(L^2) の複雑さを保持し、結果として得られる上位kの選択は連続する層間で非常に類似しているにもかかわらず、各層で独立して実行されなければなりません。我々はクロスレイヤーの冗長性を活用する IndexCache を提示します。IndexCache は層を、独自のインデクサーを走らせる少数の Full 層と、最も近い Full 層の上位kインデックスを再利用する多数の Shared 層に分割します。設定を決定し最適化するための二つの補完的アプローチを提案します。訓練不要のIndexCacheは、キャリブレーションセット上で言語モデリング損失を直接最小化する貪欲探索アルゴリズムを適用して、どの層にインデクサーを保持するかを選択します。重みの更新は不要です。訓練対応のIndexCacheは、多層蒸留損失を導入し、保持された各インデクサーを、それが担当するすべての層の平均アテンション分布に対して訓練します。これにより、単純な交互配置パターンでも全インデクサーの精度に匹敵させることを可能にします。実験は、30B の DSA モデルを用いて、IndexCache がインデクサー計算の75%を削減し、品質の低下を最小限に抑えつつ、プレフィルで最大1.82×、デコードで最大1.48×のスピードアップを達成することを示しています。これらの肯定的な結果は、図1に示す本番規模のGLM-5モデルでの予備的証拠によってさらに裏付けられています。

arXiv:2603.12201v1 Announcement Type: new

Abstract: Long-context agentic workflows have emerged as a defining use case for large language models, making attention efficiency critical for both inference speed and serving cost. Sparse attention addresses this challenge effectively, and DeepSeek Sparse Attention (DSA) is a representative production-grade solution: a lightweight lightning indexer selects the top-k most relevant tokens per query, reducing core attention from O(L^2) to O(Lk). However, the indexer itself retains O(L^2) complexity and must run independently at every layer, despite the fact that the resulting top-k selections are highly similar across consecutive layers. We present IndexCache, which exploits this cross-layer redundancy by partitioning layers into a small set of Full layers that run their own indexers and a majority of Shared layers that simply reuse the nearest Full layer's top-k indices. We propose two complementary approaches to determine and optimize this configuration. Training-free IndexCache applies a greedy search algorithm that selects which layers to retain indexers by directly minimizing language modeling loss on a calibration set, requiring no weight updates. Training-aware IndexCache introduces a multi-layer distillation loss that trains each retained indexer against the averaged attention distributions of all layers it serves, enabling even simple interleaved patterns to match full-indexer accuracy. Experimental results on a 30B DSA model show that IndexCache can remove 75% of indexer computations with negligible quality degradation, achieving up to 1.82× prefill speedup and 1.48× decode speedup compared to standard DSA. These positive results are further confirmed by our preliminary experiments on the production-scale GLM-5 model (Figure 1).