【MobileViT】携帯で使用可能な軽量ViT

画像認識 2021年11月05日

3つの要点
✔️ CNNとViTの良いところ取りをした携帯でも使える軽量モデル
✔️ ViTの学習に必要な複雑なデータ水増しを必要としない
✔️ 様々なタスクのバックボーンとして使え、全ての実験でSOTAを達成

MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer
written by Sachin Mehta, Mohammad Rastegari
(Submitted on 5 Oct 2021)
Comments: Published on arxiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV); Artificial Intelligence (cs.AI); Machine Learning (cs.LG)

code：

本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。

はじめに

Transformerは機械翻訳で提案されて以来BERTをはじめとした幅広い自然言語処理のタスクで応用されていますが、画像認識分野に応用したVision Transformer(ViT)でも様々なタスクでSOTAを生み出し続けています。

しかし、ViTはCNNのような局所的な特徴を上手く検出できる機構を持たないため、CNN以上の性能を出すためには大規模なデータセットやCNN以上のパラメータ数が必要なことが携帯端末で利用する上でネックとなっていました。その一方、ViTが持つself attention機構は大域的な関係を考慮することができる特性がありCNNに代わるアーキテクチャとして注目されています。

今回紹介するmobileViTは、局所的な検出を得意とするCNNと大域的な情報処理を得意とするViTをハイブリッドにした軽量モデルです。mobileViTはCNN以下のパラメータ数でそれ以上の性能を達成するだけでなく、基本的なデータ水増しでも精度が出るモデルとなっています。また、物体検出やセマンティックセグメンテーションなどのタスクのバックボーンとしても同程度のパラメータで最も良い結果を得られたと論文では報告されています。

論文概要

mobileViTが従来のViTと違う点は次の2つです。

パラメータ数がCNNと同程度で精度が出ること
多様なデータ水増しが必要でないこと

パラメータ数がCNNと同程度で精度が出ること

Attention機構は離れたパッチの情報も考慮できる特性を持ちますが、畳み込み層が持つ局所的な情報を集約する特性は持ってはいないので、CNNと同程度の性能を持つためにはより多くのパラメータが必要となることが知られています。例えば、セマンティックセグメンテーションでは、ViTベースのDPTでは345M、CNNベースのDeepLabv3では59Mだけで同程度の性能が出ることが挙げられます。しかし、mobileViTはCNNと同程度でそれ以上の精度を出すことができます。

多様なデータ水増しが必要でないこと

ViTは過学習をする傾向があり多様なデータ水増しが必要であるとされています。また、L2正則化の影響も受けやすく学習が困難でした。しかし、この論文では、基本的なデータ水増しであるrandom resized cropping、horizontal croppingだけで十分な結果を得られています。

モデルアーキテクチャ

mobileViTのアーキテクチャは以下の通りです。MV2はmobilenetv2ブロックで、↓2はダウンサンプリングを表しています。MV2で局所的な特徴を捉え、新規提案ブロックであるMobileViTブロックで大域的な関係性を見出しています。

MobileViTブロック

MobileViTブロックでは、3つの処理を行います。

Local representations
Transformer as Convolutions
Fusion

Local Representations

入力テンソルX∈R^H×W×Cは、n×nの通常の畳み込み層を経て、point-wiseな畳み込み層で高次元なテンソルX_L∈R^H×W×dに変換されます。dはCより大きいです。実装コードでは、n×nの通常の畳み込み層としてpaddingが1の3×3の畳み込み層が用いられています。

Transformer as Convolutions

局所的な情報を含むX_Lは、次にパッチごとに分割されX_U∈R^P^×N×dに変換されます。Nはパッチの個数でPは幅w、高さがhのパッチです。
そしてパッチのピクセルごとにTransformerに流されます。

これらのイメージを表したのが下の画像です。対象となっている赤のピクセルから他のパッチの青いピクセルに矢印が向いているがTransformerに流して大域的な関係性を見ていることに対応しています。青いピクセルから出ている青の矢印はLocal representationsを経てそのピクセルに局所的な情報が詰まっていることを表しています。

こうして得られたX_G∈R^P^×N×dは再びX_F∈R^H×W×dに変換されます。

Fusion

