💬 NLP

S $^4$ R: Selective Sampling, Subspaces, and Sparse Reconstruction for Compressed Long-Context KV Caching

The paper proposes S $^4$ R, a novel KV cache compression method that combines prompt-aware selective sampling to build low-rank subspaces with sparse reconstruction during decoding, achieving up to 5 $\times$ compression with near-full accuracy while avoiding the calibration data dependence of offline methods and the high computational cost of online full-prompt reconstruction.

Original authors: Jialong Han, You Wu, Kewei Tu

Published 2026-08-04

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Jialong Han, You Wu, Kewei Tu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to remember a massive story to tell a friend. The longer the story gets, the more mental energy it takes to hold every single detail in your head at once. In the world of artificial intelligence, specifically Large Language Models (LLMs), this "mental energy" is called memory. These models are incredibly smart, but when they try to read or write about very long documents—like entire books or hours of conversation—they run out of memory because they try to save every single word they've ever seen. To fix this, scientists have been trying to figure out how to summarize the story in their heads without losing the important parts. They've tried two main tricks: either memorizing a generic summary that works for any story (which is fast but often misses the point), or trying to summarize the specific story they are reading right now (which is accurate but takes a huge amount of time and brainpower to calculate).

Enter S4R, a new method proposed by researchers at ShanghaiTech University that tries to get the best of both worlds. Think of S4R as a super-smart librarian who doesn't just memorize the whole library, nor does she just guess what's in the books. Instead, she quickly scans a few key pages to understand the general vibe of the story, keeps the very first few sentences (which often set the tone) in perfect detail, and then only pulls out the specific pages she thinks are needed for the next sentence she's writing. This allows the AI to handle massive amounts of text without running out of memory, while still being able to answer questions accurately. The researchers tested this on popular AI models and found that it can shrink the memory needed by up to 5 times while keeping the AI's performance almost as good as if it had remembered everything perfectly.

The Problem: The "Too Much Stuff" Dilemma

Large Language Models are like students who have read the entire internet. When they answer a question, they don't just guess; they look back at everything they've read so far to make sure their answer makes sense. This "looking back" requires a special storage area called the KV Cache (Key-Value Cache). Think of the KV Cache as a whiteboard where the model writes down the most important facts about the story it's reading.

The problem is that as the story gets longer (from a few sentences to a whole novel), the whiteboard gets huge. If the story is 128,000 words long, the whiteboard takes up so much space that it can be bigger than the model's brain itself! This makes the AI slow and expensive to run.

Scientists have tried to solve this in two ways, but both have a catch:

The "One-Size-Fits-All" Approach: Some methods try to compress the whiteboard using a fixed rule that works for any story. It's fast, but if the story is weird or unique, the compression might throw away the wrong details, and the AI gets confused.
The "Analyze-Everything" Approach: Other methods try to analyze the specific story while reading it to decide what to keep. This is very accurate, but it's like trying to summarize a book while reading it for the first time—it takes so much extra time that the AI becomes incredibly slow.

The S4R Solution: The "Smart Librarian" Strategy

The S4R method (Selective Sampling, Subspaces, and Sparse Reconstruction) acts like a clever librarian who knows exactly how to manage a massive library without getting overwhelmed. It uses three main tricks:

1. The "Anchor" Pages (Sink Tokens)
The researchers noticed that the very first few sentences of a story often act as a "glue" that holds the whole thing together. No matter what happens later, these opening lines are always important. S4R treats these first few words (called "sink tokens") like precious artifacts. It keeps them in their original, high-quality form and never compresses them. This ensures the AI always remembers the beginning of the story perfectly.

2. The "Quick Scan" (Selective Sampling)
Instead of trying to read and summarize the entire 128,000-word story at once (which is slow), S4R does a quick "sniff test." It picks a small, representative sample of words from the story—some from the beginning and some from the end—to figure out the general "shape" or "subspace" of the information. It's like flipping through a few random pages of a book to get the gist of the plot without reading every word. This allows the model to build a compact, efficient summary of the story's structure without doing the heavy lifting of analyzing every single token.

