当前位置: 首页 > 科技观察

共同任务!清华提出骨干网络Flowformer实现线性复杂度|ICML2022

时间:2023-03-21 22:57:04 科技观察

任务泛化是基础模型研究的核心目标之一,也是深度学习研究通往高级智能的必由之路。近年来,得益于注意力机制的通用关键建模能力,Transformer在多个领域都有不俗的表现,逐渐呈现出通用架构的趋势。然而,随着序列长度的增加,标准注意力机制的计算呈现二次方复杂度,严重阻碍了其在长序列建模和大型模型中的应用。为此,清华大学软件学院团队针对这一关键问题进行了深入探讨,提出了一种任务通用的线性复杂度骨干网络Flowformer,在保持标准Transformer通用性的同时,将其复杂度降为线性。论文发表后被ICML2022接受。作者名单:吴海旭、吴家龙、徐杰辉、王建民、龙明生链接:https://arxiv.org/pdf/2202.06258.pdf代码:https://github.com/thuml/Flowformer与标准Transformer对比,本文提出的Flowformer模型具有以下特点:线性复杂度,可以处理数千个输入序列;不引入新的归纳偏好,保持原有注意力机制的通用建模能力;任务通用,在长序列,视觉,自然,语言,时间序列,强化学习五大任务上取得了优异的成绩。1、问题分析标准的attention机制的输入包括queries()、keys()和values()三部分,其计算方法如下:其中为attention权重矩阵,最终计算结果为通过加权融合得到。上述过程的计算复杂度为。需要注意的是,对于多项式矩阵的乘法问题,经典算法中已有很多研究。特别是对于attention机制,我们可以利用矩阵乘法的结合律来实现优化,比如可以将原来的二次复杂度降为线性复杂度。但是注意机制中的功能使得无法直接应用结合律。因此,如何去除attention机制中的function是实现线性复杂度的关键。然而,最近的许多研究表明,功能在避免琐碎的注意力学习方面起着关键作用。综上所述,我们期待一种模型设计方案能够实现以下目标:(1)去除功能;(2)避免琐碎的注意;(3)保持模型的通用性。2.Motivation针对目标(1),在之前的工作中,经常使用核方法来代替函数,即近似注意力计算(这是一个非线性函数),但直接去除它会引起普通注意力。为此,对于目标(2),之前的工作不得不引入一些归纳偏差,这限制了模型的通用性,因此不满足目标(3),例如cosFormer中的局部性假设等。Softmax中的竞争机制为了达到上述目标,我们从分析的基本性质入手。我们注意到,最初提出将“赢家通吃”最大化操作扩展为可微分形式。因此,得益于其内部的“竞争”机制,它可以区分令牌之间的注意力权重,从而避免琐碎注意力的问题。基于以上考虑,我们尝试将竞争机制引入到注意力机制设计中,以避免核方法分解带来的琐碎注意力问题。网络流量中的竞争机制我们关注的是,在图论中经典的流量网络模型中,“守恒”是一个重要的现象,即每个节点的流入等于流出。受“在资源固定的情况下,必然引起竞争”这一事实的启发,本文尝试从网络流的角度重新分析经典注意力机制中的信息流,将竞争引入注意力机制设计中通过守恒属性来避免琐碎的浓度问题。3.Flowformer3.1网络流视角下的attention机制attention机制内部:信息流可以表示为:从source(source,对应)到sink基于学习到的流容量(flowcapacity,对应注意权重)(下沉,对应)。在attention机制之外,source(v)信息来自上层网络,sink(R)信息也会提供给下层前馈层。3.2Flow-Attention基于以上观察,我们可以分别从流入和流出的角度控制注意力机制与外部网络的交互,从而实现“固定资源”,从而分别引起源和汇内部的竞争避免琐碎的注意力。不失一般性,我们将注意力机制与外部网络的交互信息量设置为默认值1。(1)sink的流入守恒(R):不难得到,在守恒之前,对于第sink,其流入的信息量为:。为了将每次流入的信息量固定为单元1,我们将引入信息流(注意力权重)的计算作为归一化。归一化后,第th个汇的流入信息量为:此时,由于汇的流入守恒,源之间存在天然竞争关系(V),我们计算每个源给出的信息(v)此时可获得的信息量:在竞争的情况下,每个来源提供的信息量也代表了每个来源的重要性。(2)源(V)的流出守恒:与上述过程类似,在守恒前,对于第th个源,流出信息量为:。为了将每个源流出的信息量固定为单元1,我们将引入信息流(注意力权重)的计算作为归一化。归一化后第j个源的输出信息量为:。此时由于source的流出量守恒,各个sink()之间存在天然的竞争关系,我们计算此时各个sink()接收到的信息量,可以得到:收到的信息。(3)总体设计基于以上结果,我们设计了如下Flow-Attention机制,具体包括三个部分:Competition,Aggregation,andAllocation:Competition引入竞争机制突出重要信息;聚合线性复杂度是基于矩阵结合律实现的;分配通过引入竞争机制来控制传递给下一层的信息量。上述过程中的所有操作都是线性复杂度。同时,Flow-Attention的设计仅仅依靠网络流中的守恒原理重新整合信息流,因此没有引入新的归纳偏好,保证了模型的普适性。通过将标准Transformer中的二次复杂度Attention替换为Flow-Attention,得到Flowformer。4.实验本文在标准数据集上进行了广泛的实验:涵盖长序列、视觉、自然语言、时间序列、强化学习五大任务;检查标准(Normal)和自回归任务(Causal)两种注意力机制类型。涵盖各种序列长度(20-4000)的输入案例。对比各领域经典模型、主流深度模型、Transformer及其变体等各种baseline方法。如下表所示,Flowformer在五项主要任务上表现出色,验证了模型的通用性。详细的实验结果可以参考论文。5.分析为了进一步说明Flowformer的工作原理,我们对ImageNet分类任务中的attention(对应Flow-Attention)进行了可视化实验,从中可以发现如果只使用核方法进行分解,比如LinearTransformer,会使模型分心,无法有效捕捉关键区域;经典的Transformer和Flowformer都可以准确捕捉图像的关键位置,但后者在计算复杂度上有优势;cosFormer在注意力机制中引入一维局部性假设,对语言任务效果显着。但在图像中(将2D数据展开为1D序列),如果不将局部性假设扩展到二维,它就无法适应视觉任务。这也印证了Flowformer中“无新归纳偏好”设计方式的优势。上面的可视化表明,通过Flow-Attention将竞争引入attention机制设计,可以有效避免trivialattention。更多可视化实验可以在论文中找到。6.总结本文提出的Flowformer将网络流中的守恒原理引入到设计中,自然地将竞争机制引入attention计算中,有效避免了普通attention问题,在保持标准Transformer的同时实现了线性复杂度。多功能性。Flowformer在长序列、视觉、自然语言、时序、强化学习五大任务上都取得了优异的成绩。此外,Flowformer中“无特殊感应偏好”的设计理念也对通用基础设施的研究有所启发。在未来的工作中,我们将进一步探索Flowformer在大规模预训练方面的潜力。