
本文收录于 NeurIPS 2021
论文链接:https://arxiv.org/pdf/2110.11945.pdf
代码链接:https://github.com/fudan-zvg/SOFT
一、背景
基于 self-attention 的 transformer 虽然取得了较好的效果,但其计算量和内存都和是输入分辨率大小的平方。
作者认为这种计算限制来源于计算概率时使用的 softmax self-attention。
一般的 self-attention 都是计算规范化的 token 特征的内积计算得到,保持这种 softmax 操作对后面的线性化操作有些挑战。
所以,本文作者提出了 softmax-free transformer, SOFT,在 self-attention 中移除了 softmax,使用高斯核函数来代替内积,能够通过低秩矩阵分解来近似得到 self-attention 矩阵。
二、方法
2.1 Softmax-free self-attention formulation
输入为 X,要进行 attention ,首先要通过映射得到 Q, K, V:
self-attention 的计算如下:
α \alpha α 为计算 self-attention map 的过程,由非线性函数 β \beta β 和 relation function γ \gamma γ 组成:
一般的形式如下:
为了简化计算,作者使用如下方式代替:
为了保持 attention matrix 的对称性,作者设定 Q 和 K 的映射函数一样,所以,本文的 self-attention matrix 如下:
2.2 通过矩阵分解来实现低秩规范化
为了降低计算量,作者参考 Nystrom[38] 来实现低秩矩阵近似,可以不用计算全部的 self-attention。
规范化的 self-attention matrix
S
^
\hat S
S^ 如下:
三、效果
SOFT 及其变体如下: