[論文レビュー] Pre-training of Lightweight Vision Transformers on Small Datasets with Minimally Scaled Images

Jen Hong Tan|arXiv (Cornell University)|Feb 6, 2024

Advanced Neural Network Applications被引用数 4

ひとこと要約

本論文は、画像スケーリングを最小限に抑え、マスク付き自己符号化器（MAE）を用いた事前学習によって、CIFAR-10 や CIFAR-100 のような小規模データセットにおいて、軽量なビジョントランスフォーマー（ViT）がResNetを上回ることを示している。学習可能な位置埋め込みと75%のマスキング率を用いたMAEを適用することで、370万パラメータ未満、計算コストも最小限に抑えながら、CIFAR-10で96.41%、CIFAR-100で96.01%の最先端の精度を達成した。

ABSTRACT

Can a lightweight Vision Transformer (ViT) match or exceed the performance of Convolutional Neural Networks (CNNs) like ResNet on small datasets with small image resolutions? This report demonstrates that a pure ViT can indeed achieve superior performance through pre-training, using a masked auto-encoder technique with minimal image scaling. Our experiments on the CIFAR-10 and CIFAR-100 datasets involved ViT models with fewer than 3.65 million parameters and a multiply-accumulate (MAC) count below 0.27G, qualifying them as 'lightweight' models. Unlike previous approaches, our method attains state-of-the-art performance among similar lightweight transformer-based architectures without significantly scaling up images from CIFAR-10 and CIFAR-100. This achievement underscores the efficiency of our model, not only in handling small datasets but also in effectively processing images close to their original scale.

研究の動機と目的

• 画像スケーリングを最小限に抑え、CIFAR-10 や CIFAR-100 のような小規模データセットにおいて、軽量なビジョントランスフォーマー（ViT）がResNetなどのCNNと同等またはそれを上回る性能を発揮できるかを調査すること。
• 自己教師付き事前学習を活用することで、低データ環境におけるViTとCNNの性能格差を是正すること。
• 畳み込み的インダクティブバイアスや画像の拡大を追加で行わない状態でも、ViTが最先端の性能を達成できることを実証すること。
• 最小限の増幅と小規模な画像（32×32）を用いたマスク付き自己符号化器（MAE）事前学習の有効性を評価すること。
• 小規模でパラメータが少ないViTモデルのベンチマークを確立し、より大きなまたは複雑なアーキテクチャと同等の性能を小規模データセットで発揮できるようにすること。

提案手法

• マスキング率が75%のマスク付き自己符号化器（MAE）フレームワークを採用し、エンコーダに入力されるパッチは全体の25%に限定する。
• 正弦波埋め込みの代わりに、エンコーダおよびデコーダの両方で学習可能な位置埋め込みを用いることで、空間認識能力と適応性を向上させる。
• エンコーダとは独立した柔軟なアーキテクチャを持つ別個のデコーダを設計し、学習可能な位置埋め込みを備えることで、再構成品質を向上させる。
• マスクされたパッチの再構成（MSEを用いて）と、未マスクパッチに対する割引損失を組み合わせた損失関数を計算し、訓練の安定性を向上させる。
• 標準的な交差エントロピー損失と標準的な最適化プロトコルを用いて、下流の分類タスクで事前学習済みViTモデルを微調整する。
• 元のデータの忠実性を保ち、計算コストを抑えるために、32×32から36×36への解像度のわずかな増加のみを適用してモデルを学習する。

Figure 1: Modified MAE: Implementation using separate learnable positional embeddings for Encoder and Decoder. The Encoder and the Decoder in this figure only consist of transformer layers, a layer of layerNorm and a linear projection layer.

実験結果

リサーチクエスチョン

• RQ1 画像の拡大なしに、軽量なビジョントランスフォーマーがCIFAR-10 や CIFAR-100 においてResNetを上回る性能を発揮できるか？
• RQ2 最小限のデータ増幅と小規模な画像解像度を用いたMAE事前学習により、低データ環境におけるインダクティブバイアスの欠如を補い、ViTが性能を発揮できるか？
• RQ3 エンコーダおよびデコーダの両方で学習可能な位置埋め込みを使用した場合、再構成品質および下流分類性能にどのような影響を与えるか？
• RQ4 あるデータセットで事前学習したViTが、別のデータセットに微調整することで効果的に一般化できるか、それともドメイン内での事前学習が不可欠か？
• RQ5 マスク済みパッチと未マスクパッチの両方の損失を組み合わせた損失関数が、訓練の安定性およびモデルの一般化性能に与える影響は何か？

主な発見

• 提案されたMAE事前学習済みViTモデル、Mae-ViT-C10 および Mae-ViT-C100 は、それぞれCIFAR-10で96.41%、CIFAR-100で96.01%のトップ-1精度を達成し、最先端の軽量ViTモデルを上回った。
• パラメータ数（372万）とMACs（0.26G）が同程度のViT-Lite-7/4モデルは、93.57%と73.94%の精度にとどまり、MAE事前学習戦略の優位性を示した。
• 畳み込み的インダクティブバイアスを追加しない状態でも、CVT-7/4（CIFAR-100で94.01%） や CCT-7/3×2（CIFAR-10で95.04%） などの他のViTバージョンを上回った。
• あるデータセットで事前学習し、別のデータセットで微調整した場合、収束が速くなったが、最終的な性能はわずかに低かった。これは、ドメイン内での事前学習がより効果的であることを示唆している。
• 訓練を重ねるごとに再構成品質が向上し、3020エポック目には画像の詳細が明確に精錬されていることが確認され、有効な特徴学習が行われていることを裏付けた。
• 未マスクパッチに対する割引損失の導入により、訓練の安定性が向上し、特に初期段階の一般化性能に寄与した。

Figure 2: The architecture of our MAE. ’B’ stands for batch size, ’N’ for number of embeddings