Transformer内部探秘：一个Token的生命周期 - Aleksa Gordić

Transformer 前向传播

作为运行示例，假设我们从数据集中采样了 2 个“文档”，其中：

• batch size = 1
• 序列长度 = 16
• 启用了文档打包

我们将追踪一个 token 如何流经 transformer，并在此过程中逐一解析每个组件。

让我们开始吧。花点时间分析以下内容：

图 1：分词阶段

我们将文档分词为整数序列，然后将两个文档打包成一个序列。

在本博客的范围内，分词器是一个黑盒组件，它接收文本并将其映射为 token 序列，每个 token 由一个整数 ID 表示。在实践中，分词器是使用 BPE 等算法在单独的文本语料库上“训练”出来的，这些算法通过反复合并最常见的字符或字节序列来学习词汇表。良好的分词器设计具有多种理想特性；例如，将数字表示为单独的 token 有助于数值推理。

与 token 一起，我们构建了两个辅助结构：

• 输入位置 - 用于位置嵌入模块 (YaRN)
• 分段掩码 - 用于注意力机制中的掩蔽

这是预处理阶段。

📝附注：

为了提高效率，数据会在训练开始前预先分块，数据加载器会将这些预处理后的结构直接送入训练循环。此时，我们永远不会处理原始字符串。（Spark）数据管道和数据加载器本身很容易各自成为一篇单独的博客文章。

接下来，我们使用输入 token 来索引嵌入表。

你可以将嵌入表视为 LLM 的词汇表。

这个索引操作将我们的整数序列转换为一个包含 16 个 4096 维的 bf16 向量的序列：

图 2：嵌入阶段

📝附注：

特殊 token 在分词过程中不会自然出现——没有任何文本映射到 token ID >= 128,000。它们是在训练期间注入（并在推理时使用）以提高性能（例如 FIM、仓库打包等）或强制执行特定行为（例如生成结束/轮次结束、工具调用）。

让我们深入探讨 FIM [3]（填空）特殊 token。

在（预）训练期间，我们将一个文档取出，将其拆分为前缀、中间（中缀）和后缀，并构造一个如下形式的序列：<FIM_PRE> 前缀 <FIM_SUF> 后缀 <FIM_MID> 中间。模型被训练为根据前缀和后缀预测中间部分。然后可以在推理时利用此能力。

例如，想象一下在你最喜欢的 IDE 中使用 Rnj 1.5 作为自动补全模型。你的光标自然地将代码拆分为前缀和后缀，中间部分缺失。通过插入 FIM token 并以 <FIM_MID> 结尾，你提示模型为缺口生成补全。这些 token 有助于向模型传达意图。

分词器本身可以成为一篇独立的博客文章，所以我就在这里打住。

现在我们准备进入第一个 transformer 层。

请注意，所有 transformer 层都具有（几乎）相同的结构，所以我将只解释一个。在实践中，我们需经过 32 个这样的层——你可以把它想象成一个 for 循环，但在 Rnj 1.5 中，每一层都有自己的可学习权重。

“几乎”是因为 Rnj-1.5 同时使用了块局部注意力和全局注意力层——唯一的区别在于掩码。在更高的抽象层次上，这个说法仍然成立。更多细节请参阅注意力部分。

另外请注意，一些 transformer 实现在层之间进行权重共享或部分权重共享（有很多变体），但在这里我们专注于 Rnj 1.5。

让我们做一次通过 transformer 块的前向传播。仔细分析以下内容：

图 3：通过 transformer 块的前向传播

在高层，该块由四个 RMSNorm 子模块、一个 MLP、一个注意力模块、两个残差连接和两个求和操作组成。残差连接只是将来自块早期部分的向量副本向前传递。

重要的是，除了注意力之外，所有子模块都操作在单个向量上。

💡额外背景：

