现代大模型笔记 01：Transformer 架构与现代改进

本篇整理Transformer架构以及它在现代 LLM 里发生了哪些变化。从图中可以看到，几个核心的步骤分别是： 1. word embedding 2. positional embedding 3. 注意力机制 4. Residual Connect 5. Normalization 6. FFN

原始 Transformer的encoder-decoder结构

原始 Transformer 是 seq2seq 架构。

Encoder 负责读完整个输入句子输出一个隐藏向量给decoder，decoder负责一步一步生成目标句子。

Encoder里是bidirectional self-attention，每个 token 可以看完整输入。Decoder 里是 causal self-attention，加上cross-attention 去看encoder输出。

那为什么现代 LLM 多用 decoder-only

GPT、LLaMA、Qwen、DeepSeek 这类通用语言模型，大多采用 decoder-only。

它没有单独的 encoder，而是把任务统一成 next token prediction：

\[ p(x_t \mid x_1, x_2, ..., x_{t-1}) \]

主要是形式简单，而且和大规模自监督预训练天然对齐。

Transformer block：attention + FFN + 残差

一个现代 decoder-only Transformer block 大概长这样：

x
-> Norm
-> Self-Attention
-> Residual Add
-> Norm
-> FFN / MLP
-> Residual Add

Self-attention 负责 token 之间的信息交互。

FFN 负责对每个 token 的表示做非线性加工。

Residual connection 让深层网络更容易训练：

\[ x_{l+1} = x_l + F(x_l) \]

如果没有 residual，几十层、上百层模型很难稳定训练。Residual path 可以理解成一条信息和梯度的主干通路，每个 block 在上面叠加修改。

注意力机制

缩放点积注意力完整计算公式： \[ \text{Attention}(Q,K,V)=\text{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right)V \]

QKV 矩阵维度完整流转过程以自注意力机制为例，先统一符号定义： - \(seq\_len\)：输入序列长度（token 个数） - \(d_{model}\)：模型整体隐藏层维度，单个token特征向量长度 - \(d_k\)：单注意力头下查询、键向量的维度 - \(d_v\)：单注意力头下数值向量维度，工程实现中默认 \(d_v=d_k\)

一般认为\(d_k = d_{model} /\) head_num

设定参数：\(seq\_len=10,\ d_{model}=512,\ d_k=64\)（对应8头注意力，均分维度）
线性投影生成 Q、K、V 原始输入特征矩阵 \(X\) 形状：\((seq\_len,\ d_{model})\) 通过三组独立线性权重矩阵 \(W_q、W_k、W_v\) 做线性映射：权重矩阵尺寸：\(W_q,W_k,W_v \in \mathbb{R}^{d_{model} \times d_k}\) \[ \begin{align*} Q &= XW_q \quad \text{尺寸：}(10,\ 64) \\ K &= XW_k \quad \text{尺寸：}(10,\ 64) \\ V &= XW_v \quad \text{尺寸：}(10,\ 64) \end{align*} \] 作用：将通用token特征，映射为专门用于检索匹配、权重计算、结果输出的三类特征向量。
点积计算相似度得分 \(QK^T\) 对键矩阵转置后与查询矩阵相乘： \(K^T\) 尺寸：\((d_k,\ seq\_len)\) 矩阵相乘维度推演： \[ (seq\_len,\ d_k) \times (d_k,\ seq\_len) = (seq\_len,\ seq\_len) \] 最终得分矩阵：\(QK^T \in \mathbb{R}^{10 \times 10}\) 矩阵元素含义：第 \(i\) 行第 \(j\) 列代表第 \(i\) 个token（查询）与第 \(j\) 个token（键）的匹配相似度。
缩放：除以 \(\sqrt{d_k}\) 对得分矩阵全体元素做标量除法，矩阵形状保持不变： \[ \frac{QK^T}{\sqrt{d_k}} \in \mathbb{R}^{10 \times 10} \]
逐行 Softmax 归一化沿着序列维度逐行执行Softmax运算，每行所有数值累加和为1，转化为注意力权重分布；矩阵尺寸依旧不变： \[ \text{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) \in \mathbb{R}^{10 \times 10} \]
加权求和得到注意力输出注意力权重矩阵与数值矩阵 \(V\) 相乘： \[ (seq\_len,\ seq\_len) \times (seq\_len,\ d_v) = (seq\_len,\ d_v) \] 本例输出尺寸：\((10,\ 64)\)

