RAT+: Train Dense, Infer Sparse -- Recurrence Augmented Attention for Dilated Inference

Xiuying Wei; Caglar Gulcehre

cs.LGFeb 20, 2026

RAT+: Train Dense, Infer Sparse -- Recurrence Augmented Attention for Dilated Inference

Authors: Xiuying Wei, Caglar Gulcehre

Organizations: CLAIRE lab at EPFL, Lausanne, Switzerland.

Abstract

Structured dilated attention has an appealing inference-time efficiency knob: it reduces the FLOPs of attention and the KV cache size by a factor of the dilation size D, while preserving long-range connectivity. While prior work studies it by training each configuration from scratch, directly sparsifying a pretrained attention model into a dilated pattern leads to severe accuracy degradation, preventing flexible reuse across inference scenarios. We introduce RAT+, a dense-pretraining architecture that augments attention with full-sequence recurrence and active recurrence learning. A single RAT+ model is pretrained densely once and can then be flexibly switched at inference time to dilated attention (optionally with local windows) or hybrid layer/head compositions, requiring only a short 1B-token resolution adaptation rather than retraining separate sparse models. At 1.5B parameters trained on 100B tokens, RAT+ closely matches dense accuracy at D = 16, and drops by about 2-3 points at D = 64 on commonsense reasoning and LongBench tasks. We further scale to 2.6B and 7.6B parameters and observe even more promising performance (e.g., a 1-point average accuracy loss with a 64x reduction in attention FLOPs and KV cache size). Code is available at https://github.com/wimh966/rat-plus.

Figures & tables

Figure 1: (a) For architectural simplicity, we adopt an extreme overlapped setting, i.e., full-sequence recurrence with $\mathtt{L=T}$ . (b) Joint training to preserve dense attention capability while enforcing active recurrence learning with desired effective length $\mathtt{L^{*}}=64$ . (c) After pretraining, the resulting model can be efficiently adapted to various sparse inference patterns including effective results on dilated attention (optional local attention). It also has better performance on top-k block attention compared to attention.

Figure 2: Efficiency results of the temporal-mixing operator on a single GH200 GPU, covering both prefilling and decoding scenarios with hidden dimension

\mathtt{H}

. Prefilling latency is measured on sequences of 262K tokens. Decoding latency is measured for 256 or 128 batches of tokens for the two hidden dimensions, respectively; the baseline runs out of memory beyond 32K tokens. We use FlexAttention ( Dong et al., 2024 ) for prefilling (with

\mathtt{D=1}

reducing to FlashAttention), and FlashAttention ( Dao et al., 2022 ) for decoding. Recurrence is implemented by associative scan in PyTorch for prefilling and a simple step update for decoding. Additional results are reported in Subsec. B.4 .

Figure 3: Results of top-k block attention with block size

\mathtt{D}

and number of selected blocks

\mathtt{K}

on the hard NIAH-MK-2 and NIAH-MK-3 tasks from the RULER benchmark with

\mathtt{T=4096}

. RAT+ (no ARL in SFT) means disabling active recurrence learning during the SFT, which further demonstrates the benefit of the recurrence. More results include remaining tasks in Table 17 , MoBA-style top-k block attention in Table 18 .

Figure 4: Maximum decoding throughput (tokens/sec) of the full 1.5B and 7B models for decoding 1024 tokens, measured at context lengths of 4096 and 16384, corresponding to prefilling lengths of 3072 and 15360 tokens, respectively.

Figure 5: Comparison with GQA/MQA using different numbers of KV heads. Joint training is also applied to GQA/MQA (

\mathtt{D^{\dagger}=1}

\mathtt{W=64}

) to match training FLOPs. RAT+ achieves lower PPL and is more flexible: it requires only a single pretraining run and allows the local KV cache size to be configured at inference time.

Figure 6: Scaling-up experiments: we report validation loss on a held-out 0.5B-token subset to illustrate the even smaller loss gap between dense and sparse variants as model size increases. The starred points refer to attention models trained with

\mathtt{D^{\dagger}=1}

and

\mathtt{W=64}

, matched in training FLOPs at the same model scale, and are included as a reference for comparison with dense RAT+.

Figure 7: Choice of

\mathtt{L=64}

\mathtt{L^{*}=16}

means

\mathtt{D=1}

and

\mathtt{D=16}

joint training to enable 16 length capability for recurrence. Different performance on

\mathtt{D=1}

comes from different training FLOPs and different recurrence abilities in capturing the corresponding length. Grey points correspond to recurrence length generalization.

Figure 8: 1B-token adaptation on two pretrained models. It is evident that various dilated patterns quickly achieve stable loss values within a few hundred million tokens. We employed a simple optimization scheme with no warmup, which may explain the slight loss increase of

\mathtt{D=1}

at first, after which it recovers.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 9: L2 norm values of recurrence outputs at different time steps. We observe that the outputs at early time steps differ significantly. The first row shows an initialized network using our simple recurrence at layers 0, 6, 18, and 23. The second row corresponds to the same initialized network but with a non-zero initial cell state provided to the recurrence. The third row shows the results of the pretrained network. We visualize the outputs of our simple recurrence here; a similar phenomenon has also been observed in standard recurrent networks, as reported in Cooijmans et al. (2016) .

SFT models	NIAH-S-1	NIAH-S-2	NIAH-S-3	NIAH-MK-1	NIAH-MK-2	NIAH-MK-3	NIAH-MV	NIAH-MQ
$\mathtt{D=1}$	100	100	100	100	100	99.6	99.85	100
$\mathtt{D=2}$	100	100	99.2	100	99.2	96	99.8	99.8
$\mathtt{D=4}$	100	100	97.8	99.8	97.4	88.8	99.75	99.5
$\mathtt{D=8}$	100	100	96	99.8	96.4	73.8	99.4	99.15
$\mathtt{D=16}$	100	100	89.2	99.6	95.2	54.6	99.3	98.15
$\mathtt{D=32}$	100	100	78.8	99.2	78	40.4	97.35	96.35

