Transformer 架构深度解析:从 Attention 到 Multi-Head Attention

Transformer 是 2017 年 Google 在论文《Attention Is All You Need》中提出的架构,彻底改变了 NLP 领域,也成为当今所有大语言模型的基础架构。

1. 自注意力机制(Self-Attention)

自注意力机制允许序列中的每个位置关注序列中的所有位置。其核心计算为:

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) * V

其中 Q(Query)、K(Key)、V(Value)是通过输入线性变换得到的三个矩阵。

2. 多头注意力(Multi-Head Attention)

多头注意力将输入投影到多个子空间,并行计算注意力,然后拼接结果。这使得模型能够关注不同位置的不同表示子空间的信息。

MultiHead(Q,K,V) = Concat(head_1, ..., head_h) * W_O
where head_i = Attention(Q*W_Q_i, K*W_K_i, V*W_V_i)

3. 位置编码(Positional Encoding)

由于 Transformer 没有循环结构,需要显式注入位置信息。原始论文使用正弦/余弦位置编码,现代模型多采用 RoPE(旋转位置编码)或 ALiBi 等改进方案。

4. PyTorch 实现示例

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        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, Q, K, V, mask=None):
        batch_size = Q.size(0)
        Q = self.W_q(Q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(K).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(V).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attn = torch.softmax(scores, dim=-1)
        return self.W_o(torch.matmul(attn, V).transpose(1,2).contiguous().view(batch_size, -1, self.num_heads * self.d_k))

5. 总结

理解 Transformer 架构是学习大模型的基础。从 GPT 到 BERT,从 LLaMA 到 DeepSeek,所有主流大模型都基于 Transformer 的 Decoder 或 Encoder-Decoder 结构。