Microsoft Research Asia's Baotong Lu: Reshaping Model Attention with Vector Retrieval | Attention

In large language models, ultra-long-context inference capability is one of the primary bottlenecks affecting performance.

This stems from the quadratic complexity of self-attention and the memory overhead of KV caching, which grows with sequence length. For instance, an 8B model processing 1M tokens of context can easily exceed 100GB of VRAM for KV cache alone — far beyond what a standard GPU can handle.

Addressing this challenge, a new installment in our Attention series interview focuses on the mechanism proposed in the paper Retrieval Attention: Accelerating Long-context LLM Inference via Vector Retrieval: a training-free, dynamic sparse attention scheme for ultra-long-context inference.

Below is Oasis Capital's interview with Baotong Lu, Senior Researcher at Microsoft Research Asia and one of the paper's core authors. Full read time: approximately 20 minutes.

Enjoy

Retrieval Attention's core insight is that each Query actually only needs strong interaction with a small subset of Keys — the rest is redundant attention; attention itself is naturally sparse.

The research team's key approach: offload most KV vectors from GPU to CPU, then during each inference step, use approximate nearest neighbor (ANN) search to identify the small number of Keys most relevant to the current Query — say, just the top 1% — and compute attention in parallel with a small set of "predictable" KVs already on GPU, merging the results.

In practice, on an RTX 4090 (24GB), an 8B-class model can generate stably at 128K context length with ~0.188s per token, achieving near-identical accuracy to full attention (verified on ∞-Bench, RULER, and other benchmarks).

The follow-up work RetroInfer further extended these performance gains, achieving 4.5× decoding throughput compared to full attention on A100 GPUs, and 10.5× throughput versus other GPU-CPU sparse attention systems at 1M token context length.

Retrieval Attention's core innovation lies in introducing vector retrieval into the attention computation path, achieving true dynamic sparsification at the computational level.

Specifically, the authors construct an attention-aware retrieval index: by learning precise Query-to-Key mappings during the prefill phase and projecting these relationships into a Key–Key graph structure, the retrieval process achieves high recall with extremely low scan rates (~1–3%). This design allows each Query to dynamically locate the most critical Keys at minimal computational cost, overcoming the inefficiency that traditional ANN search suffers when applied to attention due to mismatched Query-Key distributions.

At the system architecture level, Retrieval Attention further proposes a CPU–GPU collaborative dual-path attention mechanism: the GPU retains a small cache of "predictable" local KVs, while the CPU dynamically calls upon large-scale KV storage via retrieval. The two paths compute independently in parallel, with results fused through a numerically stable rescaling formula — achieving simultaneous reductions in memory footprint and inference latency.

Notably, the entire mechanism requires no model retraining. Retrieval Attention plugs into existing Transformers as a modular component, modifying only the attention layer's forward logic to significantly accelerate long-context inference without sacrificing accuracy. This "training-free dynamic sparse attention" approach offers a viable engineering path toward scalable long-context language models.

The following is an excerpt from Oasis Capital's conversation with Dr. Lu:

Oasis Capital: Please start by introducing your research background, particularly your work in database management systems, and what line of thinking led to this design approach.

Dr. Lu: My main research area is database management systems — how to store, retrieve, and manage massive data more efficiently. From a disciplinary perspective, database systems and machine learning may seem like separate fields, one leaning toward systems and structure, the other toward models and optimization. But they share common problems at the foundation, such as how to organize information more efficiently given limited resources.

During my PhD, I focused on how new hardware architectures — persistent memory and disaggregated memory, for example — pose new challenges for database design. This included designing new data structures and indexes for such hardware, and optimizing storage system query performance. At its core, this direction was about combining changes in underlying hardware with data management performance problems.

Oasis Capital: You mentioned that your research centered on enabling more efficient data management across different hardware architectures. How did you gradually extend from traditional database optimization problems to today's model inference-related research?

Dr. Lu: The evolution of traditional database systems is essentially a series of refinements to support more efficient data retrieval and access. The core goal is simple: make queries faster and more stable. Index structure optimization, cache mechanism improvements — these all serve to maintain query performance as data volumes grow.

But as hardware architectures change — persistent memory, disaggregated memory, tiered storage systems — these new technologies have many incompatibilities with traditional patterns. So we need to rethink: in this new architecture, how should database systems be designed to maximize hardware potential?

My PhD research mainly revolved around this.

On one hand, I focused on performance optimization under new hardware, such as reducing access latency and improving index efficiency. On the other hand, I studied how to make systems more intelligent when processing large-scale tasks. These experiences later influenced my current research direction.

