Press "Enter" to skip to content

任务通用!清华提出主干网络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在大规模预训练上的潜力。

Be First to Comment

发表评论

您的电子邮箱地址不会被公开。