深入解析Transformer模型：自注意力与多头自注意力机制的核心原理与实现

1次阅读

没有评论

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

在自然语言处理（NLP）领域，循环神经网络（RNN）和长短期记忆网络（LSTM）曾长期占据主导地位。然而，这些模型在处理长序列时存在明显缺陷：

梯度消失 / 爆炸问题 ：RNN 在反向传播时梯度会随着时间步呈指数级衰减或增长，导致难以训练深层网络。
顺序计算限制 ：必须按时间步顺序处理序列，无法充分利用现代 GPU 的并行计算能力。
长程依赖捕捉困难 ：尽管 LSTM 通过门控机制有所改善，但在超过 100 个时间步的序列中仍会丢失早期信息。

2017 年提出的 Transformer 架构通过以下设计彻底改变了序列建模：

完全基于注意力机制 ：摒弃循环结构，直接建模序列中所有位置的关系
并行化处理 ：所有时间步的计算可同时进行
恒定路径长度 ：任意两个位置的交互只需一步注意力计算

关键指标对比：

模型类型	计算复杂度	并行性	长程依赖
RNN	O(n)	❌	△
LSTM	O(n)	❌	○
Transformer	O(n²)	✔️	◎

输入表示 ：将每个词元的嵌入向量与位置编码相加
生成 QKV 矩阵 ：通过可学习权重矩阵生成查询 (Query)、键 (Key)、值 (Value)
```
Q = X @ W_Q  # [batch, seq_len, d_k]
K = X @ W_K  
V = X @ W_V
```

注意力分数计算 ：

scores = Q @ K.transpose(-2, -1) / sqrt(d_k)  # 缩放点积

Softmax 归一化 ：

attn_weights = torch.softmax(scores, dim=-1)

加权求和 ：

output = attn_weights @ V  # [batch, seq_len, d_v]

除以√d_k 是为了防止点积结果过大导致 softmax 进入梯度饱和区。假设 Q 和 K 的分量是独立随机变量，均值为 0，方差为 1，则 Q·K 的方差为 d_k。

将注意力机制并行执行多次可以捕捉不同子空间的模式：

分头处理 ：

# [batch, seq_len, num_heads, head_dim]
Q = Q.view(batch, -1, num_heads, d_k//num_heads)

独立计算注意力 ：每个头产生不同的注意力模式

合并结果 ：

# [batch, seq_len, d_model]
output = output.transpose(1,2).contiguous().view(batch, -1, d_model)

完整 PyTorch 实现示例：

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_k = d_model // num_heads
        self.num_heads = num_heads

        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)

    def forward(self, x):
        batch = x.size(0)

        # 1. 生成 QKV
        Q = self.W_Q(x).view(batch, -1, self.num_heads, self.d_k)
        K = self.W_K(x).view(batch, -1, self.num_heads, self.d_k)
        V = self.W_V(x).view(batch, -1, self.num_heads, self.d_k)

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

        # 3. 加权求和
        context = torch.matmul(attn, V)
        context = context.transpose(1,2).contiguous().view(batch, -1, self.num_heads*self.d_k)

        return self.W_O(context)

使用不同序列长度测试 GPU 计算时间（RTX 3090）：