[論文レビュー] Rethinking Query, Key, and Value Embedding in Vision Transformer under Tiny Model Constraints

Jaesin Ahn, Jiuk Hong+2 more|arXiv (Cornell University)|Nov 18, 2021

Machine Learning and ELM被引用数 2

ひとこと要約

本論文は、極めて小さなモデル制約下におけるビジョントランスフォーマー（ViT）に対して、非線形的で共有され、コードベースのクエリ（Q）、キー（K）、バリュー（V）埋め込み構造を3つ提案する。標準の線形射影を学習可能な非線形マッピングに置き換えることで、特に共有層とトレーニング可能なコードパラメータを用いることにより、パラメータ数を著しく削減しながら画像分類の精度を向上させた。本手法は、310万パラメータでImageNet-1kで71.4%のトップ1精度を達成し、元のXCiT-N12モデルを上回った。

ABSTRACT

A vision transformer (ViT) is the dominant model in the computer vision field. Despite numerous studies that mainly focus on dealing with inductive bias and complexity, there remains the problem of finding better transformer networks. For example, conventional transformer-based models usually use a projection layer for each query (Q), key (K), and value (V) embedding before multi-head self-attention. Insufficient consideration of semantic $Q, K$, and $V$ embedding may lead to a performance drop. In this paper, we propose three types of structures for $Q$, $K$, and $V$ embedding. The first structure utilizes two layers with ReLU, which is a non-linear embedding for $Q, K$, and $V$. The second involves sharing one of the non-linear layers to share knowledge among $Q, K$, and $V$. The third proposed structure shares all non-linear layers with code parameters. The codes are trainable, and the values determine the embedding process to be performed among $Q$, $K$, and $V$. Hence, we demonstrate the superior image classification performance of the proposed approaches in experiments compared to several state-of-the-art approaches. The proposed method achieved $71.4\%$ with a few parameters (of $3.1M$) on the ImageNet-1k dataset compared to that required by the original transformer model of XCiT-N12 ($69.9\%$). Additionally, the method achieved $93.3\%$ with only $2.9M$ parameters in transfer learning on average for the CIFAR-10, CIFAR-100, Stanford Cars datasets, and STL-10 datasets, which is better than the accuracy of $92.2\%$ obtained via the original XCiT-N12 model.

研究の動機と目的

• パラメータの不足に起因する小さなビジョントランスフォーマーにおける、意味的Q、K、V埋め込みの考慮不足に起因する性能低下を是正すること。
• Q、K、V埋め込み機構の再考を通じて、厳密なパラメータ制約下でもViTの性能を向上させること。
• 共有非線形層とトレーニング可能なコードパラメータを用いて、Q、K、V間の知識共有を検討すること。
• 非線形的かつ共有構造のQKV埋め込みが、ImageNet分類および転移学習性能の両方を向上させることを実証すること。

提案手法

• 入力トークンを別々の非線形空間に変換するため、2層のReLUベースの非線形埋め込みをQ、K、Vに導入する。
• Q、K、Vの全領域に共通する1層の非線形層を共有する構造を提案し、知識の転送を促進する。
• トレーニング可能なコードパラメータ（Cq、Ck、Cv）を用いた2層の共有構造を構築し、Q、K、Vの埋め込み変換を定義する。
• バックプロパゲーションを用いて、ViT分類損失を最小化するようにコードパラメータを同時に学習する。
• F-SNE可視化を用いてコードの類似性と直交性を分析し、コードが異なる、タスクに依存しない特徴を学習していることを確認する。
• パラメータ制約下でImageNet-1kおよび転移学習ベンチマーク（CIFAR-10、CIFAR-100、Stanford Cars、STL-10）における性能を評価する。

実験結果

リサーチクエスチョン

• RQ1 線形QKV射影を非線形マッピングに置き換えることで、極小モデル制約下でもViTの性能が向上するか？
• RQ2 Q、K、V間で非線形層を共有することで、特徴の学習と分類精度が向上するか？
• RQ3 Q、K、V埋め込みを同時に定義するトレーニング可能なコードパラメータを用いることで、独立した射影よりも優れた性能が得られるか？
• RQ4 提案手法のQKV埋め込み構造は、ImageNetおよび転移学習タスクにおいて、XCiT-N12などの最先端モデルと比較してどのように性能を発揮するか？

主な発見

• 提案手法は、310万パラメータでImageNet-1kで71.4%のトップ1精度を達成し、元のXCiT-N12モデル（69.9%）を上回った。
• 転移学習において、CIFAR-10、CIFAR-100、Stanford Cars、STL-10の平均精度が93.3%に達し、元のXCiT-N12モデル（92.2%）を上回った。
• トレーニング可能なパラメータを有するコードベースの共有構造は、特にCIFAR-100およびSTL-10で優れた性能を示し、タスク間での効果的な特徴学習が可能であることを示した。
• F-SNE可視化により、学習されたコード（Cq、Ck、Cv）がほぼ直交性を示しており、異なる、タスクに依存しない表現を学習していることが確認された。
• コードのl2ノルムは、ImageNet、Cars、STL-10のデータセットで一貫していたが、CIFAR-10およびCIFAR-100では異なっていたため、データセットに応じたコードの適応が行われていることが示された。
• ナノモデルではコードサイズ8が最適な性能を発揮したが、ティニー・モデルでは16が最適であった。これは、コードサイズが埋め込み次元に応じてスケーリングされるべきであることを示唆している。