Luna: Linear Unified Nested Attention
Xuezhe Ma ISI, USC、Xiang Kong∗ LTI, CMU、Sinong Wang∗ Facebook AI、Chunting Zhou LTI, CMU、Jonathan May ISI, USC、Hao Ma, Luke Zettlemoyer Facebook AI {haom, Equal contribution.
Abstract
The quadratic computational and memory complexities of the Transformer’s attention mechanism have limited its scalability for modeling long sequences. In this paper, we propose Luna, a linear unified nested attention mechanism that approximates softmax attention with two nested linear attention functions, yielding only linear (as opposed to quadratic) time and space complexity. As compared to a more traditional attention mechanism, Luna introduces an additional sequence with a fixed length as input and an additional corresponding output, which allows Luna to perform attention operation linearly, while also storing adequate contextual information. We perform extensive evaluations on three benchmarks of sequence modeling tasks: long-context sequence modeling, neural machine translation and masked language modeling for large-scale pretraining. Competitive or even better experimental results demonstrate both the effectiveness and efficiency of Luna compared to a variety of strong baseline methods including the full-rank attention and other efficient sparse and dense attention methods. The implementation of our model is available at https://github.com/XuezheMax/fairseq-apollo.
中文速览
Transformer在处理长序列时，注意力机制的计算量和内存占用都随序列长度呈平方增长，严重制约了其实际应用。为此，研究者提出了Luna（线性统一嵌套注意力机制），核心思路是引入一个长度固定的额外序列，通过"打包"和"解包"两个嵌套的线性注意力操作来近似原始的softmax全注意力，从而将时间和空间复杂度从平方降至线性。在长文本建模、神经机器翻译和大规模预训练掩码语言模型三项基准任务上的大量实验表明，Luna在速度和内存效率上相较于标准注意力及多种稀疏/低秩高效注意力方法均有显著提升，同时性能持平甚至更优，即便将固定序列长度设为仅16也能保持出色表现。这一工作的重要意义在于，Luna在保持强大建模能力的同时真正实现了线性复杂度，且支持变长序列和自回归解码，为大规模长序列Transformer提供了一个实用、通用的替代方案。
原文 arXiv:2106.01540;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2106.01540v2