Express 语言建模

arXiv:2606.10944v2 公告类型：替换 摘要：我们引入了一种新工具 Express，用于将非因果注意力近似转换为具有同等近似保证的因果近似。结合当前最先进的 Thinformer 近似方法，Express 改进了目前已知最佳的因果注意力近似保证：对于长度为 $n$ 的序列，仅需 $O(s)$ 内存和 $O(s^2 \log^2(n))$ 的压缩开销，即可实现 $\log^{3/2}(n)/s$ 的近似误差。我们为这些进展配备了高效的 I/O 感知 Triton 实现，展示了相较于 FlashAttention 2 的显著加速，并使用 Express 克服了语言建模流程中的四大资源瓶颈：长上下文预填充、KV 缓存压缩、内存受限的长文本解码，以及计算受限的长文本解码。