跳到主要内容

内容阅读

AI 底层工程阿新聊ai

Transformer 详解:自注意力、多头注意力、位置编码,逐层拆解

本文迁移自 mindcarver/91ai · 原始位置 docs/machine learning/stage5 transformer/01 transformer explained.md · 由 @阿新聊ai 整理。 <! 调研来源: 1. Vaswani et al., "Attention Is All...

本文迁移自 mindcarver/91ai · 原始位置 docs/machine-learning/stage5-transformer/01-transformer-explained.md · 由 @阿新聊ai 整理。

Transformer 详解:自注意力、多头注意力、位置编码,逐层拆解

TL;DR: Transformer 完全基于注意力机制,不使用 RNN 或 CNN。它的核心是 Scaled Dot-Product Attention:用 Query 和 Key 的点积(除以 $\sqrt{d_k}$)计算注意力权重,然后对 Value 做加权求和。多头注意力让模型同时学习多种关系模式。位置编码用正弦/余弦函数为模型注入位置信息。每个 Transformer 层 = 多头注意力 + 前馈网络 + 残差连接 + Layer Norm。整个架构可以高度并行计算,这是它取代 RNN 的关键优势。

为什么这很重要

Transformer 是现代 NLP(和越来越多的其他领域)的基础架构。BERT、GPT、T5、BLOOM、LLaMA——所有这些模型都是 Transformer 的变体。理解 Transformer 是理解大语言模型的前提。

2017 年的原始论文标题是"Attention Is All You Need"。这个标题准确地概括了 Transformer 的核心主张:只用注意力机制就够了,不需要 RNN 的递归结构,也不需要 CNN 的局部卷积。

这个主张被后来 7 年的研究充分验证了。Transformer 从机器翻译出发,征服了 NLP 的几乎所有任务,然后扩展到视觉(ViT)、语音(Whisper)、多模态(CLIP)、蛋白质结构预测(AlphaFold)——甚至数学推理和代码生成。

核心概念

Transformer 的整体架构

原始 Transformer 是一个编码器-解码器结构:

输入序列 "I love AI"
  │
  ▼
[Token Embedding + Positional Encoding]
  │
  ▼
┌──────────────────────────────┐
│  Encoder (×N 层, N=6)         │
│                                │
│  Multi-Head Self-Attention     │ ← 编码器内部的自注意力
│  + Add & Norm                  │
│                                │
│  Feed-Forward Network          │ ← 位置独立的前馈网络
│  + Add & Norm                  │
└──────────────────────────────┘
  │
  ▼ 编码器输出
  │
  ▼
┌──────────────────────────────┐
│  Decoder (×N 层, N=6)         │
│                                │
│  Masked Multi-Head Self-Attn  │ ← 解码器自注意力(带因果 mask)
│  + Add & Norm                  │
│                                │
│  Cross-Attention               │ ← 解码器关注编码器输出
│  + Add & Norm                  │
│                                │
│  Feed-Forward Network          │
│  + Add & Norm                  │
└──────────────────────────────┘
  │
  ▼
Linear + Softmax → 输出概率分布

BERT 只使用编码器部分,GPT 只使用解码器部分。

工作原理(简化的心智模型)

给 12 岁孩子的解释

想象一个会议室里坐着 10 个人,每人手里拿着一句话中的一个词。

自注意力:每个人可以问其他所有人一个问题,然后根据回答整合信息。"猫"这个词的人可能会特别关注"坐"和"垫子"——因为它们是相关的。每轮讨论后,每个人都获得了一个"综合了全局信息的理解"。

多头注意力:不是只有一轮讨论,而是同时进行 8 轮不同主题的讨论。一轮关注语法(谁做主语),一轮关注语义(谁和谁相关),一轮关注位置(相邻词的关系)等。最后把所有讨论的结果合并。

位置编码:因为所有人同时说话(并行计算),没有人知道自己在句子中的"位置"。解决方案:给每个人一个独特的"座位号徽章"(位置编码),这样模型可以通过徽章区分第一个词和最后一个词。

前馈网络:每轮讨论之后,每个人独立地"消化"刚才听到的信息——用一个小型神经网络来处理和整合。这个消化过程每个人是独立的(可以并行计算)。

多层堆叠:整个"讨论-消化"过程重复 6 次(或更多)。每一层都能在上一层的基础上建立更抽象的理解。第一层可能学到词与词的关系,更深的层可能学到短语、从句甚至段落级别的关系。

工作原理(详细机制)

1. 输入表示:Embedding + Positional Encoding

Token Embedding

输入文本首先被分词(tokenization),然后每个 token 通过一个嵌入矩阵映射为 $d_{model}$ 维向量。

$$\text{embedding}(x_i) = E_{x_i} \in \mathbb{R}^{d_{model}}$$

其中 $E \in \mathbb{R}^{V \times d_{model}}$ 是可学习的嵌入矩阵,$V$ 是词汇表大小。

Positional Encoding

Attention 机制本身是"位置无关"的——如果你把输入序列的顺序打乱,Attention 的输出只是做同样的排列变换,不会"发现"顺序变了。这意味着 Attention 不知道"猫坐垫子"和"垫子坐猫"的区别。

位置编码为每个位置注入唯一的位置信息:

$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$ $$PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$

其中 $pos$ 是位置索引,$i$ 是维度索引。

每个维度对应一个不同频率的正弦波。低维度($i$ 小)的频率高(变化快),高维度($i$ 大)的频率低(变化慢)。这就像给每个位置一个由多个不同频率的"音符"组成的独特"和弦"。

为什么选择正弦/余弦?

  1. 每个位置有唯一的编码:不同的 $pos$ 产生不同的正弦/余弦组合。
  2. 相对位置可计算:对于任意固定的偏移 $k$,$PE_{pos+k}$ 可以表示为 $PE_{pos}$ 的线性函数。这让模型可以学习相对位置关系。
  3. 可以外推到更长的序列:正弦函数的定义域是无限的,理论上可以对训练时未见过的位置生成编码。

位置编码被加到 token embedding 上(不是拼接):

$$\text{input} = \text{embedding}(x) + PE$$

为什么是相加而不是拼接?因为 $d_{model}$ 维空间足够大,可以同时容纳语义信息和位置信息而不产生严重干扰。相加也比拼接更节省参数。

其他位置编码方案

原始的正弦位置编码后来被多种方案替代:

  • 可学习位置编码:直接为每个位置学习一个向量(BERT 使用这种)。优点是灵活,缺点是不能外推到训练时未见过的长度。
  • RoPE(Rotary Position Embedding):通过旋转矩阵编码相对位置(LLaMA、PaLM 使用)。优点是自然编码相对位置,在长序列上泛化更好。
  • ALiBi:在注意力分数上直接加上距离相关的偏置。简单且支持长序列外推。

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

这是 Transformer 的核心操作。详细公式:

$$\text{MultiHead}(X) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O$$

$$\text{head}_i = \text{Attention}(X W_i^Q, X W_i^K, X W_i^V)$$

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

参数分析:

  • $W_i^Q \in \mathbb{R}^{d_{model} \times d_k}$,$d_k = d_{model} / h$
  • $W_i^K \in \mathbb{R}^{d_{model} \times d_k}$
  • $W_i^V \in \mathbb{R}^{d_{model} \times d_v}$,$d_v = d_{model} / h$
  • $W^O \in \mathbb{R}^{d_{model} \times d_{model}}$

原始 Transformer 的配置:$d_{model} = 512$,$h = 8$,$d_k = d_v = 64$。

总参数量(单层多头注意力):

  • Q/K/V 投影:$3 \times 512 \times 512 = 786,432$
  • 输出投影:$512 \times 512 = 262,144$
  • 总计:1,048,576(约 100 万参数/层)

3. 前馈网络(Feed-Forward Network)

每个 Transformer 层的第二个子层是一个位置独立的前馈网络:

$$\text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2$$

或使用 ReLU 的变体:

$$\text{FFN}(x) = W_2 \cdot \text{ReLU}(W_1 x + b_1) + b_2$$

原始配置:$d_{model} = 512$,$d_{ff} = 2048$。

这个 FFN 对每个位置独立应用同一个网络(参数共享)。它为模型提供了非线性变换能力——注意力层做"信息路由"(决定哪些信息流到哪里),FFN 做"信息处理"(对路由过来的信息做非线性变换)。

FFN 的参数量:$512 \times 2048 + 2048 \times 512 = 2,097,152$(约 200 万参数/层),是注意力层的 2 倍。这也是 Transformer 的参数主要来源。

