[論文レビュー] S$^3$-Attention:Attention-Aligned Endogenous Retrieval for Memory-Bounded Long-Context Inference

Qingsen Ma, Dianyun Wang+8 more|arXiv (Cornell University)|Jan 25, 2026

Topic Modeling被引用数 0

ひとこと要約

本論文は、Top-k Sparse Autoencoders を用いて CPU ベースの inverted index を構築し、一定の GPU メモリでほぼ全文 context 性能を維持しつつ、長文 context 推論をストリーミング・注意アラインドプロセスへと変換する内生的検索フレームワーク S3-Attention を提案する。ハイブリッド版は内生的信号と BM25 を組み合わせてロバスト性を高める。

ABSTRACT

Large language models are increasingly applied to multi-document and long-form inputs, yet long-context inference remains memory- and noise-inefficient. Key-value (KV) caching scales linearly with context length, while external retrieval methods often return lexically similar but causally irrelevant passages. We present S3-Attention, a memory-first inference-time framework that treats long-context processing as attention-aligned endogenous retrieval. S3-Attention decodes transient key and query projections into top-k sparse feature identifiers using lightweight sparse autoencoders, and constructs a CPU-based inverted index mapping features to token positions or spans during a single streaming scan. This design allows the KV cache to be discarded entirely and bounds GPU memory usage by the scan chunk size. At generation time, feature co-activation is used to retrieve compact evidence spans, optionally fused with BM25 for exact lexical matching. Under a unified LongBench evaluation protocol with fixed prompting, decoding, and matched token budgets, S3-Hybrid closely matches full-context inference across multiple model families and improves robustness in several information-dense settings. We also report an engineering limitation of the current prototype, which incurs higher wall-clock latency than optimized full-KV baselines, motivating future kernel-level optimization.

研究の動機と目的

• 外部リトリーバの意味的ミスマッチを避けつつ、 memory-efficient な長文 context 推論を動機づける。
• モデルの内部推論信号に整合する内生的検索メカニズムを開発する。
• semantic indexing をストリーミングで実行し KV キャッシュを破棄することで、プレフィル時の O(1) GPU メモリを達成する。
• 内生信号と BM25 を融合して頑健性を高める S3-Hybrid を提案する。
• LongBench 全体でほぼ無損失の性能を、モデルファミリ間で示す。

提案手法

• Dense な内部 Key/Query 投影を Top-k Sparse Autoencoders (SAE) によって離散的な sparse semantic feature に変換する。
• ストリーミングのプレフィル中にキー投影を sparse feature IDs にエンコードし、CPU ベースの inverted index を構築して GPU KV キャッシュを破棄する（O(1) GPU メモリ）。
• クエリ投影も同じ SAE でデコードしてモデルの検索意図を取得し、 feature の共活性によって spans を取得する。
• クエリと feature の共活性を用いて semantic density スコアを計算し、 rare concept を強調するために IDF で重み付け、平滑化と非極大抑制を適用して semantic-rich な spans を選択する。
• 必要に応じて endogenous signals を BM25 の lexical signals およびバイアスと融合して最終の compressed context for generation を形成する（M_final = M_S3 ∪ M_BM25 ∪ M_Bias）。
• LongBench を用いて複数のモデル（Llama-3.1、Mistral、Qwen2）を評価し、FullKV、RAG、KV-cache compression のベースラインと比較する。

Figure 1 : Overview of the $S^{3}$ -Attention framework. The framework consists of two phases connected by a Top- $k$ Sparse Autoencoder (SAE). Streaming Semantic Indexing (red flow) encodes transient key projections into sparse semantic features to build a CPU-based inverted index, enabling the den

実験結果

リサーチクエスチョン

• RQ1 内生的な attention-aligned 信号を離散的で検索可能な memory に変換して長文 context 推論を実現できるか？
• RQ2 SAE ベースの sparse feature 表現は GPU メモリを削減しつつ因果証拠を保持するか？
• RQ3 内生的検索は外部インデックスを回避しつつ RAG と競合できるか？ GPU メモリを一定に維持できるか？
• RQ4 内生信号と lexical retrieval（BM25）を融合することでタスク全般の頑健性は向上するか？
• RQ5 S3-Attention を用いると fluent ness と evidence fidelity の情報理論上のトレードオフはどうなるか？

主な発見

• S3-Hybrid は Llama-3-8B で full-context の 99.4% を維持（25.01 vs. 24.87）、Qwen2-7B では 99%以上を unified evaluation で達成。
• SAE ベースの features による内生的検索は情報密度が高いタスクで exogenous RAG より信号対雑音比が高く、いくつかの設定でノイズ低減効果を示す。
• LongBench 全体で S3-Hybrid は O(1) GPU メモリでほぼ無損失の忠実度を達成し、複数の状況で強力な外部ベースラインを上回るか同等。
• 層別のアブレーションにより、より深い semantic layer が推論タスクの性能を向上させ、複層融合は一般にタスクを跨いで頑健性を提供。
• HotpotQA に関する情報理論的分析では、S3-Hybrid は高いリコールと低い KL ダイバージェンスを達成し、 fluent さと utilty のパレート frontier に位置する。

Figure 2 : Endogenous vs. Exogenous Retrieval. Top: RAG (BGE-Small) is distracted by high lexical overlap… Bottom: S 3 -Attention (Ours) ignores the distraction… (For a larger view, please refer to Figure 4 in the Appendix.)