Attention 及各种变体

发表于 2026-11-15 分类于 machine-learning 阅读次数：本文字数： 5.2k 阅读时长 ≈ 19 分钟

本文用伪代码梳理常见 Transformer 的 Layer 与 Model 结构。伪代码中:

^N 表示重复堆叠 N 层
@ 表示矩阵乘
注释 # as (...) 标注该行表达式的形状。

#Attention

下文中的 RoPE 与 RMSNorm 见「Norm 与位置编码」一节。

#MHA

y = LinearOut(
  Dropout(
    Softmax(
      LinearQ(x)        # as (batch, head, seq, d_k)
        @ LinearK(x)^T  # as (batch, head,      d_k, seq)
      * rsqrt(d_k)
      + CausalMask()
    )
  ) @ LinearV(x)        # as (batch, head,           seq, d_k)
)                       # as (batch, head, seq,           d_k)

x: 输入序列，形状为 (batch, seq, d_model=d_k*head)
Q、K、V 都由同一个 x 投影得到，即 Self-Attention，下文记作 MHA(x)；Q 来自 x、K/V 来自另一个序列时为 CrossAttention，见下节
LinearQ, LinearK, LinearV: 线性变换层，将输入序列映射到查询、键和值的空间，d_model -> d_k*head
rsqrt(d_k): 缩放因子，防止点积过大导致 Softmax 饱和
CausalMask（可选）: 因果掩码，用于防止模型在训练过程中看到未来的信息，形状为 (seq, seq)，上三角部分为 -inf，对角线和其他部分为 0；只有 Decoder / 自回归模型需要，BERT 等 Encoder 只用 PaddingMask 屏蔽 padding
Softmax: 对注意力权重进行归一化，输出形状为 (batch, head, seq, seq)
Dropout: 随机丢弃一些注意力权重，防止过拟合
LinearOut: 线性变换层，将注意力输出映射回原始维度，d_k*head -> d_model

#CrossAttention

y = LinearOut(
  Dropout(
    Softmax(
      LinearQ(x)             # as (batch, head, seq, d_k)
        @ LinearK(memory)^T  # as (batch, head,      d_k, mem_seq)
      * rsqrt(d_k)
      + PaddingMask()
    )
  ) @ LinearV(memory)        # as (batch, head,           mem_seq, d_k)
)                            # as (batch, head, seq,               d_k)

x: Decoder 的输入序列，形状为 (batch, seq, d_model=d_k*head)
memory: Encoder 的输出，形状为 (batch, mem_seq, d_model)
没有 CausalMask：Decoder 的每个位置都可以看到完整的源序列，只需屏蔽 padding
- CausalMask 防的是目标序列的未来：推理时生成 $y_i$ 时 $y_{>i}$ 还不存在，训练时整个 $y$ 并行输入，所以要屏蔽掉；而 memory 来自源序列 $x$ ，在解码开始前就已经完整给定，没有"未来"可言，屏蔽它反而会让训练与推理不一致
- 不会泄漏目标序列的未来：位置 $i$ 的 Q 来自 Masked Self-Attention 的输出，只依赖 $y_{\le i}$ ，K、V 只依赖 $x$ ，所以输出只依赖 $y_{\le i}$ 和 $x$ ；FFN 逐位置计算，逐层往上推，整个 Decoder 的因果性完全由 Self-Attention 的 CausalMask 保证
- 注意力矩阵形状为 (seq, mem_seq)，目标第 $i$ 个位置与源第 $i$ 个位置之间没有先后关系，上三角掩码本身也没有意义
- 例外：源序列流式到达时（同声传译的 wait-k、流式语音识别），Cross-Attention 需要"目标位置 $i$ 只能看源序列前 $g(i)$ 个位置"的阶梯形掩码
与 MHA 是同一个模块，只是 K、V 的输入换成了 memory，下文记作 MHA(x, memory)
源码：PyTorch TransformerDecoderLayer.multihead_attn；HF 中 BertLayer.crossattention（add_cross_attention=True）、GPT2Block.crossattention（GPT2Attention(is_cross_attention=True)，此时 Q 用单独的 q_attn，K、V 共用 c_attn）

