ICML 2022 | Flowformer: 任务通用的线性复杂度Transformer

本文介绍本组ICML2022深度学习基础模型方向的最新工作：Flowformer: Linearizing Transformers with Conservation Flows。受网络流理论启发，本文提出任务通用的骨干网络Flowformer，实现线性复杂度，在长序列、视觉、自然语言、时间序列、强化学习五大任务上取得优秀效果。

论文链接： https://arxiv.org/pdf/2202.06258.pdf 代码链接： https://github.com/thuml/Flowformer

追求任务通用模型是基础模型研究领域的核心目标之一，也是深度学习研究通向高级智能的必经之路。

近年来，得益于注意力机制（Attention Mechanism）的通用关系建模能力，Transformer在众多领域已经展现出了优秀的效果，并且逐步呈现出通用模型的态势。但是，其核心组件注意力机制随着输入序列的增长呈现二次复杂度，严重阻碍了Transformer在长序列处理上的应用，同时也限制了其扩展至大模型（Big Model）的能力。

为解决上述困难，受到网络流理论（Flow network theory）启发，本文提出了Flowformer模型，具有以下特点：

• 线性复杂度，可以处理数千长度的输入序列；

• 没有引入新的归纳偏好，保持了原有注意力机制的通用建模能力；

• 任务通用，在长序列、视觉、自然语言、时间序列、强化学习五大任务上取得优秀效果。

问题分析2.1 Attention的二次复杂度

经典的注意力机制输入包含queries（ ），keys（ ）和values（ ）三部分， ，其计算方式如下：

其中 为计算的注意力权重矩阵（attention），最终计算结果 为将 进行加权融合所得。不难得到，上述过程计算复杂度为 。

注意到，上述三者相乘的过程可以抽象为经典的矩阵连乘优化问题。这意味着，我们可以利用矩阵乘法的结合律来实现优化，如 ，即可将原本的二次复杂度降至线性 。但是注意力机制中的 函数使得无法直接应用结合律。

2.2 Softmax中的竞争机制

近期的众多工作证明， 函数在避免平凡注意力（Trivial Attention）学习上起到了关键性作用，但它又使得模型不得不面对 相乘带来的二次复杂度。

之前的工作往往使用核方法，即通过 近似注意力（ 为非线性函数），实现可分解的注意力权重计算，进而应用结合律实现线性复杂度。此外，为了避免去掉 函数带来的平凡注意力问题，之前的工作不得不引入一些归纳偏好，这极大限制了模型的通用性，比如cosFormer中的局部性假设等。

在本文中，我们注意到， 最初被提出是用于：将“赢者通吃”的取极大值操作扩展为可微分形式。因此，得益于其内在的“竞争”机制，它可以使各个token之间的注意力权重差异化，从而避免了平凡的注意力的问题。基于以上考虑，我们试图将竞争机制引入注意力机制设计，从而避免核方法分解带来平凡注意力问题。

2.3 网络流中的竞争机制

我们关注到在图论中的经典网络流（Flow network）模型中，“守恒”（Conservation）是一个重要现象，即每个节点的流入量等于流出量。受到“固定资源情况下，必定引起竞争”的启发，在本文中，我们通过网络流视角重新分析经典注意力机制中的信息流动，并通过守恒性质将竞争引入注意力机制设计，以避免平凡注意力问题。

三、Flowformer3.1 网络流视角下的Attention

我们分别从内部和外部两个视角来分析：

注意力机制-内部视角

在注意力机制内部，信息流动可以表示为：从源（source，对应 ）基于学习到的流容量（flow capacity，对应注意力权重）汇聚至汇（sink，对应 ）。其中汇的流入信息量 ，对应着注意力权重矩阵的行和；源流出的信息量 ，对应着注意力矩阵的列和，其计算过程如下：

注意力机制-外部视角

在注意力机制外部，源（ ）的信息来自于上一层网络，汇（ ）的信息也将提供给下面的前馈层。

