---

Transformer (主要是Decoder) 算量相关分析

---

全局结构概览

• Transformer 由若干 Encoder Block 与 Decoder Block 堆叠而成。
• Encoder-only（如 BERT）与 Decoder-only（如 GPT，会去掉 cross-attention）是两种常见变体。
• 目前主流的LLM基本都是基于自回归的Decoder-only架构。

---

参数分析

• 可训练/学习参数：在反向传播中会被计算梯度并更新的参数。

  1. 词嵌入与位置嵌入

     • 词嵌入矩阵 $E\in {\mathbb{R}}^{|\text{Vocab}|\times d_{\text{model}}}$：
       将词表索引映射到 $d_{\text{model}}$ 维的向量。
     • 位置嵌入矩阵 $P\in {\mathbb{R}}^{L_{\max }\times d_{\text{model}}}$：
       对最大序列长度 $L_{\max }$ 的每个位置学习一个位置向量。
     • 参数规模示例：若 $|\text{Vocab}|=32,000$、$d_{\text{model}}=512$，则词嵌入约有 $32,000\times 512\approx 16.4$M 个参数。

  2. 自注意力模块（Self-Attention）
     每层多头注意力（Multi-Head Attention, MHA）共有 $h$ 个头，每头维度 $d_k=d_{\text{model}}/h$。

     • 查询矩阵 $W_i^Q\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$

     • 键矩阵 $W_i^K\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$

     • 值矩阵 $W_i^V\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$
       共 $3h$ 个投影矩阵。

     • 输出投影 $W^O\in {\mathbb{R}}^{h{d}_k\times d_{\text{model}}}$

     • 偏置向量（若实现中包含偏置）

     • 示例参数量（单层）

       • $d_{\text{model}}=512,h=8,d_k=64$
       • 每头 $W_i^Q$ 有 $512\times 64=32,768$ 参数，3 矩阵共 $3\times 8\times 32,768=786,432$；
       • 输出 $W^O$ 有 $(8\times 64)\times 512=262,144$；
       • 合计约 1.05M 参数。

  3. 前馈网络（Feed-Forward Network, FFN）
     包含两层全连接：

     1. 第一层
        • 权重 $W_1\in {\mathbb{R}}^{d_{\text{model}}\times d_{\text{ff}}}$
        • 偏置 $b_1\in {\mathbb{R}}^{d_{\text{ff}}}$
     2. 第二层

        • 权重 $W_2\in {\mathbb{R}}^{d_{\text{ff}}\times d_{\text{model}}}$

        • 偏置 $b_2\in {\mathbb{R}}^{d_{\text{model}}}$

        • 通常 $d_{\text{ff}}=4\times d_{\text{model}}$。

     • 示例参数量（单层）
       若 $d_{\text{model}}=512,d_{\text{ff}}=2048$，则
       [
       512\times2048 + 2048 + 2048\times512 + 512 \approx 2.1\text{M 参数}
       ]

  4. 层归一化（LayerNorm）

     • 每个 LayerNorm 对特征维度 $d_{\text{model}}$ 有一对参数：

       • 缩放 $\gamma \in {\mathbb{R}}^{d_{\text{model}}}$
       • 偏移 $\beta \in {\mathbb{R}}^{d_{\text{model}}}$

     • 通常每层包含两次 LayerNorm（注意力后与 FFN 后），共 $2\times 2\times d_{\text{model}}$ 个参数。

  5. 输出投影（Un-Embedding）

     • 映射到词表：
       • 权重 $W_{\text{out}}\in {\mathbb{R}}^{d_{\text{model}}\times |\text{Vocab}|}$
       • 偏置 $b_{\text{out}}\in {\mathbb{R}}^{|\text{Vocab}|}$
     • 权重共享：多数实现中将 $W_{\text{out}}=E^T$，与词嵌入矩阵共享同一份参数，节省参数量。

