[論文レビュー] Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet

Yuan, Li, Yunpeng Chen+7 more|arXiv (Cornell University)|Jan 28, 2021

Multimodal Machine Learning Applications被引用数 32

ひとこと要約

この論文は T2T-ViT を紹介します。Tokens-to-Token モジュールを用いて画像を段階的にトークン化し、深く狭いバックボーンを用いることで、ImageNet のスクラッチ学習をパラメータ数と FLOPs を削減しつつ ViT や強力な CNN ベースラインより高い精度を達成します。

ABSTRACT

Transformers, which are popular for language modeling, have been explored for solving vision tasks recently, e.g., the Vision Transformer (ViT) for image classification. The ViT model splits each image into a sequence of tokens with fixed length and then applies multiple Transformer layers to model their global relation for classification. However, ViT achieves inferior performance to CNNs when trained from scratch on a midsize dataset like ImageNet. We find it is because: 1) the simple tokenization of input images fails to model the important local structure such as edges and lines among neighboring pixels, leading to low training sample efficiency; 2) the redundant attention backbone design of ViT leads to limited feature richness for fixed computation budgets and limited training samples. To overcome such limitations, we propose a new Tokens-To-Token Vision Transformer (T2T-ViT), which incorporates 1) a layer-wise Tokens-to-Token (T2T) transformation to progressively structurize the image to tokens by recursively aggregating neighboring Tokens into one Token (Tokens-to-Token), such that local structure represented by surrounding tokens can be modeled and tokens length can be reduced; 2) an efficient backbone with a deep-narrow structure for vision transformer motivated by CNN architecture design after empirical study. Notably, T2T-ViT reduces the parameter count and MACs of vanilla ViT by half, while achieving more than 3.0\% improvement when trained from scratch on ImageNet. It also outperforms ResNets and achieves comparable performance with MobileNets by directly training on ImageNet. For example, T2T-ViT with comparable size to ResNet50 (21.5M parameters) can achieve 83.3\% top1 accuracy in image resolution 384$ imes$384 on ImageNet. (Code: https://github.com/yitu-opensource/T2T-ViT)

研究の動機と目的

• 純粋な Transformer アーキテクチャが ImageNet のような中規模データセットでスクラッチ学習すると CNN に劣る理由を動機づける。
• Tokens-to-Token (T2T) モジュールを提案し、局所的な画像構造を捉え、トークン長を反復的に削減する。
• 効率的な深く狭い ViT バックボーンを設計し、特徴の豊富さを改善し冗長性を減らす。
• T2T-ViT が ImageNet で同程度のサイズの CNN に対して大規模事前学習なしで上回ることを示す。
• CNN に触発されたアーキテクチャ設計が ViT バックボーンに有益であることを示す。

提案手法

• 層ごとに Re-tructurization と Soft Split を交互に行う Tokens-to-Token (T2T) モジュールを導入し、局所構造が埋め込まれたトークンへと画像を徐々に変換する。
• パラメータと MACs を減らしつつ性能を維持するため、小さめの隠れ層次とより多い層を持つ深く狭い ViT バックボーンを使用する。
• T2T モジュール内で Transformer と Performer 層を組み合わせ、メモリと計算を管理する。
• 比較対象として、同程度のモデルサイズで ImageNet 上の ViT、ResNets、MobileNets と比較する。
• T2T モジュールと深く狭いアーキテクチャの影響を定量化するアブレーションを行い、CIFAR-10/100 への転移を検討する。

実験結果

リサーチクエスチョン

• RQ1 progressive なトークン-to-トークン モジュールは、ImageNet からスクラッチで学習した ViT に対して naive なトークン化より局所的な画像構造をより上手く捉えられるか？
• RQ2 CNN にインスパイアされた深く狭いバックボーンは、標準な ViT と比較して冗長性を削減し、視覚変換器の特徴の豊富さを改善するか？
• RQ3 ImageNet でスクラッチから訓練した場合、同程度のパラメータ数と計算予算で ResNets と MobileNets に対して T2T-ViT はどのように性能差を示すか？
• RQ4 異なる T2T モジュールのバリアント（Transformer vs Performer）を使用することは、性能と効率性にどのような影響を与えるか？
• RQ5 事前学習済みの T2T-ViT モデルは CIFAR-10/100 のような下流データセットへ効果的に転移できるか？

主な発見

Models	Top1-Acc (%)	Params (M)	MACs (G)
ViT-S/16 [12]	78.1	48.6	10.1
DeiT-small [36]	79.9	22.1	4.6
DeiT-small-Distilled [36]	81.2	22.1	4.7
T2T-ViT-14	81.5	21.5	4.8
T2T-ViT-14↑384	83.3	21.5	17.1
ViT-B/16 [12]	79.8	86.4	17.6
ViT-L/16 [12]	81.1	304.3	63.6
T2T-ViT-24	82.3	64.1	13.8
T2T-ViT t-14	81.7	21.5	6.1
T2T-ViT t-24	82.6	64.1	15.0

• 21.5M パラメータと 4.8G MACs を持つ T2T-ViT は、ImageNet（224x224）をスクラッチから訓練した場合に 81.5% の top-1 精度を達成し、ViT-S/16 を上回り、同程度のサイズの ResNets を匹敵・上回る。
• 384x384 の入力で T2T-ViT-14↑ は 83.3% の top-1 精度を達成し、高解像度での強力な利得を示す。
• ResNet50（25.5M パラメータ、4.3G MACs）と比較して、 T2T-ViT-14 は 81.5% の精度（t-バリアントでは 6.1G MACs）を達成し、同等以上の計算量でより良い精度を示す。
• T2T-ViT-24 は 64.1M パラメータと 13.8G MACs で 82.3% の top-1 を達成し、より大規模なスケールで競争力のある性能を示す。
• Lite な T2T-ViT モデル（例：T2T-ViT-7/12）は MobileNets に競争力のある精度を提供する一方、MACs は高い。ただし蒸留を用いると小型モデルの性能がさらに向上する。
• 事前学習済みの T2T-ViT を CIFAR-10/100 に転移させると、ViT ベースラインより競争力のある成果を示し、転移性を示している。