#MQA / GQA

y = LinearOut(
  Dropout(
    Softmax(
      LinearQ(x)                # as (batch, head,    seq, d_k)
        @ Repeat(LinearK(x))^T  # as (batch, head_kv,      d_k, seq)
      * rsqrt(d_k)
      + CausalMask()
    )
  ) @ Repeat(LinearV(x))        # as (batch, head_kv,           seq, d_k)
)                               # as (batch, head,    seq,           d_k)

LinearK, LinearV: d_model -> d_k*head_kv，其中 head_kv < head；LinearQ 仍为 d_model -> d_k*head
Repeat: 对应 HF 的 repeat_kv，把每个 KV 头复制 n_rep = head / head_kv 份（即 num_key_value_groups），从 (batch, head_kv, seq, d_k) 扩展到 (batch, head, seq, d_k)，等价于 repeat_interleave(dim=1)；FlashAttention、SDPA（enable_gqa=True）等 kernel 内部按组索引，不会真的复制
head_kv = 1 时为 MQA，head_kv = head 时退化为 MHA
KV Cache 只需存 head_kv 个头，显存是 MHA 的 head_kv / head
HF 中由 num_key_value_heads 控制：Llama 默认为 None（等于 num_attention_heads，即 MHA），Mixtral 为 32 个 Q 头 / 8 个 KV 头

#MLA

DeepSeek-V2/V3 的 Multi-head Latent Attention（HF DeepseekV3Attention），括号内为 DeepSeek-V3 的默认配置：

c_q            = RMSNorm(LinearQA(x))     # as (batch, seq, q_lora_rank=1536)
q_nope, q_rope = Split(LinearQB(c_q))     # as (batch, head=128, seq, d_nope=128), (batch, head, seq, d_rope=64)
c_kv, k_rope   = Split(LinearKVA(x))      # as (batch, seq, kv_lora_rank=512),     (batch, 1,    seq, d_rope=64)
k_nope, v      = Split(LinearKVB(RMSNorm(c_kv)))  # as (batch, head, seq, d_nope=128), (batch, head, seq, d_v=128)

y = LinearOut(
  Softmax(
    Concat(q_nope, RoPE(q_rope))                     # as (batch, head, seq, 192)
      @ Concat(k_nope, Broadcast(RoPE(k_rope)))^T    # as (batch, head,      192, seq)
    * rsqrt(192) * mscale
    + CausalMask()
  ) @ v                                              # as (batch, head,           seq, 128)
)                                                    # as (batch, head, seq,           128)

LinearQA / LinearQB: Q 的低秩分解（q_a_proj / q_b_proj），d_model=7168 -> 1536 -> head×(128+64)，中间夹一个 q_a_layernorm
LinearKVA: kv_a_proj_with_mqa，一次性产生压缩后的 latent c_kv（512 维）和所有头共享的 k_rope（64 维，类似 MQA 的单头 K）
LinearKVB: kv_b_proj，把 latent 解压回每个头的 k_nope 和 v，512 -> head×(128+128)
解耦 RoPE：RoPE 与位置相关，如果直接作用在解压后的 K 上，kv_b_proj 就无法被吸收进 Q，所以单独拿出 64 维承载位置信息，其余 128 维（nope）不加位置编码
KV Cache 只存 c_kv 和 k_rope，每 token 每层 512+64=576 个数，而 128 头的 MHA 需要 128×(192+128)=40960 个数；HF 实现里 cache.update(kv_nope, k_rot) 也是在压缩状态下读写的
mscale: yarn_apply_mscale，使用 YaRN 扩展上下文时对 softmax scale 的修正
HF 实现每步都把 latent 解压成完整的 K、V 再做标准 attention；推理框架（vLLM、SGLang、FlashMLA）会做矩阵吸收，把 LinearKVB 合并进 Q 和 LinearOut，直接在 576 维 latent 上做 attention
num_key_value_heads = num_attention_heads = 128，所以 repeat_kv 的 n_rep=1

#LlamaAttention

