[论文解读] SOFT: Softmax-free Transformer with Linear Complexity

Jiachen Lu, Jinghan Yao+7 more|arXiv (Cornell University)|Oct 22, 2021

Advanced Neural Network Applications参考文献 45被引用 61

一句话总结

SOFT 引入了一种不含 softmax 的自注意力机制，使用高斯核和 Nyström 基于低秩近似，并通过牛顿-拉夫森法求取 Moore-Penrose 逆，达到线性时间/空间复杂度，并在 ImageNet 上相较于其他线性变换器实现更高的准确性。

ABSTRACT

Vision transformers (ViTs) have pushed the state-of-the-art for various visual recognition tasks by patch-wise image tokenization followed by self-attention. However, the employment of self-attention modules results in a quadratic complexity in both computation and memory usage. Various attempts on approximating the self-attention computation with linear complexity have been made in Natural Language Processing. However, an in-depth analysis in this work shows that they are either theoretically flawed or empirically ineffective for visual recognition. We further identify that their limitations are rooted in keeping the softmax self-attention during approximations. Specifically, conventional self-attention is computed by normalizing the scaled dot-product between token feature vectors. Keeping this softmax operation challenges any subsequent linearization efforts. Based on this insight, for the first time, a softmax-free transformer or SOFT is proposed. To remove softmax in self-attention, Gaussian kernel function is used to replace the dot-product similarity without further normalization. This enables a full self-attention matrix to be approximated via a low-rank matrix decomposition. The robustness of the approximation is achieved by calculating its Moore-Penrose inverse using a Newton-Raphson method. Extensive experiments on ImageNet show that our SOFT significantly improves the computational efficiency of existing ViT variants. Crucially, with a linear complexity, much longer token sequences are permitted in SOFT, resulting in superior trade-off between accuracy and complexity.

研究动机与目标

• 在视觉领域激发对高效变换器的需求，以在处理长令牌序列时避免二次复杂度。
• 提出一种不含 softmax 的自注意力机制，使复杂度实现线性。
• 开发基于 Nyström 的低秩近似，并通过牛顿-拉夫森计算的 Moore-Penrose 逆实现鲁棒注意力。
• 设计并在 ImageNet 上评估基于 SOFT 的骨干网络，以评估准确性与复杂度的权衡。

提出的方法

• 用高斯核注意力 S = exp(Q ⊖ K) 替代基于 softmax 的注意力，它是对称的且位于 [0,1]。
• 使用带小瓶颈 m 的 Nyström 分解来近似完整注意力矩阵，使其达到线性时间/空间复杂度，Ŝ = Pᵀ A† P。
• 通过牛顿-拉夫森迭代计算 Moore-Penrose 逆 A† 以确保数值鲁棒性（A₀ = αA, A_{k+1} = 2A_k − A_k A A_k）。
• 使用卷积、平均池化或其他采样方法对瓶颈令牌进行采样；为稳定性与效率，偏好平均池化。
• 将 SOFT 作为金字塔视觉 Transformer 骨干中的一个层实例化，使用特定超参数（如 d_e、h、n、m、sp）来构建 SOFT 变体。

实验结果

研究问题

• RQ1 在视觉变换器中，不含 softmax 的高斯核注意力是否能达到与 softmax 基于的注意力相当的准确度？
• RQ2 基于 Nyström 的低秩近似结合 Moore-Penrose 逆是否在视觉任务中提供稳定的训练和线性复杂度？
• RQ3 哪些设计选择（瓶颈大小 m、采样策略）可以优化 SOFT 的准确性与效率之间的权衡？
• RQ4 SOFT 在 ImageNet 和 NLP 基准测试中相对于其他线性/高效变换器的表现如何？
• RQ5 将 SOFT 集成到金字塔式 Transformer 骨干的视觉识别中会带来哪些架构优势？

主要发现

• SOFT 在注意力上的时间和空间复杂度达到线性（O(n)），使得可以处理更长的令牌序列。
• 通过牛顿–拉夫森实现的 Moore-Penrose 逆的 Nyström 基分解提供鲁棒的注意力近似。
• 在 ImageNet 上，基于 SOFT 的骨干在准确性–复杂度权衡方面优于若干 CNN 与 ViT 变体。
• SOFT 在类似设置下展现出与 Linformer、Performer、Nyströmformer 相当或更高的准确性。
• 消融研究显示瓶颈大小 m 约为 49 能提供良好的准确性与计算平衡；在测试方法中，平均池化采样表现最好。