並び直されたテンソルX_Fはpoint-wiseな畳み込みでC次元に落とされR^H×W×Cとなり、元のテンソルX∈R^H×W×Cと結合されR^H×W×2Cとなり、n×nの畳み込みでR^H×W×Cに縮約されます。

著者はmobileViTブロックを畳み込み層と同様の機能を果たすTransformer(Transformer as Convolution)としています。その根拠として、畳み込み層が特徴量を部分ごとに分解し、行列演算を行い、それらの結果を集積する点でmobileViTブロックと共通しているからとしています。

実験

CNNとの比較(精度)

ImageNet-1kにおいて一から学習を行い、様々なモデルと提案モデルの性能を比較した結果が以下の通りです。軽量な同程度のパラメータのCNNモデルより5%以上良く、重量なCNNモデルと比較しても少ないパラメータでそれ以上の精度を達成しています。一時期話題になった 2019年のEfficientNetと比較しても2.1%精度が良くなっています。

その他のViTとの比較(精度)

様々なViTと比較してもmobileViTは効率的なことが分かります。見ての通り、最も少ないパラメータ数であるのにも関わらず、最も良い性能を達成できています。それだけでなく、基本的なデータ水増し手法であるrandom resized cropping、horizontal croppingだけでこのような結果が得られています。
PiTはAdvancedなデータ水増し手法を用いることで5.7%精度が上がっていますが(R4とR17を参照)、それ以上の性能を基本的なデータ水増しで実現できていることから使用しやすいことが分かります。

MobileViTをBackboneとして使う

物体検出

物体検出でも最もパラメータ数が少ないにも関わらず最も良い精度を得ることができています。

実験では、SSDを軽量化したSSDLiteで比較が行われています。このSSDLiteのヘッドがseparable convolutionになっているのがSSDとの違いです。バックボーンとして様々なImageNet-1kの学習済みモデルを利用し、MS-COCOデータセットでfine tuningを行なった結果が以下の通りです。

セマンティックセグメンテーション

セマンティックセグメンテーションでも最小パラメータ数で最も良い結果が得られています。

実験は、DeepLabv3で比較が行われています。バックボーンとしてImageNet-1kの学習済みモデルを用い、PASCAL VOC2012でfine tuningが行なっています。指標はmIOU(mean intersection over union)が用いられています。以下が実験の結果です。Resnet-101の10分の1のサイズであるにも関わらず同程度の性能を得られることがmobileViTの性能の良さを物語っています。

学習効率のためのマルチスケールサンプラー

通常のViTでは様々な解像度の画像を学習する際にfine tuningが主に使われています。DeITは画像サイズが変わってしまうとpositional embeddingを補間する必要があるため、224×224の画像に固定しています。

MobileViTではpositional embeddingを用いる必要がないのでfine tuningを行う必要がなく、複数の解像度で学習を行うマルチスケールサンプラーで学習が行われています。

マルチスケールサンプラーでは、GPUを効率的に用いるために解像度によってバッチサイズを変更します。用いる画像の解像度の集合S={(H₁,W₁),⋯,(H_n,W_n)}で(H_n,W_n)が最も解像度が高いものとした時に、t回目のイタレーションで用いる解像度(H_t,W_t)を用いて以下のバッチサイズb_tにします。

携帯端末上での処理速度

実験は、CoreMLToolsを用いてCoreMLに変換されたfull precisionの学習済みmobileViTをiPhone12で100回した平均で比較されています。
MobileViTやその他のViTは、MobileNetv2より遅い推論時間となっていますが、これには2つ理由があるとしています。

GPUのTransformerに対応したCUDAのカーネルを拡張性や効率化のために用いていないこと
CNNではバッチ正則化と畳み込み層を組み合わせたハードウェアレベルの最適化があるが、Transformerにはこのようなモジュールが存在しないこと

まとめ

ViTは局所的な特徴を捉えることができないためにCNNに比べてパラメータ数が必要であると考えられてきましたが、今回提案されたmobileViTは同程度のパラメータ数でそれ以上の性能を得られます。また、基本的なデータ水増し手法で十分な性能が出たり、物体検出やセマンティックセグメンテーションのバックボーンとしても性能が優れていて、今後の研究はmobileViTがベースとなっていくのではないでしょうか。しかし、依然としてCNNに比べて推論時間が遅い問題があり、NLPのTransformerの計算量を下げる研究が応用されていくと考えられます。