[论文解读] Luna: Linear Unified Nested Attention

Xuezhe Ma, Xiang Kong+5 more|arXiv (Cornell University)|Jun 3, 2021

Topic Modeling参考文献 48被引用 49

一句话总结

Luna 通过使用两个嵌套的线性注意力函数并附带一个固定长度的输入序列 P，来实现线性时间的注意力机制，从而在保持在基准测试中的竞争性能的同时实现对长序列的高效建模。

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. Specifically, with the first attention function, Luna packs the input sequence into a sequence of fixed length. Then, the packed sequence is unpacked using the second attention function. 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

研究动机与目标

• 需要扩展 Transformer 以处理长序列，因为注意力的二次复杂度成为瓶颈。
• 提出 Luna 通过使用固定长度输入 P 的打包与解包注意力来实现线性时间/空间。
• 在保持高效性和上下文容量的同时，实现可变长度和自回归注意力。
• 展示 Luna 在长上下文建模、机器翻译以及大规模预训练/微调中的有效性。

提出的方法

• 引入打包注意力，将上下文 C 压缩到固定长度序列 P，使用查询 P。
• 引入解包注意力，使用 X 作为 Attn(X, Y_P) 的查询，将打包表示扩展回原始查询长度。
• 定义 LunaAttn(X, P, C)，同时输出 Y_X 和 Y_P，在层间传递上下文信息。
• 堆叠 Luna 层，使用 LayerNorm 和 FFN，对 Y_X、Y_P 进行归一化，对 Y_X 进行 FFN。
• 在一个层内可选地共享 W_Q、W_K、W_V 以减少参数量（W_K 与 W_V 绑定）。
• 给出 Luna 因果注意力变体，在打包阶段使用定制的激活函数（elu+1 或 softplus）以维持线性复杂度，在解包阶段使用 softmax。

实验结果

研究问题

• RQ1 如何重新表述注意力以在不牺牲长距离上下文建模的情况下实现线性时间/空间？
• RQ2 一个额外的固定长度输入序列 P 是否能够捕获足够的上下文信息以适应可变长度输入？
• RQ3 Luna 是否支持因果注意力和自回归解码同时保持高效？
• RQ4 相对于强基线，Luna 在长上下文基准、机器翻译和大规模预训练中的表现如何？

主要发现

• Luna 在长上下文序列建模、机器翻译和预训练任务中，表现与完整注意力基线及其他高效注意力方法相当甚至更好。
• 由于打包/解包注意力，Luna 展现出线性时间和内存复杂性，并与若干基线相比显示出效率提升。
• 较小的投影长度（如 l = 16）在 LRA 基准测试和机器翻译任务中提供了强大的性能。
• 编码到 P 的上下文信息相较于非上下文化的 P 在机器翻译实验中提升了性能。
• 在 WMT’14 EN→DE 的机器翻译中，使用 softplus 激活和共享键/值的 Luna 具有具有竞争力的 BLEU 得分，带上下文的 P 相较于非上下文化的 P 提供了显著提升。
• 预训练实验表明 Luna 在某些问答和分类任务上可以达到或超过基线模型，尤其是在数据规模较小的情形；全模型容量对大数据结果有影响。