XAttention

XAttention: Block Sparse Attention with Antidiagonal Scoring

• 革新Transformer推理的高效注意力机制
• 资源​​
  ​​论文链接​​：XAttention: Block Sparse Attention with Antidiagonal Scoring
  ​​代码开源​​：GitHub仓库

XAttention是韩松团队提出的一种创新的块稀疏注意力机制，旨在解决传统Transformer模型在处理长上下文时面临的计算效率瓶颈问题。该论文通过引入反斜对角线评分(antidiagonal scoring)方法，实现了在不牺牲模型性能的前提下显著加速Transformer推理过程，特别是在多模态任务中表现出色。本文将详细介绍XAttention的核心思想、技术实现、实验验证及其在长上下文Transformer模型(LCTMs)中的应用价值。

1 研究背景与动机

随着大语言模型(LLMs)和多模态模型的快速发展，长上下文Transformer模型(LCTMs)已成为处理超长序列数据的关键工具，尤其是在视频理解、视频生成等需要处理极长信息序列的任务中。然而，传统注意力机制因其二次计算复杂度(O(n²))成为模型扩展的主要瓶颈，这使得处理长序列变得极其昂贵。

当前主流的解决方案是采用块稀疏注意力(block sparse attention)机制，它通过仅计算注意力图中关键区域的值来减少计算量。然而，现有方法在平衡准确性和效率方面面临重大挑战，主要问题在于块重要性测量的高成本往往抵消了通过稀疏性获得的计算收益。具体表现为：

1. 重要性测量开销大：现有方法如Token池化或垂直斜杠检测需要大量计算资源
2. 模式捕捉不完整：统一稀疏模式难以适应不同注意力头的异质性需求
3. 精度-效率权衡：激进稀疏化常导致性能显著下降

针对这些问题，XAttention提出了反斜对角线评分这一创新方法，通过轻量级且高效的重要性评估机制，实现了更优的稀疏模式选择，从而在保持模型精度的同时大幅提升计算效率。

XAttention 的目的

XAttention 的核心目标是 在Transformer模型中实现高效的长序列注意力计算，通过 块稀疏化（Block Sparsification） 来 减少计算量，同时 尽可能保留重要的注意力模式。其关键思想是：

1. 降低计算复杂度：

   • 传统注意力计算复杂度为 (O(N^2))（N 是序列长度），而 XAttention 通过稀疏化将其降低到接近 (O(N^{1.5})) 或更低。
   • 适用于 长序列任务（如 128K+ tokens 的文本、视频帧序列等）。

2. 保持关键注意力模式：

   • 使用 反斜对角线评分（Antidiagonal Scoring） 识别重要的注意力块，避免盲目稀疏化导致的性能下降。

3. 即插即用：

   • 无需修改模型架构或重新训练，可直接应用于现有的 Transformer 模型（如 BERT、GPT、ViT 等）。

2 XAttention核心方法

XAttention框架包含三个关键组件：基于反斜对角线的块重要性预测、阈值块选择算法和动态最小阈值预测机制。这些组件共同工作，实现了高效且精确的稀疏注意力计算。

2.1 反斜对角线评分机制

• 关键思想：通过计算注意力块内反对角线（从左下到右上）元素的和作为块重要性代理，高效识别关键注意力区域。
  • 数学表达：对$B\times B$的块，以步长$S$采样反对角线元素，求和得分：
    $$\text{Score}=\sum _{i+j=k}{A}_{i,j}(k\in \text{反对角线索引})$$
  • 优势：
    • 信息保留：每个token至少贡献一个反对角线和，避免局部采样偏差。
    • 模式覆盖：反对角线天然相交于垂直/斜线依赖模式（常见于视频、文本长程关联）。

XAttention的核心创新在于发现注意力矩阵中反斜对角线值之和（从左下到右上）可以作为块重要性的强有力代理指标。这一洞察源于两个关键观察：

具体实现上，对于大小为S×S的每个注意力块，XAttention以步长S选择反斜对角线上的元素，并计算这些元素的和作为该块的重要性分数。这种方法相比传统池化方法具有明