Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Qingsen Ma, Dianyun Wang, Yaoye Wang, Lechen Ning, Sujie Zhu, Xiaohang Zhang, Jiaming Lyu, Linhao Ren, Zhenbo Xu, Zhaofeng He

TL;DR

S3-Attention reframes long-context inference as attention-aligned endogenous retrieval by compressing transient Key/Query signals into sparse semantic features with Top-$k$ Sparse Autoencoders (SAEs) and building a CPU-based inverted index. The method discards the KV cache, achieving $O(1)$ GPU memory with respect to total context length, and optionally fuses endogenous signals with BM25 to form S3-Hybrid. Across LongBench and multiple model families, S3-Hybrid attains near-full-context fidelity while remaining memory-efficient, and exhibits a denoising effect that can improve performance on information-dense tasks. The work also documents engineering latency as a current limitation and points to kernel-level optimizations as a clear path forward for production systems.

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.

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

TL;DR

S3-Attention reframes long-context inference as attention-aligned endogenous retrieval by compressing transient Key/Query signals into sparse semantic features with Top- Sparse Autoencoders (SAEs) and building a CPU-based inverted index. The method discards the KV cache, achieving GPU memory with respect to total context length, and optionally fuses endogenous signals with BM25 to form S3-Hybrid. Across LongBench and multiple model families, S3-Hybrid attains near-full-context fidelity while remaining memory-efficient, and exhibits a denoising effect that can improve performance on information-dense tasks. The work also documents engineering latency as a current limitation and points to kernel-level optimizations as a clear path forward for production systems.

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.
Paper Structure (70 sections, 5 theorems, 18 equations, 13 figures, 9 tables, 1 algorithm)

This paper contains 70 sections, 5 theorems, 18 equations, 13 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Related Works
  4. Methodology
  5. Overview: From Exogenous to Endogenous Retrieval
  6. Theoretical Framework: The Endogenous Information Bottleneck
  7. Deciphering Attention via Sparse Autoencoders
  8. Phase 1: Streaming Semantic Indexing
  9. CPU Index Size.
  10. Phase 2: Semantic Density Estimation & Retrieval
  11. Phase 3: Hybrid Fusion
  12. Experiments
  13. Experimental Setup and Methodology
  14. Qualitative Analysis: Mechanism of Endogenous Alignment
  15. Visualizing the Semantic Gap.
  16. ...and 55 more sections

Key Result

Proposition 4.1

The scoring function in Eq. eq:actual_scoring can be interpreted as a weighted soft Jaccard similarity: where $w_f^{(Q)} = a_f^{(Q)} \cdot \text{IDF}(f)$ assigns higher weight to:

Figures (13)

  • 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 dense KV cache to be discarded and maintaining an $O(1)$ GPU memory footprint. Endogenous Retrieval (blue flow) encodes query projections with the same SAE, where activated features retrieve relevant context spans via feature co-activation, which are then fed back into the LLM for answer generation.SAE is trained on K projections and reused to discretize Q.
  • Figure 2: Endogenous vs. Exogenous Retrieval.Top: RAG (BGE-Small) is distracted by high lexical overlap... Bottom: S3-Attention (Ours) ignores the distraction... (For a larger view, please refer to Figure \ref{['fig:semantic_alignment_large']} in the Appendix.)
  • Figure 3: Enlarged view of Figure \ref{['fig:semantic_alignment']}: Endogenous vs. Exogenous Retrieval.Top: RAG (BGE-Small) is distracted by high lexical overlap, ranking a generic biography (Sentence 1) higher than the true answer (Sentence 5). Bottom: S3-Attention (Ours) ignores the distraction, showing sparse activation peaks exclusively at the semantic answer anchor ("The Post") and its reasoning evidence ("Pentagon Papers").
  • Figure 4: Enlarged view of Figure \ref{['fig:semantic_alignment']}: Semantic Attention vs. Lexical Retrieval.Top: RAG (BGE-Small) is distracted by high lexical overlap, ranking a generic biography (Sentence 1) higher than the true answer (Sentence 5). Bottom: S3-Attention (Ours) ignores the distraction, showing sparse activation peaks exclusively at the semantic answer anchor ("The Post") and its reasoning evidence ("Pentagon Papers").
  • Figure 5: Entity-centric and Biochemical/Taxonomic Reasoning Samples (1/2).Sample 2. Keith Nichol (Entity-centric Question).Query: Prior to playing for Michigan State, Keith Nichol played football for a school located in what city? Description: In this experiment, we compare a traditional Retrieval-Augmented Generation (RAG) baseline with a model equipped with the S$^3$-Attention mechanism on an entity-centric background knowledge question, and the results highlight clear advantages of S$^3$ in semantic retrieval and entity-level knowledge utilization. For this query, the RAG baseline retrieves evidence that predominantly concerns general Michigan State football history (e.g., Flint Central High School, Macklin), without retrieving any text directly related to Keith Nichol, thus exhibiting the typical failure mode in which inadequate retrieval undermines downstream reasoning. By contrast, S$^3$-Attention, under the same conditions, performs activation-based ranking and assigns a high activation score to the entity "Oklahoma", despite the absence of explicit external evidence mentioning the target entity. This indicates that S$^3$ goes beyond purely document-level retrieval and leverages semantic-level attention to identify relevant entities. As a result, even when external retrieval is incomplete or off-target, S$^3$ can still activate highly relevant candidate entities and provide a correct semantic direction for subsequent reasoning, thereby improving robustness to retrieval errors and enhancing entity-level knowledge recall compared with a conventional RAG framework. Sample 26. Ribosomal Subunits (RNA Question).Query: The large subunit and small subunit that use two types of RNA are major components that make up what? Description: In Sample 26, although the baseline retriever returns generic RNA-related passages, S$^3$'s attention concentrates on subword tokens such as osomal (from "ribosomal"), which are tightly coupled to the target concept "ribosome". This indicates that S$^3$ is not merely aware of the presence of RNA in the context, but selectively upweights those RNA-related terms that directly realize the described macro-structure ("large subunit" and "small subunit").
  • ...and 8 more figures

Theorems & Definitions (8)

  • Proposition 4.1: Scoring as Weighted Jaccard Similarity
  • Remark 4.2: Connection to Mutual Information
  • Proposition 4.3: IDF as Discriminative Weighting
  • Proposition 4.4: Feature Matching Heuristic
  • Remark 4.5: On the Original Proof
  • Proposition 4.6: Memory Complexity
  • proof
  • Proposition 4.7: Time Complexity