任务通用！清华提出主干网络Flowformer，实现线性复杂度

新智元报道

作者：吴海旭

编辑：好困

【新智元导读】近年来，Transformer方兴未艾，但是其内在的二次复杂度阻碍了它在长序列和大模型上的进一步发展。清华大学软件学院机器学习实验室从网络流理论出发，提出任务通用的线性复杂度主干网络Flowformer，在长序列、视觉、自然语言、时间序列、强化学习五大任务上取得优秀效果。

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

近年来，得益于注意力机制的通用关键建模能力，Transformer在众多领域中表现优异，逐渐呈现出通用架构的趋势。但是随着序列长度的增长，标准注意力机制的计算呈现二次复杂度，严重阻碍了其在长序列建模与大模型中的应用。

为此，来自清华大学软件学院的团队深入探索了这一关键问题，提出了任务通用的线性复杂度主干网络Flowformer，在保持标准Transformer的通用性的同时，将其复杂度降至线性，论文被ICML 2022接受。

作者列表：吴海旭，吴佳龙，徐介晖，王建民，龙明盛

链接：https://arxiv.org/pdf/2202.06258.pdf

代码：https://github.com/thuml/Flowformer

相比于标准Transformer，本文提出的Flowformer模型，具有以下特点：

线性复杂度，可以处理数千长度的输入序列；
没有引入新的归纳偏好，保持了原有注意力机制的通用建模能力；
任务通用，在长序列、视觉、自然语言、时间序列、强化学习五大任务上取得优秀效果。

1. 问题分析

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

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

注意到，对于多项矩阵的连乘问题，在经典算法中已有较多研究。特别地，对于注意力机制，我们可以利用矩阵乘法的结合律来实现优化，如，即可将原本的二次复杂度降至线性。但是注意力机制中的函数使得无法直接应用结合律。

因此，如何移除注意力机制中的函数是实现线性复杂度的关键。但是，近期的众多工作证明，函数在避免平凡注意力学习上起到了关键性作用。

综上，我们期待一种模型设计方案，实现以下目标：（1）移除函数；（2）避免平凡注意力；（3）保持模型的通用性。

2. 动机

针对目标（1），在之前的工作中，往往使用核方法来替代函数，即通过近似注意力计算（为非线性函数），但直接去掉会造成平凡注意力。

为此，针对目标（2），之前工作不得不引入一些归纳偏好，这限制了模型的通用性，因此不满足目标（3），比如cosFormer中的局部性假设等。

Softmax中的竞争机制

为满足上述目标，我们从的基本性质出发进行分析。我们注意到，最初被提出是用于：将「赢者通吃」的取极大值操作扩展为可微分形式。因此，得益于其内在的「竞争」机制，它可以使各个token之间的注意力权重差异化，从而避免了平凡的注意力的问题。

基于以上考虑，我们试图将竞争机制引入注意力机制设计，从而避免核方法分解带来平凡注意力问题。

网络流中的竞争机制

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

3. Flowformer

3.1 网络流视角下的注意力机制

在注意力机制内部：信息流动可以表示为：从源（source，对应）基于学习到的流容量（flow capacity，对应注意力权重）汇聚至汇（sink，对应）。

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

3.2 Flow-Attention

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

（1）汇（）的流入守恒：

不难得到，未经过守恒之前，对于第个汇，其流入的信息量为：。

为了固定每个汇流入的信息量为单位1，我们将作为归一化引入信息流（注意力权重）的计算。经过归一化之后，第个汇的流入信息量为：

。

此时，由于汇的流入守恒，各个源（）之间存在天然的竞争关系，我们计算此时每个源（）给出的信息量，即可得到：竞争情况下，每个源所提供的信息量，这也代表着每个源的重要性。

（2）源（）的流出守恒：

与前述过程类似，未经过守恒之前，对于第个源，其流出的信息量为：。

为了固定每个源流出的信息量为单位1，我们将作为归一化引入信息流（注意力权重）的计算。经过归一化之后，第个源的流出信息量为：

。

此时，由于源的流出守恒，各个汇（）之间存在天然的竞争关系，我们计算此时每个汇（）接受的信息量，即可得到：竞争情况下，每个结果所需要最终所接受的信息量。

（3）整体设计

基于上述结果，我们设计如下Flow-Attention机制，具体包含竞争（Competition）、聚合（Aggregation）、分配（Allocation）三部分：

其中Competition将竞争机制引入中，突出重要的信息；Aggregation基于矩阵结合律实现线性复杂度；Allocation通过将竞争机制引入，控制传递到下一层的信息量。

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

4. 实验

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

覆盖了长序列、视觉、自然语言、时间序列、强化学习五大任务；
考察了标准（Normal）和自回归任务（Causal）两种注意力机制类型。
涵盖了多种序列长度的输入情况（20-4000）。
对比了各领域经典模型、主流深度模型、Transformer及其变体等多种基线方法。

如下表所示，Flowformer在五大任务上均表现优秀，验证了模型的通用性。详细实验结果请见论文。

5. 分析

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

如果仅仅使用核方法进行分解，如Linear Transformer，会造成模型注意力分散，无法有效捕捉到关键区域；
经典Transformer和Flowformer均可以准确捕捉到图像的关键位置，但是后者在计算复杂度上具有优势；
cosFormer在注意力机制中引入一维局部性假设，在语言任务上效果突出。但是在图像（将2D数据展开成1D序列）中，如果不将局部性假设扩展至二维，则无法适配视觉任务。这也印证了Flowformer中「没有引入新的归纳偏好」设计方式的优势。