引入缩放系数 \(\sqrt{d_k}\) 的底层原因

对点积结果做方差归一化，抑制Softmax进入饱和区间，防止梯度消失，保障模型训练平稳收敛，这也是「缩放点积注意力」命名的由来。

数学推导：点积方差随 \(d_k\) 线性上涨

合理初始化条件下：\(Q、K\) 内各个分量均服从均值为0、方差为1的独立分布。两个 \(d_k\) 维向量 \(q、k\) 的点积： \[ q\cdot k = \sum_{i=1}^{d_k} q_i k_i \] 方差运算性质：独立变量乘积方差 \(\text{Var}(q_ik_i)=\text{Var}(q_i)\cdot\text{Var}(k_i)=1\times1=1\)

方差有一个核心性质：独立随机变量之和的方差 = 方差之和

累加求和后总方差： \[ \text{Var}(q\cdot k) = \sum_{i=1}^{d_k}\text{Var}(q_ik_i) = d_k \]

由于\(\text{Var}(q_i\cdot k_i) = 1\), 然后有dk个1相加，所以总的方差=dk，点积标准差 \(=\sqrt{d_k}\) 向量维度 \(d_k\) 越大，点积数值波动区间越大，向量间得分差距会被持续拉大。

Softmax 饱和问题

Softmax 函数的作用是将一组任意实数（logits）映射为 0~1 之间、和为 1 的概率分布，公式为： \[ \sigma(z)_i = \frac{e^{z_i}}{\sum_{j=1}^K e^{z_j}} \]

饱和区指的是：当输入的各个 logit 之间相对差值极大时，Softmax 的输出会极度趋近于 0 或 1，此时函数曲线进入非常平坦的区域，函数自身的导数趋近于 0，梯度几乎无法有效回传。

简单理解：模型对预测结果“过度自信”，最大类别的概率接近 100%，其他类别接近 0，函数进入“几乎不再变化”的饱和状态。

二、从导数看饱和的本质

Softmax 的偏导数分两种情况： \[ \begin{aligned} 1.\ & \text{对自身输入求导}(i=j)：\frac{\partial \sigma_i}{\partial z_i} = \sigma_i(1-\sigma_i) \\ 2.\ & \text{对其他输入求导}(i\neq j)：\frac{\partial \sigma_i}{\partial z_j} = -\sigma_i\sigma_j \end{aligned} \]

饱和区的导数特征 - 当 \(\sigma_i \to 1\)（某类概率接近 1）：\(\sigma_i(1-\sigma_i) \to 0\) - 当 \(\sigma_i \to 0\)（某类概率接近 0）：\(\sigma_i(1-\sigma_i) \to 0\)

无论输出趋近 0 还是 1，Softmax 自身的导数都会无限趋近于 0，这就是“饱和”的数学本质——函数对输入的变化不再敏感，输入再怎么变，输出都几乎不动。

以二分类为例： | 输入 z = [z₁, z₂] | 输出 σ₁ | 输出 σ₂ | 导数 dσ₁/dz₁ | 饱和程度 | | :— | :— | :— | :— | :— | | [2, 0] | 0.88 | 0.12 | 0.106 | 不饱和 | | [5, 0] | 0.993 | 0.007 | 0.007 | 轻度饱和 | | [10, 0] | 0.99995 | 0.00005 | ≈ 0.00005 | 深度饱和 |