y = LinearOut(
  Softmax(
    RoPE(LinearQ(x))                # as (batch, head,    seq, d_k)
      @ Repeat(RoPE(LinearK(x)))^T  # as (batch, head_kv,      d_k, seq)
    * rsqrt(d_k)
    + CausalMask()
  ) @ Repeat(LinearV(x))            # as (batch, head_kv,           seq, d_k)
)                                   # as (batch, head,    seq,           d_k)

即 GQA + RoPE，对应 q_proj / k_proj / v_proj / o_proj
attention_bias=False：所有 Linear 都没有 bias；attention_dropout=0.0：没有 Dropout
先 RoPE，再写入 KV Cache（cache.update），再 repeat_kv

#Qwen3Attention

y = LinearOut(
  Softmax(
    RoPE(RMSNorm(LinearQ(x)))                # as (batch, head,    seq, d_k=128)
      @ Repeat(RoPE(RMSNorm(LinearK(x))))^T  # as (batch, head_kv,      d_k, seq)
    * rsqrt(d_k)
    + CausalMask()                           # 或 SlidingWindowMask()
  ) @ Repeat(LinearV(x))                     # as (batch, head_kv,           seq, d_k)
)                                            # as (batch, head,    seq,           d_k)

与 LlamaAttention 的区别是 QK-Norm：q_norm / k_norm 是作用在每个头 d_k 维上的 RMSNorm，位于 RoPE 之前，用来抑制 attention logits 过大
Qwen2 的 q/k/v_proj 带 bias，Qwen3 去掉了（attention_bias=False）
head_dim 是独立配置项（默认 128），不一定等于 d_model / head
当 use_sliding_window=True 时，第 max_window_layers 层之后的层使用 sliding_window 大小的滑动窗口掩码

#FFN

#标准 FFN

y = Dropout(
  LinearOut(
    GELU(
      LinearIn(x)  # as (batch, seq, d_ff)
    )
  )
)              # as (batch, seq, d_model)

不同实现中 Dropout 的位置略有不同：GPT-2（GPT2MLP）只在最后有一个 Dropout；PyTorch TransformerEncoderLayer._ff_block 是 Dropout(Linear2(Dropout(ReLU(Linear1(x)))))，激活后还有一个
激活函数：原始 Transformer 为 ReLU，BERT 为 gelu（erf 精确版），GPT-2 为 gelu_new（tanh 近似版）
d_ff 一般为 4×d_model

#GatedFFN / SwiGLU

y = LinearDown(
  SiLU(LinearGate(x))  # as (batch, seq, d_ff)
  * LinearUp(x)        # as (batch, seq, d_ff)
)                      # as (batch, seq, d_model)

SiLU(x) = x × sigmoid(x)（hidden_act="silu"）
有三个矩阵，为了与普通 FFN 参数量相当，d_ff 通常取 8/3×d_model 再取整，如 Llama-7B 为 4096 -> 11008
mlp_bias=False，没有 bias，也没有 Dropout

#MoE

对应 HF MixtralSparseMoeBlock：

p      = Softmax(LinearRouter(x))   # as (tokens, n_experts=8)
w, idx = TopK(p, k=2)               # as (tokens, k)
w      = w / Sum(w)
y      = Sum(w[i] * GatedFFN[idx[i]](x) for i in range(k))  # as (tokens, d_model)

x 先被展平为 (batch*seq, d_model)，每个 token 独立选专家
LinearRouter: MixtralTopKRouter，d_model -> n_experts，没有 bias；训练时可以对输入加 router_jitter_noise
每个专家都是一个 GatedFFN；HF 把所有专家的权重堆成 3D 参数 gate_up_proj：(n_experts, 2*d_ff, d_model) 和 down_proj：(n_experts, d_model, d_ff)，前向时只遍历被选中的专家，用 index_add_ 把结果累加回去
先 Softmax 再 TopK 再归一化，等价于先 TopK 再对选中的 logits 做 Softmax
训练时由 load_balancing_loss_func 计算负载均衡辅助损失，乘 router_aux_loss_coef 加到总 loss

#DeepSeekMoE

对应 HF DeepseekV3MoE，括号内为 DeepSeek-V3 的默认配置：

