稀疏注意力机制(SSA)原理剖析：如何用80%的计算量实现95%的模型精度

1次阅读

共计 2042 个字符，预计需要花费 6 分钟才能阅读完成。

当我们在 NLP 任务中使用标准 Transformer 时，全连接注意力机制的计算复杂度是 O(n²)。这意味着随着序列长度的增加，计算量会呈平方级增长。举个例子，BERT-base 模型在处理 512 长度的序列时，单层注意力需要约 235M FLOPs，而 Longformer 通过稀疏注意力将这一数字降到了 47M FLOPs——节省了 80% 的计算资源却保持了 95% 的模型精度。

数学表达：
$$A_{ij}^{local} = \begin{cases}
Q_iK_j^T & \text{if} |i-j| \leq w \
-\infty & \text{otherwise}
\end{cases}$$

只计算当前位置前后 w 个 token 的注意力
超参数选择：文本任务通常 w =32~256，蛋白质序列 w =64~512

数学表达：
$$A^{global} = Q[g_1,…,g_m]K^T$$

设计 m 个特殊 token 关注整个序列
经验值：m 通常取序列长度的 1%~5%

数学表达：
$$A_{ij}^{random} = \begin{cases}
Q_iK_j^T & \text{with probability} p \
-\infty & \text{otherwise}
\end{cases}$$

每个 token 随机关注 r 个其他位置
典型设置：p=0.1~0.3

class BlockSparseAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.local_window = config.local_window
        self.global_tokens = config.global_tokens

        # 投影层
        self.qkv = nn.Linear(config.hidden_size, 3*config.hidden_size)
        self.proj = nn.Linear(config.hidden_size, config.hidden_size)

        # 稀疏模式开关
        self.use_local = config.use_local
        self.use_global = config.use_global

    def forward(self, x, mask=None):
        B, N, C = x.shape
        qkv = self.qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=config.num_heads), qkv)

        # 局部注意力计算
        if self.use_local:
            local_mask = torch.ones(N, N, device=x.device).triu(diagonal=-self.local_window).tril(diagonal=self.local_window)
            attn_local = (q @ k.transpose(-2, -1)) * local_mask

        # 全局注意力计算
        if self.use_global:
            global_q = q[:, :, :self.global_tokens]
            attn_global = global_q @ k.transpose(-2, -1)

        # 合并注意力
        attn = attn_local + attn_global
        attn = attn.masked_fill(mask == 0, -1e9)
        attn = attn.softmax(dim=-1)

        # 梯度检查点
        if self.training:
            output = checkpoint(self._project, attn @ v)
        else:
            output = self.proj(rearrange(attn @ v, 'b h n d -> b n (h d)'))

        return output

在 PG-19 数据集（平均长度 5k tokens）上的测试结果：