差值越大，输出越靠近 0/1，导数越接近 0，饱和越严重。

Normalization：深层模型更稳

原始 Transformer 使用 Post-Norm，大概形式是：

\[ x_{l+1} = \mathrm{LayerNorm}(x_l + F(x_l)) \]

现代 LLM 更常见的是 Pre-Norm：

\[ x_{l+1} = x_l + F(\mathrm{Norm}(x_l)) \]

差别是 normalization 放在 residual branch 前面还是后面。

Pre-Norm 的好处是深层模型训练更稳定。直觉上，residual path 更像一条比较干净的通路，梯度可以更直接地传，不会每层都先被 normalization 包住。

所以在现代 LLM 里，Pre-Norm 基本成了主流选择。

RMSNorm：保留尺度控制，少做一点计算

LayerNorm 会减均值、除标准差：

\[ \mathrm{LayerNorm}(x) = \frac{x-\mu}{\sqrt{\sigma^2+\epsilon}} \odot g + b \]

RMSNorm 去掉了减均值，只保留 root mean square 的尺度归一化：

\[ \mathrm{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{d}\sum_i x_i^2+\epsilon}} \odot g \]

它关心的是向量整体尺度，而不是每一维相对均值的位置。

这有两个好处。

第一，计算更简单。单层省一点，放到几十层、几百亿参数规模上就很明显。

第二，它保留了最关键的稳定性作用：控制 activation 的尺度。

FFN：从ReLU 到 SwiGLU，再到 MoE

Transformer block 里除了 attention，还有FFN。

原始 Transformer 用的是ReLu作为激活函数，后续有相应的改进，比如Leaky ReLU，GeLU之类的，目的是让更多的残差（信息）通过梯度来更新参数。 \[ \mathrm{FFN}(x)=W_2\sigma(W_1x+b_1)+b_2 \]

随着技术进步，SwiGLU进入了视野： SwiGLU主要有两个组成部分，Swish 和GLU。

\[ \mathrm{SwiGLU}(x) = \mathrm{Swish}(xW_1) \odot (xW_2) \]

再接一个输出投影。

Gating 的直觉是：模型不只是对特征做非线性变换，还学会哪些通道该打开、哪些通道该压下去。如果把 FFN 看成每个 token 独立做一次特征加工，那么 SwiGLU 相当于给这次加工加了一个动态开关。

MoE 可以看成 FFN 的进一步扩展：不是每个 token 都走同一个大 MLP，而是路由到少数 experts。这样可以增加总参数量，但每个 token 的实际计算量不一定同比增加。后续文章中会文门提到

RoPe 位置编码

Self-attention 本身不包含顺序信息。对 attention 来说，如果不加位置编码，A B C 和 C B A 很难区分谁在前谁在后。

原始 Transformer用absolute positional encoding，直接把位置向量加到 token embedding 上。

现代 LLM 常用 RoPE，也就是 rotary position embedding。最早来自于苏剑林的Reformer工作。RoPe解读博客链接

它不是简单把位置向量加进去，而是在 query 和 key 上做旋转，使 attention score 自然带上相对位置信息。而且这种方式天然适合长度外度，可以通过一些方式让模型理解超出训练长度的文本

可以粗略理解成：

\[ q_m^\top k_n \rightarrow \text{content similarity with relative position } (m-n) \]

注意力机制的变迁

标准 multi-head attention 里，每个 head 都有自己的 \(Q, K, V\)。

这在训练和表达力上很自然，但推理时会遇到 KV cache 问题。

生成第 \(t\) 个 token 时，前面所有 token 的 key 和 value 都要保留下来。context 越长、层数越多，KV cache 越大。

所以现代模型经常改 attention 的 K/V 结构。

