点击下方卡片，关注「集智书童」公众号
导读
长上下文Transformer模型（LCTMs）对于现实应用至关重要，但由于注意力机制的二次复杂度，它们面临着高昂的计算成本。块Sparse注意力通过聚焦于关键区域来缓解这一问题，然而现有方法由于块重要性测量的成本高昂，在平衡准确性和效率方面存在困难。
在本文中，作者引入了XAttention，这是一个即插即用的框架，它通过使用Sparse注意力显著加速了Transformer模型中的长上下文推理。XAttention的关键创新在于洞察到注意力矩阵中反斜对角线值之和（即从左下到右上）为块重要性提供了一个强大的 Agent 。这允许精确识别和修剪非必要块，从而实现高Sparse性和推理速度的显著提升。
在包括语言领域的RULER和LongBench、视频理解领域的VideoMME以及视频生成领域的VBench在内的多个具有挑战性的长上下文基准测试中，XAttention实现了与全注意力相当的准确度，同时带来了实质性的计算收益。作者展示了在注意力计算中实现了高达
的加速。这些结果突显了XAttention解锁块Sparse注意力实际潜力的能力，为LCTMs在现实应用中的可扩展和高效部署铺平了道路。
- 引言
大语言模型（LLMs）的变革性影响正超越自然语言处理领域，引领着多模态能力的新时代。长上下文Transformer模型（LCTMs）正成为这一演变过程中的关键工具，尤其是在视频理解和视频生成等任务中，这些任务需要处理和生成非常长的信息序列。这些模型是解锁能够以人类方式与世界互动、在长时间内理解和生成文本以及视觉信息的卓越系统的关键。想象一下，AI Agent 进行无缝、多模态、全天候的交互，或者强大的世界模拟器生成数小时连贯的视频——这些任务都依赖于处理大量的 Token 。
然而，实现这一愿景需要克服一个重大挑战：注意力机制的计算负担。虽然对于捕捉序列中的关系至关重要，但注意力的成本与序列长度的平方成正比。这种二次扩展在预填充阶段造成了巨大的 Bottleneck ，阻碍了LCTMs在复杂、实际应用中的实际部署。
在追求更高效的Transformer的过程中，块Sparse注意力已成为一条有前景的途径。其核心思想颇具吸引力：不是计算所有 Token 对之间的注意力，而是将资源集中在注意力图的最关键区域，形成“块”的相关信息。这种选择性计算有望大幅减少计算负担，同时保留模型捕捉关键长距离依赖关系的能力。
然而，现有的块Sparse方法在充分发挥其潜力方面遇到了困难，常常在准确性和效率之间进行权衡。这源于缺乏轻量级且有效的机制来识别和优先处理真正重要的注意力块。确定块重要性的开销可能会抵消通过Sparse性获得的收益，使得这些方法在实际部署中变得不切实际。
这引出了一个问题：作者能否设计一种块Sparse注意力机制，在不牺牲准确性的情况下显著加速长上下文Transformer，从而真正释放其在现实世界应用中的潜力？
作者通过引入XAttention来回答这个问题，XAttention是一种新颖的即插即用框架，旨在显著提高长上下文中块Sparse注意力机制的效率。XAttention基于一个关键洞察：注意力矩阵中反对角值的总和可以作为块重要性的强大且计算效率高的指标。与现有方法主要依赖计算密集且损失较大的解决方案（如 Token 池化）来识别重要块不同，XAttention利用这个简单的分数提供了一种可能更加流畅和直接的方法，用于快速准确地识别关键注意力块。
该反斜对角评分算法允许XAttention积极寻找并剪枝非必要计算，在不牺牲准确性的同时实现显著的Sparse性。作者对XAttention在具有挑战性的长上下文基准测试中进行了广泛评估，包括自然语言处理的RULER和LongBench、视频理解的VideoMME以及视频生成的VBench。在这些基准测试中，XAttention在保持与全注意力相当准确性的同时，实现了显著的计算增益，在预填充阶段实现了高达
的注意力计算加速。这些结果突显了XAttention释放块Sparse注意力实际潜力的能力，为在要求严格的应用中，尤其是在多模态AI领域不断扩展的应用中，可扩展和高效的部署长上下文Transformer铺平了道路。
- 方法
在本节中，作者介绍了XAttention，XAttention。XAttention算法包含三个主要组件：
-
- 预测注意力图块的权重
-
- 选择重要的注意力块
-
- 预测注意力 Head 的最小阈值
2.1. 重要性预测
注意图的内禀Sparse性需要一种稳健的策略来预测注意块的的重要性。虽然MInference和FlexPrefill等方法结合了池化和“垂直斜杠检测”，但作者的消融研究表明，仅依赖平均或求和池化会导致预测不准确。当块中只有少数显著的垂直或斜杠模式时，池化方法尤其无效，无法捕捉这些重要性关键指标。
MInference和FlexPrefill试图通过分析输入 Query 的最后一段来识别重要的“垂直和斜杠索引”，以克服这一限制。然而，这种方法面临两个主要挑战：首先，重要的关注模式可能不会持续存在于最终的 Query 段中；其次，搜索算法本身引入了大量的计算开销（如图5所示）。
从根本上说，一个有效的块重要性预测方法应该能够自动且稳健地识别出重要模式，包括关键的垂直和斜杠模式。为了实现这一目标，作者提出了反对角线选择方法。在大小为
的每个块中，作者使用步长
（如图1所示）选择反对角线上的元素。这些所选元素的总和作为对应注意力块整体重要性的 Agent 。
该方法的有效性可以从两个角度来理解：(1) 信息保留：这种选择策略确保了所有 Token 的信息都被考虑在内，因为每个 Token 至少对一条反斜线求和有贡献。(2) 模式检测：如图2所示，反斜线与块内所有可能的垂直和斜线模式相交。XAttention的反斜线模式与块内的垂直和斜线模式相交，使得这些模式的检测变得高效，并指导有效的Sparse注意力计算。这确保了在重要性估计过程中不会错过任何关键模式。
2.2. 阈值块选择
基于反对角评分模式，作者提出了以下Sparse注意力块选择算法。令S表示步长，令B表示Sparse注意力块的大小。该过程从反对角求和开始，作者在注意力图的每个S × S块内选择反对角线上的元素，并计算这些元素的求和，得到每个反对角线的和。随后，作者通过应用softmax函数对这些反对角线的和进行softmax归一化，得到反对角线上的概率分布。最后，对于块选择，使用find blocks函数来识别反对角线概率累积和超过预定义阈值τ的最小块集。形式上，这可以表示为：
其中，A是注意力图，B是一组块，|B|表示该集合中的块数量。此过程基于反斜对角评分模式和指定的阈值，有效地确定了注意力图中最重要的块。
2.3 最小阈值预测
基于反斜对角评分模式，作者提出了以下Sparse注意力块选择算法。令
表示步长，令
表示Sparse注意力块的大小。该过程从反斜对角求和开始，作者在注意力图的每个
块内选择沿反斜对角线的元素，并计算这些元素的求和，得到每个反斜对角线的和。随后，作者通过应用softmax函数对这些反斜对角线的和进行softmax归一化，从而得到反斜对角线的概率分布。最后，对于块选择，使用find-blocks函数来识别累积反斜对角线概率超过预定义阈值
的最小块集。形式上，这可以表示为：
作者提出了一种动态规划方法来确定每个注意力 Head 的最优阈值。先前的研究表明，不同的注意力 Head 表现出不同的Sparse程度和重要性。因此，动态调整各个 Head 的阈值以优化准确性和计算效率之间的平衡是有益的。
问题表述：考虑一个具有
个注意力头的模型。作者定义一个动态规划表
，其中
代表第
个头，
表示在前
个头中进行了
次阈值调整。
存储了在前
个头中恰好进行了
次阈值调整时所能达到的最佳性能。
动态规划：作者的目标是找到每个 Head 的最优阈值，以使它们的联合贡献最大化准确率同时最小化计算量。DP表的递推关系为：
其中
是注意力图，
是一组块，
表示该集合中的块的数量。此过程基于反斜对角评分模式和指定的阈值，有效地确定了注意力图中最重要的块。
表示在调整第
个头的阈值第
次时的模型性能。这对应于在优化过程中相对于状态
将第
个头的阈值降低一步后的模型性能。
作者通过在每个步骤中将每个头的阈值减少10%来调整阈值：
这确保了计算量的逐步减少，同时保留了每个头对准确性的贡献。
请注意，此动态阈值预测方法可以进一步优化XAttention的Sparse性，但并非强制组件。作者将在消融研究中呈现详细结果。
- 实验
本节介绍了作者对XAttention有效性的实证研究。作者首先详细阐述了其实施的具体细节，随后展示了在文本和视频理解以及视频生成基准测试中对强大 Baseline 方法的评估结果。接着，作者测试了XAttention的加速性能。最后，作者提供了分析性消融研究，以进一步理解XAttention的行为。
3.1 实验设置
作者评估XAttention在三个不同的领域。对于自然语言任务，作者采用Llama-3.1-8B-Instruct。在视频理解领域，作者使用Qwen2-VL-7B-Instruct。最后，对于视频生成，作者采用HunyuanVideo模型。为了优化自然语言任务中计算效率和准确率之间的权衡，作者将作者的精确阈值预测方法应用于Llama-3.1-8B-Instruct模型。
基准线 作者将XAttention与几个强大的基准进行了比较。作者密集注意力模型的主要基准是FlashAttention，它在FlashInfer框架中实现。作者还与MInference、FlexPrefill和SeerAttention进行了比较，严格遵循它们的公开实现。对于SeerAttention，作者采用了在Gare权重上的预训练。对于MInference，作者使用了他们的官方配置，其中所有注意力头都采用“垂直斜杠”Sparse模式。对于FlexPrefill，作者将超参数设置为
和
，根据原始论文，这组参数在提供的参数集中实现了最高的准确率。
数据集作者在自然语言理解、视频理解和视频生成等多个任务上评估XAttention。对于自然语言任务，作者使用了RULER数据集，这是一个专门为评估LLMs的长期上下文能力而设计的合成基准。RULER允许自定义序列长度和任务复杂性，在扩展传统“大海捞针”测试的同时，引入了如多跳追踪和聚合等新型任务类别。作者还对来自LongBench的实际情况下的长期上下文任务进行了评估，以测试在实际场景中的性能。
对于视频理解，作者使用了Video-MME数据集，这是第一个用于评估多模态大语言模型（MLLMs）在视频分析方面的综合基准。Video-MME包含900个视频，总时长为254小时，视频时长从11秒到1小时不等，为评估长视频理解能力提供了一个强大的测试平台。
在视频生成领域，作者利用VBench提供的946个GPT增强文本 Prompt 来生成视频。然后，作者将XAttentionXAttention生成的视频与全注意力 Baseline 生成的视频进行比较，以评估XAttention在生成高质量视频内容方面的有效性。
3.2 准确性结果
在RULER基准测试中，作者应用了第3.3节中描述的动态规划方法进行最小阈值预测，使用步长
和
，最大调整次数为
。这产生了一组平均阈值为0.8的最小阈值，进一步提升了作者Sparse注意力机制的计算效率。
LongBench 表2展示了XAttention在LongBench基准测试中的真实世界任务性能，与强 Baseline 相比，使用了Llama-3.1-8BInstruct模型。在保持与RULER评估相同的配置下，作者对XAttention、MInference和FlexPrefill进行了评估。XAttention在所有任务中均取得了最高的平均分数，证明了其在实际场景中的有效性。值得注意的是，XAttention在各个任务上的性能与全注意力机制相近，表明XAttention在提高效率的同时保持了准确性。
视频理解 作者在QwenVL-2-7B模型上应用了Stride
和阈值
参数。如表3所示，在三种Sparse注意力方法中，MInference和FlexPrefill在长视频任务上未能达到最佳性能。XAttention在所有Sparse注意力方法中取得了最佳平均分数，甚至在长视频上超过了FlashAttention，以每秒1帧的帧率运行可达1小时。
视频生成作者使用HunyuanVideo模型在VBench的 Prompt 下评估了XAttention在视频生成领域的性能。HunyuanVideo模型采用Diffusion Transformer（DiT）架构，该架构采用非因果注意力。由于现有的 Baseline 模型没有实现非因果注意力，作者仅将XAttention与全注意力 Baseline 进行比较。作者的评估考虑了定量指标（PSNR、SSIM、LPIPS）和定性视觉比较。作者将DiT Backbone 网络中的所有注意力计算替换为XAttention，并使用相同的随机种子和 Prompt 与全注意力输出进行性能比较，平均结果覆盖了所有946个VBench Prompt 。生成的视频分辨率为
像素，包含129帧，并进行了50次去噪步骤。作者将XAttention配置为步长
和阈值
和
。全注意力（ Baseline ）实现了超过
的Sparse度。
一只快乐的模糊熊猫盘腿坐在噼啪作响的篝火旁，热情地弹奏着一把小原声吉他。熊猫黑白相间的毛皮与火光的温暖光辉形成鲜明对比，在周围覆盖着雪的地面上投下闪烁的阴影。熊猫身后，壮丽的雪山在暮色苍茫的天空中耸立，山顶被落日的余晖染成淡红色。熊猫在弹奏欢快的曲调时眼中闪烁着喜悦，宁静的山景和温馨的篝火营造出一个神奇、温馨的场景。
最初，在 Hunyuan Video 模型中将 XAttention 应用于去噪过程的初始阶段，与全注意力 Baseline 相比，导致输出视频的布局略有偏移，从而降低了定量评分。受扩散模型研究的启发，这些研究表明早期去噪步骤对于确定内容布局至关重要，作者引入了一个“预热”阶段。在这个阶段，作者首先使用全注意力进行前5个去噪步骤，然后切换到 XAttention。图3展示了这种预热策略的定性影响。
3.3 效率结果
进一步分析XAttention在具有不同上下文长度任务上的效率，并将其与FlashAttention、MInference和FlexPrefill进行比较。作者重点关注预填充阶段，并测量XAttention实现的关注速度提升。此外，作者将计算时间分解为模式选择和Sparse注意力组件，并将其与其他无训练模式选择方法进行对比。
注意力加速 图4展示了XAttention在不同token序列长度（从8k到
）上的预填充加速效果。作者使用步长
和
，阈值
进行实验。在较短的上下文中，由于注意力密度较高，MInference和FlexPrefill由于更广泛的模式选择而经历了增加的开销。
相比之下，XAttention保持了其加速优势。值得注意的是，对于
的上下文长度，XAttention分别以
和
的密度实现了最大预填充注意力加速，分别为
和
（见表5）。
注意力时间分解图5表明，XAttention的反对角模式，结合其高效的块选择算法，与依赖于垂直斜杠索引搜索的MInference和FlexPrefill相比，在模式选择上显著更快。具体来说，XAttention的模式选择时间分别高达MInference的
和FlexPrefill的
。此外，反对角模式的准确性使得XAttention能够实现更低的注意力密度，从而在Sparse注意力计算本身中实现大幅加速。
3.4 消融研究
为了进一步分析XAttention的组成部分，作者进行了一项消融研究，评估了反斜线模式、阈值块选择和最小阈值预测的有效性。
反对角模式
作者通过将反角模式与随机模式和正对角模式进行比较，研究了其在预测注意力块和的总和时的作用。对于随机模式，作者确保在每一个
块中选取
个元素，同时保证每行和每列至少选取一个 Token 。表6显示，反角模式在所有任务中均实现了最高的准确率，同时保持了最低的密度，从而证实了其优越性。
步长大小
作者探讨了不同步长大小
的影响。较大的步长会导致采样注意力图更Sparse，从而降低计算开销。然而，过大的步长可能会损害块选择的准确性。作者在表7中比较了4、16和64的步长。作者的结果表明，当步长过长时，它无法准确检测之前识别的斜杠注意力模式。过于Sparse的反对角线无法有效地区分从不同位置进入块的斜杠模式，从而导致性能下降。
Top-K与Top-Ratio与动态Sparse性作者评估了不同的块选择策略：Top-K、Top-Ratio以及作者的阈值块选择（动态Sparse性）。为了进行公平的比较，作者对于S=8时，将K设为8192，将Ratio设为27%；对于S=16时，将K设为16384，将Ratio设为31%，目标是实现与作者的阈值块选择相似的计算成本。表8显示，Top-K和Top-Ratio在处理具有相似计算量的多样化和动态输入序列长度时都显得力不从心。相比之下，作者的基于阈值的策略，即保留至少达到阈值 Level 的注意力块，在计算量和准确性之间实现了最佳平衡。
最小阈值预测最后，作者在RULER基准（Hsieh等，2024）上比较了作者的最小阈值预测方法与固定阈值
的性能。使用最小阈值预测，作者以
开始，并设置
，允许动态规划（DP）算法探索1,000个最优阈值组合。这导致一组更精细的阈值，平均值为0.8。表9表明，动态预测的阈值实现了更低的密度和更高的准确性，展示了该方法的有效性。
- 结论
作者提出了一种名为XAttention的新型即插即用框架，用于加速Transformer模型中的长上下文推理。通过利用注意力矩阵中反对角线求和作为块重要性的鲁棒 Agent 的见解，XAttention能够高效地识别和修剪非必要块，在不牺牲准确性的情况下实现显著的计算节省。作者在自然语言理解（RULER、LongBench）、视频理解（VideoMME）和视频生成（VBench）等具有挑战性的长上下文基准上的评估表明，XAttention在注意力计算中实现了高达13.5倍的速度提升，同时保持了与全注意力相当的性能。这些结果突出了XAttention解锁块Sparse注意力实际潜力的能力，为在现实应用中高效和可扩展地部署长上下文Transformer模型铺平了道路。
参考
[1]. XAttention: Block Sparse Attention with Antidiagonal Scoring
扫码加入👉「集智书童」交流群
（备注：方向+学校/公司+昵称）