[論文レビュー] 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正規化に基づくユニット力操作に置き換えることで、計算量とメモリ使用量を削減しながら、画像分類および密度予測タスクで最先端の性能を達成する線形複雑性を持つビジョントランスフォーマーを提案する。この手法は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²)の計算複雑性を軽減すること。
• 性能の低下を伴わずに自己注意機構におけるソフトマックス非線形性を排除し、線形時間の注意メカニズムを実現すること。
• 画像分類および密度予測を含む多様なビジョンタスクで高い精度を維持しながら、FLOPsとメモリ使用量を顕著に削減するモデルを開発すること。
• 特にオブジェクト検出やセマンティックセグメンテーションのような密度予測タスクにおいて、解像度が変動する高解像度入力への効率的なデプロイを可能にすること。

提案手法

• 自己注意機構におけるソフトマックス演算をL2正規化に置き換えることで、行列指数関数の2次的複雑性を排除する。
• 行列乗法の結合則を適用し、QK^Tを単位ベクトルに再構成することで、注意計算をO(N)の演算に要約する。
• クエリとキーを単位超球面上に射影する正規化スキームXNormを導入し、訓練の安定性を高めるとともに初期化依存性を排除する。
• ResMLPにインspiredされたアフィンスケーリングモジュールを用いて残差接続を調整し、モデルの深さにかかわらず安定性を維持する。
• 解像度に依存しないアーキテクチャを設計：モデル重みが入力解像度に依存しないため、さまざまな入力サイズで一貫した性能を発揮可能。
• 標準的なViTからのコード変更を最小限に抑え（数行のみ変更）、後方互換性と統合の容易性を確保する。

実験結果

リサーチクエスチョン

• RQ1 画像分類タスクにおいて、性能を損なわず自己注意機構の複雑性を線形化できるビジョントランスフォーマーは実現可能か？
• RQ2 自己注意機構におけるソフトマックスをL2正規化に置き換えても、標準的なトランスフォーマーや既存の線形近似手法と比較して精度が維持または向上するか？
• RQ3 提案手法は高解像度入力、特にオブジェクト検出やインスタンスセグメンテーションのような密度予測タスクに効率的にスケーリング可能か？
• RQ4 さまざまな入力解像度およびモデル容量において、メモリ消費量と推論速度が既存のトランスフォーマーやCNNと比較してどのように異なるか？

主な発見

• UFO-ViT-Sは39.7MパラメータでCOCOインスタンスセグメンテーションで44.6 mAPを達成し、ResNet50やPVT-Smallを上回る性能を示す。
• UFO-ViT-Mは56.4MパラメータでCOCOで46.0 mAPを達成し、Swin-TやXCiT-S12/16を上回るmAPを記録しながら、Swin-SのFLOPsの半分未満を実現する。
• モデルは入力解像度に比例して線形にメモリ成長を示す：標準的なトランスフォーマーとは異なり、トークン数に比例してGPUメモリ使用量が増加する。
• UFO-ViTはDeiTやSwinトランスフォーマーと比較して、1枚のV100 GPUで最大4倍のバッチサイズをサポートする。これは、優れたメモリ効率を示している。
• 推論スループットはDeiTを大きく上回り、特に高解像度ではSwinトランスフォーマーと同等またはそれ以上である。
• UFO-ViT-Bはボクシングボックス検出ではUFO-ViT-Mよりわずかに性能が低いが、小物検出およびインスタンスセグメンテーションでは優れた性能を示しており、微細なタスクに適した特徴表現を有している。