Flowformer（ICML 2022）

Flowformer：利用守恒流线性化Transformer

:triangular_flag_on_post:新闻（2024.07）Mobile-Attention，一个为移动设备量身定制的Flowformer版本，已在ICML 2024发表。注意力代码可以在这里找到。只需将标准注意力机制替换为我们的Mobile-Attention，就可以获得一个更快的模型。

Transformer在多个领域取得了令人印象深刻的成功。然而，注意力机制的二次复杂度严重阻碍了Transformer处理大量token和扩展到更大模型的能力。为了追求线性复杂度和任务通用的基础模型，我们提出了Flowformer [论文]，具有以下优点：

相对于序列长度的线性复杂度，可以处理极长序列（超过4k个token）
无特定归纳偏置，纯粹源自流网络理论
任务通用，在**$\color{red}{\text{长序列、视觉、自然语言处理、时间序列、强化学习}}$**等任务中表现出色。

流注意力设计

我们将注意力机制转化为流网络，信息流通过学习的流容量（注意力）从源（值）聚合到汇（结果）。

通过在源和汇两个方面进行守恒，我们可以将竞争引入流注意力设计，避免琐碎的注意力，体现了"固定资源会导致竞争"的精神。

图1. 具有竞争和分配机制的流注意力。

快速开始

请参考不同文件夹获取详细的实验说明。

注意：我们在为不同任务配置环境时遇到了很多困难。如果您在解决环境问题时也遇到困难，欢迎联系我们并讨论。
基准测试列表

核心代码：见Flow_Attention.py
GPT风格的Pytorch模块：见Flowformer_TorchModule
LRA中的长序列建模：见Flowformer_LRA
ImageNet-1K中的视觉识别：见Flowformer_CV
WikiText-103中的语言建模：见Flowformer_NLP
UEA中的时间序列分类：见Flowformer_TimeSeries
D4RL中的强化学习：见Flowformer_RL
CUDA加速版本

主要结果

详细结果请参见[论文]，包括近20个对比基准。

任务	指标	Flowformer	Performer	Reformer	普通 Transformer
长序列建模（LRA）	平均准确率（%）$\uparrow$	56.48	51.41	50.67	OOM
视觉识别（ImageNet-1K）	Top-1准确率（%）$\uparrow$	80.6	78.1	79.6	78.7
语言建模（WikiText-103）	困惑度$\downarrow$	30.8	37.5	33.6	33.0
时间序列分类（UEA）	平均准确率（%）$\uparrow$	73.0	71.5	71.9	71.9
离线强化学习（D4RL）	平均奖励$\uparrow$ 平均偏差$\downarrow$	73.5 $\pm$ 2.9	63.8 $\pm$ 7.6	63.9 $\pm$ 2.9	72.2 $\pm$ 2.6

普通Transformer在强化学习中指Decision Transformer。

注意力可视化

图2. 注意力可视化。Flowformer能够成功捕捉到关键部分。

引用

如果您觉得本仓库有用，请引用我们的论文。

@inproceedings{wu2022flowformer,
  title={Flowformer: Linearizing Transformers with Conservation Flows},
  author={Haixu Wu and Jialong Wu and Jiehui Xu and Jianmin Wang and Mingsheng Long},
  booktitle={International Conference on Machine Learning},
  year={2022}
}

联系方式

如果您有任何问题或想使用代码，请联系wuhx23@mails.tsinghua.edu.cn。