MuToR：为什么大语言模型在学习时需要"寄存器"？雅典娜研究中心突破性发现

在人工智能研究的前沿，来自雅典娜研究中心的Anastasios Gerontopoulos、valeo.ai的Spyros Gidaris以及克里特大学和IACM-Forth的Nikos Komodakis共同发表了一项引人注目的研究。这项名为"Multi-Token Prediction Needs Registers"（多符号预测需要寄存器）的论文于2025年5月15日发布在arXiv预印本平台上（编号：2505.10518v1），为提升大型语言模型的训练效果提出了一种全新方法。有兴趣深入了解的读者可以通过https://github.com/nasosger/MuToR访问他们的完整代码。

一、当语言模型只能"一步一步"思考时遇到的问题

想象一下，如果你在解决一个复杂的数学问题时，只能一次思考一个步骤，而且每思考一步就会立即忘记之前的思考过程，只记得当前的结果。这种方式听起来非常低效，对吧？然而，这正是当今大型语言模型（如ChatGPT背后的技术）在训练时面临的情况。

目前的大语言模型主要通过"自回归转换器"架构训练，使用"下一个词预测"的方法。简单来说，就是给模型一个句子的开头部分："今天天气真"，让它预测下一个词可能是什么（比如"好"）。这种训练方式虽然简单高效，但也存在明显的局限性：模型往往只关注短期模式，难以处理需要长远规划的复杂任务。

研究者们发现，这种"一词一词"的预测方式会导致模型走捷径学习，忽略真正有价值的信息。就像一个学生只学会了背诵公式，却不理解背后的原理一样。这种训练方式特别不适合需要推理和规划的任务。比如，当你让AI解决一个复杂的推理问题时，它可能会卡壳，因为它习惯了只看一步而不会提前规划。

二、多符号预测：让AI学会"向前看"

为了解决这个问题，研究者们发展了"多符号预测"训练方法。不同于标准的"预测下一个词"方式，多符号预测让模型同时预测多个未来的词。这就像是教一个象棋选手不仅思考下一步棋，还要预想几步之后的局面。

最近的实现主要通过额外的"头部"来实现这一功能：有些研究使用并行的头部（每个负责预测不同位置的未来词），而其他研究则采用顺序头部。这种方法只在训练时使用，目的是提供更丰富的学习信号，而不是加速推理过程。与标准的训练方法相比，多符号预测鼓励模型发展内部的"规划"表示，减轻对局部模式的过度拟合。

然而，尽管多符号预测在预训练中显示出了潜力，但其好处并未一致地推广到微调等其他设置中。这就像是学生在标准考试中表现出色，但在面对新颖的实际问题时却表现不佳。

三、MuToR：一种简单而有效的全新方法

雅典娜研究中心的研究团队提出了一种名为MuToR（多符号预测与寄存器）的新方法，它巧妙地解决了现有多符号预测方法的不足。

想象一下，如果我们在阅读一本书时，不仅关注当前的段落，还能时不时地瞥一眼几页之后的内容，这会让我们更好地理解整个故事的发展。MuToR正是基于类似的思想设计的。

与现有方法不同，MuToR没有添加额外的转换器层来预测未来的词，而是在常规词之间插入了特殊的"寄存器词"。每个寄存器词被分配一个随机采样的偏移量d，模型训练时需要预测d步之后（而不仅仅是下一个词）的词。而对于所有常规词，原始的下一个词预测目标保持不变。

这些寄存器词仅在训练期间使用，用于传播更丰富的监督信号。在推理时（即实际使用模型时），这些寄存器词会被丢弃，以保持生成速度。这是通过精心设计的注意力掩码实现的：寄存器词只能关注前面的常规词（使它们能够学习预测性表示），而常规词则完全看不到寄存器词。这确保了与标准自回归推理的完全兼容性，同时鼓励模型在训练期间内部化前瞻性、多步规划。

四、MuToR的独特优势

与那些添加输出头或转换器块的方法相比，MuToR的注册器方法提供了几个关键优势：

首先，MuToR不需要对模型架构进行任何改变。它只通过寄存器嵌入引入少量额外的可训练参数，而核心转换器层保持不变，也不需要额外的输出头。

其次，MuToR特别适合微调预训练的大型语言模型（如Llama或Gemma）。它引入的参数开销最小，为常规词保留原始注意力模式，并为寄存器词使用经过精心选择的位置ID和注意力掩码，使多符号预测更接近预训练设置。相比之下，以前的方法依赖于单独的转换器头，增加了许多必须从头开始训练的新参数，使它们在微调场景中效果较差。

最后，MuToR支持可扩展的预测范围。由于无论偏移量d如何，寄存器词的数量保持固定，因此训练成本与预测范围无关，可以任意扩展。寄存器词因此在未来词预测方面提供了更大的灵活性。例如，在自回归图像生成中，MuToR可以自然地扩展到在二维邻域中预测词，这是通过添加输出头不容易实现的功能。

五、实验验证：MuToR在各种场景下的表现

研究团队通过一系列实验验证了MuToR的有效性和通用性。实验覆盖了多种使用场景，包括监督微调、参数高效微调（PEFT）和预训练，并在语言和视觉领域的具有挑战性的生成任务上进行了测试。