后来的变体使用了不同的激活函数:

  • GELU(BERT、GPT-2 使用):$\text{GELU}(x) = x \cdot \Phi(x)$
  • SwiGLU(LLaMA 使用):$\text{SwiGLU}(x, W, V, b) = (\text{Swish}(xW + b) \otimes (xV))$
  • GLU 变体在性能上通常优于 ReLU。

4. 残差连接和 Layer Normalization

每个子层的输出经过残差连接和 Layer Norm:

$$\text{output} = \text{LayerNorm}(x + \text{Sublayer}(x))$$

残差连接($x + \text{Sublayer}(x)$)的作用和 ResNet 中一样:提供梯度直通路径,防止梯度消失。

Layer Normalization 对每个样本的每个位置独立归一化:

$$\text{LayerNorm}(x) = \frac{x - \mu}{\sigma} \cdot \gamma + \beta$$

其中 $\mu$ 和 $\sigma$ 是 $x$ 在 $d_{model}$ 维度上的均值和标准差,$\gamma$ 和 $\beta$ 是可学习的缩放和偏移参数。

为什么用 Layer Norm 而不是 Batch Norm?因为序列长度可变,Batch Norm 在序列维度上的统计量不稳定。Layer Norm 在特征维度上归一化,不受序列长度影响。

Pre-Norm vs Post-Norm:原始 Transformer 使用 Post-Norm(先加残差再做 Layer Norm),后来发现 Pre-Norm(先做 Layer Norm 再加残差)训练更稳定:

  • Post-Norm:$\text{output} = \text{LayerNorm}(x + \text{Sublayer}(x))$
  • Pre-Norm:$\text{output} = x + \text{Sublayer}(\text{LayerNorm}(x))$

GPT-2 之后的大多数模型使用 Pre-Norm。

5. 编码器(Encoder)

编码器由 $N$ 个相同的层堆叠而成(原始论文 $N=6$):

每层:
  x → LayerNorm → Multi-Head Self-Attention → Add(residual) →
  → LayerNorm → FFN → Add(residual) → output

编码器的自注意力没有 mask——每个位置可以看到所有位置。这对于理解整个输入序列是必要的(如 BERT 的双向编码)。

6. 解码器(Decoder)

解码器也由 $N$ 个相同的层堆叠,但每层有三个子层:

每层:
  x → LayerNorm → Masked Self-Attention → Add(residual) →
  → LayerNorm → Cross-Attention(Q=dec, K/V=enc) → Add(residual) →
  → LayerNorm → FFN → Add(residual) → output

Masked Self-Attention:解码器的自注意力使用因果 mask(下三角矩阵),确保位置 $i$ 只能看到位置 $1, 2, \ldots, i$。这是因为在生成任务中,模型不能"偷看"未来的 token。

Cross-Attention:Query 来自解码器,Key 和 Value 来自编码器。这让解码器的每个位置可以"回顾"输入序列的所有位置。

7. 输出层

$$P(y_t) = \text{softmax}(W_{out} \cdot \text{decoder_output}_t)$$

$W_{out} \in \mathbb{R}^{V \times d_{model}}$ 将解码器输出投影到词汇表大小的维度。在原始 Transformer 中,$W_{out}$ 与输入 embedding 矩阵共享参数(weight tying),减少参数量并可能提升效果。

8. 训练目标

原始 Transformer 使用交叉熵损失:

$$\mathcal{L} = -\sum_{t=1}^{T} \log P(y_t | y_{<t}, x)$$

训练时使用 teacher forcing:解码器的输入是真实的目标序列(shifted right,即向右移一位)。推理时使用自回归生成:每步用上一步的预测结果作为输入。

9. 完整的参数量计算

以原始 Transformer($d_{model}=512$,$h=8$,$d_{ff}=2048$,$N=6$)为例:

编码器每层

  • 多头注意力:$4 \times 512 \times 512 = 1,048,576$
  • FFN:$512 \times 2048 + 2048 \times 512 = 2,097,152$
  • Layer Norm(×2):$2 \times 512 \times 2 = 2,048$
  • 每层总计:$\approx 3.15M$

编码器 6 层:$\approx 18.9M$

解码器每层(多一个 Cross-Attention):

  • Masked Self-Attention:$1,048,576$
  • Cross-Attention:$1,048,576$
  • FFN:$2,097,152$
  • 每层总计:$\approx 4.19M$

