Training data-efficient image transformers & distillation through attention (DeiT)

どんなもの？

• Vision Transformer(ViT)は画像分類タスクでSOTAを達成しているが、大規模な学習データ(JFT-300M, 約3億枚)と計算コスト(680~2,500 TPUv3-days)を必要とする。
• DeiTはトークンベースの蒸留と学習方法の工夫によって、ViTよりも精度とスループットのトレードオフに優れたモデルの学習に成功した。トークンベースの蒸留は具体的には以下。
  • CLSトークンとは別に蒸留トークンを用意する。
  • 学習時にGTラベルを当てさせるのはCLSトークンによって行い、Teacherモデルの推論結果を当てさせるのは蒸留トークンによって行う。
  • 推論時はCLSトークンの出力と蒸留トークンの出力の平均を用いる。
• 学習データはImageNet(100万枚)のみで、計算コストはシングルPCで3日未満。

技術や手法のキモ

Distillation through attention

Distillation token

• 蒸留用のトークン(distillation token)を追加。CLSトークンと同様にtrainable。最終レイヤの蒸留トークンに対応する出力を、蒸留用のロス計算に使用する。蒸留トークンによってクラス分類用のembeddingを残しつつ、通常の蒸留のように教師モデルの出力からの学習が可能。
• 蒸留トークンとCLSトークンは異なるベクトルに収束した。二つのコサイン類似度は0.06。各レイヤの蒸留トークンとCLSトークンに対応するembeddingの類似度を見ると、レイヤが深くなるにつれて徐々に高くなり、最終レイヤでは0.93になるがそれでも1未満。各トークンは互いに似ているが、同一ではないようなembeddingを学習しようとしていると予想。