[论文解读] UFO-ViT: High Performance Linear Vision Transformer without Softmax

Jeong-geun Song|arXiv (Cornell University)|Sep 29, 2021

CCD and CMOS Imaging Sensors参考文献 46被引用 10

一句话总结

UFO-ViT 提出了一种线性复杂度的视觉变换器，通过使用基于 L2 归一化的单位力操作替代自注意力中的 softmax 非线性，实现了在图像分类和密集预测任务上的最先进性能，同时降低了计算成本和内存使用。该方法在保持 O(N) 复杂度的同时，跨不同模型尺寸维持了高精度，在 ImageNet1k 和 COCO 基准测试中优于现有的变换器和 CNN。

ABSTRACT

Vision transformers have become one of the most important models for computer vision tasks. Although they outperform prior works, they require heavy computational resources on a scale that is quadratic to $N$. This is a major drawback of the traditional self-attention (SA) algorithm. Here, we propose the Unit Force Operated Vision Transformer (UFO-ViT), a novel SA mechanism that has linear complexity. The main approach of this work is to eliminate nonlinearity from the original SA. We factorize the matrix multiplication of the SA mechanism without complicated linear approximation. By modifying only a few lines of code from the original SA, the proposed models outperform most transformer-based models on image classification and dense prediction tasks on most capacity regimes.

研究动机与目标

• 解决标准视觉变换器中自注意力机制的 O(N²) 计算复杂度问题，该问题限制了在高分辨率输入下的可扩展性。
• 在不损失性能的前提下，消除自注意力中的 softmax 非线性，实现线性时间注意力机制。
• 开发一种在多种视觉任务（尤其是图像分类和密集预测）中保持高精度的模型，同时显著降低 FLOPs 和内存使用。
• 实现在高分辨率输入上的高效部署，特别是在目标检测和语义分割等密集预测任务中，这些任务的分辨率各不相同。

提出的方法

• 用 L2 归一化替代自注意力中的 softmax 操作，消除矩阵指数运算带来的二次方复杂度。
• 通过将 QK^T 重构成单位向量，利用矩阵乘法的结合律，将注意力计算重构为 O(N) 操作。
• 引入 XNorm，一种归一化方案，将查询和键投影到单位超球面上，稳定训练过程并消除对初始化的依赖。
• 使用仿射缩放模块（受 ResMLP 启发）调整残差连接，确保在不同模型深度下保持稳定性。
• 设计架构时实现分辨率无关性：模型权重不依赖于输入分辨率，从而在不同输入尺寸下保持一致的性能表现。
• 通过极少的代码修改实现该方法——仅需对标准 ViT 修改几行代码，确保向后兼容性和易于集成。

实验结果

研究问题

• RQ1 视觉变换器是否能在不牺牲图像分类任务性能的前提下，实现自注意力的线性复杂度？
• RQ2 在自注意力中用 L2 归一化替代 softmax 是否能保持或提升与标准变换器及先前线性近似方法相比的精度？
• RQ3 所提出的方法是否能高效扩展至高分辨率输入，特别是在目标检测和实例分割等密集预测任务中？
• RQ4 在不同输入分辨率和模型容量下，该模型的内存消耗和推理速度与现有变换器及 CNN 相比如何？

主要发现

• UFO-ViT-S 在 COCO 实例分割任务上达到 44.6 mAP，仅使用 39.7M 参数，优于 ResNet50 和 PVT-Small，且参数更少。
• UFO-ViT-M 在 COCO 上实现 46.0 mAP，参数量为 56.4M，mAP 超过 Swin-T 和 XCiT-S12/16，且 FLOPs 低于 Swin-S 的一半。
• 该模型的内存消耗随输入分辨率呈线性增长：GPU 内存使用量与 token 数量成线性关系，而标准变换器则为二次方增长。
• UFO-ViT 在单张 V100 GPU 上支持的批量大小是 DeiT 和 Swin 变换器的 4 倍，表明其内存效率更高。
• 推理吞吐量显著高于 DeiT，在高分辨率下与 Swin 变换器相当或更优。
• UFO-ViT-B 在边界框检测任务上略逊于 UFO-ViT-M，但在小目标检测和实例分割任务上表现更优，表明其在细粒度任务中具有更好的特征表示能力。