SOFT:具有线性复杂度的无 Softmax Transformer

视觉 Transformer(ViTs)通过逐块图像 token 化继以自注意力,推动了多种视觉识别任务的最优性能。然而,自注意力模块的使用导致了计算与内存使用的二次复杂度。自然语言处理中已有多种将自注意力计算近似为线性复杂度的尝试。但本工作的深入分析表明,它们要么在理论上存在缺陷,要么在视觉识别中实证无效。我们进一步确认,其局限性根源在于近似过程中保留了 softmax 自注意力。具体而言,常规自注意力通过 token 特征向量间缩放点积归一化计算。保留该 softmax 操作对任何后续的线性化努力都构成挑战。基于该洞见,我们首次提出了无 softmax 的 Transformer,即 SOFT。为移除自注意力中的 softmax,使用高斯核函数替代点积相似度且不再进一步归一化。这使得完整的自注意力矩阵可通过低秩矩阵分解来近似。该近似的鲁棒性通过牛顿-拉夫森方法计算其 Moore-Penrose 逆得以实现。在 ImageNet 上的大量实验表明,我们的 SOFT 显著提升了现有 ViT 变体的计算效率。关键在于,凭借线性复杂度,SOFT 可容纳长得多的 token 序列,从而在精度与复杂度间取得更优权衡。