[論文レビュー] SimA: Simple Softmax-free Attention for Vision Transformers

Soroush Abbasi Koohpayegani, Hamed Pirsiavash|arXiv (Cornell University)|Jun 17, 2022

CCD and CMOS Imaging Sensors被引用数 8

ひとこと要約

この論文では、Vision Transformers向けに計算コストの高いSoftmaxを置き換えるsoftmaxフリーなアテンション機構であるSimAを提案する。SimAはクエリとキー行列のℓ₁正規化を導入することで、入力解像度に応じてO(ND)またはO(D²N)の計算量に低減可能な動的計算順序の決定を推奨する。SimAはImageNet、MS-COCO、自己教師ありベンチマークで最先端の精度を達成するとともに、exp(.)演算を完全に排除し、エッジデバイスにおける効率性を著しく向上させる。

ABSTRACT

Recently, vision transformers have become very popular. However, deploying them in many applications is computationally expensive partly due to the Softmax layer in the attention block. We introduce a simple but effective, Softmax-free attention block, SimA, which normalizes query and key matrices with simple $\ell_1$-norm instead of using Softmax layer. Then, the attention block in SimA is a simple multiplication of three matrices, so SimA can dynamically change the ordering of the computation at the test time to achieve linear computation on the number of tokens or the number of channels. We empirically show that SimA applied to three SOTA variations of transformers, DeiT, XCiT, and CvT, results in on-par accuracy compared to the SOTA models, without any need for Softmax layer. Interestingly, changing SimA from multi-head to single-head has only a small effect on the accuracy, which simplifies the attention block further. The code is available here: https://github.com/UCDvision/sima

研究の動機と目的

• エッジデバイスへの展開を念頭に、Vision Transformersの計算コスト、特にSoftmaxにおける高価なexp(.)演算を低減すること。
• アテンション機構からSoftmaxレイヤーを削除しても、モデルの精度を維持または向上させること。
• トークン数(N)とチャネル数(D)に基づいて、推論時における計算順序の動的再順序付けを可能にし、効率性を最適化すること。
• マルチヘッドとGELUの代わりにシングルヘッドアテンションとReLU活性化関数を用いることで、顕著な精度低下なしに代替可能かどうかを検証すること。
• 数値的に安定で、半精度推論に対応可能であり、順列不変性とトークンの重要度解釈を保持するアテンション機構を提供すること。

提案手法

• 自己アテンションにおけるSoftmax演算を、クエリとキー行列のℓ₁正規化後にそれらのドット積を取ることで置き換える。
• 単純な行列積としてアテンションを計算する：O = (Q_normalized) @ (K_normalized.T) @ Vとすることで、exp(.)とSoftmaxの必要性を排除する。
• 推論時における計算順序を動的に選択する：N > Dの場合は(Q @ K.T) @ V、D > Nの場合はQ @ (K.T @ V)を選択し、FLOPsを最小化する。
• ℓ₁正規化により数値的安定性を確保し、オーバーフローを回避しながら半精度推論を可能にする。
• トークンを順序なし集合として扱うことで順列不変性を保持し、複数のソース（複数のカメラやセンサーなど）からの入力にも適用可能にする。
• アーキテクチャの変更なしにDeiT、XCiT、CvTモデルにSimAブロックを適用し、標準的なトレーニングプロトコルを維持する。

実験結果

リサーチクエスチョン

• RQ1 softmaxフリーなアテンション機構は、視覚Transformerにおいて最先端の性能を維持できるか？
• RQ2 クエリとキーにℓ₁正規化を適用することで、モデルの精度を維持しながら計算コストを低減できるか？
• RQ3 NとDに基づいた推論時の動的計算順序付けが、顕著な効率性向上をもたらすか？
• RQ4 Softmaxが存在しない状況で、シングルヘッドアテンションとReLUはマルチヘッドとGELUに比べて性能に著しい差が出ないか？
• RQ5 正規化されたクエリとキーのベクトルの大きさは、モデルの解釈性に意味のあるサリエンシー地図として機能するか？

主な発見

• SimAはImageNet分類、MS-COCOオブジェクト検出・セグメンテーション、自己教師あり学習ベンチマークにおいて、SOTAモデルと同等の精度を達成した。
• 1536×1536の画像において、SimAはDeiTよりも58%高速であり、Softmaxのexp(.)オーバーヘッドが排除されたためである。
• モデル次元が8192の場合、SimAはXCiTと比較して推論時間を22%短縮した。これはSoftmax計算を回避したためである。
• シングルヘッドSimAは、マルチヘッドSOTAモデルと同等の精度を達成しており、ヘッド数の削減による性能低下が最小限であることが示された。
• GELUの代わりにReLUを用いた場合、精度の低下はわずかにとどまり、活性化関数の変更に対してもモデルが頑健であることが示された。
• 正規化されたクエリとキーのベクトルのℓ₂ノルムは、標準的なDeiTが示す平坦な分布とは異なり、重要な画像領域を明確に強調する意味のあるサリエンシー地図を生成した。