【Transformer】RANDOM FEATURE ATTENTION

ICLR 2021 Spotlight: https://openreview.net/forum?id=QtTKTdVrFBB

通过Random Feature Map，将高斯核转换为两个向量的内积。通过这一推论，简化self-attention的计算，降低时间和空间复杂度。

 

摘要

Transformer的核心是attention function，它建模了输入在每个时间间隔之间的关系。

由于attention的二次时间和空间复杂度，导致它不能有效处理长序列。

本文提出RFA，一种线性时间和空间复杂度的注意力机制，使用随机特征去近似softmax function，并且探索它在transformers中的应用。

RFA通过可选的门控机制（gating mechanism）提供了一种直接的学习近期偏差（recency bias）的方法，可以作为常规的softmax attention的替代

recency bias：人们倾向于以记忆中最近的历史经验和趋势判断未来，这在许多情况下也是被证明正确的。但如果只看现象不理解问题本质，情况出现转折时，有可能出现完全错误的判断而无法自拔。

在机器翻译实验中，RFA两次解码和经典transformer一样快。在3个长文本分类数据集上，RFA表现良好。我们分析认为，RFA的高效尤其体现在长序列上。

 

Background

本文的理论基础是：unbiased estimation of the Gaussian kernel

Random Fourier Features去近似一个想要的shift-invariant kernel，kernel可以看作是求两个样本之间的相似度/距离的函数，这个相似度估计函数越准确，越有利于我们对新样本的预测。

这个方法使用random feature map φ，将一对向量x和y进行非线性转换。然后用φ(x)和φ(y)之间的内积去近似x和y之间的相似度。如下所示，是d维向量，和d维x做内积，变成1维，整个并联起来是2D维向量。

Random feature methods成功的使kernel类方法变得更快，同时也被用于近似softmax。本文使用它导出exp(<.,.>)的无偏估计和对softmax attention的高效近似。

 

Model

在经典transformer中，每个query需要和M个key做内积，作为这个query对这个key的权重，时间复杂度是O(M)，N个query使得时间复杂度变为O(MN)。

首先，作者根据Theorem 1得到下式：

将上式代入attention，得到：

RFA可以作为softmax-attention的替代实现。

外积（Outer product），在线性代数中一般指两个向量的张量积，其结果为一矩阵；与外积相对，两向量的内积结果为标量。

On query and key norms和learned random feature variance。等式5假设query和key都是1范数。因此，它近似于一种softmax attention，即在将query和key相乘之前，对他们进行normalizes，然后通过将它们除以σ^2来缩放logits。从经验上来说，这个normalization step使得。这样做的一个结果是，如果没有σ， softmax输出将被flatten，σ可以设置为一个先验超参数。