[论文解读] Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention

Yunyang Xiong, Zhanpeng Zeng+5 more|arXiv (Cornell University)|Feb 7, 2021

Topic Modeling参考文献 60被引用 29

一句话总结

Nyströmformer 引入基于 Nyström 的自注意力近似，达到 O(n) 时间和内存，实现在长序列上高效的 Transformer，同时保持具有竞争力的性能。

ABSTRACT

Transformers have emerged as a powerful tool for a broad range of natural language processing tasks. A key component that drives the impressive performance of Transformers is the self-attention mechanism that encodes the influence or dependence of other tokens on each specific token. While beneficial, the quadratic complexity of self-attention on the input sequence length has limited its application to longer sequences -- a topic being actively studied in the community. To address this limitation, we propose Nyströmformer -- a model that exhibits favorable scalability as a function of sequence length. Our idea is based on adapting the Nyström method to approximate standard self-attention with $O(n)$ complexity. The scalability of Nyströmformer enables application to longer sequences with thousands of tokens. We perform evaluations on multiple downstream tasks on the GLUE benchmark and IMDB reviews with standard sequence length, and find that our Nyströmformer performs comparably, or in a few cases, even slightly better, than standard self-attention. On longer sequence tasks in the Long Range Arena (LRA) benchmark, Nyströmformer performs favorably relative to other efficient self-attention methods. Our code is available at https://github.com/mlpen/Nystromformer.

研究动机与目标

• 动机：在处理长序列的 Transformer 模型中，需要可扩展的自注意力。
• 提出基于 Nyström 的近似，用以替代完整的 softmax 自注意力，达到线性复杂度。
• 开发 landmark（Nyström）策略和一个迭代伪逆方法，以近似 softmax 矩阵。
• 在标准 NLP 基准和长序列任务上评估该方法，以将准确性和效率与基线进行比较。

提出的方法

• 将自注意力表示为 S = softmax(QK^T / sqrt(d_q))，并旨在使用 Nyström 方法将 S 近似为 O(n) 复杂度。
• 引入基于 Q 和 K 的 landmarks，形成一个小的 m x m 矩阵 A_S，通过对 landmarks 进行 softmax 形成；计算其 Moore-Penrose 逆。
• 推导出 S 的 Nyström 形式近似：hat{S} = softmax(Q ~K^T / sqrt(d_q)) A_S^+ softmax(~Q K^T / sqrt(d_q)).
• 用 Segment-means 调整 landmark 选择，以高效地在 O(n) 规模内计算 Q-landmarks 和 K-landmarks。
• 使用迭代方法计算 A_S 的 Moore-Penrose 逆 Z^*，以避免高成本的 SVD，在实践中大约稳定在 6 次迭代。
• 将自注意力输出近似为 hat{S} V，得到一个线性时间的自注意力模块，并通过在 V 上进行深度卷积实现跳跃连接。

实验结果

研究问题

• RQ1 Nyström 基于的 softmax 自注意力近似能否在 NLP 基准上达到与标准注意力相当的准确性？
• RQ2 Nyströmformer 是否在序列长度的线性时间和内存复杂度下保持性能？
• RQ3 与其他高效自注意力方法相比，Nyströmformer 在长序列任务上的表现如何？
• RQ4 在 landmark 选择（Segment-means）和伪逆近似方面对实际训练有哪些权衡？

主要发现

• Nyströmformer 在 GLUE 和 IMDB 上的表现可与基线 BERT 相近或略优，同时获得显著的效率提升。
• 在长距离任务（Long Range Arena）上，Nyströmformer 相对于其他高效自注意力方法（如 Reformer、Linformer、Performer）表现更优。
• 当 landmark 数 m 远小于 n 时，该方法在时间和内存上均按序列长度 n 线性扩展。
• Segment-means landmark 选择实现了 O(n) 的计算，开销很小。
• 实证结果表明 Nyströmformer 能在保持接近标准自注意力的准确性的同时，超越若干线性注意力变体。