DELTA: Dynamic Layer-Aware Token Attention for Efficient Long-Context Reasoning

AI summary

Overview

Research area: Efficient inference and sparse attention for large reasoning models (LLMs that generate long chains of intermediate steps), in the Natural Language Processing category.

Technical level: Advanced. The paper assumes familiarity with transformer attention, KV caching, grouped-query attention, and GPU memory-bandwidth constraints.

Scope: This paper introduces DELTA, a training-free, layer-aware sparse attention mechanism that partitions transformer layers into full-attention layers, a small set of "Δ-layers" that select salient KV pages, and sparse layers that attend only to those selected pages, in order to speed up long-context reasoning decoding without losing accuracy.

What This Paper Is About

Large reasoning models solve hard problems by generating very long chains of intermediate reasoning steps, but this makes inference slow because every newly generated token must attend to the entire growing sequence of past tokens. Existing fixes evict entries from the key-value (KV) cache to shrink the active context, but those eviction methods lose accuracy badly on reasoning tasks because token importance shifts over the course of a long derivation and selection errors accumulate. DELTA's goal is to cut the computational cost of attention while keeping the full KV cache in GPU memory, so that accuracy on reasoning benchmarks matches or beats full attention.

Key Contributions

1. A token-level analysis of attention in large reasoning models, revealing two properties: (1) strong correlation of attention patterns across consecutive layers, and (2) gradual but ongoing shifts in token importance during long generations (termed "sequential drift").

2. DELTA, a training-free sparse attention mechanism that combines a head-aware token scoring rule with a stable recency window to retain the recent context most critical for reasoning, using a three-tier layer design (initial full-attention layers, a small set of Δ-layers, and subsequent sparse-attention layers).

3. Demonstration that DELTA matches or exceeds full attention accuracy on challenging reasoning benchmarks (AIME-2024, AIME-2025, GPQA-Diamond, MATH500) while delivering up to 1.54× end-to-end speedups.

4. A comparison against state-of-the-art sparse attention methods (Quest and RaaS) showing DELTA reduces the number of attended tokens by up to 4.25× without sacrificing accuracy.

Main Findings

• Attention is strongly correlated across consecutive layers. Profiling of models such as Qwen-7B shows tokens that receive high attention in one layer tend to remain salient in the next layers, so deeper layers largely preserve the spatial configuration of attention established earlier. This means computing full attention in every layer is redundant.

• Attention sharpens with depth. In early layers, attention is diffuse and focused on nearby tokens; as depth increases, attention becomes progressively sharper and concentrates on a small set of far-away tokens.

• Attention drifts gradually during decoding ("sequential drift"). Comparing decoding steps 900 and 1000 in a reasoning sequence shows the regions of strongest attention shift across key positions, motivating query-adaptive selection rather than history-based heuristics.

• Attention dominates decode runtime. Measured decoding runtimes of FFN and attention modules show FFN cost stays nearly constant while attention latency increases almost linearly with context length; beyond 8k tokens, attention dominates total inference time, driven by repeated KV-cache memory access rather than compute.

• Accuracy parity or improvement at a 1k-token budget. With a page budget of K = 64 and page size P = 16 (1k tokens) plus a recency window of L = 8 pages (128 tokens), DELTA consistently outperforms existing sparse methods and often matches or surpasses full attention. On AIME-2024 with DS-Qwen-14B, Quest and RaaS achieve below 20% accuracy, DELTA attains nearly 50%, and full attention reaches 60%.

• Larger budgets can beat full attention, then plateau. Increasing the budget from 1k to 2k tokens often improves performance; on GPQA with DS-Qwen-7B, DELTA outperforms full attention by roughly 30%. Expanding to 4k yields marginal or no improvement and occasionally a slight decline, suggesting DELTA captures most salient context within small budgets.

• Shorter or comparable reasoning trajectories. The CDF of decoding lengths shows DELTA matches or improves over full attention and outperforms other sparse-KV baselines, meaning its sparsity does not lengthen reasoning trajectories.