• 共享参数：同一份权重在网络中被多次复用。

  1. 嵌入—输出层权重共享（Weight Tying）
     • 将输入词嵌入矩阵 $E$ 与输出投影矩阵 $W_{\text{out}}$ 设为同一份张量的转置。
     • 优势：提升性能并显著减少参数量（节省约 $|\text{Vocab}|\times d_{\text{model}}$ 个参数）。
  2. 跨层参数共享（Cross-Layer Sharing）
     • 在 ALBERT 等变体中，将所有或部分 Transformer 层（如注意力模块或 FFN 模块）复用同一套可学习参数。
     • 好处：在保持深度的同时，大幅减小模型总体参数量。
  3. Mixture-of-Experts 中的共享
     • 在 MoE（Mixture-of-Experts）模型里，专家网络的 FFN 权重各自独立，其他模块（如嵌入、注意力、LayerNorm）在所有专家间共享。

---

Encoder 与 Decoder共有结构

1. 位置编码（Positional Encoding）

• 将固定或可学习的位置向量与词向量相加，保留顺序信息。

2. 前馈网络（FFN）

• Feed-Forward Network（FFN） 是神经网络中最基本的模块之一，其特点是信息单向流动——从输入层经过若干全连接层（及激活函数）直接到输出层，中间不包含任何循环或反馈连接。
• 结构：两层全连接 + 激活（ReLU），中间维度 d_ffn = 4·d_model：

  FFN(x) = W₂(ReLU(W₁ x + b₁)) + b₂
  

• 输入/输出：
  (L, d_model) → (L, d_ffn) → (L, d_model)。

3. 残差连接 & 层归一化

• Post-LN（原始 Transformer）：子层输出 + 残差 → LayerNorm。
• Pre-LN（GPT-2/3 等）：LayerNorm → 子层 → 残差，更易深层训练。

---

Encoder

1. 多头自注意力（Multi-Head Self-Attention）

输入

• 序列长度 L，每个位置一个 d_model 维向量，整体形状 (L, d_model)。

输出

• 同样是 (L, d_model)。

计算

1. 首先通过三组线性映射，把输入投影为 Query (Q)、Key (K)、Value (V)，每个向量都是 d_model 维；
2. 在多头机制下，将 d_model 切分为 h 份，每份维度为
   [
   d_k = \frac{d_\text{model}}{h} \quad \text{(每个头的维度)}
   ]
   得到 h 组 $(Q_i,K_i,V_i)$；
3. 对每个头并行计算：
   • 打分：$Q_i{K}_i^{\top }/\sqrt{d_k}$
   • 归一化：对每行做 Softmax，得到注意力权重
   • 加权求和：将注意力权重与 $V_i$ 相乘并求和，获得该头的输出
4. 拼接与投影：

   • 将 h 个头的输出在特征维度拼回 (L, d_model)，并通过线性投影得到最终结果。

   • 在多头自注意力的计算过程中，每个头的输出维度为 $d_k$，每个头并行计算得到的输出会被拼接在一起，形成一个形状为 $(L,d_{\text{model}})$ 的矩阵。拼接操作的时间复杂度是 $O(L\cdot d_{\text{model}})$，这部分相对较轻。接下来，拼接后的结果将通过一个线性投影进行变换：

     [
     Y = X W^O + b^O,
     ]
     其中：

     • $X$ 是拼接后的输入矩阵，形状为 $(L,d_{\text{model}})$；
     • $W^O$ 是线性投影的权重矩阵，形状为 $(d_{\text{model}},d_{\text{model}})$；
     • $b^O$ 是偏置向量，形状为 $(d_{\text{model}})$。

     进行矩阵乘法 $X{W}^O$ 时的计算量为：

     • 对于每一行（共 $L$ 行），计算一个 $d_{\text{model}}$ 维的向量与一个 $d_{\text{model}}\times d_{\text{model}}$ 的矩阵的乘法，运算量为 $O(d_{\text{model}}^2)$；
     • 因此，总的计算复杂度为：
       [
       O(L \cdot d_{\text{model}}^2)
       ]

     综上，拼接和线性投影的整体复杂度为 $O(L\cdot d_{\text{model}}^2)$，其中线性投影的部分主导了计算过程。

