11 天前
Flowformer:基于保形流的Transformer线性化
Haixu Wu, Jialong Wu, Jiehui Xu, Jianmin Wang, Mingsheng Long

摘要
基于注意力机制的Transformer在多个领域取得了令人瞩目的成功。然而,注意力机制本身具有二次方复杂度,严重制约了Transformer处理大量标记(tokens)的能力,并限制了其向更大模型的扩展。此前的方法主要通过利用矩阵乘法的相似性分解与结合律,设计出线性时间复杂度的注意力机制。这些方法通常通过重新引入局部性等归纳偏置(inductive biases)来避免注意力退化为平凡分布,但这一做法以牺牲模型的通用性与表达能力为代价。本文基于流网络理论(flow network theory),提出了一种不依赖特定归纳偏置的线性化Transformer架构。我们将注意力机制建模为从源点(值,values)到汇点(结果,results)的信息流,该信息流通过学习得到的流容量(即注意力权重)进行聚合。在此框架下,我们引入流守恒(flow conservation)的性质,提出一种具有线性复杂度的Flow-Attention机制。通过分别保证汇点的输入流守恒以实现源点间的竞争,以及源点的输出流守恒以实现汇点的分配,Flow-Attention能够自然地生成具有信息量的注意力分布,而无需依赖任何特定的归纳偏置。得益于Flow-Attention机制,Flowformer在多个广泛领域中均实现了线性时间下的优异性能,涵盖长序列建模、时间序列分析、计算机视觉、自然语言处理以及强化学习等任务。相关代码与实验配置已开源,详见项目仓库:https://github.com/thuml/Flowformer。