Appendix

Table 16: Retrieval ability : Accuracy performance with exact match scoring on the Needle-in-Haystack tasks with different configurations from the RULER benchmark ( Hsieh et al., 2024 ) . We use

\mathtt{T=4096}

and bold less favorable results. We found they mainly occur when the text query/key/values are the UUID, and may incorrectly copy parts of the answers. Also, in NIAH-MK-2, larger dilation sizes also perform worse.

Quest top-k	NIAH-S-1	NIAH-S-2	NIAH-S-3	NIAH-MK-1	NIAH-MK-2	NIAH-MK-3	NIAH-MV	NIAH-MQ
attention
$\mathtt{D^{\dagger}=1}$ , All	100	100	99.8	100	100	99.6	100	99.9
$\mathtt{D^{\dagger}=16}$ , $\mathtt{K=16}$	100	99.8	99.4	99.6	71.2	5.4	93.95	94.25
$\mathtt{D^{\dagger}=16}$ , $\mathtt{K=64}$	100	100	100	100	99.6	85.2	99.9	99.6
$\mathtt{D^{\dagger}=64}$ , $\mathtt{K=8}$	100	100	85.8	98.4	20.6	0.2	87.55	88.55
$\mathtt{D^{\dagger}=64}$ , $\mathtt{K=16}$	100	100	99	99.8	63.2	15.2	97.45	96.85

Appendix

Table 17: Retrieval ability of the top-k block attention for different models. On the following three pretrained-plus-SFT models, we apply Quest ( Tang et al., 2024 ) style top-k block attention, where the critical block is determined by the Min/Max statistics of attention keys within each block.

MoBA top-k	NIAH-S-1	NIAH-S-2	NIAH-S-3	NIAH-MK-1	NIAH-MK-2	NIAH-MK-3	NIAH-MV	NIAH-MQ
attention
$\mathtt{D^{\dagger}=1}$ , All	100	100	99.8	100	100	99.6	100	99.9
$\mathtt{D^{\dagger}=16}$ , $\mathtt{K=16}$	100	91	56	84	65.2	6.2	48.1	55.9
$\mathtt{D^{\dagger}=16}$ , $\mathtt{K=64}$	100	99.8	96.2	99.6	97.8	77	95.35	96.45
$\mathtt{D^{\dagger}=64}$ , $\mathtt{K=8}$	90.2	66.8	20	54.2	29.8	5.8	30.1	34.55
$\mathtt{D^{\dagger}=64}$ , $\mathtt{K=16}$	99.2	95.8	79	91.8	66.2	22.6	70.7	76.35

Appendix

Table 18: Retrieval ability of top-k block attention across different models. For the following three pretrained-plus-SFT models, we apply MoBA-style top-k block attention ( Lu et al., 2025 ) , where the critical block is selected based on the mean-pooled attention keys within each block. Notably, we apply this mechanism to all layers and to both the prefilling and decoding stages, without any further fine-tuning. In contrast, the original MoBA paper uses continued pretraining and reports applying the method only during the prefilling stage to avoid noticeable accuracy degradation. Our goal here is primarily to demonstrate the effectiveness of RAT+ over the attention backbone.

LLM Inference Acceleration

LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.

See all 1,084 papers

Augmenting Attention with Exponentially Decaying Memory Improves Query-Aware KV Sparsity
May 27, 2026Xiuying Wei, Caglar GulcehreKV Caching Memory-Augmented Language Models
EPFL, CLAIRE Lausanne, Switzerland
RBS-Attention: Radius-Bounded Sparse Prefill for Long-Context Large Language Models
Sep 17, 2026Chuxu Song, Jiuqi Wei, Zhencan PengLong-Context Language Model Inference LLM Inference Acceleration
Department of Computer Science, Rutgers University · OceanBase, Ant Group
Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps
May 16, 2026Yanke Zhou, Yiduo Li, Hanlin Tang +6Self-Attention Long-Context Language Model Inference
Nanjing University · Alibaba Group

May 27, 2026J/K move · Enter open · S save

Augmenting Attention with Exponentially Decaying Memory Improves Query-Aware KV Sparsity

Xiuying Wei, Caglar Gulcehre

EPFL, CLAIRE Lausanne, Switzerland

KV Caching Memory-Augmented Language Models Self-Attention Sparse Attention Long-Context Language Model Inference LLM Inference Acceleration+1

Open paper AskPDF arXiv

Efficient inference is critical for long-context language models, where attention computation and KV-cache access dominate the cost. Recent work RAT+, introduces a recurrence-augmented attention backbone that enables flexible dilated attention at inference time. In this paper, we investigate whether this exponentially decaying memory can also improve existing query-aware sparse inference methods. Using representative methods including Quest, MoBA, and SnapKV, we show that RAT+ consistently improves accuracy over standard attention across sparse budgets on eight needle-in-a-haystack tasks. We validate these gains both on the released checkpoints from the RAT+ paper and on OLMo2-7B, which we continue pretraining with the added memory module for 10B tokens. Finally, we propose two hypotheses explaining why this memory module benefits query-aware sparse inference and design targeted experiments to support them.

Abstract

LLM Inference Acceleration

Explore similar work

Augmenting Attention with Exponentially Decaying Memory Improves Query-Aware KV Sparsity

RBS-Attention: Radius-Bounded Sparse Prefill for Long-Context Large Language Models

Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps