[论文解读] An Attention Free Transformer

Shuangfei Zhai, Walter Talbott+5 more|arXiv (Cornell University)|May 28, 2021

Advanced Neural Network Applications参考文献 26被引用 42

一句话总结

引入 Attention Free Transformer (AFT)，它用一个非注意力、线性复杂度的模块取代点积自注意力，并包含如 AFT-local 和 AFT-conv 的变体，在通过局部性或卷积权重共享来保持全局连通性。

ABSTRACT

We introduce Attention Free Transformer (AFT), an efficient variant of Transformers that eliminates the need for dot product self attention. In an AFT layer, the key and value are first combined with a set of learned position biases, the result of which is multiplied with the query in an element-wise fashion. This new operation has a memory complexity linear w.r.t. both the context size and the dimension of features, making it compatible to both large input and model sizes. We also introduce AFT-local and AFT-conv, two model variants that take advantage of the idea of locality and spatial weight sharing while maintaining global connectivity. We conduct extensive experiments on two autoregressive modeling tasks (CIFAR10 and Enwik8) as well as an image recognition task (ImageNet-1K classification). We show that AFT demonstrates competitive performance on all the benchmarks, while providing excellent efficiency at the same time.

研究动机与目标

• 阐明由于注意力成本的二次复杂度而需要可扩展的 Transformer 的原因。
• 提出一个可插拔的无注意力模块，在不使用点积注意力的情况下维持全局交互。
• 引入 AFT-local 和 AFT-conv，以利用局部性和空间权重共享，同时维持全局连通性。
• 在图像自回归建模、语言建模和图像分类等任务上展示具有竞争力的性能和效率。

提出的方法

• 用 Q、K、V 之间的交互替代标准的多头自注意力，其中 K 和 V 与学习到的位置偏置结合，结果再与查询逐元素相乘。
• 将 AFT 表述为 Y_t = sigma_q(Q_t) ⊙ [ sum_{t'} exp(K_{t'}+w_{t,t'}) ⊙ V_{t'} / sum_{t'} exp(K_{t'}+w_{t,t'}) ]，其中 w 是学习得到的成对位置偏置。
• 提供参数分解的参数化 w_{t,t'} = u_t^T v_{t'} 以减少参数量。
• 引入 AFT-local（在 s 有界窗口内的局部 w）和 AFT-conv（类似 CNN 的空间权重共享）变体。
• 提供 AFT-simple 作为一个高效极端变体，且无位置偏置（s=0）。
• 讨论参数化和再参数化（例如 w-分解、w 归一化）以提升训练稳定性和性能。

实验结果

研究问题

• RQ1 无注意力机制是否能够在自回归图像建模、语言建模和图像分类任务中达到甚至超过 Transformers？
• RQ2 局部性和权重共享的变体（AFT-local、AFT-conv）在保持全局连通性的同时，是否在效率和准确度上提供实际提升？
• RQ3 对位置偏置参数进行分解会如何影响模型性能和参数量？
• RQ4 与标准及其他高效 Transformer 相比，AFT 变体的复杂度和实际运行时如何？
• RQ5 AFT 变体能否作为插件集成到现有的 Transformer 架构中？

主要发现

模型	L	d	h	训练损失	测试损失	每秒迭代次数	GB/GPU
PixelCNN	-	-	-	3.08	3.14	-	-
PixelCNN++	-	-	-	-	2.92	-	-
PixelSNAIL	-	-	-	-	2.85	-	-
Sparse Transformer strided	128	256	2	-	2.80	-	-
Image Transformer local2d	12	512	4	-	2.90	1.61	22.3
Transformer	12	512	4	2.90	2.88	1.35	30.6
Transformer	24	256	2	2.90	2.86	1.36	30.4
AFT-local-256	12	512	1	2.78	2.80	1.68	11.4
AFT-local-256	24	256	1	2.75	2.74	1.67	12.8
AFT-simple	24	256	1	2.82	2.89	2.15	9.5

• AFT 的变体在 CIFAR10、Enwik8 和 ImageNet-1K 上展现出具竞争力的性能，同时具有更有利的效率。
• AFT-local 在 CIFAR10 图像自回归建模中在速度和内存使用方面均超越了若干 Transformer 基线。
• 分解的位置信偏置参数化相较于非分解的 w 大幅减少参数量并提升性能。
• AFT-simple 在效率更高的前提下提供了强结果，紧接着或甚至超过了一些视觉 Transformer 的基线。
• AFT-conv 在 ImageNet Top-1 精度上显著优于基线配置，同时通过卷积样式的权重共享维持全局连通性。
• 在各种任务中，AFT 变体相对于序列长度和特征维度表现出线性或近线性时间/空间复杂度，与标准注意力的二次成本形成对比。