AI稀疏注意力机制入门指南：从原理到PyTorch实战

1次阅读

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

传统 Transformer 的自注意力机制虽然强大，但存在一个致命缺点：计算复杂度随着序列长度呈平方级增长（O(n²)）。这意味着处理 1000 个 token 的序列时，需要计算 100 万次注意力权重！这种计算开销使得模型难以处理长文本或高分辨率图像。

稀疏注意力通过有选择地计算部分注意力权重（通常 10%-30%），将复杂度降低到 O(n√n) 甚至 O(n)。就像人类阅读时不会同时关注所有文字一样，这种机制让 AI 也能 ” 选择性聚焦 ”。

数学表达 ：$A_{ij} = \begin{cases} Q_iK_j^T & \text{if} |i-j| \leq w \ 0 & \text{otherwise} \end{cases}$
示意图 ：类似滑动窗口，每个 token 只关注左右相邻的 w 个 token
特点：保持局部上下文关系，适合连续信号处理

数学表达 ：$A_{ij} = \begin{cases} Q_iK_j^T & \text{if} i \equiv j (\text{mod} s) \ 0 & \text{otherwise} \end{cases}$
示意图 ：类似棋盘格，每个 token 固定间隔 s 关注其他 token
特点：捕获长程依赖，适合周期性模式

数学表达 ：预设少量特殊 token 参与所有注意力计算
示意图 ：某些 token 成为 ” 信息枢纽 ”
特点：平衡局部和全局信息

import torch
import torch.nn as nn
from typing import Optional, Tuple

class SparseAttention(nn.Module):
    def __init__(self, 
                 embed_dim: int, 
                 num_heads: int, 
                 window_size: int = 32,
                 stride: int = 8,
                 dropout: float = 0.1):
        super().__init__()
        self.qkv_proj = nn.Linear(embed_dim, embed_dim * 3)
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.window_size = window_size
        self.stride = stride
        self.dropout = nn.Dropout(dropout)

    def _create_local_mask(self, seq_len: int) -> torch.Tensor:
        """生成局部注意力掩码"""
        mask = torch.ones(seq_len, seq_len, dtype=torch.bool)
        for i in range(seq_len):
            start = max(0, i - self.window_size)
            end = min(seq_len, i + self.window_size + 1)
            mask[i, :start] = 0
            mask[i, end:] = 0
        return mask

    def _create_strided_mask(self, seq_len: int) -> torch.Tensor:
        """生成跨步注意力掩码"""
        return torch.eye(seq_len, dtype=torch.bool).repeat_interleave(self.stride, dim=1)

    def forward(self, 
                x: torch.Tensor,
                key_padding_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        batch_size, seq_len, _ = x.shape

        # 生成 QKV
        qkv = self.qkv_proj(x)
        q, k, v = qkv.chunk(3, dim=-1)

        # 分头处理
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 计算注意力分数
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)

        # 应用稀疏掩码
        local_mask = self._create_local_mask(seq_len).to(x.device)
        strided_mask = self._create_strided_mask(seq_len).to(x.device)
        combined_mask = local_mask | strided_mask

        # 处理 padding mask
        if key_padding_mask is not None:
            combined_mask = combined_mask & key_padding_mask.unsqueeze(1)

        # 掩码处理
        attn_scores = attn_scores.masked_fill(~combined_mask.unsqueeze(1), float('-inf'))

        # 计算注意力权重
        attn_weights = torch.softmax(attn_scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # 加权求和
        output = torch.matmul(attn_weights, v)
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)

        return self.out_proj(output)

在 NVIDIA V100 GPU 上测试不同序列长度下的表现：