s     = Sigmoid(LinearRouter(x))                    # as (tokens, n_routed=256)
group = TopK(GroupSum(Top2(s + bias)), k=4)         # 256 个专家分 8 组，选 4 组
idx   = TopK(MaskGroup(s + bias, group), k=8)       # 只在选中的组内选 8 个专家
w     = s[idx] / Sum(s[idx]) * 2.5                  # as (tokens, 8)
y     = GatedFFN_shared(x)
      + Sum(w[i] * GatedFFN[idx[i]](x) for i in range(8))

打分用 Sigmoid 而不是 Softmax（scores = router_logits.sigmoid()）
bias: e_score_correction_bias，只影响选哪些专家，不影响专家的权重 w；它在训练中根据各专家负载在梯度之外更新，从而实现无辅助损失（aux-loss-free）的负载均衡
分组路由：每组用组内最高两个分数之和打分（n_group=8、topk_group=4），用来限制每个 token 跨越的设备数
norm_topk_prob=True 时对选中的权重归一化，再乘 routed_scaling_factor=2.5
共享专家：n_shared_experts=1，所有 token 都经过，与路由专家的结果直接相加
专家的 d_ff 为 moe_intermediate_size=2048，远小于 dense FFN 的 18432（细粒度专家）

#Norm 与位置编码

#LayerNorm / RMSNorm

1 2	LayerNorm(x) = (x - Mean(x)) * rsqrt(Var(x) + eps) * weight + bias # 在 d_model 维上归一化 RMSNorm(x) = x * rsqrt(Mean(x^2) + eps) * weight

RMSNorm 去掉了减均值和 bias，计算更省；HF LlamaRMSNorm 先转 float32 计算，再转回原 dtype，与 T5LayerNorm 等价
eps: PyTorch LayerNorm 默认 1e-5，BERT 为 1e-12，Llama 的 rms_norm_eps 为 1e-6

#Sinusoidal PE

原论文的位置编码：

y = Dropout(
  TokenEmbedding(x) * sqrt(d_model)  # as (batch, seq, d_model)
  + PE(pos)                          # as (seq, d_model)
)

PE(pos)[2i]   = Sin(pos / 10000^(2i / d_model))
PE(pos)[2i+1] = Cos(pos / 10000^(2i / d_model))

加法式：只在输入处与 TokenEmbedding 相加一次，之后的层不再注入位置信息；没有可学习的参数
每一对维度 (2i, 2i+1) 用一个频率 $\omega_i = 10000^{-2i/d_{model}}$ ，波长从 $2\pi$ 到 $10000 \cdot 2\pi$ ；低维变化快，区分相邻位置，高维变化慢，编码远距离
相对位置可由线性变换得到： $(\sin, \cos)$ 在 $pos + k$ 处的值等于 $pos$ 处的值旋转角度 $k\omega_i$ ；RoPE 正是把这种旋转直接作用在 q、k 上
* sqrt(d_model)：放大 Embedding，避免被 PE 淹没；HF 中对应 embed_scale（scale_embedding=True 时启用）
原论文认为它可以外推到训练时没见过的长度，但实际外推效果并不好
源码：PyTorch nn 中没有内置位置编码（nn.Transformer 不包含），官方教程里需要手写 PositionalEncoding；HF 中如 MarianSinusoidalPositionalEmbedding，它不交错排列，而是把 sin 放在前一半、cos 放在后一半，并作为冻结的 nn.Embedding（_freeze=True）按 position_ids 查表

#可学习 PE

GPT-2、BERT 使用的可学习绝对位置编码：

y = Dropout(
  TokenEmbedding(x)         # as (batch, seq, d_model)
  + PositionEmbedding(pos)  # as (seq, d_model)
)

PositionEmbedding: 就是一个 nn.Embedding(max_position_embeddings, d_model)，与 TokenEmbedding 一样随模型训练
pos: Arange(seq)；有 KV Cache 时从已缓存的长度开始，即 GPT-2 中的 arange(seq) + past_seen_tokens
长度有硬上限：超过 max_position_embeddings 的位置没有对应的向量，GPT-2 为 1024，BERT 为 512
原论文对比过可学习 PE 和 Sinusoidal PE，两者效果几乎相同；之所以选 Sinusoidal，是希望它能外推到更长的序列
BERT 还会再加 TokenTypeEmbedding，求和之后接 LayerNorm，见 BertModel 一节
加法式 PE 的共同问题：忽略投影矩阵简化地看， $(q + p_m) \cdot (k + p_n) = q \cdot k + q \cdot p_n + p_m \cdot k + p_m \cdot p_n$ ，内容与绝对位置混在一起，模型需要自己学着从中分离出相对距离

