跳到内容

10.4 Transformer 架构与位置:并行的是训练计算,不是自回归生成顺序

注意力工作台已经能让每个位置直接检索其他位置,但单个 attention 还不是 Transformer。生产线还需要位置、残差、normalization、逐位置 FFN、mask 和任务 head;encoder 与 decoder 的信息流也不同。

Transformer 的价值不是“去掉所有序列性”。训练时能并行计算已知 token positions,decoder 生成下一个 token 时仍受自回归因果顺序约束。

本课目标

  • 组装 pre-norm/post-norm Transformer block;
  • 区分 encoder-only、decoder-only 与 encoder–decoder;
  • 理解位置表示和 causal/padding/cross masks;
  • 分析训练并行、生成 KV cache 与长序列成本;
  • 写出 shape/mask 正确的 PyTorch encoder block。

1. Token Representation = Content + Position

离散 token ID 先查 embedding:

$$ E\in\mathbb R^{|V|\times d_{model}}. $$

序列 tensor [B,T,D] 还需 position signal,否则 self-attention 不区分排列。输入可能组合:

$$ x_t=tokenEmbed(t)+position(t)+segment/type(t). $$

相加不是唯一方案,relative/rotary 方法会在 attention score/QK 关系中编码位置。

Padding token 即使 embedding 为零,叠加位置并经过 bias 后也可能非零;必须使用 padding mask,并在 loss/aggregation 处排除。

2. Position 方法的取舍

Fixed sinusoidal

不同频率 sin/cos 给出 deterministic absolute positions,不增加 position table parameters。能计算更长位置不等于模型已学会长度外推。

Learned absolute embedding

为每个 position 学 vector,简单但通常有训练最大长度/table 边界。

Relative position bias

根据 query–key 相对距离修改 attention score,更直接表达距离/方向。

Rotary position embedding

对 Q/K 分量做与位置相关旋转,使 dot product 带相对位置信号。频率缩放与长上下文扩展有多种实现,不能把 “RoPE” 当一个无参数细节。

位置方案影响 extrapolation、cache、fine-tuning 与 serving,必须和 checkpoint/config 一起版本化。

3. 一个 Transformer Block 有两类子层

Multi-head self-attention

跨 positions 交换信息。

Position-wise FFN

对每个位置独立应用相同 MLP:

$$ FFN(x)=W_2\phi(W_1x+b_1)+b_2. $$

Attention 混合 token 维,FFN 混合 feature/channel 维。二者都使用 residual 和 normalization;dropout/activation/gating 依架构不同。

4. Post-norm 与 Pre-norm

原始 post-norm 形式近似:

$$ x'=LN(x+Attention(x)), $$

$$ y=LN(x'+FFN(x')). $$

常见 pre-norm:

$$ x'=x+Attention(LN(x)), $$

$$ y=x'+FFN(LN(x')). $$

Pre-norm 通常改善很深模型的 gradient path,但两者表示/训练动态不同;不能加载 checkpoint 时随意切换。Residual branch 的 dropout、scaling 和 initialization 也属于架构定义。

5. PyTorch Pre-norm Encoder Block

python
import torch
from torch import nn

class EncoderBlock(nn.Module):
    def __init__(self, d_model, heads, d_ff, dropout=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.attention = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=heads,
            dropout=dropout,
            batch_first=True,
        )
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        self.residual_dropout = nn.Dropout(dropout)

    def forward(self, x, padding_mask=None, causal_mask=None):
        normalized = self.norm1(x)
        attended, _ = self.attention(
            normalized,
            normalized,
            normalized,
            key_padding_mask=padding_mask,  # [B,T], True 表示忽略
            attn_mask=causal_mask,
            need_weights=False,
        )
        x = x + self.residual_dropout(attended)
        x = x + self.residual_dropout(self.ffn(self.norm2(x)))
        return x

block = EncoderBlock(d_model=128, heads=4, d_ff=512)
x = torch.randn(8, 20, 128)
padding = torch.zeros(8, 20, dtype=torch.bool)
assert block(x, padding_mask=padding).shape == x.shape

当前 MultiheadAttention Boolean padding mask 中 True 表示忽略。Causal mask 的 shape/dtype 和同时传两个 masks 的规则应按固定 PyTorch 版本测试;不同 SDPA API 语义可能相反。

这个 reference block 用于理解,不代表最新模型的全部优化(RMSNorm、SwiGLU、GQA、fused kernels 等)。

6. 三种 Transformer 家族

Encoder-only

通常允许 bidirectional self-attention,得到每个 token 的上下文表示。适合分类、检索 encoding、token labeling、masked prediction。

Decoder-only

使用 causal self-attention,只允许看当前及过去 token,训练 next-token prediction。适合 autoregressive language modeling/generation。

Encoder–decoder

Encoder 读取完整 source;decoder 有 causal self-attention,再用 cross-attention 查询 encoder outputs。适合 translation、summarization 和条件 generation。

