SMYRF：基于非对称聚类的高效注意力机制

我们提出一种新颖的平衡聚类算法来近似注意力。注意力复杂度从 $O(N^2)$ 降低到 $O(N \log N)$，其中 $N$ 为序列长度。我们的算法 SMYRF 以新颖的方式使用局部敏感哈希（LSH），通过定义新的非对称变换和一种生成平衡聚类的自适应方案。SMYRF 的最大优势在于它可作为稠密注意力层的直接替代而无需任何重新训练。相反，先前的快速注意力方法施加了约束（例如查询与键共享相同向量表示）并需要从头重新训练。我们将该方法应用于预训练的最先进自然语言处理与计算机视觉模型，并报告了显著的内存与速度收益。值得注意的是，SMYRF-BERT 在 GLUE 上（略微）优于 BERT，同时使用 $50\%$ 更少内存。我们还表明 SMYRF 可在训练前后与稠密注意力互换使用。最后，我们使用 SMYRF 在高位分辨率下训练带注意力的 GAN。使用单个 TPU，我们能够在 CelebA-HQ 上的 BigGAN 中将注意力扩展到 128x128=16k 和 256x256=65k 个 token。