#RoPE

与 Sinusoidal / 可学习 PE 的加法不同，RoPE 是乘性的：它不给输入加上一个 f(pos)，而是把 q、k 旋转一个与位置有关的角度。

RoPE(x) = x * Cos(pos * inv_freq)         # x as (batch, head, seq, d_k)
        + RotateHalf(x) * Sin(pos * inv_freq)

inv_freq      = 1 / theta^(Arange(0, d_k, 2) / d_k)  # as (d_k/2,)，HF 中 Cat 两份得到 d_k 维
RotateHalf(x) = Cat(-x[..., d_k/2:], x[..., :d_k/2])

直观理解：先看 d_k=2，把 q、k 看成平面上的箭头。位置 $m$ 的 q 逆时针转 $m\theta$ ，位置 $n$ 的 k 转 $n\theta$ ，两者夹角从 $\varphi$ 变为 $\varphi + (m-n)\theta$ ，于是 $q'_m \cdot k'_n = |q||k|\cos(\varphi + (m-n)\theta)$ ，只依赖相对位置 $m-n$ ，与 $m$ 、 $n$ 各自的取值无关
每两维：旋转只在二维平面上有定义，所以把 d_k 维拆成 d_k/2 对，每一对看成一个平面各自旋转；第 $i$ 对的转速为 $\theta_i = \text{theta}^{-2i/d_k}$ （即 inv_freq），前几对转得快、区分相邻 token，后几对转得慢、编码远距离，与 Sinusoidal PE 用多个频率是同一个思路
只作用于 Q 和 K，不作用于 V：位置信息只影响"看谁"，不影响"取出什么内容"
RotateHalf 的由来：第 $i$ $i$ 对 $(x_a, x_b)$ $(x_{a}, x_{b})$ 旋转角度 $m\theta_i$ $m θ_{i}$ ，即 $x'_a = x_a \cos m\theta_i - x_b \sin m\theta_i$ $x_{a}^{'} = x_{a} cos m θ_{i} - x_{b} sin m θ_{i}$ ， $x'_b = x_b \cos m\theta_i + x_a \sin m\theta_i$ $x_{b}^{'} = x_{b} cos m θ_{i} + x_{a} sin m θ_{i}$
- 写成整个向量的矩阵乘法是一个绝大部分为 0 的 d_k×d_k 分块对角矩阵，太浪费；拆开后就是 x * cos + R(x) * sin，其中 R 把每一对 $(x_a, x_b)$ 变成 $(-x_b, x_a)$ ，即旋转 90°，只需要逐元素乘加，cos/sin 可以预先算好并在 batch、head 间广播
- HF 让第 $i$ 维与第 $i + d_k/2$ 维配对（GPT-NeoX 风格），此时 R 恰好是"后一半取负放前面，前一半放后面"，即 RotateHalf；配对的两维要用同一个角度，所以 cos/sin 要 Cat(freqs, freqs) 复制两份
- 原论文与 Meta 原始代码按相邻的 (2i, 2i+1) 配对，对应的 R 为 Stack(-x[..., 1::2], x[..., ::2])（rotate_every_two），也可以写成复数乘法 $x \cdot e^{im\theta}$ ；DeepSeek-V3 的 rope_interleave=True 即使用相邻配对
- 两种配对等价：attention 只用到 q·k 内积，对 q、k 的维度做同一个排列不改变内积，所以 HF 转换 Llama 权重时对 q_proj/k_proj 的输出维度做了一次置换；前后两半配对用的是连续切片，比 ::2 的跨步访问更适合硬件和 kernel 实现
theta: rope_theta，Llama 1/2 为 10000，Llama 3 为 500000
cos/sin 在 Model 层由 LlamaRotaryEmbedding 计算一次，作为 position_embeddings 传给每一层；KV Cache 中存的是旋转之后的 K
与 GPT-2 / BERT 的可学习绝对位置编码（只在输入处加一次）不同，RoPE 在每一层的 attention 里都要施加

