深入解析2022李宏毅自注意力机制：从理论到PyTorch实现

1次阅读

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

自注意力机制（Self-Attention）作为 Transformer 架构的核心组件，彻底改变了 NLP 领域的游戏规则。相比于传统的 RNN 和 CNN，它能够直接建模序列中任意两个位置的关系，但这种强大的能力也带来了显著的计算负担。

计算复杂度问题：
标准自注意力机制需要计算所有位置对之间的相似度，导致时间复杂度为 O(n²)，其中 n 是序列长度。当处理 512 个 token 的序列时，需要计算 262,144 次相似度。
在 PyTorch 中，即使使用优化的矩阵运算，当序列长度超过 1024 时，显存占用会急剧上升。
内存瓶颈：
每个注意力头的 QKV 矩阵存储需要 3×n×d 的内存（d 是特征维度）
注意力权重矩阵 (n×n) 在 float32 精度下，处理 2048 长度的序列就需要 16MB 显存

import torch
import torch.nn as nn
import math

class SelfAttention(nn.Module):
    def __init__(self, embed_size):
        super(SelfAttention, self).__init__()
        self.embed_size = embed_size
        # 同时生成 Q /K/ V 的线性变换
        self.qkv = nn.Linear(embed_size, embed_size * 3)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x, mask=None):
        batch, seq_len, _ = x.shape
        # 并行计算 Q /K/V [batch, seq_len, embed_size*3] -> 各[batch, seq_len, embed_size]
        qkv = self.qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: t.view(batch, seq_len, -1), qkv)

        # 缩放点积注意力
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.embed_size)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        attention = self.softmax(scores)
        out = torch.matmul(attention, v)
        return out

维度变换技巧：
将 embed_size 拆分为 num_heads × head_dim

使用 einops 库简化 reshape 操作：

from einops import rearrange
q = rearrange(q, 'b n (h d) -> b h n d', h=self.num_heads)

注意力掩码处理：
对于 padding 部分，使用 -inf 填充
解码器的因果掩码需要结合 triu 和 expand 操作

class SparseAttention(nn.Module):
    def __init__(self, block_size=64):
        self.block_size = block_size

    def forward(self, q, k, v):
        # 将序列分块计算
        q_blocks = q.split(self.block_size, dim=1)
        k_blocks = k.split(self.block_size, dim=1)
        v_blocks = v.split(self.block_size, dim=1)

        outputs = []
        for q_block in q_blocks:
            block_attn = []
            for k_block, v_block in zip(k_blocks, v_blocks):
                # 只计算局部注意力
                attn = torch.matmul(q_block, k_block.transpose(-1,-2))
                block_attn.append(torch.matmul(attn, v_block))
            outputs.append(torch.cat(block_attn, dim=1))
        return torch.cat(outputs, dim=1)

def print_memory_usage(module):
    print(f"Allocated: {torch.cuda.memory_allocated()/1024**2:.2f}MB")
    print(f"Cached: {torch.cuda.memory_reserved()/1024**2:.2f}MB")

梯度爆炸预防：
在注意力层后立即添加 LayerNorm
使用 AdamW 优化器并设置 weight decay
warmup 学习率调度策略

注意力可视化：

import matplotlib.pyplot as plt

def plot_attention(attention_weights):
    plt.matshow(attention_weights.detach().cpu().numpy())
    plt.colorbar()
    plt.show()