[論文レビュー] Less is More: Pay Less Attention in Vision Transformers

Zizheng Pan, Bohan Zhuang+3 more|arXiv (Cornell University)|May 29, 2021

Advanced Neural Network Applications被引用数 8

ひとこと要約

本稿では、計算コストを削減しながら性能を維持する階層的ビジョントランスフォーマーである Less is More: Pay Less Attention in Vision Transformers (LIT) を提案する。LIT は、初期の自己注意層を多層パーセプトロン (MLP) に置き換えることで、高解像度特徴マップにおける自己注意の高コストを回避する。深層部でのみ自己注意を適用し、学習可能な可変的トークンマージング (DTM) モジュールを導入することで、画像分類、物体検出、インスタンスセグメンテーションの分野で、従来のビジョントランスフォーマーと比較して顕著に少ない FLOPs とメモリ使用量で最先端の精度を達成する。

ABSTRACT

Transformers have become one of the dominant architectures in deep learning, particularly as a powerful alternative to convolutional neural networks (CNNs) in computer vision. However, Transformer training and inference in previous works can be prohibitively expensive due to the quadratic complexity of self-attention over a long sequence of representations, especially for high-resolution dense prediction tasks. To this end, we present a novel Less attention vIsion Transformer (LIT), building upon the fact that the early self-attention layers in Transformers still focus on local patterns and bring minor benefits in recent hierarchical vision Transformers. Specifically, we propose a hierarchical Transformer where we use pure multi-layer perceptrons (MLPs) to encode rich local patterns in the early stages while applying self-attention modules to capture longer dependencies in deeper layers. Moreover, we further propose a learned deformable token merging module to adaptively fuse informative patches in a non-uniform manner. The proposed LIT achieves promising performance on image recognition tasks, including image classification, object detection and instance segmentation, serving as a strong backbone for many vision tasks. Code is available at: https://github.com/zhuang-group/LIT

研究の動機と目的

• 階層的ビジョントランスフォーマー (HVT) の初期段階における自己注意の高い計算コストとメモリ使用量を軽減すること、特に高解像度特徴マップにおいて。
• 自己注意の計算量が二次関数的であることを考慮すると、HVT の初期自己注意層が最終性能に有意義に寄与しているかどうかを調査すること。
• 標準的な均一なパッチマージングよりも、より適応的で情報量の多いトークンマージング機構を設計することで、階層的表現学習を向上させること。
• 密度予測タスクの性能を損なわずに、モデルの精度と推論効率のバランスを改善すること。

提案手法

• 高解像度特徴における高コストな自己注意を回避するため、階層的ビジョントランスフォーマーの最初の2段階を純粋な MLP ブロックに置き換えることで、局所的なパターンを捉える。
• シーケンス長が短くなった深い段階でのみ、標準的なマルチヘッド自己注意 (MSA) を適用し、長距離依存関係の効率的モデリングを可能にする。
• 物体の形状や幾何変換に適応する空間オフセットを学習する、可変的トークンマージング (DTM) モジュールを提案する。これにより、幾何的変換のモデリングが向上する。
• 固定された均一なパッチマージングとは異なり、顕著な空間領域に注目する学習可能な非均一なマージング戦略を採用し、標準的な畳み込み的または線形パッチマージングを凌駕する。
• PVT や Swin Transformer と公平な比較が行えるよう、パッチサイズやチャネル次元を含む一貫したアーキテクチャハイパーパrameterを LIT に設計する。
• ImageNet、COCO、ADE20K のベンチマークで LIT を訓練・評価し、性能と効率の向上を検証する。

実験結果

リサーチクエスチョン

• RQ1 高解像度特徴マップにおける自己注意の計算コストが二次関数的であることを考えると、ビジョントランスフォーマーの初期段階での自己注意が、性能向上に顕著な寄与をしているのか？
• RQ2 初期の自己注意を MLP に置き換えることで、FLOPs やメモリ使用量を大幅に削減しながら、性能を維持または向上させることができるのか？
• RQ3 学習可能な可変的トークンマージング機構は、階層的ビジョントランスフォーマーにおける固定された均一なパッチマージングと比較して、特徴表現の質を向上させることができるのか？
• RQ4 自己注意を深層部でのみ適用することで、性能と効率の最適なバランス（最適な性能-効率の最適点）を達成できるのか？
• RQ5 提案された LIT アーキテクチャは、多様なビジョンベンチマークにおいて、精度、FLOPs、メモリ消費量の観点で最先端のビジョントランスフォーマーと比較してどうなるのか？

主な発見

• LIT は、1.8G FLOPs のみで ImageNet-1K 分類で 85.7% のトップ-1精度を達成し、Swin-Tiny や PVT-S を上回る性能を示した。パラメータ数も少ない。
• COCO 物体検出およびインスタンスセグメンテーションにおいて、LIT-B はそれぞれ 46.8 AP および 42.3 マスク AP を達成し、FLOPs を低く抑えながら Swin-B や PVT-B と同等またはそれ以上の性能を示した。
• 可変的トークンマージング (DTM) モジュールは、物体の形状や幾何的変換に適応する空間オフセットを学習し、標準的なマージング手法よりも特徴表現の質が向上した。
• 可視化の結果、標準的な HVT における初期 MSA レイヤーは、わずかな局所領域にのみ注目していることが判明し、初期段階での自己注意の恩恵は限定的であることが示された。
• Swin-Tiny や PVT-S と比較して、LIT は最大 30% の FLOPs 削減を達成しながら、精度を維持または向上させた。これは、優れた効率-精度のトレードオフを示している。
• アブレーションスタディの結果、初期の自己注意を MLP に置き換えることで、性能の低下がほとんどないまま顕著な効率的向上が得られ、LIT のコア設計原理が妥当であることが確認された。