MHA：每个 query head 都有自己的 K/V，表达力强，但 cache 大
MQA：所有 query heads 共享一组 K/V，cache 小，但表达力可能受影响
GQA：一组 query heads 共享一组 K/V，在效果和成本之间折中
MLA：把 K/V 压到 latent 表示里缓存，需要时再恢复，进一步降低 KV cache

下面给出一个最小的 MHA 实现。代码同时覆盖 prefill 和带 KV cache 的 decode；True 表示对应 key 可以被当前 query 看见。

import torch
import torch.nn as nn

class MHA(nn.Module):
    def __init__(self, hidden_size, num_heads):
        super().__init__()
        assert hidden_size % num_heads == 0

        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads
        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)

    def split_heads(self, x):
        # [B, L, D] -> [B, H, L, Dh]
        batch, seq_len, _ = x.shape
        x = x.view(batch, seq_len, self.num_heads, self.head_dim)
        return x.transpose(1, 2)

    def forward(self, x, allowed_mask=None, past_kv=None, use_cache=False):
        batch, query_len, hidden_size = x.shape

        q = self.split_heads(self.q_proj(x))
        k = self.split_heads(self.k_proj(x))
        v = self.split_heads(self.v_proj(x))

        past_len = 0
        if past_kv is not None:
            past_k, past_v = past_kv
            past_len = past_k.size(2)
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)

        key_len = k.size(2)
        scores = q @ k.transpose(-1, -2) / (self.head_dim ** 0.5)
        # scores: [B, H, query_len, key_len]

        if allowed_mask is None:
            # decode 时 query 的绝对位置要加上 past_len
            query_position = torch.arange(query_len, device=x.device) + past_len
            key_position = torch.arange(key_len, device=x.device)
            allowed_mask = key_position[None, :] <= query_position[:, None]
            allowed_mask = allowed_mask[None, None, :, :]

        scores = scores.masked_fill(~allowed_mask, float("-inf"))
        attention_prob = torch.softmax(scores, dim=-1)
        context = attention_prob @ v
        # context: [B, H, query_len, Dh]

        context = context.transpose(1, 2).contiguous()
        context = context.view(batch, query_len, hidden_size)
        output = self.o_proj(context)

        new_kv = (k, v) if use_cache else None
        return output, new_kv

MHA 只负责 token 间的信息交互。把它放进 Pre-Norm decoder block 时，residual、normalization 和 FFN 应该在 MHA 外部组织：

class DecoderBlock(nn.Module):
    def __init__(self, hidden_size, num_heads, ffn_size):
        super().__init__()
        self.norm_1 = nn.LayerNorm(hidden_size)
        self.attention = MHA(hidden_size, num_heads)
        self.norm_2 = nn.LayerNorm(hidden_size)
        self.ffn = nn.Sequential(
            nn.Linear(hidden_size, ffn_size),
            nn.GELU(),
            nn.Linear(ffn_size, hidden_size),
        )

    def forward(self, x, allowed_mask=None, past_kv=None, use_cache=False):
        attention_output, new_kv = self.attention(
            self.norm_1(x),
            allowed_mask=allowed_mask,
            past_kv=past_kv,
            use_cache=use_cache,
        )
        x = x + attention_output
        x = x + self.ffn(self.norm_2(x))
        return x, new_kv

这里的两条 residual path 都保持原始 \(x\) 直接通过，Norm 则放在 attention / FFN 之前，对应前面介绍的 Pre-Norm。

参考资料

Vaswani et al., Attention Is All You Need
Shazeer, GLU Variants Improve Transformer
Zhang and Sennrich, Root Mean Square Layer Normalization
Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding
Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
Qwen Team, Qwen3 Technical Report
DeepSeek-AI, DeepSeek-V3 Technical Report
苏剑林的科学空间，RoPe解读博客链接

现代大模型与前沿论文笔记

#笔记 #大模型 #Transformer #RoPE #KV cache

现代大模型笔记 03：LLM 推理与效率优化上一篇

深度学习笔记 05：注意力机制下一篇