RAT+: 训练稠密、推理稀疏 —— 基于循环增强注意力的稀疏注意力方法

结构化稀疏注意力在推理时间效率方面具有引人注目的调节能力：它通过膨胀因子 D 减少注意力计算量和 KV 缓存大小，同时保持长程连通性。虽然已有工作通过从头训练每种配置来研究稀疏注意力，但直接将预训练注意力模型稀疏化为稀疏模式会导致严重的精度下降，阻碍了在不同推理场景下的灵活复用。我们引入 RAT+，一种稠密预训练架构，通过引入全序列循环和主动循环学习来增强注意力机制。一个 RAT+ 模型只需一次稠密预训练，即可在推理时灵活切换为稀疏注意力（可选带局部窗口）或混合层/头构数组合，仅需进行短暂的 10 亿 token 分辨率适应，而无需为每个稀疏模型单独重新训练。在 15 亿参数、训练 1000 亿 token 的规模下，RAT+ 在 D=16 时几乎匹配稠密模型的精度，在 D=64 时在常识推理和 LongBench 任务上仅下降约 2-3 分。我们进一步扩展至 26 亿和 76 亿参数，观察到更具前景的性能（例如，在注意力计算量和 KV 缓存大小降低 64 倍的情况下，仅损失 1 分的平均准确率）。代码已公开于 https://github.com/wimh966/rat-plus。