提出Express工具，让非因果注意力转为因果注意力并提升效率。
Express Language Modeling
- 将非因果注意力近似转换为因果近似，保持精度保证
- 实现log³ᐟ²(n)/s误差，仅需O(s)内存和O(s² log²(n))开销
- 显著提速，解决长文本建模中的四大资源瓶颈
我们提出一种新工具Express，可将非因果注意力近似转化为因果近似，且保持相同的近似精度保障。结合当前最先进的Thinformer近似方法，Express在序列长度为n时，实现了log³ᐟ²(n)/s的近似误差，仅需O(s)内存和O(s² log²(n))的压缩开销。我们进一步设计了高效的I/O感知Triton实现，展示出相比FlashAttention 2的显著加速效果，并利用Express克服语言建模流水线中的四个资源瓶颈：长上下文预填充、键值缓存压缩、内存受限的长文本解码，以及计算受限的长文本解码。
原文摘要 · Abstract (English)
We introduce a new tool, Express, for converting a non-causal attention approximation into a causal approximation with matching approximation guarantees. When combined with the state-of-the-art Thinformer approximation, Express improves upon the best known causal attention guarantees, delivering $\log^{3/2}(n)/s$ approximation error with only $O(s)$ memory and $O(s^2 \log^2(n))$ compression overhead for a sequence of length $n$. We pair these developments with an efficient I/O-aware Triton implementation, demonstrate substantial speedups over FlashAttention 2, and use Express to overcome four resource bottlenecks in the language modeling pipeline: long-context prefill, KV cache compression, long-form memory-constrained decoding, and long-form compute-constrained decoding.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。