3.2 Flow-Attention

基于上述观察，我们可以通过分别控制注意力机制与外部网络的交互，来实现“固定资源”，从而分别引起源和汇内部的竞争，以避免平凡注意力。不失一般性，我们将注意力机制与外部网络的交互信息量设置为默认值1。

具体地，我们可以通过以下归一化，分别实现源的流出守恒与汇的流入守恒：

上述权重归一化操作会影响计算所得的流容量，从而实现守恒，可以通过以下公式验证：

至此，通过将守恒引入注意力机制，我们得到了引入竞争之后的源、汇分别的流出、流入量：

注意到，此时 表示汇（ ）流入信息量一定情况下，每个源（ ）提供的信息量，因此代表了源（ ）的重要性； 表示源（ ）流出信息量一定情况下，竞争之后，每个汇（ ）的所得信息量，代表了汇（ ）所需要获取的信息量。

Flow-Attention

基于此，我们设计如下Flow-Attention机制，自然地将竞争机制引入其中，具体包含竞争、聚合、分配三部分：

Flow-Attention的伪代码如下：

Flow-Attention伪代码

上述过程中的所有操作均为线性复杂度。同时，Flow-Attention的设计仅仅依赖于网络流中的守恒原理，对信息流的重新整合，因此并没有引入新的归纳偏好，保证了模型的通用性。将标准Transformer中的二次复杂度Attention替换为Flow-Attention，即得到了Flowformer。

四、实验

我们在标准数据集上进行了广泛的实验：

• 覆盖了长序列、视觉、自然语言、时间序列、强化学习五大任务；

• 考察了标准（Normal）和自回归任务（Causal）两种注意力机制类型。

• 涵盖了多种序列长度的输入情况（20-4000）。

各个数据集的具体情况如下：

实验数据集4.1 主要结果

主要实验结果

Flowformer在五大任务上均表现优秀，验证了模型的通用性。

详细实验结果请见论文，包含20多个基线模型，其中有各领域的经典模型、主流深度模型、Transformer及其变体们。

4.2 注意力可视化

注意力可视化

为了进一步说明Flowformer的工作原理，我们对ImageNet分类任务中的注意力（对应Flow-Attention中的 ）进行了可视化实验，从中可以发现：

• 经典Transformer和Flowformer均可以准确捕捉到图像的关键位置，但后者更加高效；

• Linear Transformer仅仅使用核方法进行分解，失去Softmax之后会造成模型注意力分散，无法有效捕捉到关键区域；

• cosFormer在注意力机制中引入局部性假设，在语言任务上效果突出。但是在图像（将2D数据展开成1D序列）中，如果不将局部性假设扩展至二维，则无法适配视觉任务。这也印证了Flowformer中“没有引入新的归纳偏好”设计方式的优势。

本文深入研究了注意力机制存在的二次复杂度问题，通过将网络流中的守恒原理引入设计，自然地将竞争机制引入到注意力计算中，有效避免了平凡注意力问题。

我们提出的任务通用的骨干网络Flowformer，实现了线性复杂度，同时在长序列、视觉、自然语言、时间序列、强化学习五大任务上取得优秀效果。

在长序列建模应用上，如蛋白质结构预测、长文本理解等，Flowformer具有良好的应用潜力。此外，Flowformer中“无特殊归纳偏好”的设计理念也对通用基础架构的研究具有良好的启发意义。

公众号：【THUML】 作者：吴海旭

Illustration b y Alena Kosareva from icon s8

-The End-

7.6 周三 19:00

快来直播预约，或扫码报名吧！

扫码观看！

本周上新！

关于我“门”

将门是一家以专注于发掘、加速及投资技术驱动型创业公司的新型创投机构，旗下涵盖将门创新服务、将门技术社群以及将门创投基金。

将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”:

bp@thejiangmen.com

点击右上角，把文章分享到朋友圈

⤵一键送你进入TechBeat快乐星球