[论文解读] Stable, Fast and Accurate: Kernelized Attention with Relative Positional Encoding

Shengjie Luo, Shanda Li+7 more|arXiv (Cornell University)|Jun 23, 2021

Advanced Neural Network Applications被引用 10

一句话总结

本文提出了一种新型的核化注意力机制，结合相对位置编码（RPE），利用快速傅里叶变换（FFT）实现 O(n log n) 的复杂度，从而在 Transformer 模型中实现稳定、快速且准确的长序列建模。该方法通过利用 RPE 矩阵的托普利茨（Toeplitz）结构加速计算，同时通过降低注意力分数的方差来提高训练稳定性。

ABSTRACT

The attention module, which is a crucial component in Transformer, cannot scale efficiently to long sequences due to its quadratic complexity. Many works focus on approximating the dot-then-exponentiate softmax function in the original attention, leading to sub-quadratic or even linear-complexity Transformer architectures. However, we show that these methods cannot be applied to more powerful attention modules that go beyond the dot-then-exponentiate style, e.g., Transformers with relative positional encoding (RPE). Since in many state-of-the-art models, relative positional encoding is used as default, designing efficient Transformers that can incorporate RPE is appealing. In this paper, we propose a novel way to accelerate attention calculation for Transformers with RPE on top of the kernelized attention. Based upon the observation that relative positional encoding forms a Toeplitz matrix, we mathematically show that kernelized attention with RPE can be calculated efficiently using Fast Fourier Transform (FFT). With FFT, our method achieves $\\mathcal{O}(n\\log n)$ time complexity. Interestingly, we further demonstrate that properly using relative positional encoding can mitigate the training instability problem of vanilla kernelized attention. On a wide range of tasks, we empirically show that our models can be trained from scratch without any optimization issues. The learned model performs better than many efficient Transformer variants and is faster than standard Transformer in the long-sequence regime.

研究动机与目标

• 解决标准自注意力机制在长序列中因二次方复杂度导致的效率低下问题。
• 实现使用相对位置编码（RPE）的 Transformer 模型中高效注意力计算，而 RPE 是当前最先进模型中的标准配置。
• 克服先前近似方法的局限性，这些方法无法处理 RPE，因为 RPE 的表达能力超出了标准的点积后指数化注意力机制。
• 稳定核化注意力模型的训练过程，这类模型在实践中通常因优化不稳定性而表现不佳。
• 在不依赖微调的前提下，同时实现计算效率和长序列场景下的高性能。

提出的方法

• 通过利用相对位置编码矩阵的托普利茨（Toeplitz）结构，重新表述了结合 RPE 的核化注意力机制。
• 使用快速傅里叶变换（FFT）加速 RPE 矩阵与核化特征矩阵之间的矩阵-向量乘法，将时间复杂度降低至 O(n log n)。
• 该方法用核化注意力机制替代标准的 softmax 注意力机制，通过特征映射来近似注意力函数。
• 引入归一化注意力（NPRF）以降低方差，提升训练稳定性，尤其在查询和键的范数较大时表现更优。
• 该方法将核化注意力机制应用于编码器和解码器模块，实现与 Transformer 架构的完整集成。
• 支持多种特征映射函数（例如：PRF、TRF、Sphere-PRF、ORF），展示了在不同选择下的泛化能力。

实验结果

研究问题

• RQ1 能否在标准的点积后指数化公式之外，高效计算结合相对位置编码的核化注意力？
• RQ2 在核化注意力中使用 RPE 是否相比原始核化注意力能提升训练稳定性？
• RQ3 能否利用 RPE 矩阵的托普利茨结构，在注意力计算中实现亚二次方时间复杂度？
• RQ4 所提出方法在长序列任务上的性能与标准 Transformer 及高效变体相比如何？
• RQ5 所提出方法在核化注意力机制中对不同特征映射维度和类型是否具有鲁棒性？

主要发现

• 通过在托普利茨结构的 RPE 矩阵上应用 FFT，所提方法实现了 O(n log n) 的时间复杂度，从而支持高效的长序列建模。
• 采用归一化核化注意力与 RPE 的模型能够从零开始稳定训练，无优化问题，而原始核化注意力则常出现训练不稳定现象。
• 在 ImageNet 数据集上，采用 RPE 的 NPRF-Transformer 模型达到 80.9% 的 top-1 准确率，优于标准 DeiT-base（81.2%）和 PRF 转换模型（79.5%）。
• 在 IWSLT14 德语到英语翻译任务中，采用 RPE 和归一化注意力的模型达到 32.1 的 BLEU 分数，与经过充分训练的标准 Transformer（32.3）相当，同时保持线性复杂度。
• 该方法对特征映射维度具有鲁棒性，即使在 16 维时性能下降极小，并且在不同特征映射函数（PRF、TRF、Sphere-PRF、ORF）之间具有良好泛化能力。
• RPE 的使用可降低从标准注意力到核化注意力转换时的近似误差，尤其在结合归一化时，消融实验显示性能下降更小。