在实践中，你会找到 transformer 块的多种变体。设计选择包括归一化层的位置、类型和数量，确切的 MLP 结构（门控 vs. 非门控、门控函数的选择等），残差连接结构（恒等映射、注意力残差 [4] 等），尤其是注意力模块。

广义上讲，注意力机制是序列长度的二次方（例如 MLA [5]、缩放点积注意力等）或线性（例如 Kimi Linear [6]）关系，每种方法都在建模能力（尤其是在长上下文下）和效率之间进行权衡。

一旦向量离开最终的 transformer 块，它们通过矩阵乘法被投影到一个 128,256 维的空间中。这会产生 logits，通过 softmax 转换为概率分布。我们在推理时从中采样，并在训练时用于交叉熵损失。

图 4：

接下来，让我们深入了解各个子层。这次我将按相反顺序进行，这恰好使我们从最简单到最复杂：

1. RMSNorm（均方根层归一化）
2. GeGLU MLP（多层感知器）
3. 注意力（缩放点积注意力）

RMSNorm（均方根层归一化）

RMSNorm [7] 是一种用于稳定深度神经网络训练的归一化技术。

如前所述，RMSNorm 在单个向量上操作，因此我们将专注于一个单一的 bf16 4096 维向量（所有其他向量都以相同方式并行处理）。输出具有相同的形状和数据类型：

图 5：RMSNorm

GeGLU MLP（多层感知器）

MLP 是一个简单的、逐点的前馈神经网络，用于学习输入和输出向量之间的非线性关系。

我们的变体是 GeGLU（GELU 门控线性单元 [8]），其中门控机制使用 GELU，形式为 W2 @ GELU(W0@X)*(W1@X)：

图 6：GeGLU MLP

对于 ReLU，“门”更字面化，因为门控向量是非负的，所以它只抑制或缩放特征。对于 GELU，门控值可以是负数，所以门也可以反转特征的符号，这使得“门”成为一个较宽松的历史术语。

MHA（多头注意力）

MHA 是一种自注意力机制，用于建模序列中不同 token 之间的关系。我们使用一种特殊的 MHA 变体，称为 GQA，即分组查询注意力（K/V 头的数量比 Q 头少，因此多个查询（分组）关注同一个键）。

首先，我会给出高层概述——然后我们将深入探讨两个最有趣的组件：YaRN 和核心注意力。

我们首先将每个向量独立地映射为查询、键和值向量。然后我们重塑它们，对查询和键进行归一化，并应用 YaRN（通过旋转注入位置信息）。接下来是核心注意力，它在不同位置之间混合信息。最后，我们应用一个线性投影来产生输出。

图 7：MHA - 多头注意力

YaRN

YaRN [9] 以巧妙的方式修改了 RoPE [10]（旋转位置嵌入），从而更好地外推到更长的上下文长度。

但首先，为什么我们需要位置嵌入？

图 8：位置嵌入背后的原因

现在我们已经理解了原因，让我们看看 RoPE 是如何工作的：

图 9：YaRN 频率表

下面是一个可视化，展示了不同 YaRN 频率的行为。请注意，我们的最慢频率每 1.088M 个位置才完成一个周期！

图 10：YaRN 频率

有了这些，我们就可以看到位置嵌入在前向传播过程中是如何注入的：

图 11：YaRN - 前向传播

以上就是 YaRN 的前向传播。

现在我们理解了其机制，你可能仍然想知道：YaRN 是如何通过对查询和键向量进行成对坐标旋转，然后进行点积来编码相对位置信息的？

图 12：RoPE 如何编码相对位置信息？

RoPE/YaRN 就是这样工作的！:)

核心注意力

最后，让我们分析核心注意力机制。在实践中，我们使用 FlashAttention，它值得一篇单独的博客文章（我实际上在 2023 年写过一篇，点此查看 [11]）。在这里，我将讲解 vanilla attention。

核心注意力是建模序列中 token 之间关系的机制。花些时间分析这个：

图 13：计算 (seqlen, seqlen) 注意力分数矩阵