#Layer

#Token Mixing 与 Channel Mixing

一个层的输入是 (seq, d_model) 的矩阵，可以沿两个方向混合信息，各种 Layer 都可以看成这两种子层的组合：

1 2	z = x + TokenMixer(Norm(x)) # 沿 seq 方向：在不同 token 之间交换信息 y = z + ChannelMixer(Norm(z)) # 沿 d_model 方向：每个 token 独立加工自己的特征

术语来自 MLP-Mixer（2021），它用一个沿 seq 方向、权重固定的 MLP 代替 attention 做 token mixing；MetaFormer（PoolFormer 的论文，2022）把 TokenMixer 抽象成一个可替换的槽位，上式就是 MetaFormer 结构
Attention 也不全是 token mixing：Q/K/V/O 四个投影对每个 token 单独计算，其实是 channel mixing；只有 Softmax(QK^T) @ V 在 token 之间混合，混合权重由输入实时算出（数据相关），同一个头内所有通道共用一组权重，不同的头则对不同的通道组使用不同的权重
FFN 就是 ChannelMixer：position-wise，所有位置共享同一套权重，因为文本中每个位置都是同一种东西

不同的 TokenMixer：

TokenMixer	代表模型	特点
Softmax Attention	Transformer	数据相关，O(n²)
固定权重 MLP	MLP-Mixer	权重可学习但与输入无关，序列长度必须固定，难以做因果掩码
傅里叶变换	FNet	没有参数
平均池化	PoolFormer	没有参数
线性 Attention / SSM	RWKV、Mamba	循环或卷积形式，O(n)
多头重排	RankMixer	没有参数，见下文

TokenMixer 有多重要取决于任务：视觉分类中相邻 patch 高度相关，PoolFormer 只用池化就能达到与 ViT 相当的精度，说明通用结构本身就很强；语言建模需要从很远的上下文中精确检索并复制 token，softmax attention 很难被替代，SSM 在关联回忆（associative recall）任务上明显较弱，这也是如今常见 attention 与 SSM 混合架构的原因之一
CNN 中早有同样的分解：MobileNet 的深度可分离卷积中，depthwise conv 对应 token mixing（每个通道单独在空间上卷积），1×1 conv 对应 channel mixing（每个像素单独在通道间混合）
RankMixer 把这一结构用在推荐排序模型上，并针对推荐特征的**异构**性改了两处：
- TokenMixer 是一个没有参数的重排：输入的 T 个 token 是按语义分组的特征，每个 token 切成 H 份，新的第 h 个 token 为 Concat(x[1][h], x[2][h], ..., x[T][h])，取 H=T 以保持 token 数不变；不同 token 处于不同的语义空间，内积相似度没有意义，Self-Attention 的效果反而略差，计算成本还更高
- ChannelMixer 是 Per-token FFN：每个 token 有自己的一套 FFN 参数，而不是像 Transformer 那样所有位置共享
- 重排后的第 h 个 token 已不再对应原来的第 h 个 token，残差直接相加会造成语义错位；后续的 TokenMixer-Large 先把排列还原，再做残差相加

#TransformerEncoderLayer

对应 PyTorch nn.TransformerEncoderLayer：

# norm_first=False（默认，Post-LN，即原论文结构）
z = LayerNorm(x + Dropout(MHA(x)))
y = LayerNorm(z + Dropout(FFN(z)))

# norm_first=True（Pre-LN）
z = x + Dropout(MHA(LayerNorm(x)))
y = z + Dropout(FFN(LayerNorm(z)))

MHA: nn.MultiheadAttention，其 dropout 参数作用在注意力权重上；默认没有 CausalMask，需要通过 src_mask / is_causal 传入
FFN: Linear2(Dropout(ReLU(Linear1(x))))，默认 d_model=512、nhead=8、dim_feedforward=2048、dropout=0.1
batch_first=False 是默认值，此时输入形状为 (seq, batch, d_model)
Post-LN 的残差路径上有 LayerNorm，层数深时梯度不稳定，需要 warmup；Pre-LN 的残差路径是恒等映射，更容易训练，GPT-2 之后的 LLM 基本都采用 Pre-Norm

