语境化机器

本文信息来源：stochasm

引言

本文旨在阐述我构建 Transformer 心智模型的过程，某种程度上也是对我过去几个月各种思考与想法的整合。阅读本文需要事先了解 Transformer 架构的相关知识。

我看到很多人说 Transformer 是下一词元预测器（next-token predictors）。虽然这确实描述了 LLMs 和其他模型的工作方式，但我感觉这实际上并没有建立一个关于 Transformer 如何运作的良好心智模型。随着时间的推移，我建立了一个心智模型来帮助我理解 Transformer 的行为：我从根本上将它们视为情境化机器（contextualization machines）。毕竟，下一词元预测是一个学习目标，而不是一种架构。话虽如此，由于我主要与 LLMs 打交道，所以这篇文章仍将重点关注 LLMs（仅解码器架构等）。在接下来的内容中，我们将通过特定的视角来审视 Transformer 的每个组件 contextualization. 我将通过展示这个心智模型如何帮助我们理解不同的研究结果和论文，来阐述它的作用。

当我谈到“语境化”（contextualization）时，我指的是对 Token 和隐状态（hidden states）的语境化。对于仅解码器（decoder-only）架构的 Transformer，有一个我觉得很有用的视角：将残差连接链（residual chain）视为模型的主干，而将各层视为增量变换（additive transformations）；而不是将各层之间的状态流向视为的主干，把残差视为某种锚定状态。下图展示了这一观点的含义。

从某种意义上说，每一层对隐状态的变换都可以看作是对嵌入向量（embedding）的一种语境化操作，然后这种语境化信息会被加回到 Token 表示中。如果你绘制各层之间隐状态相关性（以余弦相似度衡量）的图表，你会发现每一层输出的隐状态与输入时的隐状态非常相似，只存在细微的差异（也就是我所称的额外语境化信息）。以下是一张 Llama-3.2-1B 中不同层隐状态之间余弦相似度的图表。

所以，将语境化操作想象成一种通过更多信息来丰富 Token 嵌入或隐状态的过程，让我们来看看我是如何构建 Transformer 的框架的。

分词器（The Tokenizer）与嵌入矩阵（The Embedding Matrix）

我将分词器（tokenizer）和嵌入矩阵（embedding matrix）概念化为“语境化操作”的方式非常简单。分词器的词汇表大小决定了最终模型在开始前向传播时，将使用多少个经过“预语境化”的片段；而嵌入矩阵则在训练过程中累积信息，从而使这些片段得以被预先语境化。

嵌入包含了构成特定 token 的原子（指不可进一步分解的信息原子）的语境。一般来说，一个 token 中包含的这类原子越多，该 token 中存储的信息语境就越丰富。例如，“hello”比“h”包含更多的信息。你可以说，包含在 token “hello” 中的信息原子具有更高程度的语境化，拥有更丰富的语境含义，因为它们以特定的方式组合在一起，其含义在嵌入阶段就已经被结合了。

现在，把这个概念延伸到词汇表大小上：对于给定的语料库，词汇表越大，最终生成的 token 就越大——因此，原子组也会越大。这意味着在开始前向传播时，每个嵌入都可以拥有更具体的含义。

过度分词的 Transformer：扩大词汇表通常是值得的

这也恰好引出了我们的第一篇相关论文，即 Over-Tokenized Transformer paper。我直接将摘要复制在下方：

分词（Tokenization）是大型语言模型（LLMs）的一个基本组成部分，但其对模型缩放和性能的影响尚未被完全探索。在本文中，我们介绍了“过度分词 Transformer”（Over-Tokenized Transformers），这是一种新颖的框架，通过解耦输入和输出词表来提升语言建模的性能。具体而言，我们的方法通过扩充输入词表来利用多克（multi-gram）标记。通过大量的实验，我们揭示了输入词表大小与训练损失之间存在对数线性关系，证明了无论模型大小如何，更大的输入词表都能持续提升模型性能。通过使用较大的输入词表，我们在不增加额外成本的情况下，达到了与两倍大小基线模型相当的性能。我们的发现凸显了分词在缩放定律中的重要性，并为分词器设计提供了实用的见解，为更高效、更强大的 LLMs 铺平了道路。

