机器之心报道
编辑：泽南、杜伟

谷歌的新研究 ， 让 transformer 模型的效率有了巨大提升 ， 该方法的核心在于减少注意力机制 。

Transformer 是目前深度学习领域最热门的技术 ， 在语言和视觉领域都取得了很多突破 。 尽管因为模型体量不断增加而受到质疑 ， 但大多数 Transformer 因为输入长度的二次复杂度问题 ， 仍然限于短上下文大小——由于记忆容量有限 ， 不得不抛弃较早的信息 。 这一限制阻止了 Transformer 模型处理长期信息的能力 ， 而后者是许多应用程序的关键属性 。
在学界 ， 已经有很多技术试图通过更有效的注意力机制来扩展上下文能力 。 尽管其中一些方法具有线性理论复杂性 ， 但常规 Transformer 仍然是最先进系统中的主要选择 。 从实际的角度来研究 ， 现有的有效注意力方法至少存在以下缺点之一：

• 性能低下 。 实验表明 ， 通过几个简单的调整来增强普通 Transformer 可以比文献中使用的常见基线强得多（见图 1） 。 与增强的 Transformer 相比 ， 现有的高效注意力方法通常会导致质量显着下降 ， 而这种性能下降超过了它们的效率优势 。
• 计算开销 。 由于有效的注意力方法通常会使 Transformer 层复杂化并需要大量的数据 / 内存格式化操作 ， 因此其理论复杂性和 GPU 或 TPU 等加速器的经验速度之间可能存在不小的差距 。
• 低效率的自回归训练 。 大多数注意力线性化技术在推理中享受快速解码 ， 但在语言建模等自回归任务上训练可能非常慢 。 这主要是由于它们在大量步骤中的 RNN 式顺序状态更新 ， 使其无法在训练期间充分利用现代加速器的优势 。

文章图片

来自康奈尔大学、谷歌大脑的研究人员近日提出了一个新模型 FLASH（Fast Linear Attention with a Single Head） ， 首次不仅在质量上与完全增强的 Transformer 相当 ， 而且在现代加速器的上下文大小上真正享有线性可扩展性 。 与旨在逼近 Transformers 中的多头自注意力 (MHSA) 的现有高效注意力方法不同 ， 谷歌从一个新层设计开始 ， 自然地实现更高质量的逼近 。 FLASH 分两步开发：
首先设置一个更适合有效近似的新层 ， 引入门控机制来减轻自注意力的负担 ， 产生了下图 2 中的门控注意力单元 (Gated Attention Unit, GAU) 。 与 Transformer 层相比 ， 每个 GAU 层更便宜 。 更重要的是 ， 它的质量更少依赖于注意力精度 。 事实上 ， 小单头、无 softmax 注意力的 GAU 与 Transformers 性能相近 。
虽然 GAU 在上下文大小上仍存在二次复杂度 ， 但它削弱了注意力的作用 ， 允许稍后以最小的质量损失进行近似 。
随后作者提出了一种有效的方法来逼近 GAU 中的二次注意力 ， 从而导致在上下文大小上具有线性复杂度的层变体 。 其思路是首先将标记分组为块 ， 然后在一个块内使用精确的二次注意力和跨块的快速线性注意力（如下图 4 所示） 。 在论文中 ， 研究者进一步描述了如何利用此方法自然地推导出一个高效的加速器实现 ， 在实践中做到只需更改几行代码的线性可扩展能力 。
在大量实验中 ， FLASH 在各种任务、数据集和模型尺度上均效果很好 。 FLASH 在质量上与完全增强的 Transformer (Transformer++) 相比具有竞争力 ， 涵盖了各种实践场景的上下文大小 (512-8K) ， 同时在现代硬件加速器上实现了线性可扩展 。
例如 ， 在质量相当的情况下 ， FLASH 在 Wiki-40B 上的语言建模实现了 1.2 倍至 4.9 倍的加速 ， 在 Transformer++ 上 C4 上的掩码语言建模实现了 1.0 倍至 4.8 倍的加速 。 在进一步扩展到 PG-19 (Rae et al., 2019) 之后 ， FLASH 将 Transformer++ 的训练成本降低了 12.1 倍 ， 并实现了质量的显着提升 。