自适应稀疏自注意力(ASSA)机制入门指南：从原理到PyTorch实现

1次阅读

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

在处理长序列任务时，传统自注意力机制的计算复杂度是 O(n²)，其中 n 是序列长度。这意味着当序列长度增加时，计算和内存需求呈平方级增长。例如，处理 1024 长度的序列需要约 100 万次计算，而 2048 长度的序列则需要约 400 万次计算。

自适应稀疏自注意力 (ASSA) 机制通过动态学习稀疏模式，仅计算最重要的注意力权重，显著降低了计算开销。与传统注意力机制相比，ASSA 具有以下优势：

计算复杂度降低到 O(n log n)
内存消耗大幅减少
能够自动学习序列中最相关的部分

ASSA 通过引入一个可学习的稀疏模式矩阵，动态决定哪些位置的注意力权重需要计算。这个矩阵通过以下公式生成：

S = softmax(QK^T / sqrt(d_k))

其中 Q 和 K 分别是查询和键矩阵，d_k 是键的维度。

对于每个查询位置，ASSA 只保留 top- k 最大的注意力权重：

A_ij = {S_ij if S_ij in top-k(S_i) 
         0 otherwise }

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

class ASSA(nn.Module):
    def __init__(self, embed_dim, num_heads, k_ratio=0.25):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.k_ratio = k_ratio

        # 线性变换层
        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
        k = max(1, int(seq_len * self.k_ratio))  # 计算 top- k 值

        # 1. 线性变换
        q = self.q_proj(x)  # [B, L, E]
        k = self.k_proj(x)  # [B, L, E]
        v = self.v_proj(x)  # [B, L, E]

        # 2. 多头切分
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)  # [B, H, L, D]
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)  # [B, H, L, D]
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)  # [B, H, L, D]

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

        # 4. top- k 稀疏化
        topk_values, topk_indices = torch.topk(attn_scores, k=k, dim=-1)  # [B, H, L, k]
        sparse_attn = torch.zeros_like(attn_scores).scatter_(-1, topk_indices, topk_values)  # [B, H, L, L]

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

        # 6. 注意力加权
        output = torch.matmul(attn_weights, v)  # [B, H, L, D]

        # 7. 合并多头输出
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)

        # 8. 最终线性变换
        return self.out_proj(output)