Transformer架构解析：自注意力机制如何实现高效并行计算

1次阅读

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

在 NLP 领域，Transformer 的出现彻底改变了序列建模的方式。传统的 RNN/LSTM 虽然能够处理序列数据，但存在两个致命缺陷：

顺序依赖性 ：必须严格按时间步计算，无法并行化
长程依赖衰减 ：随着距离增加，早期时间步信息会逐渐丢失

而 Transformer 通过自注意力机制完美解决了这两个问题。下面我们通过三个维度深入解析其核心设计。

自注意力机制的核心是计算每个位置与其他位置的关联程度，公式表示为：

Attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d_k}})V

其中：

查询 (Query)：当前关注的位置向量 $Q = XW^Q$
键 (Key)：被比较的位置向量 $K = XW^K$
值 (Value)：实际使用的信息向量 $V = XW^V$

计算步骤分解：

相似度计算：$QK^T$ 得到位置间点积分数
缩放处理：除以 $\sqrt{d_k}$ 防止梯度消失
权重归一化：softmax 转换为概率分布
信息聚合：加权求和得到最终表示

与传统 RNN 的对比：

graph LR
  RNN[RNN 计算] --> step1[t= 1 计算]
  step1 --> step2[t= 2 计算]
  step2 --> step3[...]

  Transformer[自注意力] --> 矩阵乘法 [QKV 矩阵运算]
  矩阵乘法 --> 并行 [所有位置同步计算]

关键优势：

所有时间步的 QKV 矩阵可一次性计算
注意力权重通过矩阵乘法并行获取
摆脱了 RNN 的串行计算链

import torch
import torch.nn.functional as F

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_k = d_model // n_heads
        self.n_heads = n_heads
        # 线性变换层
        self.wq = torch.nn.Linear(d_model, d_model)
        self.wk = torch.nn.Linear(d_model, d_model)
        self.wv = torch.nn.Linear(d_model, d_model)

    def forward(self, x):
        bs, seq_len, _ = x.shape
        # 1. 计算 QKV [batch_size, seq_len, d_model]
        q = self.wq(x)
        k = self.wk(x)
        v = self.wv(x)

        # 2. 多头拆分 [batch_size, n_heads, seq_len, d_k]
        q = q.view(bs, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        k = k.view(bs, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        v = v.view(bs, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # 3. 缩放点积注意力
        scores = (q @ k.transpose(-2, -1)) / (self.d_k ** 0.5)
        attn = F.softmax(scores, dim=-1)

        # 4. 信息聚合与拼接
        output = (attn @ v).transpose(1, 2).contiguous()
        return output.view(bs, seq_len, -1)