Sparse Sinkhorn Attention
Yi Tay Dara Bahri Liu Yang Donald Metzler Da-Cheng Juan
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.
中文速览
自注意力机制（self-attention）的计算量和显存消耗随序列长度呈二次方增长，严重制约了Transformer在长序列任务上的应用，而简单的局部窗口注意力又因感受野受限而难以捕捉长程依赖。为此，本文提出了稀疏Sinkhorn注意力（Sparse Sinkhorn Attention），核心思路是在注意力计算前，先用一个可学习的"排序网络"（SortNet）对输入序列的分块进行可微排列重组——通过Sinkhorn迭代归一化生成双随机矩阵作为软置换，将原本相隔遥远但语义相关的块排到相邻位置，再只在局部窗口内计算注意力，从而以局部计算近似全局感知。此外，文中还提出了适用于自回归解码的因果Sinkhorn平衡（Causal Sinkhorn Balancing）和进一步将复杂度降至线性的动态截断方案SortCut，将显存复杂度从O(ℓ²)大幅压缩至近似线性。在语言建模、图像逐像素生成、文档分类、自然语言推理及序列排序等多项大规模任务上，Sinkhorn Transformer与标准全连接注意力性能相当，并持续超越Sparse Transformer等高效Transformer基线，同时也是可微神经排序首次在大规模真实任务中得到成功应用的验证。
原文 arXiv:2002.11296;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2002.11296v1