MAGE让块扩散模型提前锁定关键记忆，大幅提速长文本生成。
MAGE: All-[MASK] Block Already Knows Where to Look in Block Diffusion LLM
- 利用初始步的全掩码块预判后续所有步骤需关注的键值对
- 在128K上下文下实现6.82倍端到端加速，精度几乎无损
- 无需训练，适合长序列生成场景的推理优化
块扩散语言模型是并行生成语言的新范式，但其键值缓存导致长上下文推理中内存访问成为主要瓶颈。稀疏注意力通过仅关注少量键值对可降低延迟且精度损失小。但在块扩散中，每个块的B个令牌必须共享同一组键值对，这一约束使现有稀疏键值估计器召回率下降高达25%。我们发现块扩散训练目标自然产生一种特性：去噪步骤中块平均查询保持一致，因此首个[掩码]块在第一步就已揭示整个轨迹所需的键值子集。基于此，我们提出MAGE（[掩码]-引导稀疏注意力），一种无需训练的方法：仅在第一步执行一次精确注意力，然后复用其top-k索引集用于块内所有后续步骤。在LongBench上的三个块扩散模型家族中，MAGE在k=512时达到与精确注意力近乎无损的精度，128K上下文下端到端速度提升最高达6.82倍，并比专为自回归模型设计的Quest和全双向扩散模型设计的SparseD分别快3.35倍和2.28倍。
原文摘要 · Abstract (English)
Block diffusion LLMs are an emerging paradigm for parallel language generation, but their KV caching makes memory access the dominant bottleneck in long-context inference. Sparse attention, which attends only to a small KV subset per query, can reduce this latency with minimal accuracy loss. In block diffusion, however, the B tokens of each block must share a single KV subset, and we show this per-block constraint degrades existing sparse KV estimators by up to 25% in recall. We address this challenge by exploiting a property that emerges from the block-diffusion training objective: it aligns the block-average query across denoising steps, so the All-[MASK] block at the first step already reveals the per-block KV subset for the entire trajectory. We exploit this in MAGE ([MASK]-Guided Sparse Attention), a training-free method that runs one exact attention pass at the first step and reuses its top-k index sets for all remaining steps within the block. Across three block-diffusion families on LongBench, MAGE matches Exact Attention at k=512 with near-lossless accuracy, achieves up to 6.82x end-to-end speedup at 128K context, and runs up to 3.35x and 2.28x faster than Quest and SparseD, designed for AR LLMs and fully bidirectional diffusion LLMs, respectively.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。