Express 语言建模
Express Language Modeling
Albert Gong, Annabelle Michael Carrell, Raaz Dwivedi, Lester Mackey
AI 速览：提出 Express 工具，将非因果注意力近似转化为因果近似，并结合 Thinformer 实现更优的因果注意力保证与高效实现。
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.