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

任务通用是基础模型研究的核心目标之一，同时也是深度学习研究通向高级智能的必经之路。近年来，得益于注意力机制的通用关键建模能力，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，对应）。

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

3.2 Flow-Attention

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

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

不难得到，未经过守恒之前，对于第个汇，其流入的信息量为：。为了固定每个汇流入的信息量为单位1，我们将作为归一化引入信息流（注意力权重）的计算。经过归一化之后，第个汇的流入信息量为：

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

（2）源（V）的流出守恒：与前述过程类似，未经过守恒之前，对于第个源，其流出的信息量为：。为了固定每个源流出的信息量为单位1，我们将作为归一化引入信息流（注意力权重）的计算。经过归一化之后，第j个源的流出信息量为：。此时，由于源的流出守恒，各个汇（）之间存在天然的竞争关系，我们计算此时每个汇（）接受的信息量，即可得到：竞争情况下，每个结果所需要最终所接受的信息量。

（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中「没有引入新的归纳偏好」设计方式的优势。

上述可视化表明，通过Flow-Attention将竞争引入注意力机制设计可以有效避免平凡注意力。更多可视化实验可见论文。

6. 总结

本文提出的Flowformer通过将网络流中的守恒原理引入设计，自然地将竞争机制引入到注意力计算中，有效避免了平凡注意力问题，在实现线性复杂度的同时，保持了标准Transformer的通用性。Flowformer在长序列、视觉、自然语言、时间序列、强化学习五大任务上取得优秀效果。此外，Flowformer中「无特殊归纳偏好」的设计理念也对通用基础架构的研究具有一定的启发性。在未来工作中，我们将进一步探索Flowformer在大规模预训练上的潜力。