模型对哪些类型的 token 预测得好,哪些又预测得不好?这个问题对于混合架构来说尤其引人入胜——这是一种已经开始挑战标准 Transformer 架构的语言模型架构,我们一直在通过 Olmo Hybrid 对其进行研究。
混合模型在标准评测基准上可以媲美甚至超越 Transformer 模型,但那些头条数字并不能揭示混合模型相比 Transformer 具体有哪些优势。
为了阐明这些 token 级别的行为,我们最近进行了实验,将我们自己最强的 7B Transformer 模型 Olmo 3 与混合模型 Olmo Hybrid 进行直接对比。具体来说,我们以细粒度方式比较了模型在不同类型 token(即作为大语言模型输入的信息单元)上的预测差异。
由于 Olmo 3 和 Olmo Hybrid 在架构之外的设计尽可能保持一致——在数据、分词器和训练方案上高度匹配——因此它们预测结果的任何差异主要反映了架构本身。在 token 层面审视这些差异,使我们能够深入了解混合模型相对于 Transformer 的具体优势。
我们的结果表明,混合模型在众多 token 上的优势是真实存在的,但并非对所有 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 工作中的三个 10 亿参数模型:一个 Transformer 架构、一个混合架构,以及一个完全没有注意力机制的纯循环模型。
在非重复的有意义 token 上,混合架构和纯循环模型超越了 Transformer,其中混合架构表现最佳。而在重复 token 上,纯循环模型——由于缺乏用于回溯复制的注意力机制——落后于混合架构和 Transformer。
因此,这些过滤后的 token 损失揭示了架构之间更细粒度的差异,包括复制能力以及在内容词上的差异,这些差异在训练早期原本是无法显现的。
结论与展望
过滤 token 损失在 10 亿参数预训练过程中揭示了架构差异。图中展示了 Transformer、混合架构和纯循环神经网络(RNN)在 WSD 退火检查点上的 token 损失曲线。
这项工作带来了两点启示。
第一,单一的总体损失——模型在所有 token 上的平均误差——过于粗糙,不足以比较 Transformer 和混合架构。仅针对测试特定模型能力的 token 计算损失,才能揭示关键差异。
第二,具体到混合模型,我们发现其在开放类 token 上具有特定优势,这可能与 RNN 层的状态追踪能力有关。
作为下一步,我们正在将这些发现应用到正在进行的混合模型研究工作中。我们相信,最佳的混合架构将来自于逐 token 理解模型每个组件的优势所在。我们希望此类研究能帮助整个 AI 社区加深这种理解。
我们鼓励您阅读完整报告、探索 Olmo 3、尝试 Olmo Hybrid,并深入研究其相关的开放成果。