#TransformerDecoderLayer

对应 PyTorch nn.TransformerDecoderLayer：

# norm_first=False（默认，Post-LN）
z1 = LayerNorm(x  + Dropout(MHA(x)))           # Self-Attention，需要 CausalMask
z2 = LayerNorm(z1 + Dropout(MHA(z1, memory)))  # CrossAttention
y  = LayerNorm(z2 + Dropout(FFN(z2)))

# norm_first=True（Pre-LN）
z1 = x  + Dropout(MHA(LayerNorm(x)))
z2 = z1 + Dropout(MHA(LayerNorm(z1), memory))
y  = z2 + Dropout(FFN(LayerNorm(z2)))

memory: Encoder 的输出
Self-Attention 的 CausalMask 需要调用者传入 tgt_mask（可由 nn.Transformer.generate_square_subsequent_mask 生成）或设置 tgt_is_causal=True
三个子层分别对应 self_attn / multihead_attn / linear1+linear2，以及 norm1/2/3、dropout1/2/3

#BertLayer

1 2	z = LayerNorm(x + Dropout(MHA(x))) # BertAttention = BertSelfAttention + BertSelfOutput y = LayerNorm(z + Dropout(FFN(z))) # BertIntermediate + BertOutput

结构与 PyTorch TransformerEncoderLayer(norm_first=False) 相同，即 Post-LN
MHA 没有 CausalMask，只有 PaddingMask，所以是双向注意力；is_decoder=True 时才使用因果掩码
MHA 的 LinearOut 和 FFN 的 LinearOut 分别放在 BertSelfOutput.dense 和 BertOutput.dense 中，与残差相加、LayerNorm 一起完成
FFN: BertIntermediate（Linear + GELU）+ BertOutput（Linear）；BERT-base 为 d_model=768、head=12、d_ff=3072

#GPT2Layer

1 2	z = x + MHA(LayerNorm(x)) y = z + FFN(LayerNorm(z))

对应 HF GPT2Block，采用 Pre-LN
MHA 的 LinearOut 之后还有一个 resid_dropout，即 Dropout(LinearOut(...))
c_attn 一次性算出 QKV 再 split；GPT-2 的 Linear 层用的是 Conv1D，本质是权重转置存储的 Linear（weight 形状为 (in, out)）
FFN 为 GPT2MLP，激活函数为 gelu_new，d_ff=4×d_model

#LlamaLayer

1 2	z = x + LlamaAttention(RMSNorm(x)) y = z + GatedFFN(RMSNorm(z))

对应 HF LlamaDecoderLayer，两个 RMSNorm 分别为 input_layernorm 和 post_attention_layernorm
Pre-Norm、RMSNorm、RoPE、SwiGLU、无 bias、无 Dropout，是之后大多数开源 LLM 的基础结构

#Qwen3Layer

1 2	z = x + Qwen3Attention(RMSNorm(x)) y = z + GatedFFN(RMSNorm(z))

与 LlamaLayer 相同，只是 attention 加了 QK-Norm 和可选的滑动窗口

#MixtralLayer

1 2	z = x + LlamaAttention(RMSNorm(x)) # GQA: head=32, head_kv=8 y = z + MoE(RMSNorm(z)) # 8 选 2

把 LlamaLayer 的 GatedFFN 换成 MoE，其他不变

#DeepseekV3Layer

1 2	z = x + MLA(RMSNorm(x)) y = z + DeepSeekMoE(RMSNorm(z)) # layer_idx < 3 时为 GatedFFN(d_ff=18432)

first_k_dense_replace=3：前 3 层用 dense FFN，之后的层用 DeepSeekMoE

#Model

#Transformer

对应 PyTorch nn.Transformer，即原论文的 Encoder-Decoder 结构：

