摘要

基于注意力机制的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。