Transformer 架构详解：核心模块、变体与推理成本

Transformer 架构是一组重复堆叠的神经网络 Block。每个 Block 使用 Attention 在位置之间传递信息，再用逐位置 Feed-Forward Network 变换特征；Residual Path 维持共享表示流，Normalization 控制数值尺度，Position Signal 编码顺序，Mask 则决定哪些位置可以交换信息。

Attention 只是架构中的一个运算。可训练、可部署的 Transformer 还必须明确输入表示、Block 顺序、归一化位置、输出头、训练目标和推理状态。这些契约共同决定信息如何流动、模型可以预测什么，以及系统需要多少显存与延迟预算。

核心结论

• Transformer Block 通常保持 [Batch, Token, Model Dimension] 的外部形状不变：Attention 跨位置混合信息，FFN 对每个位置独立变换。
• 2017 年原始 Transformer 是采用 Post-LayerNorm、正弦位置编码、ReLU FFN 和多头注意力的编解码器模型；这些具体选择不是所有 Transformer 的必要条件。
• 仅编码器、仅解码器和编解码器模型的核心差异是可见范围、训练目标与输出流，而不是简单的「理解模型」和「生成模型」标签。
• 因果模型训练时可以在 Mask 约束下并行处理已知目标位置；自回归生成仍然串行，因为下一 Token 依赖刚刚选出的 Token。
• 稠密 Attention 不是全部成本。Projection 和 FFN 随序列长度线性增长、随模型宽度快速增长；KV Cache 则由 Token 数、层数、KV Head、Head Dimension、Batch 和精度共同决定。
• 正确性验证必须覆盖 Mask、Position、缓存解码、Tokenizer/Checkpoint 契约、有限数值与任务表现。

什么构成 Transformer

Transformer 把跨位置的信息路由和逐位置的特征变换放进残差堆栈。原始论文 Attention Is All You Need 面向机器翻译，使用一个 Encoder 和一个自回归 Decoder。后续模型保留了 Block 结构，却改变了可见范围、训练目标、归一化、位置信息、激活函数和 Attention Head 布局。

对形状为 [B, T, D] 的 Hidden State，常见的 Pre-Norm Block 可写成：

H = X + Attention(Norm(X), mask, position)
Y = H + FFN(Norm(H))

原始论文使用 Post-Norm：

H = Norm(X + Attention(X, mask, position))
Y = Norm(H + FFN(H))

二者都保持 [B, T, D] 外部形状，因此可以堆叠，但优化行为并不相同。Xiong 等人的研究分析了归一化位置对初始化阶段梯度的影响；结论来自论文设定，不能简化为某种布局对所有模型都更好。

本文把注意力机制视为已经定义的基础运算。独立的注意力机制指南详细解释 QKV Shape、Mask 极性、稳定 Softmax、MHA/MQA/GQA 与 Kernel 行为。

一个 Transformer Block 的内部结构

理解 Block 的实用方式是：从 Residual Stream 读取表示，计算本子层的增量，再把增量写回。Attention 在允许的位置之间搬运信息；FFN 使用同一组参数独立变换每个位置。

Attention 子层

单个 Head 的缩放点积注意力为：

Attention(Q, K, V, M) = softmax((QK^T / sqrt(d_head)) + M) V

Multi-Head Attention 使用多组独立投影，再拼接 Head 输出并做 Output Projection。每个 Head 没有预先规定的人类可读职责。「第三个 Head 负责语法」之类结论需要干预实验或受控分析，单张 Attention Map 不能证明因果作用。

Mask 属于架构契约：

• Bidirectional Self-Attention 允许读取全部合法输入位置；
• Causal Self-Attention 屏蔽未来位置；
• Cross-Attention 的 Query 来自目标流，Key/Value 来自来源流；
• Padding Mask 与业务 Mask 排除无效位置。

Feed-Forward 子层

Feed-Forward Network（FFN）对每个位置独立应用同一变换：

FFN(x) = activation(x W1 + b1) W2 + b2

原始模型使用 ReLU，后续设计也使用 GELU 或门控变体。GLU Variants Improve Transformer在其序列到序列实验中报告多种门控 FFN 的收益，但这不能证明某个激活函数在所有模型上都最优。

Attention 经常获得最多讨论，但 FFN 可能占据稠密模型的大量参数与矩阵乘法。容量和延迟规划必须同时计算两个子层。

残差连接与归一化

Residual Connection 允许子层只写入一个增量，而不必替换整份表示。Normalization 的位置会改变优化行为、Checkpoint 兼容性和真实前向函数。

LayerNorm 对特征做中心化与尺度调整；RMSNorm 根据均方根调整尺度。RMSNorm 论文在特定模型上评估了效率与任务表现。LayerNorm 与 RMSNorm、Pre-Norm 与 Post-Norm、Epsilon 和仿射参数都应视为权重契约，不能随意互换。