3. The "Just-in-Time" Retrieval (Sparse Reconstruction)
This is the magic trick. When the AI needs to write the next word, it doesn't try to reconstruct the entire compressed story. That would be too slow. Instead, it looks at the compact summary and asks, "Which parts of the story are actually relevant to what I'm writing right now?"

It always keeps the most recent few words (the "local window") because those are usually the most important.
It then scans the summary to find a few other "globally important" words from deep in the past that might be needed.
It only "reconstructs" (brings back to full detail) those specific words and the recent ones. It ignores the rest of the story for that specific moment.

What the Results Show

The researchers tested S4R on two major challenges: LongBench (a test of how well AI understands long documents) and RULER (a test of how well AI can find specific needles in a haystack of text). They used popular AI models like Llama and Qwen.

Here is what they found:

Massive Memory Savings: S4R was able to shrink the memory needed for the KV cache by up to 5 times. This is a huge deal because it means the AI can run on smaller computers or handle much longer stories.
Accuracy Stays High: Even with all that compression, the AI's accuracy remained very close to the "full memory" version. On the LongBench test, S4R scored nearly as high as the uncompressed models, beating other compression methods that tried to be too aggressive.
Speed Wins: When compared to other methods that try to analyze the whole story on the fly (like a method called xKV), S4R was much faster. It reduced the time it takes to start generating an answer (from about 80 seconds down to 27 seconds in one test) and made the overall writing speed about 4 to 5 times faster than those slower, heavy-duty methods.

The Bottom Line

S4R suggests that you don't need to remember everything perfectly, nor do you need to guess blindly. By keeping the "anchors" of the story safe, doing a smart, quick scan to understand the big picture, and only bringing back the specific details needed for the next step, AI models can become much more efficient. The researchers showed that this approach works well across different types of AI models and tasks, offering a practical way to make long-context AI faster and cheaper without losing its smarts. While the method isn't perfect (it still struggles slightly with very specific types of "needle-in-a-haystack" tasks compared to full memory), it represents a significant step forward in making long-document AI usable for everyone.

Technical Summary: S4R - Selective Sampling, Subspaces, and Sparse Reconstruction for Compressed Long-Context KV Caching

1. Problem Statement

The expansion of context windows in Large Language Models (LLMs) to hundreds of thousands or millions of tokens has created a critical memory bottleneck. The Key-Value (KV) cache, required to store intermediate computations for autoregressive generation, grows linearly with context length, often surpassing the memory footprint of the model parameters themselves.

Existing solutions for KV cache compression face a fundamental trade-off between efficiency and adaptivity:

Offline/Fixed Methods: Approaches like Palu, LoRC, and Eigen Attention perform low-rank decomposition (e.g., SVD) on calibration data or weights before inference. While efficient at decoding time, they are sensitive to distribution shifts and rely on fixed compression designs that may not align with specific prompts.
Online/Prompt-Dependent Methods: Approaches like ShadowKV and xKV estimate low-rank representations from the current prompt during inference. While more robust to context shifts, they incur prohibitive computational costs, particularly during the prefilling stage where decomposing the entire long prompt is expensive, and often require full-cache reconstruction during decoding, hurting throughput.

The paper identifies a gap: a method that combines the low prefilling overhead of fixed compression with the context adaptivity of prompt-dependent methods, without incurring the cost of full-prompt decomposition or full-cache reconstruction.

2. Methodology: S4R

S4R (Selective Sampling, Subspaces, and Sparse Reconstruction) addresses this dilemma through a two-stage process: Prefilling and Decoding.

2.1 Prefilling Stage: Selective Subspace Construction

Instead of applying Singular Value Decomposition (SVD) to the entire long-context sequence, S4R constructs low-rank subspaces from a selectively sampled subset of tokens.

