arXiv:2601.17702cs.CLcs.LG2026-01

用稀疏特征检索取代传统缓存,实现长文本推理的内存高效处理

S$^3$-Attention:Attention-Aligned Endogenous Retrieval for Memory-Bounded Long-Context Inference

  • 通过稀疏自编码器提取关键特征标识,构建基于CPU的倒排索引
  • 在固定预算下,性能接近全上下文推理,信息密集场景更鲁棒
  • 适合需要低内存、长文本生成的部署场景,如文档摘要与问答

大语言模型在多文档和长文本输入中的应用日益广泛,但长上下文推理仍面临内存和噪声效率低的问题。键值(KV)缓存随上下文长度线性增长,而外部检索方法常返回语义相似但因果无关的内容。本文提出S3-Attention,一种以内存为中心的推理时框架,将长上下文处理视为注意力对齐的内生检索。S3-Attention通过轻量级稀疏自编码器将解码过程中的临时键和查询投影转化为top-k稀疏特征标识,并在单次流式扫描中构建基于CPU的倒排索引,将特征映射到标记位置或片段。该设计允许完全丢弃KV缓存,且GPU内存使用量受扫描块大小限制。生成时,通过特征共激活检索紧凑证据片段,可选地融合BM25实现精确词法匹配。在统一的LongBench评估协议下,固定提示、解码和匹配令牌预算,S3-Hybrid在多个模型族上接近全上下文推理表现,并在多个信息密集场景中提升鲁棒性。我们也报告了当前原型的工程局限:相比优化后的全KV基线,存在更高壁钟延迟,推动未来内核级优化。

原文摘要 · Abstract (English)

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.

长文本推理内存优化检索增强稀疏编码

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。