自注意力机制（Self-Attention Mechanism）原理解析与实战应用

1次阅读

没有评论

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

自注意力机制是 Transformer 架构的核心，彻底改变了自然语言处理领域。然而，许多开发者在理解其工作原理时常常遇到困惑。主要痛点包括：

QKV 矩阵的物理意义不直观
缩放因子（scaling factor）的作用难以理解
多头注意力的实现细节复杂
位置编码的必要性和选择

这些难点导致许多开发者只是简单地调用现成的 Transformer 库，而无法根据实际需求进行定制化修改或优化。

自注意力机制的核心是让序列中的每个元素都能关注到序列中的所有元素（包括自己），从而捕捉长距离依赖关系。其计算过程可分为以下步骤：

查询 (Query)、键(Key)、值(Value) 矩阵的计算：
$$Q = XW^Q, K = XW^K, V = XW^V$$
其中 X 是输入序列，W 是可学习参数矩阵。
注意力分数的计算：
$$\text{Attention}(Q, K, V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V$$
这里 $d_k$ 是键向量的维度，缩放因子 $\frac{1}{\sqrt{d_k}}$ 用于防止点积过大导致 softmax 梯度消失。

多头注意力将上述过程并行执行多次（h 次），然后将结果拼接：
$$\text{MultiHead}(Q, K, V) = \text{Concat}(head_1, …, head_h)W^O$$
其中每个 head 的计算为：
$$head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$$

下面是完整的 PyTorch 实现，包含多头注意力、残差连接和层归一化：

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # 线性变换层
        self.Wq = nn.Linear(d_model, d_model)
        self.Wk = nn.Linear(d_model, d_model)
        self.Wv = nn.Linear(d_model, d_model)
        self.Wo = nn.Linear(d_model, d_model)

    def forward(self, x):
        batch_size, seq_len, _ = x.size()

        # 线性变换得到 Q,K,V
        Q = self.Wq(x)  # [batch_size, seq_len, d_model]
        K = self.Wk(x)
        V = self.Wv(x)

        # 分割为多头
        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)

        # 计算缩放点积注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))
        attn_weights = F.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V)

        # 合并多头
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)

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

class TransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.attention = MultiHeadAttention(d_model, num_heads)
        self.norm1 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(nn.Linear(d_model, 4 * d_model),
            nn.ReLU(),
            nn.Linear(4 * d_model, d_model)
        )
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x):
        # 自注意力 + 残差连接 + 层归一化
        attn_output = self.attention(x)
        x = self.norm1(x + attn_output)

        # 前馈网络 + 残差连接 + 层归一化
        ffn_output = self.ffn(x)
        return self.norm2(x + ffn_output)

自注意力机制的时间和空间复杂度都是 $O(n^2)$，其中 n 是序列长度。这对于长序列处理是一个瓶颈。