[論文レビュー] Fast Vision Transformers with HiLo Attention

Zizheng Pan, Jianfei Cai+1 more|arXiv (Cornell University)|May 26, 2022

Advanced Neural Network Applications被引用数 112

ひとこと要約

tldr: 論文は HiLo 注意機構を用いた高速 Vision Transformer のバックボーン LITv2 を提案し、局所の高周波と全球の低周波の相互作用を分離して高速スループットと競争力の精度を実現する。

ABSTRACT

Vision Transformers (ViTs) have triggered the most recent and significant breakthroughs in computer vision. Their efficient designs are mostly guided by the indirect metric of computational complexity, i.e., FLOPs, which however has a clear gap with the direct metric such as throughput. Thus, we propose to use the direct speed evaluation on the target platform as the design principle for efficient ViTs. Particularly, we introduce LITv2, a simple and effective ViT which performs favourably against the existing state-of-the-art methods across a spectrum of different model sizes with faster speed. At the core of LITv2 is a novel self-attention mechanism, which we dub HiLo. HiLo is inspired by the insight that high frequencies in an image capture local fine details and low frequencies focus on global structures, whereas a multi-head self-attention layer neglects the characteristic of different frequencies. Therefore, we propose to disentangle the high/low frequency patterns in an attention layer by separating the heads into two groups, where one group encodes high frequencies via self-attention within each local window, and another group encodes low frequencies by performing global attention between the average-pooled low-frequency keys and values from each window and each query position in the input feature map. Benefiting from the efficient design for both groups, we show that HiLo is superior to the existing attention mechanisms by comprehensively benchmarking FLOPs, speed and memory consumption on GPUs and CPUs. For example, HiLo is 1.4x faster than spatial reduction attention and 1.6x faster than local window attention on CPUs. Powered by HiLo, LITv2 serves as a strong backbone for mainstream vision tasks including image classification, dense detection and segmentation. Code is available at https://github.com/ziplab/LITv2.

研究の動機と目的

• ViT の評価を FLOPs のみでなく、対象ハードウェア上の直接的な速度測定で行う動機付け。
• 高周波の局所的な詳細と低周波のグローバル構造を個別にモデル化する HiLo 注意を提案。
• HiLo ベースの LITv2 が ImageNet、COCO、ADE20K で高速なスループットと競争力のある精度を達成することを示す。
• 先進的な ViT や一部の CNN に対してメモリ使用量と効率性の利点を示す。

提案手法

• HiLo 注意を導入し、注意ヘッドを高周波（Hi-Fi）と低周波（Lo-Fi）グループに分割。
• Hi-Fi は高解像度マップ上で局所ウィンドウ自己注意を使用（例：2x2 ウィンドウ）。
• Lo-Fi は各ウィンドウに平均プーリングを適用して低周波信号を得て、これらのプーリング信号からダウンサンプルしたキー/値で注意を実行。
• ヘッドを割合 α で分割し、（1−α）N_h を Hi-Fi、αN_h を Lo-Fi に割り当て、出力を結合。
• 固定相対位置エンコーディングを FFN のディープウェイ 3x3 畳み込みに置換して位置情報を注入し受容野を拡大。
• LITv2 の三つのモデルサイズ（S/M/B）で HiLo を採用し、ImageNet-1K で訓練、下流タスクで評価。

実験結果

リサーチクエスチョン

• RQ1 HiLo 注意は標準的な MSA および他の効率的なアテンションと比較して、CPU/GPU 上で速度アップを提供しつつ競争力の精度を維持できるか。
• RQ2 HiLo の高周波/低周波の分離は高解像度および密な予測タスクの性能にどう影響するか。
• RQ3 畳み込み FFN や RPE の除去などのアーキテクチャ変更は、視覚タスク全般で速度と精度をさらに改善するか。
• RQ4 異なるワークロードと解像度に対して最適なヘッド分割 α とウィンドウサイズ s はどれか。
• RQ5 HiLo ベースのバックボーンは下流の検出/分割タスクで SoTA ViTs と比較してスループットとメモリ使用を改善するか。

主な発見

Model	Param (M)	FLOPs (G)	Throughput (imgs/s)	Train Mem (GB)	Test Mem (GB)	Top-1 (%)
ResNet-50 resnet_back	26	4.1	1,279	2.8	80.4
ConvNext-Ti convnext	28	4.5	1,079	1.7	82.1
PVT-S pvt	25	3.8	1,007	1.3	79.8
Swin-Ti swin	28	4.5	961	1.5	81.3
CvT-13 cvt	20	4.5	947	1.5	81.6
Focal-Tiny yang2021focal	29	4.9	384	3.3	82.2
Twins-PCPVT-S chu2021Twins	24	3.8	998	1.2	81.2
LITv1-S lit	27	4.1	1,298	1.2	81.5
LITv2-S	28	3.7	1,471	1.2	82.0
ResNet-101 resnet_back	45	7.9	722	3.0	81.5
ConvNext-S convnext	50	8.7	639	1.8	83.1
PVT-M pvt	44	6.7	680	1.5	81.2
Twins-SVT-B chu2021Twins	56	8.3	621	1.9	83.2
Swin-S swin	50	8.7	582	1.7	83.0
LITv1-M lit	48	8.6	638	1.4	83.0
LITv2-M	49	7.5	812	1.4	83.3
ResNet-152 resnet_back	60	11.6	512	2.9	82.0
ConvNext-101 convnext	89	15.4	469	2.9	83.8
Twins-SVT-L chu2021Twins	99	14.8	440	3.1	83.7
Swin-B swin	88	15.4	386	2.4	83.3
LITv1-B lit	86	15.0	444	2.1	83.4
LITv2-B	87	13.2	602	2.1	83.6
DeiT-B ↑384 deit	86	55.4	159	2.5	83.1
Swin-B ↑384 swin	88	47.1	142	OOM	6.1	84.5
LITv2-B ↑384	87	39.7	198	4.6	84.7

• HiLo を用いた LITv2 は FLOPs が少なくとも LITv1 よりも精度が上で、スループットも速い。
• HiLo は GPU での推論速度を ImageNet-1K 実測値で LITv1 と比較して S: 13%、M: 27%、B: 36% の速度向上をもたらす。
• 比較可能な FLOPs の場合、LITv2-S は ImageNet-1K で PVT-S および Twins-PCPVT-S よりも高速。
• COCO の物体検出/インスタンスセグメンテーションで、LITv2 バックボーンは ResNet バックボーンより高い AP を示し、SoTA ViTs と競合または優位、FPS と FLOPs の点で有利。
• ADE20K のセマンティックセグメンテーションで、LITv2 系は複数のベースラインより高い mIoU と FPS を達成（例：LITv2-S/LITv2-M/LITv2-B が Swin-Ti/Swin-S/Swin-B を mIoU で上回り FPS が増加）。
• HiLo は ImageNet-1K および下流タスクで他の効率的なアテンション機構（SRA、W-MSA、T-MSA）に対して FLOPs を一貫して削減し速度を改善する。