在语言建模方面，研究团队专注于数学推理任务（带有思维链）和摘要生成任务，这两类任务都是对模型能力的严格测试。他们微调了两个预训练的仅解码器语言模型：Gemma 2B和Llama 3 8B。测试基准包括GSM8K、MATH500和AQUA-RAT等数学推理基准，以及SAMSum和DialogSum等摘要生成基准。

对于数学推理任务，MuToR始终优于基线方法和Multi-Token方法，而后者引入了大量额外的可训练参数。此外，MuToR的性能提升在不同的训练集大小下都能保持，表明即使在高质量微调数据的设置中也能保持有效性。相比之下，Multi-Token的好处似乎随着更大的模型或更多的数据而减弱。

在摘要生成任务中，MuToR提高了所有ROUGE分数，在序列到序列生成任务中展示了广泛的适用性。

研究团队还验证了MuToR在参数高效微调（LoRA）场景中的有效性。在这些实验中，MuToR结合LoRA不仅优于标准LoRA微调，甚至匹配或超过了全微调的性能，展示了它在PEFT设置中的实用性。

此外，团队还将MuToR应用于自回归图像生成，在预训练设置中提高了性能，进一步突显了其跨不同领域和训练设置的广泛潜力。

六、深入理解MuToR的工作原理

要理解MuToR的工作原理，想象一下学习弹钢琴的过程。传统的"下一个词预测"就像初学者只能一个音符一个音符地学习，而无法理解整体的乐句。而MuToR则像是在学习过程中，教师不时地引导学生预测几个音符之后应该弹什么，帮助学生建立对音乐全局的感知。

具体来说，MuToR通过以下方式工作：

首先，将特殊的"寄存器词"插入到输入序列中。每个寄存器词被分配一个随机抽样的偏移量d，模型需要预测当前位置之后d步的内容。

其次，通过精心设计的注意力掩码，寄存器词只能看到前面的常规词，而常规词完全看不到寄存器词。这确保了常规词的表示不受影响，同时让寄存器词学习更丰富的预测表示。

第三，为寄存器词设置特殊的位置ID。虽然常规词保持其自然位置t，但每个插入在xt之后的寄存器词（预测xt+d）接收位置t+d-1。这匹配了在标准"下一个词预测"下通常会预测xt+d的常规词的位置ID。

最后，在训练过程中，模型同时优化标准的"下一个词预测"损失和辅助的寄存器损失，通过加权和组合这两个目标。

在推理时，寄存器词被完全丢弃，不会影响模型的计算图和延迟。这使得MuToR在保持标准自回归过程不变的同时，通过更丰富的训练信号改善了模型的学习能力。

七、MuToR在不同领域的特殊应用

MuToR不仅在标准语言建模任务中表现出色，还展示了在多领域的适应性和灵活性。

在自回归图像生成方面，MuToR展示了独特的优势。传统的图像生成模型将图像视为一维令牌序列，按照光栅扫描顺序（从左到右，从上到下）生成图像。MuToR通过二维偏移采样修改了这一方法，考虑了图像的二维结构。

具体来说，对于每个序列，模型采样一对二维偏移量(dh,dw)，然后计算光栅化偏移量。这种二维扩展通过捕获视觉数据中固有的空间依赖性来丰富训练信号，同时只需最小的架构更改。不同于之前的多符号预测方法（需要为每个可能的二维偏移量添加多个额外的预测头），MuToR以极小的参数开销实现了这一功能。

实验结果表明，在ImageNet图像生成任务上，使用二维偏移的MuToR-2D在各种评估指标上都显著优于基线方法和使用一维偏移的MuToR-1D。研究团队还发现，即使使用较少的随机放置寄存器（每张图像仅80个寄存器），MuToR-2D也能实现与全寄存器设置（256个寄存器）非常相似的性能，同时减少计算量。

此外，研究团队还在合成数据上测试了MuToR，特别是在星形图路径查找问题上。这个问题突显了"下一个词预测"目标的局限性。在这个任务中，MuToR成功地克服了标准教师强制失败的"作弊现象"，在各种图配置上有效地解决了任务，恢复了有价值的训练信号。

八、研究结论与未来展望

研究团队的工作表明，MuToR提供了一种简单而有效的方法来实现多符号预测，通过插入可学习的寄存器令牌来预测未来目标。与现有方法相比，MuToR引入的参数极少，不需要架构改变，与预训练目标保持一致，特别适合监督微调场景。

研究结果一致地表明，MuToR在各种生成任务中提高了性能，包括数学推理、摘要生成和图像生成。特别值得注意的是，即使在使用高质量微调数据的情况下，MuToR的优势仍然保持，表明它提供了真正的模型改进，而不仅仅是弥补数据不足。

然而，值得注意的是，MuToR目前使用均匀插入或随机放置的寄存器词，这些策略可能不一定与特定任务的结构或语义最佳对齐。通过学习或适应寄存器词的放置—可能由模型不确定性或特定任务的先验引导—MuToR可能会提供更有针对性的监督，使用更少的辅助词进一步提高效率和性能。

这项研究为使用基于令牌的前瞻机制在训练期间传播更丰富的监督信号奠定了有希望的基础。未来的工作可能会探索更复杂的寄存器放置策略、动态偏移量选择或领域特定的适应，进一步推动语言和视觉模型的性能边界。