arXiv:cs.LG· Albert Gong, Annabelle Michael Carrell, Raaz Dwivedi, Lester Mackey·· 7 小时前AI 评分45
Express:将非因果注意力近似转换为因果近似,改进 Thinformer 注意力保证
Express Language Modeling
AI 导读
研究者推出 Express 工具,可将非因果注意力近似转换为因果近似并保持近似保证。与 Thinformer 结合后,对长度为 n 的序列实现 log^{3/2}(n)/s 的近似误差,仅需 O(s) 内存和 O(s^2 log^2(n)) 压缩开销。
正文
Abstract: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.
| Subjects: | Machine Learning (cs.LG); Data Structures and Algorithms (cs.DS); Statistics Theory (math.ST); Methodology (stat.ME); Machine Learning (stat.ML) |
| Cite as: | arXiv:2606.10944 [cs.LG] |
| (or arXiv:2606.10944v2 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2606.10944 arXiv-issued DOI via DataCite |
Submission history
From: Albert Gong [view email]
[v1]
Tue, 9 Jun 2026 14:48:56 UTC (285 KB)
[v2]
Mon, 5 Oct 2026 21:36:43 UTC (1,559 KB)
来源:arXiv:cs.LG · arxiv.org