• Latency and throughput gains. On DS-Qwen-1.5B with batch size 64 and a maximum decoding length of 18k tokens, full attention latency grows from 7.5 ms to about 30 ms, whereas DELTA rises to only 13 ms, roughly 4× smaller growth. Overall decoding time decreases from 403 to 261 seconds, and throughput increases from 2,921 to 4,517 tokens/s, a 55% improvement.

• More Δ-layers cost more. Runtime consistently increases as the number of Δ-layers grows, because each Δ-layer runs full attention to refresh selected pages; as the count approaches the total number of layers, DELTA reduces to full attention.

• Δ-layer placement matters. Evaluating 10 different Δ-layer configurations on Mixed120 shows overall accuracy can vary noticeably, more so for the larger model, though relative trends across datasets remain largely consistent within a model.

• Recency window is a tuning trade-off. Accuracy is sensitive to L with differences up to 10%: at K = 256 and K = 512 pages (4k and 8k tokens), L = 8 performs best, while at the smallest budget K = 64 pages (1k tokens), a larger recency window performs better.

• Calibrated model-specific Δ-layers. DS-Qwen-1.5B uses layers [2, 14, 23] out of [0-27], DS-Qwen-7B uses [2, 14, 22] out of [0-27], DS-Qwen-14B uses [2, 6, 42] out of [0-47], and DS-Llama-8B uses [2, 8, 31] out of [0-31]. Layers [0, 1] always use full attention, and layer [2] is always the first Δ-layer.

Methodology in Plain English

The researchers began by measuring how attention behaves inside large reasoning models. They visualized attention maps layer by layer at different decoding steps and separately timed the FFN and attention components during generation. Two patterns stood out: attention patterns look very similar in adjacent layers, and attention focus moves gradually as generation proceeds.

Based on that, they designed DELTA to divide transformer layers into three groups. The first few layers (layers 0 and 1) run ordinary full attention, because early attention is diffuse and has no stable sparse structure. A small number of intermediate "Δ-layers" also run full attention, but their job is to score which tokens matter: for each token they take the maximum attention weight across heads, group tokens into fixed-size pages, sum the per-token scores within a page, keep the most recent pages as a recency window, and fill the rest of the budget with the highest-scoring older pages. Every remaining layer then runs sparse attention, attending only to the tokens in those selected pages. Because the selections are refreshed at each Δ-layer and recomputed for every generated token, the reduced context keeps up with the evolving query.

Crucially, DELTA never removes tokens from the KV cache. It reduces computation, not memory, so a token that seemed unimportant early can still be recovered later. The implementation uses the FlashInfer Just-In-Time module to extract attention logits directly from the decoding kernel and PyTorch's native topk for page selection, with a fixed page size of 16. The Δ-layers themselves are chosen by a lightweight calibration procedure that runs full-attention decoding on a small set, computes the average inter-layer shift between consecutive layers using 1 − cosine similarity averaged over decoding steps and samples, and picks layers with large shifts that are well distributed across depth.

Experiments used four distilled DeepSeek-R1 variants — DeepSeek-R1-Distill-Qwen-1.5B, 7B, and 14B, and DeepSeek-R1-Distill-Llama-8B — evaluated on AIME-2024, AIME-2025, GPQA-Diamond, and MATH500, with the baselines Quest and RaaS reproduced in Hugging Face Transformers. All experiments ran on a single node with eight NVIDIA A100 (SXM4, 40GB) GPUs. All 30 problems from AIME-2024 and AIME-2025 were used, along with the first 100 problems from GPQA-Diamond and MATH500.

Why This Matters

Impact on research. DELTA reframes the efficiency problem for reasoning models: instead of discarding KV entries permanently (which risks losing tokens that become important later), it exploits the observation that adjacent layers carry nearly the same attention structure, letting a few "refresh" layers pay the full-attention cost on behalf of many. This layer-aware reuse is presented as a promising direction for efficient reasoning-time inference, and the work is training-free and plug-and-play, so it does not require retraining or post-training.

Real-world applications.

• Serving reasoning assistants that generate tens of thousands of tokens of step-by-step mathematics or science reasoning, where decoding latency dominates user-perceived response time.
• Long-context inference for question answering over large documents, where the KV cache grows linearly with sequence length and batch size.
• Any deployment of small-to-mid-sized distilled reasoning models (1.5B to 14B) on cost-constrained GPU fleets, where the reported latency and throughput improvements translate directly to serving capacity.
• Reasoning-heavy agentic or tutoring pipelines that issue many sequential reasoning requests and are bottlenecked by per-request decode time.

Industry relevance. The paper notes that with a 32K-token context and batch size 128, the KV cache of Llama-3-8B in float16 already exceeds 500 GB, and that using full attention in HuggingFace, DeepSeek-R1-Distill-Llama-8B requires more than 15 minutes on a single NVIDIA A100 GPU to generate 32K tokens for one AIME problem. Reducing decoding bandwidth pressure is therefore a direct lever on serving cost, and DELTA's training-free integration makes it attractive for existing inference stacks. A key caveat for industry: DELTA does not reduce peak KV memory, so it does not address out-of-memory failures at extreme context lengths or on smaller GPUs.

Future Directions

• Combining DELTA with memory-saving techniques. The limitations section proposes integrating DELTA with quantization, eviction under guarantees, or offloading, so that it reduces peak memory as well as compute while maintaining high selection recall.
• Extending beyond distilled DeepSeek-R1 models. Generalization to other architectures, modalities, and workloads such as open-ended conversation and code generation remains unverified and may require re-tuning of the Δ-layer schedule and context budgets.
• Reducing sensitivity to hyperparameters. Performance depends on Δ-layer placement and the (K, L) budget pair; adaptive per-sample scheduling or lightweight learned selectors are suggested as fixes, along with addressing the small overhead of max-attention scoring when attention drift is fast.
• Reducing the selection-path overhead. The appendix breakdown shows DELTA introduces higher overhead in the collect-pages and planning stages because it supports two attention paths rather than one, leaving room for optimizing the selection machinery itself.

Target Audience

This paper is most useful to machine learning systems researchers and engineers working on LLM inference efficiency, KV-cache management, and sparse attention, as well as practitioners deploying long-chain reasoning models in production who need to reduce decoding latency without retraining. Readers with a working knowledge of transformer attention, grouped-query attention, and GPU memory bandwidth will get the most from it; the empirical sections on accuracy budgets and Δ-layer configuration are also accessible to readers focused on benchmarking rather than kernel implementation.

Authors’ abstract

Large reasoning models (LRMs) achieve state-of-the-art performance on challenging benchmarks by generating long chains of intermediate steps, but their inference cost is dominated by decoding, where each new token must attend to the entire growing sequence. One approach to reduce this latency is to evict entries from the key-value (KV) cache, thereby reducing the active context used in attention computation. However, such sparse attention methods suffer from severe accuracy degradation on reasoning tasks due to cumulative selection errors and the evolving importance of tokens over long derivations. We present \textbf{DELTA}, a training-free sparse attention mechanism that improves computational efficiency without sacrificing model accuracy. DELTA partitions transformer layers into three groups: initial layers that use full attention, a small set of \emph{$Δ$-layers} that identify salient tokens via aggregated head-level attention scores, and subsequent sparse-attention layers that attend only to the selected subset. This design preserves the full KV cache in GPU memory for accuracy, while avoiding expensive full-attention computation over many layers. On reasoning benchmarks such as AIME and GPQA-Diamond, DELTA matches or surpasses full attention in accuracy, while reducing the number of attended tokens by up to $4.25\times$ and delivering $1.54\times$ end-to-end speedup. Our results show that selective reuse of intermediate attention maps offers a robust path toward efficient long-context reasoning. The code is available at https://github.com/hoenza/DELTA.