• 缩放因子 $\frac{1}{\sqrt{d_k}}$ 的目的：
  1. 数值稳定：计算自注意力时，我们计算的是 Query $Q$ 和 Key $K$ 的点积 $Q{K}^{\top }$。如果 $d_k$ 很大，那么 $Q{K}^{\top }$的值也可能非常大。控制打分方差在 $O(1)$ 级别，避免 Softmax 输入过大导致梯度消失或爆炸。
  2. 统一尺度：确保不同头之间打分尺度一致，便于拼接和后续处理。
• 复杂度推导：
  • 每个头的打分计算为 $O(L^2\cdot d_k)$；
  • h 个头并行共 $O(L^2\cdot d_k\cdot h)=O(L^2\cdot d_{\text{model}})$；
  • 加上最后的拼接与投影 $O(L\cdot d_{\text{model}}^2)$。
  • 当序列长度 L 远大于 d_model 时，注意力部分主导，总体约为 $O(L^2\cdot d_{\text{model}})$。

---

Decoder

1. Masked Self-Attention

• 类似于 Self-Attention，但对未来时刻的部分进行掩码（$-\infty$），确保位置 $t$ 仅能关注 $t$ 之前的内容。

具体流程

1. 计算 QKᵀ 得分矩阵；
2. 对上三角（未来位置）加入掩码（值 = $-\infty$）；
3. 再做 Softmax，得到只关注过去和当前位置的注意力权重；
4. 用该权重乘以 V 并求和，得到输出。

理解掩码的妙处

• 虽然语义上我们要“看不到”未来 token，但在实现上，所有位置的 attention 计算流程是完全一致、可并行的——掩码只是通过矩阵运算一次性“把”对未来位置的注意力权重置为 $-\infty$，让 softmax 输出时对应权重为零。所以Decoder在训练时，每个token计算量是一致的。但是在推理时，由于自回归的限制，掩码名存实亡，除了初始输入也就是生成首token的时候，掩码真切地发挥了作用，后续的自回归过程，每个token都确实只计算和前文的注意力，造成每个token计算量的差异。严格来说，掩码在推理阶段并未完全失效，而是通过自回归的序列生成方式隐式替代了显式掩码机制

2. Cross-Attention（Encoder–Decoder Attention）

• Decoder 的 Query 来自前一子层的输出，Key/Value 来自 Encoder 的输出。Decoder 在生成时，不仅自回看自身序列，还要“交叉”关注 Encoder 的输出，上下文对齐。

计算流程

1. Query 来自 Decoder 子层输出；Key/Value 来自 Encoder 最终输出；
2. 计算注意力：$\text{softmax}(Q_{\mathrm{d}\mathrm{e}\mathrm{c}}{K}_{\mathrm{e}\mathrm{n}\mathrm{c}}^{\top }/\sqrt{d_k})V_{\mathrm{e}\mathrm{n}\mathrm{c}}$；
3. 无掩码——因为 Encoder 可见全序列；（Query 数量 = Decoder 端的序列长度，Key/Value 数量 = Encoder 端的序列长度。）
4. 结果与残差相加并归一化。

3. 多层堆叠

(Embed + PosEncode)
    ↓
[Masked Self-Attention → Add&Norm]
    ↓
[Cross-Attention     → Add&Norm]
    ↓
[FFN                 → Add&Norm]
    ↓
(重复 N 层 Decoder Block)
    ↓
[线性 + Softmax 输出概率分布]

• 从 (L×d_model) 到 (L×V) 的全流程
  1. Embedding：词 ID → d_model 维向量，加位置编码 → (L, d_model)。
  2. Decoder Stack：N 层 Decoder Block，始终 (L, d_model)。
  3. 输出投影：线性层映射 (L, d_model) → (L, V)，V 为词表大小。
  4. Softmax：对每行做 Softmax，得到 (L, V) 的概率矩阵。
     • 第 i 行是位置 i 的“下一个 token”预测分布。
     • 在自回归推理中，只取最后一行（第 L 行）来生成第 L+1 个 token，再将其拼入序列继续推理。
  5. 权重共享：Embedding 矩阵与输出投影矩阵转置共享，节省参数并提高泛化。
• 堆叠层输出差异
  • 中间层：每层输出为 (L, d_model)，用于下一层继续上下文融合。
  • 最后一层：同维度表示，经输出投影后才变预测概率分布。
