一个模型既能全量训练又能稀疏推理，灵活适配不同效率需求。
RAT+: Train Dense, Infer Sparse -- Recurrence Augmented Attention for Dilated Inference
- 用全序列循环增强注意力，支持训练一次、推理时动态切换稀疏模式。
- 在1.5B参数下，稀疏度64时仅损失2-3点准确率，推理效率提升64倍。
- 适合需要高效推理且场景多变的部署应用，如边缘设备或实时系统。
结构化稀疏注意力在推理时具有理想的效率调节能力：通过稀疏度D可将注意力计算量和键值缓存大小降低D倍，同时保持长程连接。以往方法需对每种配置从头训练，而直接将预训练稠密模型稀疏化会导致严重精度下降，难以灵活复用。本文提出RAT+，一种稠密预训练架构，通过引入全序列递归与主动递归学习，使单个模型可在推理时灵活切换至稀疏注意力（可选局部窗口）或混合层/头结构，仅需约10亿词元的微调即可实现，无需重新训练多个稀疏模型。在1000亿词元上训练的15亿参数模型，在稀疏度为16时接近稠密模型性能，稀疏度达64时在常识推理与LongBench任务上仅下降2-3点；进一步扩展至26亿和76亿参数后表现更优，例如在稀疏度64下注意力计算量与缓存大小减少64倍，平均准确率仅损失1点。代码已公开于https://github.com/wimh966/rat-plus。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。