位置信息是模型契约

只看 Q/K/V 的 Attention 不包含序列顺序，因此 Transformer 必须在表示或 Attention 计算中加入 Position Information。

原始 Transformer 给出正弦和可学习绝对位置两种选择；BERT 使用 Token、Segment 与可学习 Position Embedding；T5 使用相对位置 Bias。RoFormer提出 RoPE，通过旋转编码绝对位置，并让相对位移进入 Attention 关系。

运行时更换位置方案不是无损优化。Position ID、缩放规则、Cache Offset 和训练时最大上下文必须保持一致。模型能接收更长 Tensor，并不证明它能可靠使用新增上下文。

Encoder、Decoder 与 Encoder-Decoder

三类主要 Transformer 具有不同的信息流。只有把架构名称与实际 Mask、Objective 和 Output Head 放在一起，分类才有意义。

BERT证明了使用 Masked Token 和历史 NSP 目标训练双向 Encoder 的可行性，但不能代表所有 Encoder-Only 模型都采用 BERT 目标。T5展示了 Encoder-Decoder 的 Text-to-Text 框架和 Span Corruption，同样不是所有条件任务的通用配方。

Decoder-Only 模型也可以在目标与输出处理匹配时执行分类、抽取或 Embedding；Encoder-Only 模型也能作为检索器或编码器参与生成系统。选型应依据所需信息流和实测系统行为，而不是固定任务标签。

训练、Prefill 与 Decode 是三种负载

因果 Transformer 的训练与生成会执行相同的学习 Block，但两者拥有不同的已知信息。

Teacher Forcing 训练时，完整目标序列已经存在。Causal Mask 阻止位置读取未来目标，但 Kernel 仍可并行计算多个位置。推理阶段则分为：

1. Prefill：处理已知 Prompt，并为各层生成初始 Key 与 Value；
2. Decode：处理刚选出的 Token，复用历史 K/V，再输出下一位置的 Logit；
3. 采样或解码策略选出 Token，形成下一步必须等待的数据依赖。

KV Cache避免重复计算历史 Token 的 K/V Projection，但不会让自回归 Token 选择并行化，还会引入显存分配、Position、淘汰、复用和租户隔离问题。Hugging Face 的缓存文档指出，Cache Position 和 Attention Mask 长度必须同时覆盖历史与当前 Token，且缓存面向推理而非训练。

计算量、参数量与 KV Cache

Transformer 成本同时取决于序列长度和模型宽度。只说「Transformer 是 O(n²)」会隐藏重要开销，并可能导致错误容量估计。

对序列长度 T、模型宽度 D、FFN 宽度 F 的稠密 Block：

不同实现的常数项和融合算子不同。FlashAttention 减少中间结果与显存访问，但不会消除稠密 Attention 的成对算术；Sparse、Local、Recurrent 或 State Space 方案改变了运算本身，必须在各自质量和硬件条件下评估。

可运行的 Go 架构审计程序

下面的无依赖 Go 程序会验证 Decoder 配置，并估算无 Bias 稠密 Block 参数量和 KV Cache 容量。程序假设 SwiGLU 风格 FFN 使用三个矩阵，不包含 Embedding、输出头、Bias、MoE、额外 Norm、Optimizer State、Activation 和内存分配器开销。

package main

import (
	"errors"
	"fmt"
)

type Config struct {
	Layers     int64
	ModelDim   int64
	FFNDim     int64
	QueryHeads int64
	KVHeads    int64
}

func (c Config) validate() error {
	if c.Layers <= 0 || c.ModelDim <= 0 || c.FFNDim <= 0 {
		return errors.New("layers and dimensions must be positive")
	}
	if c.QueryHeads <= 0 || c.KVHeads <= 0 {
		return errors.New("head counts must be positive")
	}
	if c.ModelDim%c.QueryHeads != 0 {
		return errors.New("model dimension must be divisible by query heads")
	}
	if c.QueryHeads%c.KVHeads != 0 {
		return errors.New("query heads must be divisible by KV heads")
	}
	return nil
}

func (c Config) headDim() int64 {
	return c.ModelDim / c.QueryHeads
}

func (c Config) blockParameters() int64 {
	kvWidth := c.KVHeads * c.headDim()
	attention := c.ModelDim*c.ModelDim +
		2*c.ModelDim*kvWidth +
		c.ModelDim*c.ModelDim
	swigluFFN := 3 * c.ModelDim * c.FFNDim
	twoRMSNorms := 2 * c.ModelDim
	return attention + swigluFFN + twoRMSNorms
}