这篇论文的主要结论是：无论模型规模大小，增加分词器（tokenizer）的大小总是能持续提升模型性能。研究还发现，词表大小（vocabulary size）与训练损失（training loss）之间存在着标准的、符合缩放定律（scaling-law）的对数线性关系。

如果从 Transformer 本质上是语境化机器这一角度来看，这一点对我来说立刻就变得非常合情合理。实际上，这是将部分语境化工作从模型本身（各层的语境化操作）转移出来，通过允许更高的precontextualization.，将其直接融入分词器（tokenizer）和嵌入矩阵中。更细粒度的信息原子可以被打包进一个嵌入向量里，这样该嵌入向量就不再需要依赖相邻的嵌入向量来丰富自身，就能达到同等水平的contextualization.。

我再尝试用一个例子来说明这一点：假设较小的词表将“he”和“llo”作为独立的词元（token），而较大的词表将“hello”作为一个完整的词元。词表较小的模型必须至少经过一次注意力操作，这两个嵌入（embedding）才能意识到它们是“hello”这个词的一部分；但在较大的模型中，从嵌入矩阵中取出的初始嵌入本身就能识别出它是“hello”这个词。

分词问题与越狱

人们在 LLMs 身上发现了很多有趣的怪癖，比如数单词里的字母，结果发现这些都与分词器（tokenizer）有关。许多越狱攻击手段似乎也依赖于这些分词器的怪癖。在阅读大量越狱案例时，我发现它们通常依赖于拼写错误或类似拼写错误的概念，比如黑客语（leetspeak）。我认为这个例子很好地说明了我的意思。那么，究竟是拼写错误，或者确切地说是轻微的拼写错误，导致了越狱呢？

再一次，我认为将 transformer 视为一台“语境化机器”能让这件事变得直观易懂。让我们考虑一下单词“hello”和它的错误拼写版本“hwllp”。虽然这个词有两个拼写错误，但我们仍然能弄清楚它本来是什么意思。

“hello”这个词被视为一个独立的 token（词元）。相反，“hwllp”则会被 token 化为“hw”+“ll”+“p”（以 GPT-4o tokenizer 为例）。由于“hello”是一个单一的 token，经过嵌入层（embedding layer）后，它的含义是“预语境化”的：模型已经知道它是什么意思。然而，“hwllp”是 3 个 token，模型必须至少使用一个注意力层（attention layer）来弄清楚到底发生了什么，搞明白这一串符号原本想表达什么含义。“hw”、“ll”和“p”各自本身并没有包含太多预语境化的含义。拼写错误还为含义增加了一些额外的噪声，因为模型必须意识到这些 token 组合在一起代表的是“hello”的拼写错误版本。

现在，将这个想法延伸到“越狱”（jailbreaks）：如果模型必须花费额外的层数通过语境化来有效地对输入进行“去噪”，那么输入可能无法足够快地被去噪，导致模型来不及将整个请求语境化认定为有害内容。这将导致当我们到达语言模型头（LM-head）时，得到的是一个“仅做了足够语境化而尚不足以促发拒绝机制”的嵌入，这实际上绕过了模型意识到需要拒绝请求所需的内部计算。这就是我对于为什么拼写错误有助于越狱的直觉判断。在下一节中，我将结合一篇论文详细阐述这一点。

大语言模型中的拒绝能力由单一方向调控

通过这篇论文 ，我认为其中关于越狱（jailbreaks）的观点得到了进一步支持。为了方便参考，我将摘要复制如下。

