XAttention 计算步骤详解及示例

XAttention 计算步骤详解及示例

XAttention 是一种高效的块稀疏注意力机制，通过 反对角线评分（Antidiagonal Scoring） 和 动态阈值选择 来优化长序列 Transformer 模型的推理效率。以下是其核心计算步骤及具体示例。

---

1. XAttention 的核心步骤

Step 1: 计算原始注意力分数

输入：

• Query $Q\in {\mathbb{R}}^{n\times d}$
• Key $K\in {\mathbb{R}}^{m\times d}$
• Value $V\in {\mathbb{R}}^{m\times d_v}$

计算未缩放的注意力分数：
$S=Q{K}^T$

Step 2: 反对角线评分（Antidiagonal Scoring）

1. 分块计算：将 $S$ 划分为 $B\times B$ 的块（如 $8\times 8$）。
2. 反对角线求和：对每个块，计算反对角线（从左下到右上）元素的和，作为块的重要性分数：
   $\text{Score}={\sum }_{i+j=k}{A}_{i,j}$
   • 其中 $k$ 是反对角线索引，例如 $k=0,1,...,2B-2$。

Step 3: 阈值块选择

1. 归一化：对块分数进行 softmax 归一化：
   $P=\text{softmax}(\text{Score})$
2. 选择关键块：保留累积概率超过阈值 $\tau$ 的最小块集合 $B^{\ast }$：
   $B^{\ast }=\arg \min |B|\text{s.t.}{\sum }_{(i,j)\in B}{P}_{i,j}>\tau$

Step 4: 稀疏注意力计算

仅计算选中的关键块 $B^{\ast }$ 的注意力权重，并加权聚合 $V$：
$\text{Output}={\sum }_{(i,j)\in B^{\ast }}{A}_{i,j}{V}_j$

---

2. 计算示例

输入数据

假设 $d=2$，输入如下：

• Query (Q)：
  $Q=\left[\begin{array}{cc}1.0& 2.0\\ 3.0& 4.0\end{array}\right]$
• Key (K)：
  $K=\left[\begin{array}{cc}5.0& 6.0\\ 7.0& 8.0\\ 9.0& 10.0\end{array}\right]$
• Value (V)：
  $V=\left[\begin{array}{ccc}1.0& 0.0& 1.0\\ 0.0& 1.0& 0.0\\ 1.0& 1.0& 0.0\end{array}\right]$

Step 1: 计算原始注意力分数 $S=Q{K}^T$

$S=\left[\begin{array}{ccc}1\cdot 5+2\cdot 6& 1\cdot 7+2\cdot 8& 1\cdot 9+2\cdot 10\\ 3\cdot 5+4\cdot 6& 3\cdot 7+4\cdot 8& 3\cdot 9+4\cdot 10\end{array}\right]=\left[\begin{array}{ccc}17& 23& 29\\ 39& 53& 67\end{array}\right]$

Step 2: 反对角线评分（假设块大小 $2\times 2$）

• 块 1（$S_{1:2,1:2}$）：
  反对角线元素 = { 17 , 53 } 和 = 17 + 53 = 70 \text{反对角线元素} = \{17, 53\} \quad \text{和} = 17 + 53 = 70 反对角线元素={17,53}和=17+53=70
• 块 2（$S_{1:2,2:3}$）：
  反对角线元素 = { 23 , 67 } 和 = 23 + 67 = 90 \text{反对角线元素} = \{23, 67\} \quad \text{和} = 23 + 67 = 90 反对角线元素={23,67}和=23+67=90

Step 3: 阈值块选择（假设 $\tau =0.6$）

1. 归一化：
   $P=\text{softmax}([70,90])\approx [0.27,0.73]$
2. 选择关键块：
   • 累积概率：$0.27+0.73=1.0>\tau$，因此选择 块 2（$S_{1:2,2:3}$）。

Step 4: 稀疏注意力计算

仅计算块 2 的注意力权重：
$A_{\text{selected}}=\text{softmax}\left(\frac{[23,67]}{\sqrt{2}}\right)\approx [0.0001,0.9999]$
加权聚合 $V$：
$\text{Output}=0.0001\cdot [0.0,1.0,0.0]+0.9999\cdot [1.0,1.0,0.0]\approx [1.0,1.0,0.0]$

---

3. 关键优势

1. 计算高效：反对角线评分的计算复杂度低，仅需 $O(B^2)$ 而非 $O(N^2)$。
2. 模式保留：反对角线能捕捉垂直/斜线依赖（如视频帧间的时空关联）。
3. 动态适应性：通过阈值调整可平衡计算量与精度。

XAttention 在长文本和视频任务中可加速 13.5 倍，同时保持全注意力的精度。