深入解析 AssaNet 自适应稀疏自注意力机制：从原理到新手实践

1次阅读

没有评论

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

传统注意力机制（如 Transformer 中的自注意力）虽然能有效建模长距离依赖关系，但其计算复杂度随序列长度呈平方级增长（O(n²)）。这导致以下问题：

计算资源消耗大 ：处理长序列（如文档、高分辨率图像）时需要极高的显存和算力
推理延迟高 ：实时应用场景难以满足性能要求
信息冗余 ：并非所有 token 之间的交互都有实际意义

自适应稀疏自注意力通过动态选择最相关的注意力连接，将计算复杂度降低到 O(n√n) 甚至 O(nlogn)，同时保持模型性能。

常见的稀疏化方法及其特点：

固定模式稀疏
优点：实现简单，计算可预测
缺点：无法适应不同输入特性
示例：局部窗口注意力、带状稀疏模式
基于内容的稀疏
优点：动态适应输入特征
缺点：需要额外计算相似度
示例：Top- k 选择、聚类注意力
混合稀疏策略
结合固定模式和动态选择
AssaNet 采用此类策略

自适应稀疏注意力得分计算：

A_{ij} = \begin{cases}
\frac{Q_iK_j^T}{\sqrt{d_k}} & \text{if} j\in S_i \\
-\infty & \text{otherwise}
\end{cases}

其中 S_i 是通过稀疏策略选择的邻居集合。

graph TD
    A[输入序列] --> B[计算 Q,K,V]
    B --> C[稀疏模式选择]
    C --> D[稀疏注意力计算]
    D --> E[输出加权和]

import torch
import torch.nn as nn
import torch.nn.functional as F

class SparseSelfAttention(nn.Module):
    def __init__(self, dim, num_heads, sparse_ratio=0.3):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.sparse_ratio = sparse_ratio

        # 线性变换层
        self.qkv = nn.Linear(dim, dim*3)
        self.proj = nn.Linear(dim, dim)

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C//self.num_heads)
        q, k, v = qkv.unbind(2)  # [B, N, H, C/H]

        # 计算原始注意力分数
        attn = (q @ k.transpose(-2, -1)) * (1.0 / torch.sqrt(torch.tensor(q.size(-1))))

        # 稀疏化处理
        k = int(N * self.sparse_ratio)
        topk = torch.topk(attn, k=k, dim=-1)
        sparse_attn = torch.full_like(attn, float('-inf'))
        sparse_attn.scatter_(-1, topk.indices, topk.values)

        # softmax 归一化
        attn = F.softmax(sparse_attn, dim=-1)

        # 加权求和
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        return self.proj(x)

关键实现说明：
1. 通过 topk 操作动态选择每个 token 最相关的 k 个连接
2. 使用 scatter 操作构建稀疏注意力矩阵
3. softmax 仅在非负无穷的值上计算

假设序列长度 n =1024，稀疏比 s =0.3：