对话式大语言模型既经过指令遵循微调，也经过安全微调，从而形成既能服从良性请求，又会拒绝有害请求的模型。尽管这种拒绝行为在各类聊天模型中普遍存在，但其潜在机制仍鲜为人知。在这项工作中，我们在多达 13 个流行的开源聊天模型（最大参数量达 72B）中展示了拒绝能力是由一个一维子空间调控的。具体而言，对于每个模型，我们都找到了一个单一方向，只要从模型的残差流激活中消除该方向，就能阻止其拒绝有害指令；反之，加入该方向则会诱发其对无害指令的拒绝。基于这一洞察，我们提出了一种新颖的白盒越狱方法，该方法可以在对其他能力影响最小的情况下，精准地移除拒绝机制。最后，我们从机理上分析了对抗性后缀是如何抑制拒绝调控方向传播的。我们的发现凸显了当前安全微调方法的脆弱性。更广泛地说，我们的工作展示了如何利用对模型内部的理解来开发控制模型行为的实用方法。

这就很有意思了，我特别欣赏这篇论文的一点是：既然作者确定了“拒绝”（refusal）对应嵌入空间中的某一特定方向，他们就能逐层测量模型趋向拒绝的程度。他们确实这么做了，最后得出的图表就是这张。

该图表显示，对于不安全输入，其与拒绝方向的相关性通常出现在解码器堆栈的大约 75% 处。如果拼写错误越狱（typo-jailbreaks）导致更多的层被用于语境化，那么图表中的曲线（可能）会向右偏移，使得当嵌入到达语言模型头（LM head）时，与拒绝方面的相关性还不够强，从而导致越狱成功。

通过分词器和嵌入矩阵，模型以预先语境化的含义为起点，但这些静态表征需要动态机制来理解它们在序列上下文中的角色。

注意力机制

提到 Transformer 背景（上下文）下的“语境化（contextualization）”一词时，注意力机制可能是脑海中浮现的核心内容，因此我认为这里不需要过多阐述其思维模型。它是 Transformer 中 token 嵌入/隐藏状态相互交互并共享信息的唯一方式。

不过，我要强调的一点是，我倾向于将其想象为一种局部或上下文内的语境化形式（我现在说了多少次这个词了？）。嵌入向量通过与其他嵌入向量的交互接收新信息，这些新信息被用来丰富每个嵌入向量中包含的意义。然而，这些新信息具体来自于其他嵌入向量中所包含的意义：例如，“ball（球）”的嵌入向量可能会看到在其之前的 token 序列中出现了“black（黑）”和“white（白）”这两个词。这可能表明这个球是黑白相间的，而这一信息可能会被整合进“ball”的嵌入向量中。此时，“ball” token 位置上的隐藏状态可能就包含了类似“黑白相间的球”这样的信息。

（警告：此处纯属推测）然而，没有任何迹象表明仅凭注意力机制的变换就能看出这个球可能是足球，因为在这个序列中没有任何地方提到这个球是足球。因此，我倾向于将注意力视为一种局部语境化的形式，用于处理序列中已经包含的信息。这一点我也可能搞错，也许另一个注意力层中添加到嵌入向量（embedding）的值，可能会根据键（Key）在查询（Query）中看到的内容，包含一些关于足球的方向性信息。不过，鉴于注意力机制与上下文学习（in-context learning）的联系如此紧密，我认为这是一个可靠的心智模型。

虽然注意力机制允许 token 能够根据局部上下文共享并结合其含义，但模型仍然需要一种途径去访问输入序列之外的更广博的知识。

前馈层

（这部分直觉很大程度上来自 3blue1brown 关于 LLMs 如何存储事实的视频，我强烈建议观看。其中的可视化效果非常棒。）现在，与充当局部语境化形式的注意力机制相反，前馈层/多层感知机块充当的是一种全局语境化形式contextualization.。我倾向于将其理解为：通过一个定向查找表，在模型见过的整个训练数据分布的全景图中，找到当前嵌入向量所处的位置。