Sink Token Preservation: The first $s$ tokens (identified as "sink tokens" which act as global pivots in attention) are excluded from compression and stored in full precision to maintain stable long-range dependencies.
Hybrid Sampling: For the remaining non-sink tokens, S4R samples a representative subset $X[I]$ $X [I]$ using a hybrid strategy:
1. Uniform sampling from early positions.
2. A contiguous block of the most recent tokens.
Normalization & SVD: The sampled keys and values are normalized along the feature dimension to prevent high-magnitude vectors from dominating the decomposition. Truncated SVD is then performed on this subset to derive low-rank bases ( $B_K, B_V$ ).
Coefficient Storage: Non-sink KV states are projected onto these learned bases, storing only the low-rank coefficients ( $C_K, C_V$ ) rather than full-dimensional vectors.
Query Projection: Local queries are also projected into the same low-rank subspace to enable efficient relevance estimation later.

2.2 Decoding Stage: Sparse Reconstruction

During generation, S4R avoids reconstructing the entire cache. Instead, it reconstructs only a sparse set of tokens deemed relevant.

Latent-Space Relevance Estimation: Token relevance is estimated directly in the latent coefficient space. The projected query coefficients are multiplied with the cached key coefficients to compute attention logits without full reconstruction.
Sparse Selection: A subset of tokens is selected for reconstruction based on:
1. Local Window: The most recent $W$ tokens are always retained.
2. Global Importance: The top- $k$ tokens from the remaining history are selected based on the averaged relevance scores derived from the latent space.
Sparse Reconstruction: Only the selected local and global tokens are reconstructed from their coefficients using the learned bases. These are concatenated with the full-precision sink tokens to perform the final attention computation.

3. Key Contributions

Prompt-Aware Low-Rank Compression: S4R constructs key/value bases from selectively sampled prompt tokens rather than external calibration data or the full prompt, trading a small amount of prefilling adaptation for better context alignment.
Sparse Reconstruction Strategy: By retaining only informative positions (local window + globally important tokens) during decoding, S4R avoids the computational bottleneck of full-cache reconstruction, significantly improving throughput.
Sink Token Preservation: The method explicitly preserves the first few "sink" tokens in full precision, addressing the instability often caused by compressing critical global attention pivots.
Efficiency-Accuracy Trade-off: The design successfully combines the low prefilling overhead of fixed methods with the adaptability of prompt-dependent methods.

4. Experimental Results

The authors evaluated S4R on LongBench and RULER benchmarks using Llama (1B, 8B) and Qwen (4B, 14B) model families.

Accuracy:
- LongBench: S4R achieves up to 5× KV cache compression while maintaining accuracy close to the full-cache baseline. On Qwen3-4B, it achieves an average score of 47.79 (vs. 47.79 for full cache), outperforming fixed baselines like Palu and competitive prompt-dependent methods like ShadowKV.
- RULER: Under a 2.5× compression ratio with a 128k context, S4R achieves an average score of 93.61, nearly matching the full KV cache (93.97) and slightly outperforming xKV (93.47). It shows particularly strong performance on multi-key and multi-query retrieval tasks.
Efficiency:
- S4R substantially reduces the overhead of prompt-dependent methods. Compared to xKV, S4R reduces Time to First Token (TTFT) from 79.90s to 27.65s and 512-token generation latency from 260.76s to 61.54s.
- While slightly slower than fixed methods like Palu in unoptimized implementations, S4R offers a superior accuracy-efficiency trade-off under aggressive compression ratios.

5. Significance and Claims

The paper claims that S4R provides a practical path toward adaptive KV cache compression for long-context LLM inference. Its primary significance lies in resolving the "dilemma" of existing low-rank methods:

It avoids the calibration-data dependence of offline methods.
It avoids the prohibitive prefilling and reconstruction costs of online methods.

By demonstrating that the dominant KV subspace can be approximated from a carefully selected subset of tokens and that full reconstruction is unnecessary at every decoding step, S4R enables high-fidelity long-context inference with significantly reduced persistent memory and manageable computational overhead. The authors note that their implementation is not system-optimized, suggesting that reported efficiency metrics reflect algorithmic overhead rather than a fully deployed system, and that further gains are possible with kernel-level optimizations.