[論文レビュー] Separable Self-attention for Mobile Vision Transformers

Sachin Mehta, Mohammad Rastegari|arXiv (Cornell University)|Jun 6, 2022

Advanced Neural Network Applications被引用数 182

ひとこと要約

論文は、MobileViT における標準のマルチヘッド自己注意を置換する線形 O(k) 複雑さを持つ分離可能自己注意を導入し、モバイルデバイス上で推論を高速化しつつ精度はほぼ同等。

ABSTRACT

Mobile vision transformers (MobileViT) can achieve state-of-the-art performance across several mobile vision tasks, including classification and detection. Though these models have fewer parameters, they have high latency as compared to convolutional neural network-based models. The main efficiency bottleneck in MobileViT is the multi-headed self-attention (MHA) in transformers, which requires $O(k^2)$ time complexity with respect to the number of tokens (or patches) $k$. Moreover, MHA requires costly operations (e.g., batch-wise matrix multiplication) for computing self-attention, impacting latency on resource-constrained devices. This paper introduces a separable self-attention method with linear complexity, i.e. $O(k)$. A simple yet effective characteristic of the proposed method is that it uses element-wise operations for computing self-attention, making it a good choice for resource-constrained devices. The improved model, MobileViTv2, is state-of-the-art on several mobile vision tasks, including ImageNet object classification and MS-COCO object detection. With about three million parameters, MobileViTv2 achieves a top-1 accuracy of 75.6% on the ImageNet dataset, outperforming MobileViT by about 1% while running $3.2\times$ faster on a mobile device. Our source code is available at: \url{https://github.com/apple/ml-cvnets}

研究の動機と目的

• モバイルデバイス向けの視覚トランスフォーマーにおけるマルチヘッド自己注意（MHA）の待ち時間のボトルネックを動機づけ、解決する。
• 線形計算量と要素ごとの演算を特徴とする分離可能な自己注意メカニズムを提案する。
• この分離可能な自己注意を MobileViT に組み込み、MobileViTv2 を形成する。
• ImageNet-1k、MS-COCO、およびセグメンテーションのベンチマークで、精度を維持または改善しつつ推論速度の改善を実証する。

提案手法

• MobileViT の MHA を、潜在トークン L に対する文脈スコアを計算する分離可能な自己注意に置換し、複雑さを O(k^2) から O(k) に削減する。
• W_I を用いる単一の潜在射影 I によって文脈スコアを計算し、次に softmax で c_s を得て、W_K を用いた K 投影されたトークンの加重和として文脈ベクトル c_v を得る。
• ReLU(xW_V) とのブロードキャスト付き要素ごとの乗算を介して文脈情報を V に伝搬し、最終的な線形射影 W_O で出力 y を生成する。
• 分離可能な自己注意を y = ( sum( sigma(xW_I) * (xW_K) ) * ReLU(xW_V) ) W_O として記述し、要素ごとの演算を強調する。
• MobileViT の MHA を置換して分離可能な自己注意を MobileViTv2 に組み込み、モデル幅を {0.5, 2.0} の乗数 α を用いて探索する。

実験結果

リサーチクエスチョン

• RQ1 モバイルデバイス向けの視覚タスクで性能を維持しつつ、自己注意を線形計算量に再定式化することは可能か。
• RQ2 分離可能な自己注意は、モバイルデバイス上で待ち時間を大幅に削減しつつ、競争力のある精度を維持できるか。
• RQ3 分離可能な自己注意を用いた MobileViTv2 は、分類・検出・セグメンテーションタスク全般で MobileViTv1 および他のモバイル視覚モデルとどのように比較されるか。

主な発見

Attention unit	Latency (ms)	Top-1 (%)
Self-attention in Transformer (Fig. 3(a))	9.9	78.4
Self-attention in Linformer (Fig. 3(b))	10.2	78.2
Separable self-attention (Ours; Fig. 3(c))	3.4	78.1

• MobileViTv2 はモバイルデバイス上で MobileViTv1 より約3倍推論速度が向上し、ImageNet-1k での精度はほぼ同等を維持。
• ImageNet-1k では、分離可能な自己注意を用く MobileViTv2 は MobileViTv1 の精度（約0.1% 程度の差内）に一致し、アテンション単位の待ち時間を 3.4 ms に改善（ベースラインは 9.9–10.2 ms）。
• 分離可能な自己注意バックボーンを使用した場合、MS-COCO の物体検出および ADE20k/PASCAL VOC のセグメンテーションで競争力またはそれ以上の性能を示す。
• 実験は、モバイルハードウェア上で Transformer/Linformer のベースラインと比較して、Top-1 精度のほぼ損失なしに注意関連遅延を低減することを示している。
• タスクを超えて、MobileViTv2 はモバイルデバイス上で CNN と ViT の間の待ち時間のギャップを縮めつつ、精度を維持または向上させる。