[論文レビュー] FastViT: A Fast Hybrid Vision Transformer using Structural Reparameterization

Pavan Kumar Anasosalu Vasu, James Gabriel+3 more|arXiv (Cornell University)|Mar 24, 2023

Advanced Neural Network Applications被引用数 45

ひとこと要約

FastViTはRepMixerを導入します。再パラメータライズ可能なトークンミキサーで、高速なハイブリッドビジョントランスフォーマーを実現し、モバイルとGPUプラットフォーム全体で優れた遅延-精度のトレードオフを達成し、複数のビジョンタスクをサポートします。

ABSTRACT

The recent amalgamation of transformer and convolutional designs has led to steady improvements in accuracy and efficiency of the models. In this work, we introduce FastViT, a hybrid vision transformer architecture that obtains the state-of-the-art latency-accuracy trade-off. To this end, we introduce a novel token mixing operator, RepMixer, a building block of FastViT, that uses structural reparameterization to lower the memory access cost by removing skip-connections in the network. We further apply train-time overparametrization and large kernel convolutions to boost accuracy and empirically show that these choices have minimal effect on latency. We show that - our model is 3.5x faster than CMT, a recent state-of-the-art hybrid transformer architecture, 4.9x faster than EfficientNet, and 1.9x faster than ConvNeXt on a mobile device for the same accuracy on the ImageNet dataset. At similar latency, our model obtains 4.2% better Top-1 accuracy on ImageNet than MobileOne. Our model consistently outperforms competing architectures across several tasks -- image classification, detection, segmentation and 3D mesh regression with significant improvement in latency on both a mobile device and a desktop GPU. Furthermore, our model is highly robust to out-of-distribution samples and corruptions, improving over competing robust models. Code and models are available at https://github.com/apple/ml-fastvit.

研究の動機と目的

• モバイル機器とデスクトップで精度と遅延のバランスを取る効率的なビジョンモデルを動機づける。
• 畳み込みとトランスフォーマーの長所を活用するハイブリッドアーキテクチャを開発する。
• スキップ接続の構造的再パラメータライズによりメモリアクセスコストを削減する。
• 学習時のオーバーパラメータライズと大きなカーネルを用いて遅延ペナルティなしにモデル容量を増加させる。
• 分類、検出、セグメンテーション、3D手メッシュ推定などのタスクで頑健性と一般化を示す。

提案手法

• 推論時に再パラメータライズ可能で、スキップ接続を取り除くトークンミキサーRepMixerを導入する。
• 密なk×k畳み込みを、ファクタライズされた深さ方向畳み込みとポイントワイズ畳み込みに置換し、線形な学習時オーバーパラメータライズによって補助する。
• FFNとパッチ埋め込み層に大きなカーネル畳み込みを組み込み、初期の自己注意を代替する。
• 深さ方向畳み込みによって生成される条件付き位置エンコーディングを使用する。
• ステム、パッチ埋め込み、および投影層でオーバーパラメータライズを用いて容量を増強して訓練する。

実験結果

リサーチクエスチョン

• RQ1 再パラメータライズ可能なトークンミキサーは、精度を犠牲にせずメモリアクセスコストと遅延を低減できるか。
• RQ2 因子分解された畳み込み設計の下で、線形な学習時オーバーパラメータライズは精度を改善しますか。
• RQ3 初期段階の大きなカーネル畳み込みは、ハイブリッドアーキテクチャにおける自己注意と比べて遅延に配慮した精度向上をもたらしますか。
• RQ4 実世界の遅延制約の下で、FastViTは画像分類、検出、セグメンテーション、3D手メッシュ推定の各タスクでどのように性能を発揮しますか。
• RQ5 競合アーキテクチャと比較して、ノイズや分布外入力に対してモデルは頑健ですか。

主な発見

• FastViTは、モバイル（iPhone 12 Pro）とデスクトップGPU（RTX-2080Ti）で優れた遅延-精度トレードオフを達成しつつ、競争力のある精度を維持します。
• FastViT-S12はImageNet-1kでTop-1 83.9%で、同等のFLOPsでモバイル上のEfficientNet-B5より4.9×、ConvNeXt-Bより1.9×速く、GPU上のEfficientNetV2-Sより1.6×速い。
• FastViT-S12はiPhoneでMobileOne-S4より26.3%、GPUでは同等精度で26.9%速い。FastViT-MA36は、パラメータ数とFLOPsを削減しつつ、いくつかの最先端モデルに匹敵するかそれを上回る。
• RepMixerはスキップ接続を取り除くことでメモリアクセスコストを削減し、特に高解像度入力（例: 384×384、1024×1024）で低遅延を実現します。
• ステム、パッチ埋め込み、投影層での学習時オーバーパラメータライズは精度向上をもたらし（例: ImageNetでTop-1最大0.9%）、訓練時間オーバーヘッドは控えめです。
• FFNとパッチ埋め込み層の大きなカーネル畳み込みは、遅延の影響を適度に抑えつつ、頑健性と精度を向上させます。