[論文レビュー] Learning to Merge Tokens in Vision Transformers

Cédric Renggli, André Susano Pinto+4 more|arXiv (Cornell University)|Feb 24, 2022

Multimodal Machine Learning Applications被引用数 24

ひとこと要約

PatchMerger を導入し、ViT および V-MoE モデルのトランスフォーマーブロック間のトークン数を削減する学習可能なモジュールを提示。精度をほぼ維持しつつ、計算量を 40-60% 削減。

ABSTRACT

Transformers are widely applied to solve natural language understanding and computer vision tasks. While scaling up these architectures leads to improved performance, it often comes at the expense of much higher computational costs. In order for large-scale models to remain practical in real-world systems, there is a need for reducing their computational overhead. In this work, we present the PatchMerger, a simple module that reduces the number of patches or tokens the network has to process by merging them between two consecutive intermediate layers. We show that the PatchMerger achieves a significant speedup across various model sizes while matching the original performance both upstream and downstream after fine-tuning.

研究の動機と目的

• Vision Transformers の大規模化に伴う計算コストの増大を動機づけ、対応する。
• エンコーダーブロック間でトークンを統合する PatchMerger を提案し、トークン数を削減する。
• 出力トークン数を不変に保ちつつ、異なる解像度でファインチューニングを可能にする。
• PatchMerger が性能を維持しつつ、FLOPs および実行時間の大幅な節約を実現することを示す。

提案手法

• PatchMerger は二つのエンコーダーブロックの間に学習可能な統合レイヤーを配置し、N 個の入力トークンから固定の小さなトークン数 M を出力する（例: 8）。
• 層正規化後に学習可能な W ∈ R^{D×M} および入力 X ∈ R^{N×D} を用いて Y = softmax(W^T X^T) X を計算し、ボトムアップ型のアテンション風の統合を実行する。
• 出力トークン数 M は入力トークン数 N に依存せず固定され、解像度を跨いだ柔軟なファインチューニングを可能にする。
• PatchMerger に続くレイヤー正規化は、入力/出力トークン数が異なる場合のスケールを安定化する。
• ViT および V-MoE バックボーンを JFT でトレーニングし、ImageNet への下流転移を評価する（Few-shot およびフルファインチューニング）。
• 設計選択の評価: 配置場所（どのエンコーダーブロックに置くか）と出力トークン数の目標値（8 対他の値）。

実験結果

リサーチクエスチョン

• RQ1 Vision Transformers で、計算コストを削減しつつ精度を損なわずにトークンを統合できる軽量モジュールは存在するか？
• RQ2 コストと性能のトレードオフの観点から、ネットワーク内のどこでトークン統合を行うべきか？
• RQ3 PatchMerger はどの程度の出力トークン数を生成すべきか、効率と精度のバランスはどうか？
• RQ4 大規模データセットで前学習した後、下流タスク（例: ImageNet）への転移性能は保たれるか？
• RQ5 標準の ViT と mixture-of-experts（V-MoE）バックボーンに対して、PatchMerger のバリエーションはどのような性能になるか？

主な発見

TPU-Days	ExoFLOPs	JFT Prec@1	INet/1shot	INet/5shot	INet/10shot	INet/finetuned
ViT-H/14	2387.71	4275.92	56.56	61.35	76.76	79.12	87.97
Merger ViT-H/11	1920.82	3464.17	58.15	62.77	77.53	79.84	88.24
Merger ViT-H/14	1518.49	2207.33	57.27	62.07	77.25	79.26	87.90
ViT-L/16	651.18	1572.74	53.37	59.04	73.92	76.51	86.84
Merger ViT-L/16	388.44	819.29	53.71	59.99	74.99	77.01	87.01
ViT-L/32	97.29	196.11	46.52	50.34	66.59	69.58	84.45
Merger ViT-L/32	66.49	114.38	46.61	51.29	67.61	70.32	84.08
ViT-B/16	95.04	224.45	44.58	48.21	63.50	66.94	–
Merger ViT-B/16	54.22	117.22	43.75	49.49	64.35	67.15	–
ViT-B/32	27.62	56.07	39.48	40.66	55.47	58.93	81.39
Merger ViT-B/32	17.76	32.99	38.62	40.53	57.23	60.02	79.96
ViT-S/32	7.72	12.27	29.05	29.37	43.21	46.38	76.34
Merger ViT-S/32	4.83	7.34	27.82	26.42	40.27	43.75	73.77
V-MoE-H/14	2735.70	4750.73	60.12	62.95	78.08	80.10	88.41
Merger V-MoE-H/14	1520.59	2220.96	59.84	63.21	77.56	79.97	88.10
V-MoE-L/16	761.27	1666.10	56.76	61.46	76.53	78.64	87.47
Merger V-MoE-L/16	388.19	823.54	56.14	61.72	75.86	78.03	87.04
V-MoE-L/32	110.65	207.94	49.68	54.52	69.90	72.80	85.05
Merger V-MoE-L/32	67.11	116.70	48.75	52.03	68.98	71.63	84.46
V-MoE-B/16	130.86	250.70	48.31	54.92	68.84	71.81	–
Merger V-MoE-B/16	54.78	118.59	46.65	53.33	67.53	70.05	–
V-MoE-B/32	36.80	62.75	43.17	48.04	62.45	65.72	83.27
Merger V-MoE-B/32	18.51	34.61	41.77	44.88	60.24	63.62	80.57
V-MoE-S/32	12.40	14.40	33.26	35.49	50.90	54.16	78.63
Merger V-MoE-S/32	5.33	7.95	32.85	33.59	49.23	52.74	75.07

• PatchMerger は、 backbone ViT および V-MoE モデルと比較して、upstream および downstream の性能をほぼ維持しつつ、計算量を大幅に削減する（大規模モデルで 40-60%）。
• Merger ViT バリアントは FLOPs 対精度のパレートフロンティアで ViT の他の手法を上回り、特に大規模スケールで顕著な利得を示す（例: ViT-H/14 vs Merger ViT-H/14）。
• ImageNet の完全ファインチューニング時において、Large Merger モデル（L/32, L/16, H/14）はバックボーンの性能と同等で、FLOPs/実行時間の大幅な節約を実現する（例: ViT-H/14: 51.6% FLOPs, 63.5% runtime）。
• ViT-H/11 など初期トークン数が多い場合でも、Merger バリアントは標準バックボーンと比較して upstream、few-shot、ファインチューニング精度で優れており、コストは抑えられる。
• 大型 V-MoE Merger モデルは、ImageNet での下流結果が同等または改善されつつ、計算量を substantial に削減する（約 53% FLOPs、44% runtime）。