func (c Config) kvCacheBytes(batch, tokens, bytesPerScalar int64) int64 {
	return 2 * batch * tokens * c.Layers *
		c.KVHeads * c.headDim() * bytesPerScalar
}

func gib(bytes int64) float64 {
	return float64(bytes) / (1 << 30)
}

func main() {
	config := Config{
		Layers:     32,
		ModelDim:   4096,
		FFNDim:     11008,
		QueryHeads: 32,
		KVHeads:    8,
	}
	if err := config.validate(); err != nil {
		panic(err)
	}

	fmt.Printf("head dimension: %d\n", config.headDim())
	fmt.Printf("parameters per block: %d\n", config.blockParameters())
	fmt.Printf(
		"KV cache: %.2f GiB\n",
		gib(config.kvCacheBytes(4, 8192, 2)),
	)
}

预期输出：

head dimension: 128
parameters per block: 177217536
KV cache: 4.00 GiB

结果是容量规划估算，不是 Profiler 实测值。真实峰值内存还包括权重、临时 Workspace、Activation、碎片、Runtime Metadata，以及可能重复或量化的状态。

架构图不会暴露的失败模式

结构图看似正确的 Block，仍可能生成错误训练结果或错误 Token：

1. 未来 Token 泄漏：Causal Mask 缺失、错位或极性相反；
2. Padding 泄漏：Padding Key 仍可见，或 Loss 包含 Padding Target；
3. Position 漂移：缓存解码重新从零计算 Position ID，或 RoPE Offset 错误；
4. Checkpoint 不匹配：Tokenizer Vocabulary、Special Token、Norm、Head 数或 Weight Tying 与权重不一致；
5. 训练/评估模式不一致：Dropout 仍开启、Norm 行为变化，或训练阶段错误启用 Cache；
6. 缓存解码漂移：整段计算与逐 Token 缓存计算的 Logit 超出预期数值误差；
7. 整行 Mask 或数值不稳定：Softmax 没有合法 Key，或低精度下溢出；
8. 目标函数不匹配：Output Head、Label Shift 和 Loss 没有表达目标任务。

框架 API 还有额外契约。当前 PyTorch TransformerEncoderLayer暴露 batch_first 与 norm_first，并明确它是功能有限的基础参考实现；文档还警告，错误的 Causal Hint 可能导致错误执行。

一套可执行的验证阶梯

应从局部不变量逐步验证到用户可见行为：

1. Shape 测试：断言每个 Projection、Residual Add、Head Split 与输出维度；
2. Mask 测试：扰动未来或 Padding Token，证明更早的合法输出不变；
3. 数值测试：拒绝 NaN/Inf，与稳定参考 Attention 比较并定义混合精度容差；
4. Cache Parity：逐位置比较整段计算与缓存解码的 Next-Token Logit；
5. 极小 Batch 过拟合：用干净小样本暴露接线、Loss 与 Label Shift 错误；
6. Checkpoint 契约：固定 Tokenizer、Special Token ID、架构配置、权重版本、Dtype 和位置设置；
7. 任务评测：在代表性切片上与简单 Baseline 比较，而不是只看一个总体分数；
8. 服务评测：在真实并发下测量首 Token 时间、Token 间延迟、吞吐、峰值内存、长上下文退化和失败行为。

架构选型应遵循这些证据。Encoder、Decoder、Encoder-Decoder、循环、卷积、State Space、Retrieval 和 Hybrid 系统需要在同一任务与资源预算下比较。Mamba 与状态空间模型指南介绍另一类序列处理架构，但不表示任一方案在所有任务上占优。

一手资料

• Vaswani 等：Attention Is All You Need：原始编解码器 Transformer、Attention、FFN、Residual、Normalization 与正弦位置设计。
• Devlin 等：BERT：双向 Encoder 与 Masked Token 预训练。
• Raffel 等：T5：Encoder-Decoder Text-to-Text 框架与相对位置 Bias。
• Xiong 等：On Layer Normalization in the Transformer Architecture：Pre-LN 与 Post-LN 的梯度行为分析。
• Su 等：RoFormer：Rotary Position Embedding。
• Ainslie 等：Grouped-Query Attention：KV Head 共享与 Decoder 推理权衡。
• PyTorch TransformerEncoderLayer 文档：当前参考层的 Shape 与 API 契约。

总结

Transformer 不只是 Self-Attention，而是由 Mask、Normalization、Position System、FFN、Objective 和 Inference State 共同定义的 Residual Architecture。学习时应先理解原始 Encoder-Decoder，再把现代 Decoder-Only 选择视为显式、与 Checkpoint 绑定的变体。

实现阶段应让 Shape 与 Visibility Contract 可执行；生产阶段应区分 Prefill 和 Decode，在压测前计算 KV Cache 容量，并以缓存解码一致性和任务级证据作为可信门槛。