机器之心 | 预训练无需注意力，扩展到4096个token不成问题，与BERT相当

news2024/11/15 9:43:38

本文来源公众号“机器之心”，仅用于学术分享，侵权删，干货满满。

原文链接：预训练无需注意力，扩展到4096个token不成问题，与BERT相当

本文提出了双向门控 SSM（BiGS）模型，结合基于状态空间模型（SSM）的 Routing 层和基于乘法门的模型架构，在不使用注意力的情况下能够复制 BERT 预训练结果，并可扩展到 4096 个 token 的长程预训练，不需要近似。

Transformer 作为 NLP 预训练模型架构，能够有效的在大型未标记的数据上进行学习，研究已经证明，Transformer 是自 BERT 以来 NLP 任务的核心架构。

最近的工作表明，状态空间模型（SSM）是长范围序列建模有利的竞争架构。SSM 在语音生成和 Long Range Arena 基准上取得了 SOTA 成果，甚至优于 Transformer 架构。除了提高准确率之外，基于 SSM 的 routing 层也不会随着序列长度的增长而呈现二次复杂性。

本文中，来自康奈尔大学、 DeepMind 等机构的研究者提出了双向门控 SSM （BiGS），用于无需注意力的预训练，其主要是将 SSM routing 与基于乘法门控（multiplicative gating）的架构相结合。该研究发现 SSM 本身在 NLP 的预训练中表现不佳，但集成到乘法门控架构中后，下游准确率便会提高。

实验表明，在受控设置下对相同数据进行训练，BiGS 能够与 BERT 模型的性能相匹配。通过在更长的实例上进行额外预训练，在将输入序列扩展到 4096 时，模型还能保持线性时间。分析表明，乘法门控是必要的，它修复了 SSM 模型在变长文本输入上的一些特定问题。

论文地址：https://arxiv.org/pdf/2212.10544.pdf

方法介绍

SSM 通过以下微分方程将连续输入 u (t) 与输出 y (t) 联系起来：

对于离散序列，SSM 参数被离散化，其过程可以近似为：

这个方程可以解释为一个线性 RNN，其中 x_k 是一个隐藏状态。y 也可以用卷积计算：

Gu 等人展示了一种在神经网络中使用 SSM 的有效方法，他们开发了参数化 A 的方法，称为 HiPPO，其产生了一个稳定而高效的架构，称为 S4。这保留了 SSM 对长期序列建模的能力，同时比 RNN 训练更有效。最近，研究人员提出了 S4 的简化对角化版本，它通过对原始参数更简单的近似实现了类似的结果。在高层次上，基于 SSM 的 routing 为神经网络中的序列建模提供了一种替代方法，而无需二次计算的注意力成本。

预训练模型架构

SSM 能取代预训练中的注意力吗？为了回答这个问题，该研究考虑了两种不同的架构，如图 1 所示的堆叠架构（STACK）和乘法门控架构（GATED）。

具有自注意力的堆叠架构相当于 BERT /transformer 模型，门控架构是门控单元的双向改编，最近也被用于单向 SSM。带有乘法门控的 2 个序列块（即前向和后向 SSM）夹在前馈层中。为了进行公平比较，门控架构的大小保持与堆叠架构相当。

图 1：模型变量。STACK 是标准 transformer 架构，GATED 为基于门控单元。对于 Routing 组件（虚线），该研究同时考虑双向 SSM（如图所示）和标准自注意力。门控（X）表示逐元素乘法。

实验结果

预训练

表 1 显示了 GLUE 基准测试中不同预训练模型的主要结果。BiGS 在 token 扩展上复制了 BERT 的准确率。这一结果表明，在这样的计算预算下，SSM 可以复制预训练 transformer 模型的准确率。这些结果明显优于其他基于非注意力的预训练模型。想要达到这个准确率，乘法门控是必要的。在没有门控的情况下，堆叠 SSM 的结果明显更差。为了检查这种优势是否主要来自于门控的使用，本文使用 GATE 架构训练了一个基于注意力的模型；然而，结果显示该模型的效果实际上低于 BERT。