Flowformer:基于守恒流的 Transformer 线性化方法
机器学习
2022-06-17 v2 人工智能
摘要
基于注意力机制的 Transformer 在多个领域取得了令人瞩目的成功。然而,注意力机制具有二次复杂度,严重阻碍了 Transformer 处理大量 token 并扩展到更大模型。以往方法主要利用相似度分解和矩阵乘法的结合律来设计线性时间注意力机制。它们通过重新引入局部性等归纳偏置来避免注意力退化为平凡分布,但代价是牺牲了模型的通用性与表达能力。本文基于流网络理论,在不依赖特定归纳偏置的情况下对 Transformer 进行线性化。我们将注意力视为通过学习的流容量(注意力)从源(值)聚合到汇(结果)的信息流。在此框架下,我们将流守恒性质应用于注意力,提出具有线性复杂度的 Flow-Attention 机制。通过分别对汇的入流进行源竞争守恒、对源的出流进行汇分配守恒,Flow-Attention 无需使用特定归纳偏置即可天然生成信息丰富的注意力。借助 Flow-Attention,Flowformer 在长序列、时间序列、视觉、自然语言与强化学习等广泛领域的线性时间内均取得强劲性能。代码与设置可见于该仓库:https://github.com/thuml/Flowformer。
引用
@article{arxiv.2202.06258,
title = {Flowformer: Linearizing Transformers with Conservation Flows},
author = {Haixu Wu and Jialong Wu and Jiehui Xu and Jianmin Wang and Mingsheng Long},
journal= {arXiv preprint arXiv:2202.06258},
year = {2022}
}