📊 statistics

Express Language Modeling

本文介绍了 Express，这是一个将非因果注意力近似转换为因果注意力的工具，它具有改进的理论保证和高效的 Triton 实现，能够在克服长文本语言建模中关键资源瓶颈的同时，实现相对于 FlashAttention 2 的显著加速。

原作者： Albert Gong, Annabelle Michael Carrell, Raaz Dwivedi, Lester Mackey

发布于 2026-06-10

📖 1 分钟阅读☕ 轻松阅读

原作者： Albert Gong, Annabelle Michael Carrell, Raaz Dwivedi, Lester Mackey

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在试图通过阅读一本非常长的书来回答一个问题。在阅读过程中，你需要记住之前遇到的每一个重要细节，才能理解当前的句子。在人工智能语言模型的世界里，这种“记忆”过程被称为注意力（Attention）。

问题在于，随着书变得越来越长，记忆所有内容的精力需求会呈爆炸式增长。这就像是在一个挤满了人的房间里进行对话，每当你开口说一个新词，你都必须向房间里的每一个人重新介绍自己，并重新阅读你们之前的整个对话历史。这会让过程变得极其缓慢且耗费内存，以至于无法处理非常长的故事或复杂的推理任务。

这篇论文介绍了一个名为 Express（以及一个特定版本 Thinformer Express）的新工具，它充当了这些 AI 模型的超级高效图书管理员。以下是它的工作原理，我们使用简单的类比来解释：

1. 问题所在：“二次方”瓶颈

通常情况下，为了理解一个新句子，AI 会查看之前所有的单词。如果你有 1,000 个单词，它要做 1,000 次检查。如果你有 100 万个单词，它就必须做 1 万亿次检查。这就是论文中提到的“二次方（quadratic）”成本。这就像是在干草堆里寻找一根特定的针，你必须一遍又一遍地逐一检查每一根干草。

2. 解决方案：智能总结者（稀疏化/Thinning）

作者意识到，你并不需要完美地记住每一个单词。你只需要一个关于整个故事的、高质量的“摘要”或“核心集”，来代表全文。

他们构建了一个名为 Express 的工具，它能将一个“非因果（non-causal）”总结器（一种可以先读完一整本书然后挑选出最精彩的 100 页的工具）转化为一个“因果（causal）”总结器。

非因果 就像是先读完一整本书，然后再挑选精华。
因果就像是实时地逐页阅读书籍，在阅读过程中立即决定哪些精华值得放入口袋，而永远不会向前看（预知未来）。

Express 是一个神奇的技巧，它允许 AI 在不损失准确性的情况下实现这种实时总结。它确保了即使 AI 只关注过去极小的一部分内容（一个“稀疏化/thinned”后的版本），它仍然能得到正确的答案。

3. “膨胀”技巧

论文描述了一种巧妙的机制，让 AI 的记忆像气球一样“膨胀”和“收缩”。

精确阶段（Exact Phase）： 在最开始，AI 会完美地记住一切。
稀疏阶段（Thin Phase）： 随着更多单词的到来，AI 将它们分组。它不再保留所有单词，而是使用一种特殊的算法将这一组单词“压缩”成一个更小的、带有权重的摘要。
减半阶段（HALVE Phase）： 当摘要变得太大时，AI 会执行“减半”操作，在仔细保留最重要的信息的同时，将大小减半。

这个过程确保了无论故事有多长，内存大小都能保持在较小且恒定的水平。这就像是一个会自动缩小体积以适应空间的背包，只保留最核心的物品，这样你就永远不会耗尽空间。

4. 现实世界的结果：加速 AI

作者不仅做了数学推导，还使用了一种名为 Triton 的编程语言（类似于计算机芯片的高速引擎）构建了一个快速版本的工具。他们在四个特定场景下进行了测试：

长上下文预填充（阅读书籍）： 当给 AI 一个巨大的文档进行阅读前进行处理时，对于超长文本（512,000 个单词），Express 比当前最优秀的方法（FlashAttention 2）快了 82 倍。
压缩内存（KV 缓存）： AI 模型会将过去的单词存储在“缓存”中。Express 帮助它使用其他流行的方法压缩此缓存，在不损失准确性的情况下使整个过程更快。
内存高效解码（逐步思考）： 在解决需要长链推理的难题时，Express 让 AI 仅使用标准方法所需内存的 61%，却能得到同样的正确答案。
计算高效解码（加速思考）： 对于同样的数学问题，Express 让 AI 完成任务的时间仅为通常时间的 56%，同时仍能得到正确答案。

总结

简而言之，Express 是一种让 AI 处理长对话和复杂任务的新方法。它充当了一个智能过滤器，不断总结过去的内容，使 AI 的记忆保持精简、思考速度保持高效，同时不会让 AI “忘记”重要的细节。它解决了 AI 在处理长文本或复杂推理时变得过慢或内存耗尽的问题。

技术摘要：Express 语言建模

问题陈述

大型语言模型（LLM）中的注意力机制随序列长度 ( $n$ ) 呈二次方增长，这在长上下文预填充（prefilling）、顺序解码（sequential decoding）以及键值（KV）缓存管理方面造成了显著的计算和内存瓶颈。虽然目前已存在最先进的亚二次方（sub-quadratic）注意力近似方法（例如 Thinformer [4]），但它们主要解决的是非掩码（unmasked）注意力问题。对于语言建模而言，更关键的挑战在于因果掩码（causal masking），即每个查询（query）必须仅关注不断演进的先前键值对上下文。现有的因果近似方法要么缺乏强有力的理论保证，要么面临高内存开销，或者无法提供适用于令牌生成增量性质的高效流式更新。

