Surpassing RAG with 20% Memory: ASM Latency Optimization on the NBA Benchmark

by
Table of contents

🌏 閱讀中文版本

The challenge of long-context inference often comes down to the physical constraints of attention computation.

In a standard Transformer architecture, processing a prefix during inference faces two structural constraints. First, prefix-attention overhead expands linearly with prefix length. Second, as generation progresses, the prefix gradually has less influence on model behavior. Existing approaches usually balance the high latency of retaining the full prefix against the high maintenance cost of internalizing information through gradient descent.

Attention-State Memory (ASM) offers a training-free alternative. It neither compresses the prefix nor modifies model weights. Instead, it turns the prefix into a lightweight, lookup-based external memory. The design aims to reduce inference overhead while maintaining, or even exceeding, the performance of traditional In-Context Learning (ICL) and RAG.

Two Structural Constraints: Linear Cost and Information Decay

To understand ASM’s design motivation, we first need to clarify two physical constraints in long-context inference.

The first is the linear expansion of attention computation. In a standard Transformer architecture, KV Cache memory usage and the cost of attending to the existing prefix at each generated token, or decode step, both grow linearly with prefix length. This affects not only latency in the initial prefill stage, but also the speed of subsequent decoding.

The second is the decay of prefix influence. Research indicates that as generation progresses, the model’s attention spreads across more tokens, gradually reducing the influence of the initial prefix on model behavior. In long-context settings, this means simply extending the prefix may not improve instruction following or reasoning ability in direct proportion.

Existing paths each work well in specific situations:

  1. Compression methods, such as KV Cache Compression: These approaches suit settings that allow approximation and are primarily constrained by memory. However, experiments show that compression approaches such as KVZip perform below ASM or ICL in some configurations. This suggests that aggressive KV cache compression can reduce key label information in ICL prompts.

  2. Parameterized methods, or training-intensive approaches: These internalize prefix information into model weights through gradient descent. This route requires frequent training cycles and is less flexible when prefixes need frequent updates, especially for applications that adjust prefixes dynamically.

Core Mechanism: Attention-State Memory (ASM)

ASM provides a training-free approach built around a simple idea: transform the prefix into a lightweight, lookup-based external memory.

Construction: From Queries to Centroids

ASM does not precompute every possible token. Instead, it builds its structure through the following steps:

  1. Collect representative queries: Run a set of representative forward passes and collect their attention outputs over the prefix for a target query set.

  2. Cluster and compress: Cluster these attention outputs and extract centroids as memory entries.

Inference: Hierarchical Lookup and Lossless Merging

During inference, ASM no longer runs full attention over the original prefix. It improves performance through the following mechanisms:

1. Hierarchical Lookup

By indexing centroids, ASM reduces lookup cost to O(logK). This decoupling changes the relationship between retrieval cost and the number of memory entries from linear to logarithmic growth. It allows memory to grow without adding latency linearly. Even as memory expands, inference latency does not rise proportionally.

2. Online-Softmax Identity

This is the mathematical foundation that enables ASM’s lossless merge. Using this identity, ASM can merge a query’s attention state with the attention states of precomputed centroids. The process can recover the full attention output precisely without actually performing full attention over the prefix. One distinction matters: “lossless” here applies only to the online-softmax merge step. Centroid clustering and lookup remain approximate selections, so ASM does not precisely reconstruct the original full-prefix attention for every possible query.

Empirical Results and Performance Analysis

The research evaluated LLaMA-3.1-8B on a single NVIDIA RTX Ada 4500 GPU with batch size 1 and 512-token queries. The following are the key experimental findings:

Accuracy and Memory Efficiency

  • ManyICLBench: Across memory budgets from 1K to 8K, ASM showed higher accuracy than In-Context Learning (ICL) on LLaMA-3.1-8B. At an 8K memory budget, the paper reports a 1.36× improvement in attention latency, using the paper’s definition of that ratio.

  • NBA Benchmark: ASM exceeded the performance of full-attention RAG while using only 20% of its memory footprint.

  • Banking77: Across different memory sizes, ASM matched or exceeded ICL baselines at the same budget.

Key Technical Observations

  1. Non-monotonicity: Accuracy does not increase linearly as the number of memory entries grows. Instead, it peaks at an intermediate codebook size. The best number of memory entries is therefore a task-specific hyperparameter that needs tuning, rather than a value determined by prefix length alone.

  2. Whitening variant: On tokenization-sensitive reasoning tasks such as gpqa_cot and nba, the whitening variant performs better. The paper attributes this to leading-space tokens playing a key role in distinguishing answer choices.

  3. Peak memory for long-prefix calibration: The paper also notes that ASM’s composable structure can complete long-prefix calibration with substantially lower peak memory while retaining the performance of full-prefix construction. This adds an important deployment distinction: the 20% memory result comes from the NBA Benchmark, while the long-prefix calibration claim concerns peak memory during construction. They are not the same measurement and should not be interpreted together.

Tradeoffs & Boundaries

ASM is not an unrestricted optimization. Its performance is bounded by several dimensions:

1. Offline Cost vs. Online Efficiency

ASM trades additional offline construction, including forward passes for clustering, for logarithmic lookup cost during online inference. In the single NVIDIA RTX Ada 4500 GPU test environment, this strategy of using precomputation to reduce latency has potential when the same long prefix is reused frequently. However, the source does not quantify the cost of memory construction, clustering, or updates. It is worth validating end-to-end value through reuse frequency, construction cost, online latency, and accuracy together.

2. Applicability Boundary: Query-Distribution Representativeness

ASM’s performance rests on the assumption that representative queries cover the target query distribution. If online queries differ substantially from the distribution used during offline construction, the precomputed centroids may not represent attention states precisely. This makes ASM suitable for settings with relatively static or predictable prefixes, such as fixed rule bases. In highly dynamic, unpredictable prefix settings, offline-built centroids may no longer apply.

3. Applicability Boundary: Prefix Updates and Reconstruction Cost

ASM’s performance is also affected by dynamic prefix updates. When online prefixes change frequently, precomputed centroids cannot be applied directly. The system may need to repeat offline construction or rebuild memory. This expands the operational surface of deployment and bounds its real-time advantage in highly dynamic settings.

4. Task Sensitivity and Hyperparameter Tuning

ASM’s results vary across tasks. For reasoning tasks that require fine-grained distinctions between token features, such as gpqa_cot, selecting the right variant is important. Because accuracy is non-monotonic with respect to the number of memory entries, the optimal codebook size also needs task-specific tuning.

Conclusion

ASM offers a technical path for transforming long prefixes from a computational burden into a lookup task. Through lossless online merging and hierarchical indexing, it substantially reduces memory use in specific settings such as the NBA Benchmark while maintaining or improving accuracy across several benchmarks. Its performance, however, is bounded by the representativeness of the query distribution and hyperparameter tuning for codebook size.

This is a direction worth considering: before deployment, validate whether query-distribution stability is sufficient to support offline precomputation. For highly dynamic settings, the boundary of this offline approach may require validation with more real traffic to clarify.

Sources