前向传播

Transformer Block 前向传播

x
↓
LayerNorm
↓
Attention
↓
Residual Add
↓
LayerNorm
↓
FFN / MoE
↓
Residual Add
↓
输出

输入：

\[x \in \mathbb{R}^{\text{batch\_size} \times \text{seq\_len} \times \text{hidden\_dim}} \]

假设输入的 \(x\) 有 \(m\) 个文本，每条文本被分为 \(n\) 个 token，每个 token 有 \(p\) 个维度。

实际应用时，常通过补足 [PAD] 到该 batch 中最长文本、补足到固定长度，或者截断来确保 seq_len 一致。

---

LayerNorm

输入 \(x\) 时，先进行 LayerNorm 层归一化，对 hidden_dim 做层归一化，使其 token 数值标准化。

\[\mu = \operatorname{mean}(x_{\text{token}}) \]

\[\sigma^2 = \operatorname{variance}(x_{\text{token}}) \]

\[\hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \varepsilon}} \]

其中，\(\varepsilon\) 是一个很小的常数，用来防止分母为 \(0\)。

---

Attention

经过 LayerNorm 的 \(x\) 进行 Attention。在 Attention 内部，每个 token 将会融合上下文信息。

\(x\) 进行运算之后得到 \(Q\)、\(K\)、\(V\) 三个矩阵。其中：

• \(Q\) 为 Query，查询矩阵，代表当前 token 想要查询的信息；
• \(K\) 为 Key，键矩阵，代表其他 token 被打上的索引标签；
• \(V\) 为 Value，值矩阵，代表其他 token 携带的信息。

\[Q = xW_q \]

\[K = xW_k \]

\[V = xW_v \]

Attention 公式为：

\[\operatorname{Attention}(Q, K, V)= \operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} + \operatorname{mask} \right)V \]

---

Multi-Head Attention

在多头注意力架构中，每个 head 负责分析不同关系，如指代关系、语义关系、位置关系、时间关系、语法关系等等。

每个 head 有不同的 \(W_q\)、\(W_k\)、\(W_v\)。

\[\operatorname{MultiHeadAttention}(x)= \operatorname{Concat} \left( \operatorname{head}_0, \operatorname{head}_1, \operatorname{head}_2, \dots, \operatorname{head}_h \right)W_o \]

输出形状还是：

\[x.\operatorname{shape}=(\text{batch\_size}, \text{seq\_len}, \text{hidden\_dim}) \]

这样方便之后和原始输入 \(x\) 相加。

---

Residual Add

记：

\[\operatorname{attn_out}=\operatorname{Attention}(\operatorname{LayerNorm}(x)) \]

残差连接：

\[x = x + \operatorname{attn_out} \]

不使用：

\[x = \operatorname{attn_out} \]

是为了提高稳定性。

---

FFN / MoE

经过 Attention 之后，每个 token 已经拥有了上下文信息。这个时候进入 FFN 层，也就是 MLP。对于使用 MoE 的模型，则是进入 MoE。

注意，这里的 FFN / MoE 是对每个 token 单独起作用。

\[\operatorname{FFN}(x)= W_2 , \operatorname{activation}(W_1x) \]

结构是：

hidden_dim -> intermediate_dim -> hidden_dim

比如：

4096 -> 11006 -> 4096

记：

\[\operatorname{FFN_out}= \operatorname{FFN}(\operatorname{LayerNorm}(x)) \]

对于 MoE 架构，则是会路由到对应专家：

\[\operatorname{MoE}(x)= \sum_i \operatorname{score}_i \cdot \operatorname{expert}_i(x) + \operatorname{shared_expert}(x) \]

\[\operatorname{MoE_out} = x + \operatorname{MoE}(\operatorname{LayerNorm}(x)) \]