[论文解读] 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

一句话总结

本文提出 S3-Attention，一种内生检索框架，将长上下文推理转化为流式、与注意力对齐的过程，使用 Top-k Sparse Autoencoders 构建基于 CPU 的倒排索引，在保持近全上下文性能的同时实现恒定的 GPU 内存。Hybrid 变体将内生信号与 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.

研究动机与目标

• 在不产生外部检索器引发的语义错配的前提下，激发记忆高效的长上下文推理。
• 开发与模型内部推理信号对齐的内生检索机制。
• 通过流式语义索引并丢弃 KV 缓存，在预填阶段实现 O(1) GPU 内存。
• 提出 S3-Hybrid，将内生信号与 BM25 融合以提升鲁棒性。
• 在 LongBench 上对多种模型家族（Llama-3.1、Mistral、Qwen2）展示近乎无损的性能，并对比 FullKV、RAG 与 KV-cache 压缩基线。

提出的方法

• 通过 Top-k Sparse Autoencoders (SAEs) 将密集的内部 Key/Query 投影转化为离散的稀疏语义特征。
• 在流式预填阶段，将 Key 投影编码为稀疏特征 ID，构建基于 CPU 的倒排索引以丢弃 GPU KV 缓存（O(1) GPU 内存）。
• 使用相同的 SAE 对查询投影进行解码，获得模型的检索意图并通过特征共激活检索片段。
• 利用与查询的特征共激活并结合 IDF 的权重来计算语义密度分数，强调罕见概念，随后进行平滑和非极大值抑制以选择语义丰富的片段。
• 可选地将内生信号与 BM25 词汇信号及偏置融合，形成最终压缩后的生成上下文（M_final = M_S3 ∪ M_BM25 ∪ M_Bias）。
• 在 LongBench 上对多模型（Llama-3.1、Mistral、Qwen2）进行评估，并与 FullKV、RAG 和 KV-cache 压缩基线进行对比。

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 内生的、与注意力对齐的信号是否可以离散化为一个轻量、可搜索的记忆用于长上下文推理？
• RQ2 基于 SAE 的稀疏特征表示是否在降低 GPU 内存使用的同时保留因果证据？
• RQ3 内生检索是否能够在保持恒定 GPU 内存且避免外部索引的情况下与 RAG 竞争？
• RQ4 将内生信号与基于词汇的检索（BM25）融合是否能提升跨任务的鲁棒性？
• RQ5 在使用 S3-Attention 时，流畅性与证据保真之间的信息论权衡如何？

主要发现

• S3-Hybrid 在 Llama-3-8B 的全上下文性能保留率为 99.4%（25.01 对 24.87），在 Qwen2-7B 的统一评估中也超过 99%。
• 基于 SAE 的内生特征在信息密集型任务上相较外源性 RAG 显示出更高的信号与噪声比，并在某些设置中展现去噪效果。
• 在 LongBench 的多种设置中，S3-Hybrid 以 O(1) GPU 内存实现近损失的保真度，在若干情景下优于或匹配强力的外源基线。
• 层级消融显示更深的语义层对推理任务性能有提升，多层融合通常在跨任务上提供鲁棒性。
• HotpotQA 的信息理论分析表明 S3-Hybrid 具备有利的召回率和更低的 KL 散度，使其在流畅性与实用性之间的帕累托前沿上。

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.)