因此，在 Transformer 中，前馈网络（FFN）（通常）由两层组成：一个上投影（up-projection）和一个下投影（down-projection），两者都是矩阵乘法（中间夹一个非线性激活函数）。第一个矩阵乘法可以看作是计算矩阵的每一行与嵌入向量之间的点积，这意味着中间的隐藏状态是一个相似度向量，编码了上投影矩阵的每一行与输入隐藏状态的相似程度。你可以将其视为对嵌入空间中语义方向的一种加权索引，其中语义方向是由上投影矩阵的行定义的。“在这一组我已知的含义语境中，这个 token 处于什么位置？”

随后，向下投影（down-projection）可以被看作是根据这些相似度检索相关信息，并将这些信息汇集到一个大小为 embedding_dim 的方向向量中，该向量随后会被（通过残差连接）加到之前的隐藏状态上。因此，我认为这是以一种“全局”的方式对隐藏状态进行语境化处理，即与网络在训练过程中见过的所有含义进行比对。例如，如果原始隐藏状态中存在表示“黑白球”的向量，并且与“足球”方向存在关联，这可能会增加一些朝向“足球”含义的信息。

这样，我将主解码器栈的前向传递过程视为不断为 token 增添额外含义并精炼其语境，直到它们的含义变得清晰透彻。

下一 token 预测的学习目标与语言模型（LM）头

我知道，我说过下一 token 预测是目标，而不是架构，但我认为这里仍有一点值得讨论，即我如何将其纳入我的心智模型。我是这样看的：既然下一 token 预测是目标，那么关于下一个 token 的信息就必须在某个时刻被纳入隐藏状态中。所以，我在脑海中建立的模型是：注意力语境化操作不仅纳入了关于先前 token 的信息，同时也纳入了关于未来语境可能是什么样子的推测性信息。我觉得，你对一个 token 在特定语境下的确切含义了解得越多，你对接下来会出现什么的推测也就越清晰。

推测越强烈，隐藏状态分布就开始越像最终输出分布。在我看来，解码器的后几层最终不得不更多地专注于这种推测性的语境化，因为到了这个阶段，大多数隐藏状态的含义理应已经得到了充分的语境化处理。

语言模型头（Language Model Head），或称去嵌入矩阵（Unembedding Matrix）

我想象中的语言模型头（或称去嵌入矩阵）是这样的：它从每个词元（token）经过充分语境化的表示中提取出对“下一个词元”的推测信息。在经历了所有的注意力层和前馈操作之后，每个词元的嵌入（embedding）都已经富含了局部和全局的语境信息，同时也包含了关于“接下来的内容”的推测信息。此时，去嵌入矩阵的作用就是提取出这些推测信息，并将其转化为覆盖整个词汇表的概率分布。

通过多词元预测实现更好、更快的大型语言模型

这也让多 token 预测（multi-token prediction）显得更加直观。让我们以 Meta 的多 token 预测设置为例。他们的架构保持了 Transformer 的主干部分不做修改，然后为多 token 预测添加了多个 LM 头（Language Model heads）。我在下面再次引用了摘要内容。论文链接在这里 。

像 GPT 和 Llama 这样的大型语言模型通常是使用下一个词元（next-token）预测损失进行训练的。在这项工作中，我们提出训练语言模型一次性预测未来多个词元可以带来更高的样本效率。具体而言，在训练语料库的每一个位置，我们通过位于共享模型主干之上的 $n$ 个独立输出头，要求模型预测接下来的 $n$ 个词元。我们将多词元预测视为一种辅助训练任务，结果显示在不增加训练时间开销的情况下，代码和自然语言模型的下游能力均得到了提升。该方法对于更大规模的模型越来越有效，并且在多轮次（epoch）训练中依然保持其优势。在诸如代码生成等生成式基准测试中，性能提升尤为显著，我们的模型始终比强大的基准模型高出几个百分点。我们拥有 130 亿（13B）参数的模型在 HumanEval 上多解决了 12%的问题，在 MBPP 上多解决了 17%的问题，表现优于同类的下一个词元预测模型。在小型算法任务上的实验表明，多词元预测有利于归纳头（induction heads）和算法推理能力的发展。作为一个额外的优势，即便在大批量（batch sizes）处理时，使用 4 词元预测训练的模型在推理速度上也快了 3 倍。