解码器 6 层:$\approx 25.2M$

Embedding + 输出层:$\approx V \times 512 \times 2$

总参数量(不含 embedding):$\approx 44M$。加上 embedding($V = 37000$):$\approx 82M$。

代码示例(完整可运行的 Python)

从零实现完整 Transformer

"""
完整 Transformer 实现:编码器-解码器架构
运行要求: pip install torch
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import copy


def clones(module, N):
    """克隆 N 个相同的层"""
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


class PositionalEncoding(nn.Module):
    """正弦/余弦位置编码"""
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        # 预计算位置编码矩阵
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数维度
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数维度
        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)


class MultiHeadAttention(nn.Module):
    """多头注意力"""
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_k = d_model // n_heads
        self.n_heads = n_heads
        self.linears = clones(nn.Linear(d_model, d_model), 4)  # Q, K, V, O
        self.attn = None
        self.dropout = nn.Dropout(dropout)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)

        # 投影并分头: (batch, seq, d_model) → (batch, heads, seq, d_k)
        Q, K, V = [
            lin(x).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
            for lin, x in zip(self.linears[:3], (query, key, value))
        ]

        # Scaled Dot-Product Attention
        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_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        attn_output = torch.matmul(attn_weights, V)
        self.attn = attn_weights

        # 拼接多头并输出投影
        attn_output = attn_output.transpose(1, 2).contiguous().view(
            batch_size, -1, self.n_heads * self.d_k
        )
        return self.linears[-1](https://raw.githubusercontent.com/mindcarver/91ai/main/docs/machine-learning/stage5-transformer/attn_output)


class PositionwiseFeedForward(nn.Module):
    """位置独立前馈网络"""
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))


class EncoderLayer(nn.Module):
    """编码器层"""
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ff = PositionwiseFeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Pre-Norm 架构
        x_norm = self.norm1(x)
        attn_out = self.self_attn(x_norm, x_norm, x_norm, mask)
        x = x + self.dropout1(attn_out)

        x_norm = self.norm2(x)
        ff_out = self.ff(x_norm)
        x = x + self.dropout2(ff_out)
        return x


class DecoderLayer(nn.Module):
    """解码器层"""
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ff = PositionwiseFeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self, x, enc_output, src_mask=None, tgt_mask=None):
        # Masked Self-Attention
        x_norm = self.norm1(x)
        attn_out = self.self_attn(x_norm, x_norm, x_norm, tgt_mask)
        x = x + self.dropout1(attn_out)

        # Cross-Attention
        x_norm = self.norm2(x)
        cross_out = self.cross_attn(x_norm, enc_output, enc_output, src_mask)
        x = x + self.dropout2(cross_out)

        # FFN
        x_norm = self.norm3(x)
        ff_out = self.ff(x_norm)
        x = x + self.dropout3(ff_out)
        return x


class Transformer(nn.Module):
    """完整 Transformer 模型"""
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model=512,
                 n_heads=8, d_ff=2048, N=6, dropout=0.1, max_len=5000):
        super().__init__()
        self.d_model = d_model

        # Embeddings
        self.src_embed = nn.Embedding(src_vocab_size, d_model)
        self.tgt_embed = nn.Embedding(tgt_vocab_size, d_model)
        self.pos_enc = PositionalEncoding(d_model, max_len, dropout)

        # Encoder &#x26; Decoder
        self.encoder = clones(EncoderLayer(d_model, n_heads, d_ff, dropout), N)
        self.decoder = clones(DecoderLayer(d_model, n_heads, d_ff, dropout), N)

        # 输出层
        self.norm = nn.LayerNorm(d_model)
        self.output_proj = nn.Linear(d_model, tgt_vocab_size)

        self._init_weights()

    def _init_weights(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def encode(self, src, src_mask=None):
        x = self.pos_enc(self.src_embed(src) * math.sqrt(self.d_model))
        for layer in self.encoder:
            x = layer(x, src_mask)
        return self.norm(x)

    def decode(self, tgt, enc_output, src_mask=None, tgt_mask=None):
        x = self.pos_enc(self.tgt_embed(tgt) * math.sqrt(self.d_model))
        for layer in self.decoder:
            x = layer(x, enc_output, src_mask, tgt_mask)
        return self.norm(x)

    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        enc_output = self.encode(src, src_mask)
        dec_output = self.decode(tgt, enc_output, src_mask, tgt_mask)
        logits = self.output_proj(dec_output)
        return logits

    @staticmethod
    def make_causal_mask(size):
        """创建因果 mask(下三角)"""
        mask = torch.tril(torch.ones(size, size)).unsqueeze(0).unsqueeze(0)
        return mask


# ============================================================
# 验证代码
# ============================================================
if __name__ == "__main__":
    torch.manual_seed(42)

    # 小型 Transformer 配置
    src_vocab = 1000
    tgt_vocab = 1000
    d_model = 128
    n_heads = 4
    d_ff = 512
    N = 2  # 2 层(用于测试)
    batch_size = 4
    src_len = 15
    tgt_len = 12

    print("=" * 60)
    print("完整 Transformer 测试")
    print("=" * 60)

    model = Transformer(src_vocab, tgt_vocab, d_model, n_heads, d_ff, N)
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\n模型配置:")
    print(f"  d_model={d_model}, n_heads={n_heads}, d_ff={d_ff}, N={N}")
    print(f"  总参数量: {total_params:,}")

    # 创建输入
    src = torch.randint(1, src_vocab, (batch_size, src_len))
    tgt = torch.randint(1, tgt_vocab, (batch_size, tgt_len))

    # 创建 mask
    tgt_mask = Transformer.make_causal_mask(tgt_len)

    # 前向传播
    logits = model(src, tgt, tgt_mask=tgt_mask)

    print(f"\n前向传播:")
    print(f"  源序列: {list(src.shape)}")
    print(f"  目标序列: {list(tgt.shape)}")
    print(f"  输出 logits: {list(logits.shape)}")
    print(f"  输出范围: [{logits.min().item():.2f}, {logits.max().item():.2f}]")

    # 位置编码可视化
    print("\n" + "=" * 60)
    print("位置编码分析")
    print("=" * 60)

    pe = PositionalEncoding(d_model)
    pe_values = pe.pe[0, :20, :8].detach()  # 前 20 个位置,前 8 个维度

    print("\n前 10 个位置的前 6 个维度:")
    print(f"  {'Pos':>4}", end="")
    for d in range(6):
        print(f"  {'dim'+str(d):>8}", end="")
    print()
    for pos in range(10):
        print(f"  {pos:>4}", end="")
        for d in range(6):
            print(f"  {pe_values[pos, d]:>8.4f}", end="")
        print()

    # 验证位置编码的唯一性
    cos_sim = F.cosine_similarity(
        pe_values.unsqueeze(0), pe_values.unsqueeze(1), dim=-1
    )
    print(f"\n位置编码余弦相似度(部分):")
    print(f"  Pos 0 vs Pos 1: {cos_sim[0, 1]:.4f}")
    print(f"  Pos 0 vs Pos 5: {cos_sim[0, 5]:.4f}")
    print(f"  Pos 0 vs Pos 10: {cos_sim[0, 10]:.4f}")
    print(f"  Pos 5 vs Pos 6: {cos_sim[5, 6]:.4f}")

    # 注意力权重分析
    print("\n" + "=" * 60)
    print("注意力权重分析")
    print("=" * 60)

    # 获取最后一层编码器的注意力权重
    with torch.no_grad():
        enc_output = model.encode(src)
        print(f"  编码器输出: {list(enc_output.shape)}")
        print(f"  输出范数: {enc_output.norm(dim=-1).mean().item():.4f}")

    # 不同大小的 Transformer 参数量对比
    print("\n" + "=" * 60)
    print("不同配置的参数量")
    print("=" * 60)

    configs = [
        ("Tiny", 128, 2, 2, 512),
        ("Small", 512, 8, 6, 2048),
        ("Base", 768, 12, 12, 3072),
        ("Large", 1024, 16, 24, 4096),
    ]

    for name, d, h, n, d_ff in configs:
        m = Transformer(30000, 30000, d, h, d_ff, n)
        params = sum(p.numel() for p in m.parameters())
        print(f"  {name:>8}: d={d}, h={h}, N={n}, d_ff={d_ff}{params/1e6:.1f}M 参数")

评论

0
登录后可以参与评论和讨论。
💬

还没有评论

欢迎留下第一条评论,帮助这篇内容更快形成讨论。