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
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 问题。
练习
- 为 encoder-only、decoder-only、encoder–decoder 画信息流和 masks。
- 实现 sinusoidal position,检查 odd/even dimension 和长度外推。
- 在 pre-norm block 中改变 future token,验证 causal invariance。
- 比较 full-prefix 与 KV-cache incremental logits。
- 估算一个模型 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 只是改变行为的不同接口。