自适应稀疏自注意力机制：如何优化Transformer的长序列处理能力

1次阅读

没有评论

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

Transformer 模型在自然语言处理等领域表现出色，但其核心的自注意力机制在处理长序列时面临两个主要问题：

计算复杂度高：传统自注意力机制的计算复杂度为 O(n²)，这意味着输入序列长度增加一倍，计算量会增加四倍。对于长文档、高分辨率图像等任务，这会导致训练和推理变得极其缓慢。
内存消耗大：自注意力需要存储 n×n 的注意力矩阵，当序列长度达到数千甚至数万时，这会消耗大量 GPU 内存，使得模型难以在普通硬件上运行。

目前主要有几种解决自注意力计算复杂度的方法：

局部注意力(Local Attention)：只计算每个 token 附近窗口内的注意力，复杂度降为 O(n×k)，k 为窗口大小。缺点是难以捕捉长距离依赖。
Reformer 的 LSH 注意力 ：使用局部敏感哈希(LSH) 将相似 token 分到同一桶中，只在桶内计算注意力。实现复杂，且哈希质量影响效果。
ASSA(自适应稀疏自注意力)：动态决定每个 token 应该关注哪些其他 token，既能降低计算量，又能保留重要的长距离交互。这是本文重点介绍的方法。

ASSA 的核心思想是为每个查询 token 选择最相关的 k 个键 token，而不是计算所有可能的组合。具体步骤：

对每个查询 qi，计算其与所有键 kj 的初始相关性分数 sij = qi·kj/√d
只保留每个 qi 对应的 top- k 个 sij，其余置为负无穷
对筛选后的分数做 softmax 得到最终注意力权重

数学表达式：

Attention(Q,K,V) = softmax(top_k(QK^T/√d))V

其中 top_k(·)操作保留每行最大的 k 个元素。

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

class ASSA(nn.Module):
    def __init__(self, embed_dim, num_heads, sparsity_ratio=0.1):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scaling = self.head_dim ** -0.5
        self.sparsity_ratio = sparsity_ratio

        # 初始化 QKV 投影矩阵
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        """
        输入 x: (batch_size, seq_len, embed_dim)
        输出: (batch_size, seq_len, embed_dim)
        """
        batch_size, seq_len, _ = x.shape

        # 计算 Q,K,V
        q = self.q_proj(x)  # (batch, seq_len, embed_dim)
        k = self.k_proj(x)
        v = self.v_proj(x)

        # 多头切分
        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.scaling  # (batch, heads, q_len, k_len)

        # 动态稀疏化：每行只保留 top- k 个元素
        k = max(1, int(seq_len * self.sparsity_ratio))
        topk_values, topk_indices = torch.topk(attn_scores, k, dim=-1)

        # 创建稀疏注意力矩阵
        sparse_attn = torch.full_like(attn_scores, float('-inf'))
        sparse_attn.scatter_(-1, topk_indices, topk_values)

        # 计算注意力权重
        attn_weights = F.softmax(sparse_attn, dim=-1)

        # 应用注意力
        output = torch.matmul(attn_weights, v)
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)

        return self.out_proj(output)

在实际应用中，稀疏模式的选择需要考虑：