Pre-LN(层归一化在子层输入之前)
① LayerNorm①(子层前)
x_n=(x-μ)/√(σ²+ε)·γ+β
维度保持(B,N,D)
↓
② 多头自注意力(7步完整展开)
步骤1-整体QKV投影:
Q=X·W_Q,K=X·W_K,V=X·W_V
W_*∈ℝ^(D×D);Q/K/V∈(B,N,D)
步骤2-来源:自注意力
Q/K/V来自同层输入X∈(B,N,D)
步骤3-多头拆分:d_k=D/h
Q_i=Q[:,:,i*d_k:(i+1)*d_k]
Q_i/K_i/V_i∈(B,N,d_k)
reshape等价:
Q.reshape(B,N,h,d_k).transpose(1,2)
→(B,h,N,d_k)
步骤4-△Q·Kᵀ产生注意力分数
∈(B,h,N,N);每token对其他的关注度
V_i不参与分数计算
只被分数加权聚合,作为输出值
步骤5-单头计算:
head_i=Softmax(Q_i·K_iᵀ/√d_k)·V_i
∈(B,N,d_k);√d_k缩放防梯度消失
步骤6-多头拼接:
Concat(head_1..head_h)
→(B,N,h*d_k)=(B,N,D)
步骤7-输出投影:
·W_O+b_O,W_O∈ℝ^(D×D)
out:(B,N,D)
↓
③ 残差加法①
x←x+attn_out
维度保持(B,N,D)
↓
④ LayerNorm②(第二子层前)
x_n=(x-μ)/√(σ²+ε)·γ+β
↓
⑤ Linear升维
out=x·W+b,W∈ℝ^(D×4D)
out:(B,N,4D)
↓
⑥ GELU
GELU(x)=x·Φ(1.702x)
↓
⑦ Linear降维
W∈ℝ^(4D×D)
out:(B,N,D)
↓
⑧ 残差加法②
x←x+ffn_out
维度(B,N,D)
x = x + Attn(LN(x))
x = x + FFN(LN(x))
ViT-Base/Large
DeiT
Swin-T
MAE
BEiT
DINO
CaiT
GPT-2
LLaMA