如果我们就此打住，会出现以下情况：

1. 文档 1 的 token 可以关注文档 2 的 token（反之亦然）
2. token i 可以关注 token i+1（未来 token），这破坏了因果性

为了防止这种情况，我们需要引入掩码！

图 14：注意力掩码与值向量聚合

现在想象一下，我们的序列长度是 32,768 而不是 16。为简单起见，假设是单个文档且没有填充。掩码会是什么样子？

图 15：混合注意力：块局部 + 全局

这是另一种可视化布局的方式，聚焦于位置 9,000 和 10,000 的 token：

图 16：混合注意力布局

你可以看到，在大多数层（块局部）中，这两个 token 不能关注超过 4,096 的位置。在其余的八层（全局）中，它们可以一直回溯到位置 0。

Transformer 数学

最后，我想简要谈谈 KV 缓存，因为它是理解推理的一个极其重要的概念。到目前为止，我们一直在看训练期间的前向传播。

在推理期间，transformer 是自回归的——我们一次生成一个 token。在每一步重新计算所有先前 token 的键和值将极其低效。幸运的是，没有必要：在因果 transformer 中，它们保持不变。相反，我们计算它们一次并将其存储在缓存中。

让我们来看一下基本的 KV 缓存存储需求：

图 17：KV 缓存计算

现在让我们也计算一下 Rnj 1.5 有多少可学习参数。我们只需要遍历架构并统计所有可学习的权重。

图 18：可学习参数数量计算

对于快速心算，请注意你只需要计算 MLP 中的 3 个矩阵和注意力中的 4 个矩阵，其他的可以忽略。

现在让我们计算每个 token 需要多少计算量（FLOPs）。这在规划集群规模时非常有价值——更多内容请参阅本节之后的部分。

图 19：每个 token 的 FLOPs 计算

值得记住 6N 这个公式。也值得记住这个公式成立的设置（即 seqlen << 内部模型维度）。

最后，让我们看看如何将上述公式用于集群规模规划：

图 20：集群规模计算

你现在可以去找孙正义（Masayoshi Son）要 10 亿美元的种子轮融资了。

图 21：利润

结语

我们已经看到了单个 token 如何流经 transformer，以及所有子组件如何协同工作。

我们深入探讨了 YaRN 和注意力，并推导出了一些最重要的 transformer 公式。

在后续文章中，我将更深入地探讨 MoE、Muon（优化器 [12]）以及一些架构创新，例如 MLA（DeepSeek）、MTP（多 token 预测）和 DSA（稀疏注意力 [13]）。

参考文献

1. "Attention Is All You Need", https://arxiv.org/abs/1706.03762
2. RNJ 1.0, https://essential.ai/research/rnj-1
3. "Efficient Training of Language Models to Fill in the Middle", https://arxiv.org/abs/2207.14255
4. "Attention Residual Learning", https://arxiv.org/abs/2603.15031
5. "DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model", https://arxiv.org/abs/2405.04434
6. "Kimi Linear: An Expressive, Efficient Attention Architecture", https://arxiv.org/abs/2510.26692
7. "Root Mean Square Layer Normalization", https://arxiv.org/abs/1910.07467
8. "GLU Variants Improve Transformer", https://arxiv.org/abs/2002.05202
9. "YaRN: Efficient Context Window Extension of Large Language Models", https://arxiv.org/abs/2309.00071
10. "RoFormer: Enhanced Transformer with Rotary Position Embedding", https://arxiv.org/abs/2104.09864
11. "Eli5 Flash Attention", https://learnblockchain.cn/article/26360
12. Muon, https://kellerjordan.github.io/posts/muon/
13. "Dissecting Sparsity in Large Language Models: Intrinsic Data-Aware Sparse Attention", https://arxiv.org/abs/2512.02556

• 原文链接： aleksagordic.com/blog/tr...
• 登链社区 AI 助手，为大家转译优秀英文文章，如有翻译不通的地方，还请包涵～