“BERT/GPT/T5”不只是 mask 不同,objective、tokenizer、normalization、position、data 与训练方案也不同。

7. Decoder Training 能并行,Generation 仍串行

训练时完整 target sequence 已知,可把 inputs/labels shift 并用 causal mask,一次 matrix computation 计算所有 positions 的 next-token losses。这是 Transformer 相对 recurrent training 的重要并行优势。

Autoregressive generation 时第 $t+1$ token 不存在,必须等第 $t$ token 采样/选择后再继续。Batch、heads、layers 内仍并行,但 token steps 有依赖。

“训练快数百倍”不是架构保证;速度取决于 length、batch、hardware、kernel、RNN baseline、communication 和 memory。

8. KV Cache 避免重复投影历史

Decoder generation 若每步重算整个 prefix 的 K/V 很浪费。KV cache 为每层保存历史 keys/values,新 token 只计算新 Q/K/V,再让 query attend 到 cache。

收益与代价:

  • 避免历史 K/V projection 重算;
  • 每 token attention 仍随当前 context length 增长;
  • cache memory 约随 layers × sequence × KV heads × head dim × dtype 增长;
  • beam/batch、GQA/MQA、offload/quantization 改变占用;
  • position index、mask 和 cache eviction 必须正确。

Cache 是 serving state,tenant isolation、TTL 和 model version 不匹配都可能造成严重错误。

9. 长序列成本

Dense self-attention 的 score/weight 通常 $O(T^2)$,projection/FFN 常为 $O(TD^2)$ 量级。瓶颈随 $T,D$ 和 hardware 变化。

长上下文策略:

  • Flash/memory-efficient exact kernels:减少 memory I/O/materialization;
  • local/sliding-window/block-sparse attention:限制连接;
  • low-rank/kernel/linear approximations:近似 attention;
  • recurrence/memory/compression/retrieval:改变信息来源;
  • chunking:需处理跨 chunk 依赖。

宣称支持 128k context 只表示接口/训练配置允许,不证明模型能在所有位置可靠检索、推理。需要 needle、multi-hop、position、lost-in-the-middle 与真实任务评估。

10. Mask 是信息安全边界

至少测试:

  • causal:改变未来 token 不影响过去 logits;
  • padding:改变 padded token ID 不影响有效 outputs;
  • cross-attention:source padding 不被读取;
  • loss:只统计目标有效 positions;
  • packed/batched:不同长度样本互不污染;
  • cache:incremental logits 与 full-prefix logits 在容差内一致。

Mask 只要错一位,模型可能训练得更快、loss 更低,因为它看到了答案。

11. 训练稳定性

Transformer 仍依赖第 9 章原则:

  • residual/norm layout 与 initialization;
  • optimizer/weight decay parameter groups;
  • warmup 与 schedule;
  • gradient clipping、mixed precision 和 loss scaling;
  • dropout 与 mode;
  • validation、checkpoint、seed 和数据顺序。

LayerNorm 不依赖 batch statistics,但不能自动防止 activation outlier、attention logit overflow 或深层 residual accumulation。

12. Attention 模型仍需要任务 Bias

Transformer 的连接更灵活,不等于完全没有先验:tokenization、mask、position、context window、weight sharing 和 objective 都是 bias。

图像 Transformer 需要 patch/position/augmentation;时间序列需要 causal cutoff、scale 与 calendar;集合数据可能无需 absolute position。架构应匹配数据生成与部署过程。

常见误区

  • Transformer 去掉了顺序:它仍需要 position;decoder 生成还有因果顺序。
  • 所有 token 可并行生成:并行的是 teacher-forced training positions,不是 autoregressive steps。
  • KV cache 让长上下文成本恒定:attention 与 cache memory 仍随 context 增长。
  • 长 context 等于会利用长 context:需要任务级位置/检索验证。
  • LayerNorm 解决全部训练稳定性:仍有初始化、精度、优化和 residual 问题。

练习

  1. 为 encoder-only、decoder-only、encoder–decoder 画信息流和 masks。
  2. 实现 sinusoidal position,检查 odd/even dimension 和长度外推。
  3. 在 pre-norm block 中改变 future token,验证 causal invariance。
  4. 比较 full-prefix 与 KV-cache incremental logits。
  5. 估算一个模型 attention matrix、FFN activations 与 KV cache 内存。

小结

Transformer block 把 attention、position、FFN、residual 与 normalization 组合成可堆叠结构。Encoder、decoder 和 encoder–decoder 由信息流与 objective 区分。训练可并行计算已知 positions,自回归生成仍逐 token 推进;长上下文还受 attention、cache 和实际利用能力限制。

下一章进入预训练:tokenizer 决定模型看到的离散单元,objective 决定从海量语料学什么,fine-tuning 与 prompt 只是改变行为的不同接口。

Built with VitePress | Software Systems Atlas