若从语境化机器的视角来看：得益于改进后的学习目标，模型被激励去对更长远的未来进行推测性语境化，从而允许语言模型（LM）头从词元嵌入中提取出多个推测词元。

这篇论文的有趣之处在于，模型（在更大规模下）展现出了性能的提升和更丰富的内部表征，进而改善了下游任务的表现。这表明多 token 预测从根本上改进了模型在前向传播过程中处理和情境化信息的方式。这对我来说非常符合直觉：如果模型是利用来自下游多个 token 位置的信号进行训练的，那么情境化操作就会被训练去推测更远的未来。因此，更强的情境化操作会带来模型内部对概念更强的表征，这种想象并不牵强：Transformer 本质上就是一台情境化机器。

大语言模型是否在潜在地执行多跳推理？

我想提到的最后一篇论文是 Deepmind 的这篇论文 。其摘要如下：

我们研究了大语言模型 (LLMs) 在面对诸如“‘Superstition’这首歌演唱者的母亲是”这类复杂提示词时，是否在潜在地执行多跳推理。我们寻找一种潜在推理路径的证据，即 LLM (1) 潜在地将“‘Superstition’的演唱者”识别为 Stevie Wonder（即桥接实体），并且 (2) 利用其关于 Stevie Wonder 母亲的知识来补全提示词。我们单独分析这两个跳跃步骤，并将它们的共现视为潜在多跳推理的标志。对于第一跳，我们测试如果将提示词改为间接提及桥接实体而非其他实体，是否会增加 LLM 对桥接实体的内部召回。对于第二跳，我们测试增加这种召回是否会导致 LLM 更好地利用其关于桥接实体的知识。我们在某些关系类型的提示词中发现了潜在多跳推理的有力证据，这种推理路径在超过 80% 的提示词中被使用。然而，这种利用具有高度的情境依赖性，在不同类型的提示词中表现各异。此外，平均而言，关于第二跳和完整多跳遍历的证据相当温和，只有第一跳的证据较为充分。而且，我们发现随着模型规模的增加，第一跳推理呈现出明显的扩展趋势，但第二跳则没有。我们的实验发现揭示了 LLMs 未来发展和应用所面临的潜在挑战与机遇。

在本文中，他们提出了一种名为 ENTREC 的指标，该指标通过模型的隐藏层表示来衡量模型对桥接实体的内部召回情况。在附录 C 中，他们对该指标进行了验证，结果表明，当提高“Stevie Wonder”在中后层的内部召回率时，模型在描述性短语后的逗号之后，生成“Stevie”作为下一个 token 的可能性会增加。有趣的是，他们绘制了在特定层提升该指标所产生的影响，得到了下图：

我认为这可以看作是对我们思维模型的间接支持。随着层深度的增加，这种效应缓慢增强，随后维持在平台期一段时间，最后当我们接近解码器的末端时开始减弱。在语境化框架下，这意味着随着隐藏状态（hidden states）获得更多的语境化处理，它们逐渐收集关于该 token 的信息，从而更有可能实现检索跳跃。这似乎是合乎逻辑的。到了解码器的末端，这种效应之所以下降，是因为隐藏状态已经完成了充分的语境化，转而进入了推测阶段。

他们发现了单跳推理的有力证据，但对完整的多跳遍历只发现了较为温和的证据，这一事实也表明，那种认为“首先尽可能对 Token 进行语境化，然后推测性地对其进行语境化以收集有关潜在下一个 Token 的信息”的心智模型可能是准确的。我猜想某种层块循环式架构（blocks-of-looped-layers）可能有助于解决这个问题，不过当然，实现这种功能的方法有很多，我目前还远不能确定哪种方法才是最佳途径。

结论

希望这篇博文能帮助我传达出自己关于 Transformers 运作机制的心智模型，也希望它能让你对此建立起更强的直觉。