• 多层堆叠 Transformer 中，Token Embedding + Positional Embedding 只在模型的最底部做一次（也就是对原始输入序列进行一次），不会随每一层都重复。

4. 训推计算对比

• 自回归推理时，一般会Cache一些中间结果，用于加速推理，即KV Cache。Cache的原因是，自回归时，未来的token是真的没来，不像训练时用掩码遮蔽的情况。
  • 推理时的缓存（KV Cache）

    • 在自回归生成（auto-regressive）场景下，为了避免对已经生成的前缀重复计算“键”“值”：

      1. 第一次推理：对输入序列每层分别计算 Q, K, V；

      2. 缓存 K/V：把每层的 K 和 V 存入缓存；

      3. 下一个 token：只需计算新的 Q，以及把旧的 K/V 拼接到新的 K/V 后面，并做注意力计算。

    • 这样做能将时间复杂度从 O(L²) 降到 O(L)（L 为已生成长度），极大加速生成。

    • 通常只缓存 K/V，而不会缓存中间的 Query-Key 点积或 Attention 输出；中间的激活（hidden states）也不全量保留，只保留最后一层用于下一次输出前的输入。

• 训练或者处理初始输入时，上面的中间结果直接在并行计算时被利用，谈不上Cache。

5. Decoder 单层每token计算量粗略估算

Embedding

• 每个 token 都要做一次嵌入（token embedding + position embedding）, 每个token $O(\cdot d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}})$
• token embedding：在实现里，大部分深度学习框架（PyTorch、TensorFlow）都会用跳转地址的方式（维护一个嵌入矩阵，查表）：给定一个 token id，立刻算出其embedding在内存中的起始地址，然后拷贝连续 $d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}$ 个数的出来，所以是 $O(\cdot d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}})$。
• position embedding：位置嵌入有多种计算方式，比如可以做成一个可学习的查表矩阵，总之基本都是$O(\cdot d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}})$。
• 加法开销：token embedding 和 position embedding相加，需要$O(\cdot d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}})$ 次加法。

QKV计算

• 每个token的嵌入都要计算QKV，固定开销 $O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}^2)$（QKV 权重矩阵维度为 [d_model, d_model]，且每个注意力头的维度为 $d_k=d_{\text{model}}/h$）

$Q{K}^{\top }$与加权和

假设当前是第t个token

1. $Q{K}^{\top }$:

   • 需计算与 ​前 $t$ 个 Key​ 的点积：
   • 单头计算量：$t\times 2d_k$（每个 Key 对应 $2d_k$ 次运算，更准确说是$2d_k-1$）
   • 总计算量（h个头）：$h\times t\times 2d_k=2t{d}_{\text{model}}$

2. 缩放、掩码与 Softmax​

   1. ​缩放​：对 ​前 $t$ 个 Key​ 的注意力分数除以 $\sqrt{d_k}$，运算量为 $t$ 次除法。
   2. ​掩码​：通过矩阵加法一次性处理未来位置，运算量可忽略。
   3. ​Softmax​：
      • 指数运算：$t$ 次
      • 分母求和：$t-1$ 次加法
      • 归一化：$t$ 次除法

   ​单个 Token 合计​：
   $$t(\text{缩放})+t(\text{指数})+(t-1)(\text{求和})+t(\text{归一化})=4t-1\approx 4t$$

3. ​加权（α×V）​​：

   • 需加权 ​前 $t$ 个 Value​：
   • 单头计算量：$t\times d_k$（乘法）
   • 总计算量：$h\times t\times d_k=t{d}_{\text{model}}$

4. ​加权结果求和​：

   • 单头计算量：$t\times d_k$（加法）
   • 总计算量：$h\times t\times d_k=t{d}_{\text{model}}$

​单个 Token 合计​：
$$2t{d}_{\text{model}}+4t+t{d}_{\text{model}}+t{d}_{\text{model}}=4t{d}_{\text{model}}+4t$$

Multi‐Head Attention 输出拼接后的线性映射

