[论文解读] Skyformer: Remodel Self-Attention with Gaussian Kernel and Nystr\\"om Method

Yifan Chen, Qi Zeng+2 more|arXiv (Cornell University)|Oct 29, 2021

Neural Networks and Applications参考文献 24被引用 7

一句话总结

Skyformer 提出了一种核化注意力机制，用高斯核替代标准自注意力中的 softmax，以提升训练稳定性，并通过改进的 Nystrom 方法对非半正定矩阵实现高效近似。该方法在 Long Range Arena 基准测试中实现了与完整自注意力相当或更优的性能，同时显著降低了计算成本并提升了收敛稳定性。

ABSTRACT

Transformers are expensive to train due to the quadratic time and space complexity in the self-attention mechanism. On the other hand, although kernel machines suffer from the same computation bottleneck in pairwise dot products, several approximation schemes have been successfully incorporated to considerably reduce their computational cost without sacrificing too much accuracy. In this work, we leverage the computation methods for kernel machines to alleviate the high computational cost and introduce Skyformer, which replaces the softmax structure with a Gaussian kernel to stabilize the model training and adapts the Nystr\\"om method to a non-positive semidefinite matrix to accelerate the computation. We further conduct theoretical analysis by showing that the matrix approximation error of our proposed method is small in the spectral norm. Experiments on Long Range Arena benchmark show that the proposed method is sufficient in getting comparable or even better performance than the full self-attention while requiring fewer computation resources.

研究动机与目标

• 解决标准自注意力在 Transformer 中计算成本高和训练不稳定的挑战。
• 利用核方法通过用高斯核替代 softmax 结构来稳定训练。
• 针对非半正定矩阵开发一种高效的核化注意力近似方法，基于 Nystrom 方法。
• 为所提方法在谱范数下的矩阵近似误差提供理论保证。
• 证明所提方法在长序列任务中可实现与资源消耗更低的性能竞争。

提出的方法

• 用基于高斯核的注意力机制替代标准 softmax 注意力，以稳定训练并改善数值行为。
• 将 Nystrom 方法改进以近似核化注意力矩阵，通过将其提升为更大的半正定矩阵，实现低秩近似。
• 使用随机化压缩（sketching）降低核矩阵的维度，加速计算同时保持精度。
• 理论分析表明，该方法的矩阵近似误差在谱范数下较小，提供了近似质量的正式保证。
• 采用对称化技术处理自注意力计算中由独立查询和键矩阵产生的非 PSD 核矩阵。
• 将该方法应用于 Long Range Arena 基准，评估其在多样化 NLP 任务中的性能与效率。

实验结果

研究问题

• RQ1 将自注意力中的 softmax 结构替换为高斯核，是否能提升训练稳定性并降低对超参数的敏感性？
• RQ2 Nystrom 方法能否有效适配自注意力中出现的非半正定核矩阵，以实现高效近似？
• RQ3 所提方法在谱范数下的近似误差理论边界是什么？
• RQ4 所提方法是否在降低计算与内存成本的同时，实现与标准自注意力相当或更优的性能？
• RQ5 该方法在序列长度增加时如何扩展？序列长度与近似效率之间存在何种关系？

主要发现

• 与标准自注意力相比，Skyformer 在文本分类和文档检索任务中实现了近 4 倍的加速，训练时间与内存使用显著减少。
• 在文本分类任务上，准确率较标准自注意力基线提升 2.75%；在文档检索任务上提升 1.37%。
• 使用 Skyformer 训练时，不同随机种子下的表现更稳定，收敛失败率更低，且对学习率变化的敏感性更小。
• 理论分析证实，Skyformer 的矩阵近似误差在谱范数下较小，支持其可靠性。
• 该方法在更长序列上表现更优，因为核矩阵的统计维度相对于序列长度更小，使得近似更有效。
• 实验表明，对标准自注意力直接应用 Nystrom 近似会导致数值不稳定，而核化注意力变体则能稳定训练。