在人工智能领域,理解语言最强大的工具依赖于一种被称为“Transformer”的结构。想象一位读者正在扫描一篇长文档:为了理解一个特定的词,他们必须回顾出现在该词之前的词汇,以把握上下文语境。标准的AI模型通过两个截然不同的步骤来完成这一过程。首先,它们扫描整个文本的历史,以收集来自过去词汇的相关信息,这个过程类似于图书管理员快速定位出与某个主题相关的所有书籍。其次,它们将收集到的信息通过每一个位置上独立的、单独的过滤器进行处理,以精炼其含义。这种“两步舞步”在领域内一直表现良好,但计算量巨大,因为随着文本变得越来越长,模型在生成每个新词时都需要进行大量的计算。研究人员长期以来一直在寻求一种合并或简化这些步骤的方法,希望能够构建出既更聪明又更高效的模型,使其能够处理海量信息,而不会被处理这些信息所需的庞大数学运算所拖累。
印第安纳大学布鲁明顿分校的一位研究人员引入了一种全新的方法,从根本上重塑了这些模型与语言交互的方式。他们将这一创造物称为 TANGO,代表“标记聚合非线性门控算子”(Token-Aggregated Nonlinear Gating Operators)。TANGO 不再将信息收集和信息精炼这两个步骤分开,而是将它们合并为一个统一的动作。在这种新设计中,句子中的每一个词都充当了控制源。当模型阅读时,每个词都会生成一组特定的指令,即“门控”(gates),这些指令决定了在不同文本特征上投入多少权重。当模型读到一个新词时,它不仅仅是回顾过去以寻找最相关的词,它还会回顾那些过去的词准备好了什么样的指令。随后,它会对这些指令进行平均,并利用它们来缩放(或调整)当前词的特征,然后再将其添加到模型的记忆中。这意味着,过去一个词的影响力不仅在于它说了什么,还在于它如何告诉当前的词去如何解读自身。
研究人员开发了两个版本的系统,以测试准确性与速度之间的不同权衡。第一个版本是 TANGO,它会查看当前词之前出现的每一个词,无论这些词在文本中离得有多远。这种“全前缀”(full-prefix)方法允许模型汲取整个对话或文档的历史信息,从而实现对上下文的高度准确理解。然而,由于它必须将当前的词与之前每一个词进行比较,其工作量会随着文本变长而迅速增长。第二个版本被称为 WANGO,它针对极长文本采取了一种更务实的方法。它会密切关注最近出现的词,并以对待第一个模型那样详尽的全历史方式来处理它们。对于落在近期窗口之外的较旧词汇,WANGO 使用了一种更高效的方法来总结它们的影响力。它维护着一个由这些旧词生成的指令的运行总计,从而能够在不重新计算与每一个词的关系的情况下,吸收它们的智慧。这种修改使得 Wango 模型的负载随文本长度增加呈直线增长,而不是在复杂度上爆炸式增长,从而使其在处理长序列时更加快速。
为了观察这些新设计在现有技术面前的表现如何,研究人员使用完全相同的数据、相同数量的参数以及相同的训练计划训练了六个不同的模型。他们在三个非常不同的挑战任务上测试了这些模型:一个大规模的教育类网络文本集合、一个用名为 Lean 的编程语言编写的正式数学证明库,以及来自 DeepMind 的一系列数学问题。结果显示,拥有全历史视图的 T多少 TANGO 模型在所有三项测试中都实现了最高的准确度,在预测序列中的下一个词时表现出了最可靠的预测能力。它的表现优于所有其他模型,包括那些使用传统方法在层间共享参数的模型。Wango 模型同样表现出色,特别是在旨在追求效率的模型类别中。在那些能够处理长文本且计算成本不会失控的架构中,Wego 产生了最准确的结果。它甚至击败了一个计算量相近但依赖于较旧标准技术的模型。
这项研究强调了构建此类系统的一种重要转变。通过允许词汇通过这些聚合门控来控制其他词汇的处理,研究人员发现了一种创建既强大(在 Wango 的情况下也是高效)的模型的方法。TANGO 模型证明,即使在模型被迫进行更多数学运算的情况下,通过单一的、统一的步骤来收集和精炼信息,也能超越传统的两步处理过程。与此同时,Wango 模型证明了这种高水平的性能并不一定以牺牲速度为代价;通过智能地总结旧信息,它在保持高准确度的同时,使计算成本保持在可控范围内。研究人员谨慎地指出,这些发现是基于模型预测序列中下一个词的能力(这是衡量语言理解能力的标准指标),并不意味着这些模型一定能够完美地解决数学问题或编写证明。然而,结果表明,这种新的信息门控方式为构建能够以更高精度和效率理解复杂语言及长上下文的 AI 提供了一条充满前景的路径。
技术摘要:TANGO 与 WANGO
问题陈述
标准的 Transformer 架构将跨 Token 交互(由自注意力处理)与逐位置非线性变换(由前馈网络处理)分离。这种分离使得每个模块需要两个不同的子层。虽然这种设计行之有效,但可能在效率和参数利用率方面并非最优。此外,现有的线性复杂度序列模型往往依赖于某种近似,这可能会牺牲跨 Token 交互的丰富性,或者需要复杂的门控机制,无法将非线性变换与注意力机制充分整合。
方法论
核心架构:Token 聚合非线性门控算子 (TANO)
本文引入了 TANGO,这是一种解码器模块,它用单个跨 Token 门控残差更新取代了标准的自注意力和逐位置前馈子层。
源条件 SwiGLU 门控:
- 不同于传统的 SwiGLU(Swish-Gated Linear Unit,其中门控和特征均由同一个 Token 计算),TANGO 允许由源位置(xj)计算出的门控来调制目标位置(xi)的投影特征。
- 每个位置 i 生成一个门控向量 gi=SiLU(Wgxi) 和一个投影特征向量 vi=Wvxi。
- 对于目标位置 i,模型计算所有因果可见源(j≤i)的门控向量的基于内容的加权平均值,记作 gˉi。
- 残差更新计算如下:
hi+=hi+Wo(gˉi⊙vi)
- 此处,聚合后的门控 gˉi 作为目标位置投影特征 vi 的对角缩放算子,作用于最后的输出投影 Wo 之前。
权重聚合机制:
论文提出了两种变体,其区别在于如何计算用于聚合源门控的权重:
- TANGO (Full-Prefix/全前缀): 基于每个目标与每个因果可见源之间的缩放余弦相似度计算 Softmax 权重。相对于序列长度 T,其复杂度为二次方(O(T2))。
- WANGO (Windowed Aggregation/窗口聚合): 一种线性复杂度变体(对于固定的窗口/特征维度,复杂度为 O(T))。
- 近期源(窗口): 对位于末尾窗口 Li 内的源,使用与 TANGO 相同的指数余弦相似度得分。
- 较旧的源(前缀): 使用确定性的正特征映射 ϕ(u) 来计算得分。这使得模型能够维护变换后的 Key 及其与门控向量外积的运行累加和。较旧源的贡献通过这些运行累加和进行计算,从而避免了需要重新评估所有过去的配对。
实验设置
作者将 TANGO 和 WANGO 与四个基准模型进行了对比:
- Recurrent Transformer++: 重复使用四次的标准 Transformer 模块(共享参数)。
- Untied Transformer++: 四个独立参数化的 Transformer 模块。
- Full-attention GAU (Gated Attention Unit/全注意力门控单元): 使用全成对注意力的四个独立模块。
- FLASH: 结合了局部成对注意力和全局线性注意力摘要的四个独立模块。
约束条件:
- 所有模型的非嵌入参数量均匹配(约 44.3M 个独立参数)。
- 模型使用相同的训练数据顺序、初始化种子和训练预算进行训练。
- 实验在 FineWeb-Edu(自然语言)、Lean(形式化证明与源代码)以及 DeepMind Mathematics 上进行。
- 比较过程涵盖了 4 块和 8 块序列应用的场景。
核心贡献
- TANGO 模型: 引入了一种统一的模块,通过跨 Token 门控更新,将自注意力和前馈网络合二为一,其中源 Token 生成的 SwiGLU 门控用于调制目标特征。
- WANGO 模型: 开发了 TANGO 的线性复杂度变体,在保持与 TANGO 相同参数化的同时,保留了局部窗口内的精确成对交互,并使用特征映射前缀和处理较旧的源。
- 严谨对比: 通过严格匹配参数量、训练目标和数据顺序,提供了对六种架构的受控比较,从而隔离了架构设计与训练方差的影响。
结果
基准测试表现
- FineWeb-Edu:
- TANGO 取得了整体最低的平均验证负对数似然(NLL)(3.293)。
- WANGO 在所有具有线性序列长度缩放能力的架构中取得了最低的 NLL(3.355),超越了 FLASH (3.465) 和 Recurrent Transformer++ (3.392)。
- 值得注意的是,在接近相同的解析前向传递乘加运算(MAC)计数下(对于 8,192 个 Token 的序列,分别为 1.90T vs 1.94T),WANGO 的 NLL 低于 Recurrent Transformer++。
- Lean (源代码与证明补全):
- TANGO 获得了最低的联合 NLL (1.719),以及在源代码建模和证明补全任务中各自最低的 NLL。
- WANGO (2.271) 优于 FLASH (2.377),但在 4 块实验中并未优于 Recurrent Transformer++ (2.530);然而,在 8 块实验中,WANGO (2.231) 确实优于 Recurrent Transformer++ (2.364)。
- DeepMind Mathematics:
- TANGO 实现了最低的验证 NLL (0.473)。
- WANGO (0.487) 的表现优于 Recurrent Transformer++ (0.537) 和 Untied Transformer++ (0.556)。
计算复杂度
- TANGO: 随序列长度呈二次方增长(O(T2))。对于 8,192 个 Token 的序列,需要约 9.44T MACs。
- WANGO: 对于固定维度,随序列长度呈线性增长(O(T))。对于同样的序列,需要约 1.90T MACs。
- Recurrent Transformer++: 二次方复杂度,但由于参数共享且缺乏完整的跨 Token 门控聚合,其 MACs 较低(约 1.94T)。
意义与声明
论文声称,在所评估的参数规模下,TANGO 在自然语言、形式化代码和数学推理任务中,实现了测试架构中的最优验证似然度。
作者强调,WANGO 证明了在匹配解析运算量和参数大小时,线性复杂度模型可以超越二次复杂度基准模型(如 Recurrent Transformer++),特别是在 FineWeb-Edu 上。具体而言,WANGO 通过运行累加和聚合 SwiGLU 门控的能力,使其能够在保持线性缩放的同时维持高性能。
作者明确说明了其声明的局限性:
- 结果是基于教师强制(teacher forcing)下的验证 NLL,并未衡量生成的答案准确率、证明接受率或实际训练速度。
- 实验并未建立超出测试参数量(约 70M 总参数)之外的缩放法则(scaling laws)。
- 不应将 TANGO 较低的 NLL 解释为速度或内存优势,因为它会产生显著更高的计算成本。
总而言之,本文提出了一种新型的架构模块,通过跨 Token 门控统一了注意力和前馈机制,证明了这种方法能带来更优的预测似然度,而其线性复杂度变体(WANGO)为长上下文应用提供了一个具有竞争力的替代方案。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。