方法论

EXPRESS 元程序

作者引入了 EXPRESS，这是一个元程序，旨在将任何高质量的、基于稀疏化（thinning）（即对输入点进行子采样）的非因果（非掩码）注意力近似算法，转换为高质量的、具有流式保证的因果（掩码）近似算法。

EXPRESS 作为一个可更新的加权缓存对象运行，该对象通过目标缓存大小 $n_{out}$ 、膨胀因子 $m$ 以及基础减半算法（如 Kernel Halving）进行初始化。它通过三个不同的阶段来管理缓存：

精确阶段（Exact Phase）： 前 $n_{out}$ 个点直接存储，不进行压缩。
稀疏阶段（Thin Phase）： 新进入的点被分批处理，并在进入主缓存之前，通过结合均匀子采样和一种广义压缩算法（COMPRESS2）被增量稀疏化至 $n_{out}$ 大小。
减半阶段（HALVE Phase）： 当缓存大小达到阈值（ $4n_{out}$ ）时，通过应用两次基础减半算法将整个缓存压缩回 $n_{out}$ ，并增加膨胀因子 $m$ 。

这种结构确保了无论序列长度 $n$ 为多少，缓存大小始终保持在有界范围内（最多为 $6n_{out}$ ），同时维持亚高斯（sub-Gaussian）近似保证。

Thinformer Express

作者使用 Kernel Halving (KH)（最先进的非因果稀疏化算法）来实例化 EXPRESS，从而创建了 Thinformer Express。该系统利用键值注意力核 $\kappa_{att}$ 来确保近似过程能够保留注意力计算所需的特定平均值。

关键技术组件包括：

流式保证： 理论证明表明，对于长度为 $n$ 、缓存大小为 $s$ 的序列，Thinformer Express 实现了 $O(\log^{3/2}(n)/s)$ 的近似误差，且仅需 $O(s)$ 的内存和 $O(s^2 \log^2(n))$ 的压缩开销。
面向 I/O 的实现： 受 FlashAttention 2 的启发，作者使用 Triton 内核实现了 Thinformer Express。该实现利用分块（tiling）技术来最小化高带宽内存（HMB）的读写，在批次（batch）和头（head）维度上进行并行化，并避免了中间核矩阵的实例化。

核心贡献

因果转换框架： 引入了 EXPRESS，它弥合了非掩码与掩码注意力近似之间的鸿隙，将离线稀疏化算法转化为高效的流式因果近似算法。
改进的理论边界： Thinformer Express 比现有的因果近似方法（如 HyperAttention [15], BalanceKV [19]）提供了更紧密的近似保证。具体而言，它提供了更快的误差衰减率（ $n^{-a}$ 对比 $n^{-a/6}$ ）以及对值矩阵范数更好的依赖关系。
高效实现： 基于 Triton 的实现最大化了 GPU 利用率，证明了其相对于精确注意力和现有近似方法具有显著的速度提升。
全面评估： 在语言建模流水线的四个关键资源受限场景下进行了实证验证。

实验结果

作者在标准基准测试中，针对 FlashAttention 2 和其他领先的 KV 缓存压缩方法（如 SnapKV, StreamingLLM, PyramidKV 等），对 Thinformer Express 进行了评估：

长上下文预填充：
- 在序列长度为 512k 个 token 时，Thinformer Express 相比 FlashAttention 2 实现了 82 倍的加速。
- 与 HyperAttention 相比，它在困惑度（perplexity）和加速比方面均实现了统一的提升。
- 与 HyperAttention 不同，后者在 A6000 GPU 上处理 512k token 时会出现显存溢出（OOM），而 Thimer Express 则能顺利完成任务。
加速 KV 缓存压缩：
- 当与 SnapKV、StreamingLLM 和 PyramidKV 配对使用时，Thinformer Express 在不牺牲准确性的情况下，降低了 LongBench-E 任务的总注意力运行时间。
- 它有效地将生成初始键值（K/V）的二次方成本与压缩步骤解耦。
内存高效的长文本解码：
- 在 MATH-500 推理问题上，Thinformer Express 在准确率与缓存大小的权衡中占据优势。
- 它在使用仅 61% 缓存大小的情况下，达到了与精确注意力相当的准确率。
加速长文本解码：
- 在同样的 MATH-500 任务中，它在准确率与运行时间的权衡中占据优势。
- 它仅用 56% 的计算时间就达到了与精确注意力相当的准确率。

重要性与主张

论文声称，Thifer Express 代表了在资源受限条件下实现长上下文语言建模的重要进展。通过提供逐 token 准确率保证、常数级内存使用（独立于序列长度）以及亚二次方查询时间，它解决了 LLM 推理中的主要瓶颈：长上下文预填充、KV 缓存管理以及内存/计算受限的解码。

作者强调，他们的方法不仅仅是一种启发式方法，而是植根于严谨的亚高斯稀疏化理论，提供了比以往因果近似方法更强的理论保证。此外，高效的 Triton 实现证明了这些理论收益可以转化为现实世界的性能提升，使得处理高达 512k token 的序列成为可能，且速度远超当前最先进的水平（FlashAttention 2）。

论文最后指出了一些局限性，包括目前主要关注英文和数学推理，以及未来利用尚未在当前实现中充分利用的新型 GPU 特性（如 Tensor Memory Accelerator, FP8 支持）进行进一步优化的潜力。