《Born》第15章：Transformer 架构深度解析

Transformer 是现代 NLP 和视觉模型的基石。这一章拆解它的每一个组件，对照 Born 的源码实现。

Transformer Block 结构

Input
  ├──► RMSNorm ──► GQA ────────────► Add ──┐
  │                                          │
  └──────────────────────────────────────────┘
  ├──► RMSNorm ──► SwiGLU FFN ─────► Add ──┐
  │                                          │
  └──────────────────────────────────────────┘
Output

多头注意力（MHA）

mha := nn.NewMHA[float32](dModel, nHead, backend)

// Q, K, V 投影
Q := mha.Wq.Forward(x)  // [batch, seq, dModel]
K := mha.Wk.Forward(x)
V := mha.Wv.Forward(x)

// 分头
Q = Q.Reshape([batch, seq, nHead, headDim])

// 注意力得分
scores := Q.MatMul(K.T()) / sqrt(headDim)
scores = scores.Softmax(-1)
output := scores.MatMul(V)

RoPE：旋转位置编码

// RoPE 把位置信息编码进 Q/K 向量
func applyRoPE(x Tensor, pos int) Tensor {
    // x: [seq, headDim]
    // 对每个位置 m，应用旋转矩阵
    for i := 0; i < headDim/2; i++ {
        theta := math.Pow(10000, -2*float64(i)/float64(headDim))
        angle := float64(pos) * theta
        cosVal := math.Cos(angle)
        sinVal := math.Sin(angle)
        // 旋转二维子空间
        x1, x2 := x[i], x[i+headDim/2]
        x[i] = x1*cosVal - x2*sinVal
        x[i+headDim/2] = x1*sinVal + x2*cosVal
    }
}

RoPE 的优势：外推性好，长序列表现稳定。

GQA：分组查询注意力

// MHA: 每个头有自己的 K/V
// GQA: nHead 个 Q 共享 nKVHead 个 K/V
gqa := nn.NewGQA[float32](dModel, nHead, nKVHead, backend)
// nKVHead < nHead，减少 KV-Cache 内存

LLaMA 2/3 使用 GQA，内存占用减少 ~30%。

📘 《Born》连载技术书，第 15/22 章。