[论文解读] Sparse Sinkhorn Attention

Yi Tay, Dara Bahri+3 more|arXiv (Cornell University)|Feb 26, 2020

Natural Language Processing Techniques参考文献 24被引用 77

一句话总结

Sparse Sinkhorn Attention 引入了一种可微分的排序机制，创建内存高效的稀疏注意力，通过局部窗口实现准全局上下文，并在与原生 Transformer 相当的性能方面具有竞争力。

ABSTRACT

We propose Sparse Sinkhorn Attention, a new efficient and sparse method for learning to attend. Our method is based on differentiable sorting of internal representations. Concretely, we introduce a meta sorting network that learns to generate latent permutations over sequences. Given sorted sequences, we are then able to compute quasi-global attention with only local windows, improving the memory efficiency of the attention module. To this end, we propose new algorithmic innovations such as Causal Sinkhorn Balancing and SortCut, a dynamic sequence truncation method for tailoring Sinkhorn Attention for encoding and/or decoding purposes. Via extensive experiments on algorithmic seq2seq sorting, language modeling, pixel-wise image generation, document classification and natural language inference, we demonstrate that our memory efficient Sinkhorn Attention method is competitive with vanilla attention and consistently outperforms recently proposed efficient Transformer models such as Sparse Transformers.

研究动机与目标

• 在学习稀疏注意力输出的同时，推动降低点积注意力的二次内存成本。
• 提出一种基于可微分排序的注意力机制，通过重新排列输入块在局部计算中实现准全局上下文。
• 引入因果变体（因果 Sinkhorn 平衡）用于自回归解码，以及用于截断序列的动态编码（SortCut）。
• 探索将 Sinkhorn 注意力与标准注意力结合的混合模型以提升性能。
• 展示在语言建模、图像生成、文档分类和自然语言推断等任务中的有效性。

提出的方法

• 引入一个带参数的元排序网络（SortNet），输出分块置换矩阵。
• 使用可微分的 Sinkhorn 平衡将排序矩阵投影到双随机（Birkhoff）多胞体。
• 对块而非单个标记进行排序，以在保持局部邻域的同时实现准全局注意力。
• 在已排序的块内计算注意力，并将共享的排序算子应用于键和值。
• 加入 Gumbel 噪声和温度以近似离散置换（Gumbel-Sinkhorn）。
• 可选地将 Sinkhorn 注意力与标准的原生注意力混合（混合模型）。
• 扩展到因果设置，包括 (i) 基于累计和的因果排序网络，以及 (ii) 带未来屏蔽的因果 Sinkhorn 平衡。
• 提供 SortCut，在排序后对序列进行截断以达到预算（N_k），进一步降低复杂性。

实验结果

研究问题

• RQ1 Sparse Sinkhorn Attention 是否在降低内存复杂度的同时达到与 vanilla 注意力相竞争的性能？
• RQ2 所提出的基于排序的稀疏注意力变体（因果 Sinkhorn 平衡、SortCut）对编码/解码效率和准确性有何影响？
• RQ3 在多种任务中，Sinkhorn 与 vanilla 注意力的混合是否能超过单独使用的每种方法？
• RQ4 在实际中，Sparse Sinkhorn Attention 的内存复杂度与密集注意力和 Sparse Transformers 相比如何？
• RQ5 排序超参数（温度、Sinkhorn 迭代次数）对性能的影响是什么？

主要发现

• Sparse Sinkhorn Attention 将自注意力的内存从 O(ell^2) 降低到 O(B^2 + N_B^2)，在某些设置下，结合 SortCut 可降至线性时间 O(ell N_k)。
• 在算法排序、语言建模、逐像素生成、文档分类和自然语言推断等任务中，Sinkhorn Transformers 的表现与 vanilla Transformers 相当或更好，且常常超过 Sparse Transformers 和 Local Attention。
• 在 LM1B 子词建模中，Sinkhorn Transformers 在若干块大小下优于局部注意力和 Sparse Transformer，其中 Sinkhorn 混合模型取得最佳结果。
• 在 CIFAR-10 的图像生成中，Sinkhorn Transformer 在测试的高效模型中获得了最佳的每字节维数（Bpd）。
• SortCut 编码器在显著降低内存使用的同时，能够达到与 vanilla Transformers 相当的性能。
• 消融研究表明，排序网络在简单线性形式下效果最好，且应用 Sinkhorn 归一化（N_k > 0）对性能至关重要。