[論文レビュー] Discrete Representations Strengthen Vision Transformer Robustness

Chengzhi Mao, Lu Jiang+4 more|arXiv (Cornell University)|Nov 20, 2021

Advanced Neural Network Applications参考文献 41被引用数 10

ひとこと要約

本論文では、標準的な連続的ピクセルトークンをベクターライズドVQ-GANエンコーダーからの離散的トークンに置き換えることで、耐性を向上させるVision Transformerの変種、Dr. ViTを提案する。離散的トークンと残留ピクセル埋め込みを組み合わせることで、より不変で形状に敏感な表現を学習し、ImageNetの精度を損なわずに7つのImageNet耐性ベンチマークで最大12%の性能向上を達成する。

ABSTRACT

Vision Transformer (ViT) is emerging as the state-of-the-art architecture for image recognition. While recent studies suggest that ViTs are more robust than their convolutional counterparts, our experiments find that ViTs trained on ImageNet are overly reliant on local textures and fail to make adequate use of shape information. ViTs thus have difficulties generalizing to out-of-distribution, real-world data. To address this deficiency, we present a simple and effective architecture modification to ViT's input layer by adding discrete tokens produced by a vector-quantized encoder. Different from the standard continuous pixel tokens, discrete tokens are invariant under small perturbations and contain less information individually, which promote ViTs to learn global information that is invariant. Experimental results demonstrate that adding discrete representation on four architecture variants strengthens ViT robustness by up to 12% across seven ImageNet robustness benchmarks while maintaining the performance on ImageNet.

研究の動機と目的

• Vision Transformerの局所的なテクスチャへの過剰依存と分布シフトに対する一般化性能の低さを是正すること。
• ImageNetにおけるドメイン内性能を損なわずに耐性を向上させること。
• 離散的トークン化が、ViTにおけるグローバルで形状不変な特徴の学習を促進するかどうかを検証すること。
• さまざまなViTバリアントに適用可能な、プラグアンドプレイ式の入力層変更を構築すること。
• 生成タスクで使用されてきた離散的表現が、分類タスクにおける一般化性能の向上に寄与するかどうかを示すこと。

提案手法

• 標準的な連続的ピクセル埋め込みを、VQ-GANエンコーダーから得られる離散的トークンに置き換える。このエンコーダーは画像パッチをコードブックインデックスに量子化する。
• 学習されたコードブックを用いて、画像パッチを離散的視覚語彙にマッピングし、構造的および形状情報の保持を図る。
• 離散的トークン埋め込みと元のピクセル埋め込みを連結することで、特に小さな物体の局所的詳細を保持する。
• 標準的な最適化（Adam、コサインスケジューリング）とデータオーグメンテーション（RandAug、Mixup）を用いて、両方の表現を併用してViTを訓練する。
• 小規模モデルには軽量なVQエンコーダー、ViT-Bにはより大きなエンコーダーを用い、コードブックサイズは1,024（ImageNet）および8,192（ImageNet-21K）とする。
• ViT-Bでは最終特徴量に対して平均プーリングを適用し、より高い解像度の入力（384×384、512×512）で微調整することで、一般化性能を向上させる。

実験結果

リサーチクエスチョン

• RQ1 離散的トークン化は、Vision Transformerの分布シフトに対する耐性を向上させることができるか？
• RQ2 連続的ピクセルトークンを離散的表現に置き換えることで、ViTがより形状不変な特徴を学習するよう促されるか？
• RQ3 離散的および連続的トークンの組み合わせが、ドメイン内および分布外ベンチマークにおける性能に与える影響は何か？
• RQ4 この手法は、アーキテクチャの変更なしに、さまざまなViTアーキテクチャに適用可能か？
• RQ5 離散的トークンの使用は、局所的テクスチャではなくグローバルなコンテキストを反映するような、より良い注意メカニズムをもたらすか？

主な発見

• Dr. ViTは、Stylized-ImageNetで最大12%、ImageNet-SketchおよびImageNet-Cで最大10%の耐性向上を達成し、標準的なViTを上回る。
• 本手法は、データセット固有のデータオーグメンテーションを用いない状態で、7つのImageNet耐性ベンチマークのうち4つで最先端の性能を達成した。
• 離散的トークンの追加により、注意メカニズムのパターンが改善され、モデルがグローバル構造に敏感になり、局所的テクスチャに依存しなくなることが、可視化比較で示された。
• 離散的トークンのみのバージョンは、小規模モデル（例：ViT-Ti）では性能が低いが、ピクセル埋め込みと組み合わせることで顕著に向上する。
• ViT-Bを300エポックではなく800エポック訓練することで、ImageNetで0.1〜1%の追加精度向上が得られ、離散的トークンを用いた長期訓練の利点が示された。
• 本手法は、既存の耐性向上技術と直交しており、最小限の変更で任意のViTベースアーキテクチャに統合可能である。