memory = LayerNorm(
  TransformerEncoderLayer^6(
    src                                     # as (src_seq, batch, d_model=512)
  )
)
y = LayerNorm(
  TransformerDecoderLayer^6(
    tgt,                                    # as (tgt_seq, batch, d_model)
    memory
  )
)                                           # as (tgt_seq, batch, d_model)

nn.Transformer 只包含 Encoder 与 Decoder 主体，不包含 Embedding、位置编码和输出层，需要调用者自己添加；原论文中为 Dropout(TokenEmbedding(x) * sqrt(d_model) + SinusoidalPE(pos))，输出层 Linear 与 Embedding 共享权重，最后接 Softmax
即使 norm_first=False，nn.Transformer 也会在 Encoder、Decoder 末尾各加一个 LayerNorm（encoder_norm / decoder_norm）
默认 batch_first=False，形状为 (seq, batch, d_model)

#BertModel

y = BertLayer^12(
  Dropout(
    LayerNorm(
      TokenEmbedding(x)          # as (batch, seq, d_model=768)
      + PositionEmbedding(pos)
      + TokenTypeEmbedding(seg)
    )
  )
)                                # as (batch, seq, d_model)
pooled = Tanh(LinearPool(y[:, 0]))  # as (batch, d_model)

三种 Embedding 对应 BertEmbeddings 中的 word_embeddings / position_embeddings / token_type_embeddings，位置编码是可学习的绝对位置
TokenTypeEmbedding: 区分句子 A / B（type_vocab_size=2）
由于是 Post-LN，最后一层输出已经过 LayerNorm，所以末尾没有额外的 LayerNorm
pooled: BertPooler，取第一个 token（[CLS]）经过 Linear + Tanh，用于句子级分类
MLM 头（BertLMPredictionHead）为 LinearDecoder(LayerNorm(GELU(Linear(y))))，LinearDecoder 与 TokenEmbedding 共享权重，vocab_size=30522

#GPT2Model

y = LMHead(
  LayerNorm(
    GPT2Layer^12(
      Dropout(
        TokenEmbedding(x)        # as (batch, seq, d_model=768)
        + PositionEmbedding(pos)
      )
    )
  )
)                                # as (batch, seq, vocab=50257)

TokenEmbedding / PositionEmbedding: wte / wpe，可学习的绝对位置编码，最大长度 n_positions=1024
由于是 Pre-LN，最后一层的输出没有归一化，所以需要末尾的 ln_f
LMHead: 属于 GPT2LMHeadModel，tie_word_embeddings=True，与 TokenEmbedding 共享权重
GPT-2 small / medium / large / xl 分别为 12 / 24 / 36 / 48 层

#LlamaModel

y = LMHead(
  RMSNorm(
    LlamaLayer^32(
      TokenEmbedding(x)          # as (batch, seq, d_model=4096)
    )
  )
)                                # as (batch, seq, vocab=32000)

输入处没有位置编码，也没有 Dropout；位置信息全靠每层 attention 中的 RoPE
rotary_emb 和 create_causal_mask 在 Model 层只算一次，然后传给每一层
LMHead: 属于 LlamaForCausalLM，默认 tie_word_embeddings=False，不与 TokenEmbedding 共享权重
Qwen3Model、MixtralModel、DeepseekV3Model 的顶层结构与此相同，只是把 LlamaLayer 换成对应的 Layer；DeepSeek-V3 为 61 层、d_model=7168，HF 实现中不包含 MTP（Multi-Token Prediction）模块

#总结对比

Model	Norm	位置	位置编码	Attention	FFN	激活
Transformer	LayerNorm	Post	Sinusoidal 绝对	MHA	FFN	ReLU
BERT	LayerNorm	Post	可学习绝对	MHA（双向）	FFN	GELU
GPT-2	LayerNorm	Pre	可学习绝对	MHA	FFN	GELU (tanh)
Llama	RMSNorm	Pre	RoPE	MHA / GQA	GatedFFN	SiLU
Qwen3	RMSNorm	Pre	RoPE	GQA + QK-Norm	GatedFFN / MoE	SiLU
Mixtral	RMSNorm	Pre	RoPE	GQA	MoE	SiLU
DeepSeek-V3	RMSNorm	Pre	解耦 RoPE	MLA	DeepSeekMoE	SiLU