• 在计算完所有头的加权和并将 $h$ 个头的输出拼接成 ${\mathbb{R}}^{d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}}$ 之后，还要再做一次线性映射，把它投回到 $d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}$ 维空间：
  [
  \text{Concat}([\text{head}1, \dots, \text{head}h]);\bigl(\in \mathbb R^{d{\rm model}}\bigr);\times W_O\in \mathbb R^{d{\rm model}\times d_{\rm model}}.
  ]
• 复杂度：每个 token $O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}^2)$。

Encoder–Decoder Cross‐Attention（仅对完整 Decoder 而言）

• 如果实现的是标准的 “Self‐Attention → Cross‐Attention → FFN” 结构，那么在自注意力之后，还会有一段用 Decoder 隐藏态去 attend Encoder 输出的交叉注意力：
• QKV 投影：与自注意力完全类比，Decoder 的 hidden 做一次 $O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}^2)$ 的 Q 投影，Encoder 的输出做一次 $O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}^2)$ 的 K、V 投影。
• 点积与加权和：对第 $t$ 个 token，与 Encoder 序列长度 $m$ 做点积与加权求和，总量级是 $O(m{d}_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}})$。
• 输出线性映射：同样还要一个 $O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}^2)$ 的 $W_O$ 映射。

LayerNorm 与残差连接

• 几乎每个子层（Self‐Attn、Cross‐Attn、FFN）前后都会插入 LayerNorm 和残差（元素级加法）：
• LayerNorm：对 $d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}$ 维向量做均值／方差计算再归一化，近似 $O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}})$。
• 残差加法：把子层输入和输出相加，也是 $O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}})$。

Feed-Forward Network (FFN)

• 两层全连接，开销 $O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}\cdot d_{\mathrm{f}\mathrm{f}\mathrm{n}})=O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}^2)$（因为 $d_{\mathrm{f}\mathrm{f}\mathrm{n}}=4d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}$）。

小结

组成	复杂度（每 token）
Embedding	$O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}})$
Self-Attention QKV	$O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}^2)$
Self-Attention 点积等	$O(t{d}_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}})$
Self-Attention 输出映射	$O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}^2)$
Cross-Attention QKV	$O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}^2)$
Cross-Attention 点积	$O(m{d}_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}})$
Cross-Attention 输出映射	$O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}^2)$
FFN	$O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}^2)$
LayerNorm × 3 + Residual × 3	$O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}})$ each ≈ $O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}})$ total

其中，$t$ 是当前已生成长度，$m$ 是 Encoder 序列长度。最重的几项依然是那些 $O(d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}^2)$ 和 $O(\max (t,m)d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}})$ 项。

核心结论

1. 实际上由于掩码机制，在训练阶段，掩码注意力层的相关计算的复杂度和t是不相关的，而是和序列长度L相关。因为掩码只是模拟看不到后面的token，实际上还是能看到的，只是通过把注意力权重搞成0，让其不起作用。这样每个token的计算行为是相同的，有利于并行化。
2. 而在基于自回归的推理场景中，除了初始输入序列的计算，后续基于自回归生成的过程中，注意力相关计算的复杂度的确是和token的位置t相关的，两个连续的由自回归生成的token，后者比前者多出的运算数，如果只看单层Decoder的话，大概是 $4d_{\text{model}}+4$ 。

---

相关图示

• 用GPT找的，忘了让标注出处了。

Transformer架构

完整Transformer架构与位置编码

Decoder-Only示意

多层Decoder结构

Decoder的输入输出与Softmax

多头注意力机制架构

Softmax前后的掩码示例

---

Referencess

1. https://papers.neurips.cc/paper/7181-attention-is-all-you-need.pdf
2. https://en.wikipedia.org/wiki/Transformer_%28deep_learning_architecture%29
3. https://medium.com/%40hunter-j-phillips/position-wise-feed-forward-network-ffn-d4cc9e997b4c
4. https://ifwind.github.io/2021/08/17/Transformer%E7%9B%B8%E5%85%B3%E2%80%94%E2%80%94%EF%BC%887%EF%BC%89Mask%E6%9C%BA%E5%88%B6/#self-attention%E4%B8%AD%E7%9A%84padding-mask

---