无 Softmax 的线性 Transformer
计算机视觉与模式识别
2024-03-18 v3 人工智能
机器学习
摘要
视觉 Transformer (ViT) 已推动了视觉感知任务的最先进水平。支撑 ViT 优势的自注意力机制在计算和内存使用上均具有二次复杂度。这促使了以线性复杂度近似自注意力的研究。然而,本工作的深入分析揭示,现有方法在理论上存在缺陷或在视觉识别上实证无效。我们指出,它们的局限性根源于近似过程中继承了基于 softmax 的自注意力,即使用 softmax 函数对 token 特征向量间的缩放点积进行归一化。由于保留 softmax 操作对任何后续的线性化努力都构成挑战。基于这一洞见,我们提出了一系列无 Softmax Transformer (SOFT)。具体而言,采用高斯核函数替代点积相似度,使得完整的自注意力矩阵可在低秩矩阵分解下被近似。为计算鲁棒性,我们仅在前向过程中使用迭代 Newton-Raphson 方法估计 Moore-Penrose 逆,而在反向过程中仅计算一次其理论梯度。为进一步扩展适用性(例如密集预测任务),引入了高效的对称归一化技术。在 ImageNet、COCO 和 ADE20K 上的大量实验表明,我们的 SOFT 显著提升了现有 ViT 变体的计算效率。凭借线性复杂度,SOFT 允许长得多的 token 序列,从而在精度与复杂度之间实现了更优的权衡。代码与模型见 https://github.com/fudan-zvg/SOFT。
引用
@article{arxiv.2207.03341,
title = {Softmax-free Linear Transformers},
author = {Jiachen Lu and Junge Zhang and Xiatian Zhu and Jianfeng Feng and Tao Xiang and Li Zhang},
journal= {arXiv preprint arXiv:2207.03341},
year = {2024}
}
备注
Accepted by IJCV. arXiv admin note: substantial text overlap with arXiv:2110.11945