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

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 64× reduction in attention FLOPs and KV cache size). Code is available at https://github.com/wimh966/rat-plus.

Large language models are powerful but expensive to run, especially when they process long contexts. Many efficient model designs require choosing one fixed efficiency setting during training, which makes them hard to adapt to different deployment needs. This paper introduces RAT+, a model architecture that is trained once in a dense setting but can flexibly switch at inference time to different sparse attention patterns and efficiency levels. RAT+ adds a lightweight recurrent memory that helps preserve information when the model attends to fewer tokens. Experiments show that the flexible dilated attention in RAT+ can reduce computation and memory use while maintaining strong performance across short-context, long-context, and retrieval-heavy tasks. This flexibility makes RAT+ a promising direction for building language models that are both accurate and easier to deploy.