Now our team is more focused on how to migrate the "retrieval" logic from traditional database systems to the model level. For example, when models face long-context inference, the bottleneck they encounter is very similar to databases: how to quickly locate key information within massive storage. This led us to think that perhaps we could combine database retrieval thinking with model attention mechanisms.

Our approach attempts to transplant mature vector retrieval methods from databases into the language model inference process, allowing the model to access only the "most relevant" information during generation. In other words, rather than retraining the model, we hope to make it utilize existing memory more efficiently through system-level design. This was essentially the starting point for our current paper.

Oasis Capital: Could you elaborate on how this hardware-oriented research approach took shape? How did you think about and integrate hardware or system-level considerations at the time?

Dr. Lu: This idea actually originated from a very concrete problem.

We realized that the attention mechanism is essentially a vector retrieval process — when generating, the model needs to continuously "look up" relevant information from existing context. And in this process, memory (VRAM) becomes the primary constraint.

Modern GPUs have very limited memory. We found that during long-context inference, models accumulate massive KV caches, and the way these caches are managed is very similar to the "index + cache" logic in traditional database systems. So we wondered: since databases can efficiently manage massive data through tiered storage and retrieval mechanisms, could model attention mechanisms borrow similar ideas?

So we began experimenting with offloading part of the cache from GPU to larger storage layers, such as CPU memory, then dynamically calling needed portions through retrieval. This "tiering + retrieval" structure allows the system to remain efficient in long-context scenarios without incurring additional training costs.

Oasis Capital: So essentially, you've migrated the "storage—index—retrieval" thinking from database systems to model attention computation?

Dr. Lu: Yes, the thinking is very similar. Databases aim to make data access faster; we focus on making the process of "finding information" faster during model generation. The difference is that we face unstructured, dynamically changing context information, so retrieval and caching strategies must be redesigned.

You could say we're reapplying the performance optimization accumulated in traditional database systems to model attention mechanisms, making their computation closer to real-world "storage and access" processes.

This was also the original motivation behind Retrieval Attention.

Oasis Capital: Could we start from more fundamentals to help readers understand unfamiliar concepts? In database management systems, what is the core logic of "data management mechanisms"? Can it be understood as the core of a "large database"?

Dr. Lu: Let me briefly explain. In machine learning or cognitive science contexts, the attention mechanism originally emerged as a simulation of human attention — when processing a sequence of information, how does the model determine which content is more important and which can be ignored?

For example, when given a question, the model doesn't treat all input information equally; instead, it automatically distributes "attention weights," concentrating on portions more relevant to the question. This is the core idea of attention mechanisms: within a sequence, based on task objectives, focus computational resources on that small subset of truly important elements.

Early sequence models like RNNs or Seq2Seq already possessed this "selective focus" characteristic, though they relied more on sequential structures.

The advent of Attention enabled models to establish connections between arbitrary positions: no longer constrained by time or position, it could directly capture global dependencies. This is why Transformer became a milestone architecture.

From an implementation perspective, attention can be understood as a "query—match—weight" process. The model decomposes input into three vector types: Query, Key, and Value.

• Query represents the information currently being processed — "what am I looking for right now";
• Key represents all potentially matchable clues;
• Value is the content corresponding to these clues.

The model computes similarity between Query and each Key to obtain a set of weights, then uses these weights to compute a weighted sum of all Values. The result is what the model considers "most worth attending to" — the final attention output.

If we place attention mechanisms within a database system framework, it can be understood as a dynamic information retrieval system. Each time the model generates a new token, it needs to "query" the most relevant information within existing semantic space — very similar to how databases execute query requests.

In this analogy, attention weights represent the "matching degree" of a query. Higher weights indicate that the currently generated token should "attend to" that portion of input more, extracting more information from the corresponding Value. In other words, every time the model generates a word, it's performing a "find and summarize" operation.

Mapping further to database semantics: Query is the user's query statement, Key corresponds to fields in an index table, and Value is the stored content. The model computes similarity between Query and Key to find the best-matching Value, then computes a weighted sum for output. Each generation step resembles a dynamic, multi-field fuzzy query within a database.

From a systems perspective, this mechanism resembles an "implicit retriever" — it completes information indexing and matching inside the model, but without explicit structuring. Our research starts from here, hoping to make this implicit, invisible retrieval logic explicit, and optimize it through systems approaches.

Simply put, we want to make the model's attention mechanism more like a "controllable database":

Each time it generates, rather than passively traversing all context, the model can actively query, filter, and call upon the information it truly needs.

Oasis Capital: Retrieval Attention — this is also your paper's core proposition. When you started this project, what specific problem did you most want to solve? Were there any unexpected discoveries or gains during the research process?

Dr. Lu: What we wanted to solve was the unavoidable performance bottleneck in long-context inference.

Traditional attention mechanisms see quadratic growth in computational complexity with context length, while simultaneously requiring storage of large numbers of Key–Value (KV) vectors, consuming enormous VRAM. Taking Llama-3 8B as an example: at 1 million token context, KV cache alone requires over 125GB VRAM, nearly exceeding single-GPU limits.

To reduce memory footprint, many systems transfer part of KV cache to CPU, but new problems emerge: at each generation step, which critical information should the model extract from thousands of cached tokens? These "important tokens" change dynamically and cannot be predicted in advance.

Retrieval Attention proposes a solution precisely for this problem. We store all KVs in larger CPU memory, while GPU retains only a small "active cache." Each time a new token is generated, the system uses vector retrieval to find the most relevant KVs in CPU memory, then dynamically loads them back to GPU for computation. This reduces memory footprint to roughly 1/10 of original, with almost no accuracy loss.

The real challenge lies in determining "which KVs are most valuable." To this end, we designed an approximate nearest neighbor (ANN)-based dynamic retrieval mechanism, and built a CPU–GPU collaborative scheduling system enabling parallel retrieval and computation. Ultimately, rather than passively traversing all context during generation, the model can actively "query" the information it needs, achieving efficient inference in long-context scenarios.

Oasis Capital: This is also one of this work's main contributions?

Dr. Lu: Yes, I think this work's main contributions can be divided into two levels.

The first is theoretical. We propose a new perspective: attention mechanisms can essentially be viewed as retrieval systems. Traditional vector database queries typically assume queries and data share the same distribution, whereas in attention mechanisms, queries and keys naturally lie in different distributions. Therefore, we redesigned attention's "retrieval structure," enabling the model to more precisely find that portion of truly important context information. In our follow-up work RetroInfer, we further observed that sparsity degrees (amount of important data) vary across different model layers and attention heads; our subsequent vector index design efficiently covers these varying sparsity degrees, achieving accuracy nearly identical to full attention.

The second level is systems and engineering implementation. Retrieval Attention is not merely an algorithmic improvement, but also a system optimization for CPU–GPU resource scheduling. We restructured traditional "linear cache" into a dynamic allocation structure, enabling the model to efficiently switch between different storage layers, significantly reducing memory and computation overhead in large-scale inference scenarios.

Additionally, our follow-up work RetroInfer made improvements for online service availability. Traditional systems often incur heavy overhead when building indexes, while our work achieves online dynamic indexing and querying through lightweight retrieval structures, substantially reducing construction and maintenance costs.

Overall, this work's value lies in: it not only proposes a new theoretical framework, but also makes attention mechanisms operationally scalable at the systems level.

Oasis Capital: Following up — what is Retrieval Attention's approximate theoretical error bound? From your perspective, what directions remain for future improvement?

Dr. Lu: This is an excellent question. To be frank, in this paper and current research, we remain primarily at the engineering and experimental exploration stage, without yet forming a complete theoretical framework. Retrieval Attention validates feasibility mainly through extensive empirical testing. For example, testing model performance at ultra-long context across different tasks, seeing to what extent this method maintains accuracy.

We discovered an interesting phenomenon: compared to original full attention mechanisms, dynamic retrieval-based attention performs almost equally on most tasks, but shows slight accuracy degradation on some high-precision tasks. This indicates an as-yet-unquantified "theoretical error bound" between computational efficiency and accuracy.

From a research perspective, this is actually quite interesting. Traditional sparse attention work often assumes this approximation is lossless, but our experiments show this assumption doesn't always hold. An important future direction is to establish more rigorous theoretical models to explain and bound the conditions of this approximation.

I personally believe subsequent research should achieve complementarity between theory and systems: on one hand continuing to optimize dynamic retrieval computation paths, on the other hand mathematically defining clear error upper bounds for sparse attention. This is both an engineering challenge and a theoretical topic worth deep exploration.

Oasis Capital: We'd like to ask an extended question about models relearning and optimizing attention at more fundamental levels, such as computation and systems. How do you see optimization opportunities in the long-context era?

Dr. Lu: We indeed focused subsequent research on systems-level optimization. Retrieval Attention's implementation essentially depends on a heterogeneous architecture, with CPU and GPU each assuming different tasks. GPU has strong computation but limited memory; CPU has relatively weaker computation but much larger memory space. These strengths are naturally complementary.

Therefore we chose to let GPU focus on forward computation, offloading main Key–Value cache to CPU memory, thereby supporting long-context inference on lower-cost hardware. But this also brought new challenges: data transfer bandwidth between CPU and GPU is small, becoming the system's primary bottleneck. To address this, our follow-up work RetroInfer made multiple optimization designs at the systems level.

One key idea draws on database system caching mechanisms. During inference, we observed that model attention distributions exhibit temporal locality — when continuously generating a text segment, the model tends to focus on a relatively small window.

Based on this characteristic, we designed a hot-cold data tiered caching strategy: storing high-activity "hot data" temporarily in GPU memory; placing less frequently accessed "cold data" in CPU memory; through dynamic cache updates, reducing frequent data exchange between the two, thereby significantly reducing inference latency.

Furthermore, we optimized CPU-side cache management software, reducing additional scheduling overhead. Experiments show this heterogeneous architecture enables the system to complete ultra-long-context inference at lower cost without sacrificing performance — for example, achieving 128K token inference on consumer-grade CPUs. Meanwhile, large-capacity CPU memory can support more concurrent requests, significantly improving system throughput.

Overall, we believe future large model inference frameworks should no longer be "GPU-only," but rather hybrid architectures that fully exploit heterogeneous hardware advantages. Retrieval Attention provides one validation of this: it enables cheaper, more scalable systems to achieve performance approaching mainstream GPU clusters.

Oasis Capital: In previous expert interviews on other attention mechanisms, we mentioned NSA, a dynamic absorptive learning mechanism, while Retrieval Attention leans more toward a training-free, engineering optimization path. Is there possibility for combining these two directions? Or could Retrieval Attention complement methods like NSA in the future?

Dr. Lu: I believe these two paths are not fundamentally opposed, but mutually complementary. Retrieval Attention and RetroInfer contain substantial long-context systems optimization, while NSA emphasizes the model's own dynamic learning and structural adaptability. The former improves inference efficiency through systems design; the latter hopes to let models learn to autonomously regulate attention distribution, obtaining an "index" through training.

Currently Retrieval Attention can stably handle context lengths above 64K, which is already considerable engineering-wise. But from a longer-term evolution perspective, dynamic attention (like NSA) may become the next breakthrough point, enabling models to autonomously "remember" which information deserves long-term retention without explicit retrieval.

From a systems perspective, we're also thinking about how to more efficiently implement this capability within GPU-CPU collaborative structures. Retrieval Attention's optimization occurs mainly at the boundary between storage and computation: we distribute attention computation load through CPU–GPU hybrid architecture, improving throughput and energy efficiency while maintaining accuracy.

Overall, these two research directions may converge in the future — model-level adaptive learning (NSA) on one side, systems-level efficient retrieval and scheduling (Retrieval Attention) on the other. When combined, NSA becomes an index, and Retrieval Attention can efficiently extend it to GPU-CPU architecture. We may see a new attention system that can both actively learn and self-manage "memory."

Oasis Capital: A more macro-level question: what position do you think this type of research will occupy long-term? From a broader perspective, what does this "long-context attention" approach mean for AI development overall?

Dr. Lu: I think the long-term significance of this research lies mainly in enabling models to possess true "long-term memory" capability. In the past, large models were often constrained by window size when processing information — they could only remember local context. With long-context attention, we can begin to let models maintain semantic consistency across enormous ranges, and build new application scenarios accordingly.

For example, a model with long-term memory can not only understand complete legal texts, source code, or entire books, but also maintain logical coherence across multi-turn interactions. This transforms the model from a "local understander" to a "systematic reasoner" — a qualitative change in capability.

Of course, this also brings new challenges.

As context grows longer, we need to rethink the coordination between system architecture and algorithms, how to let models dynamically allocate attention within massive information flows, and how to find balance between computational cost and inference efficiency. This is no longer a pure algorithm problem, but a comprehensive problem of systems design, data management, and computation optimization.

From a farther perspective, this work may push us to rethink how "knowledge" is organized.

Models' long-term memory will eventually require independent structures for management and updating, much like databases. Perhaps in the future, we'll see AI systems with autonomous knowledge management capabilities, capable of retaining information long-term, continuously learning, and achieving true scalability at both logical and physical levels.

About the Interviewee and Paper

This issue's interviewee: Dr. Baotong Lu

Currently Senior Researcher at Microsoft Research Asia, PhD from The Chinese University of Hong Kong (2023). Received the ACM SIGMOD Research Highlight Award in 2021; two works (Dash and APEX) selected for arXiv's most influential papers list.

Paper link: https://arxiv.org/pdf/2409.10516

Personal Github: https://baotonglu.github.io/

Thanks to Dr. Lu for participating in this conversation.

Oasis Capital sincerely invites frontline researchers, developers, and thinkers worldwide to jointly advance this exploration of attention.

If you're working on attention-related research — whether in algorithms, systems, cognition, neuroscience, or product design and content construction — and would like to explore and dialogue with us, please contact us through the QR code at the end!