混合模型在哪些 token 上预测得更好？

📄 Tech report: https://arxiv.org/abs/2606.20936

模型对哪些类型的 token 预测得好，对哪些预测得不好？这个问题对于混合模型来说尤其引人入胜——混合模型是一种语言模型架构，已经开始挑战标准的 Transformer，我们一直在用 Olmo Hybrid 对它进行深入研究。

混合模型在标准基准测试中可以媲美甚至超越 Transformer，但这些 headline 数字并不能揭示混合模型相对于 Transformer 具体有哪些优势。

为了阐明这些 token 层面的行为，我们最近进行了实验，将自己最强的 7B Transformer 模型 Olmo 3 与混合模型 Olmo Hybrid 进行直接对比。具体来说，我们以细粒度的方式比较了不同类型 token（即作为大语言模型输入的信息单元）上的模型预测差异。

由于 Olmo 3 和 Olmo Hybrid 在架构以外的方面尽可能保持一致——数据、分词器和训练方案都高度匹配——因此它们预测结果的任何差异主要反映了架构本身。从 token 层面观察这些差异，使我们能够深入了解混合模型相对于 Transformer 的具体优势。

我们的结果表明，混合模型在很多 token 上确实有优势，但并非所有。Olmo Hybrid 在承载意义的 token（如名词、动词和形容词）上表现最强，在只能通过跟踪上下文来预测的 token（如代词指代的是谁）上也是如此。但混合模型在那些只是简单重复输入中已有内容的 token 上优势几乎消失——即从前面逐字复制的单词或短语——答案就在那里，可以查得到。这正是 Transformer 的强项所在。

注意力机制与循环机制，以及测量差异

语言模型由一堆重复的层构建而成，每一层都利用周围的 token 来精炼每个 token 的表示。

Transformer 架构在每一层都使用了注意力机制。模型可以同时直接访问每一个更早的 token，评估每个 token 与当前预测的相关性。这使得注意力机制能够精确地回忆起某一个较早的 token，即使该 token 出现在输入的较远位置。代价是每个 token 都要与所有更早的 token 进行比较，因此随着输入增长，注意力的计算成本会急剧上升。此外，虽然注意力在回忆和聚合信息方面表现出色，但它也难以表示随时间顺序演变的信息。

混合模型保留少量注意力层，其余层替换为循环层。与注意力层不同，循环层从左到右读取 token，并携带一个固定大小的记忆，将每个新 token 折叠到记忆中，因此无论输入多长，处理每个 token 的成本都保持恒定。这种记忆是压缩且有损的，因此循环层无法像注意力层那样回溯寻找精确的早期 token。但它非常适合持续记录模型读取 token 时发生的任何变化，提供了对注意力机制的互补优势。

为了分离注意力层和循环层的优势与劣势领域，我们向 Olmo 3 和 Olmo Hybrid 输入了文本段落：文章、维基百科条目、书籍和科学论文，以及 Python、HTML 和 LaTeX 等结构化文本。我们根据每个模型在给定样本中从前面的 token 预测每个 token 的准确度进行了评分。

两个模型都看到了相同的更早 token，并为每个可能的下一个 token 分配了一个概率。我们记录了每个模型对实际出现的 token 所赋予的概率。然后，我们通过计算损失差距（即两个模型之间的损失差异）来逐 token 总结两个模型之间的差异。正差距表示混合模型对真实下一个 token 的预测更好。负差距则表示 Transformer 模型做得更好。

为了找出损失差距可能集中在哪里，我们进行了多项分析。首先，我们将每个 token 归类，并计算这些类别内的平均损失差距。由于原始平均值可能受到其他因素的影响，例如某个类别的稀疏性或 token 在文本样本中的重复频率，我们通过回归方法重新检验了每个模式，该方法在控制其他变量不变的情况下，估算该类别本身的影响。

真实文本所呈现的结果

我们发现，在大多数类型的 token 上，Olmo Hybrid 的损失低于 Olmo 3，但不同 token 的降低幅度并不相同。

在散文文本中，最明显的分界在于实义词——即承载意义的名词、动词和形容词——与功能词（如“the”、“of”、“is”）之间。混合模型对实义词的预测效果优于 Transformer，损失差距约为 [原文此处未给出具体数值]，而对功能词的损失差距则接近 [原文此处未给出具体数值]。

特别是在副词和形容词这类实义词类别中，混合模型的优势尤为显著；不过，某些功能词类别（例如表示存在的词，如“there”）也显示出混合模型具有较大优势。简而言之，混合模型的优势在那些决定句子主要内容的词汇上最大，而在那些任何模型几乎都能根据句法猜测出来的语法词汇上最小。

相比之下，我们发现了一些特定场景，混合模型相对于 Transformer 的优势消失了。第一个场景是闭合括号（而非开放括号）——这一模式在语言、代码和标记语言的各种括号中均稳定存在。为什么？因为已知注意力机制足以表示括号匹配，这表明仅凭注意力机制就能预测闭合括号。

混合模型优势几乎消失的第二个场景是：当下一个 token 仅仅是重复前文中已有的内容时。我们通过查找重复的 n-gram 来识别这些情况：即一段文本中，能够补全某个序列的 token 此前在同一段落中已经逐字出现。重复的片段越长，混合模型的领先优势就越小，直至趋近于零。

最后，受这些发现启发，我们探索在特定类型 token 上使用过滤损失作为评估手段，以便在预训练实验中更好地比较不同架构。我们使用了来自早期 Olmo Hybrid 工作的三个 1B 参数模型：一个 Transformer、一个混合模型，以及一个完全没有注意力机制的纯循环模型。

在非重复的有含义 token 上，混合模型和纯循环模型超越了 Transformer，其中混合模型表现最佳。而在重复 token 上，纯循环模型——由于没有注意力机制来向前追溯复制——则落后于混合模型和 Transformer。

因此，这些过滤后的 token 损失揭示了不同架构之间更细粒度的差异，包括复制能力以及实词上的差异，这些差异在训练早期原本是无法察觉的。

这给我们带来了什么

过滤后的 token 损失在 1B 预训练过程中揭示了架构差异。图中展示了 Transformer、混合模型和一个纯循环神经网络（RNN）在 WSD 退火检查点处的 token 损失曲线。

这项工作带来了两个经验教训。

首先，单一的整体损失——模型在所有 token 上的平均误差——过于粗糙，无法用于比较 Transformer 和混合架构。仅对测试特定模型能力的 token 计算损失，才能揭示关键差异。

其次，具体到混合模型，我们发现了其在开放类 token 上具有特定优势的证据，这可能与 RNN 层的状态追踪能力有关。

作为下一步，我们正在将这些发现应用于我们持续的混合建模工作中。我们相信，最佳混合架构将来自逐 token 理解模型中每个组件的优势。我们希望这样的研究能帮助整个 AI 社区增进这种理解。

我们鼓励您阅读我们的完整报告，探索 Olmo 3，尝试 Olmo Hybrid，并深入研究其相关的开放成果。