【Transformer】SOFT: Softmax-free Transformer with Linear Complexity

本文收录于 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 $\hat{S}$ 如下：

三、效果

SOFT 及其变体如下：