[論文レビュー] Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers

John Guibas, Morteza Mardani+4 more|arXiv (Cornell University)|Nov 24, 2021

CCD and CMOS Imaging Sensors被引用数 78

ひとこと要約

この論文は Adaptive Fourier Neural Operators (AFNO) を vision transformer の効率的なトークンミキサーとして提案し、フーリエベースのグローバル畳み込みと適応的でスパースなブロック構造のチャネル混合を組み合わせることで、実質的にリニアに近い計算量と下流パフォーマンスの向上を実現する。特に高解像度および few-shot のセマンティックセグメンテーションタスクで効果を発揮する。

ABSTRACT

Vision transformers have delivered tremendous success in representation learning. This is primarily due to effective token mixing through self attention. However, this scales quadratically with the number of pixels, which becomes infeasible for high-resolution inputs. To cope with this challenge, we propose Adaptive Fourier Neural Operator (AFNO) as an efficient token mixer that learns to mix in the Fourier domain. AFNO is based on a principled foundation of operator learning which allows us to frame token mixing as a continuous global convolution without any dependence on the input resolution. This principle was previously used to design FNO, which solves global convolution efficiently in the Fourier domain and has shown promise in learning challenging PDEs. To handle challenges in visual representation learning such as discontinuities in images and high resolution inputs, we propose principled architectural modifications to FNO which results in memory and computational efficiency. This includes imposing a block-diagonal structure on the channel mixing weights, adaptively sharing weights across tokens, and sparsifying the frequency modes via soft-thresholding and shrinkage. The resulting model is highly parallel with a quasi-linear complexity and has linear memory in the sequence size. AFNO outperforms self-attention mechanisms for few-shot segmentation in terms of both efficiency and accuracy. For Cityscapes segmentation with the Segformer-B3 backbone, AFNO can handle a sequence size of 65k and outperforms other efficient self-attention mechanisms.

研究の動機と目的

• 高解像度のビジョン・トランスフォーマーにおける二次項的な自己注意計算量を超える、効率的なトークン混合を動機づける。
• 離散的不連続性と高解像度を持つ画像に対応するために Fourier Neural Operators を適用して AFNO を開発する。
• 効率性と一般化を改善するためにアーキテクチャの修正（ブロック対角ウェイト、適応的ウェイト共有、スパース性）を導入する。
• アップサンプリング/インペインティング pretrained、few-shot segmentation、Cityscapes/高解像度セグメンテーションで AFNO の有効性を実証する。
• 最先端ミキサーとのコードと経験的比較を提供する。

提案手法

• トークン混合を演算子学習による連続グローバル畳み込みとしてフレーム化する。
• FFTベースの離散実装を用いてトークン混合を行う Fourier Neural Operators (FNO) を採用する。
• パラメータ数を削減するためにチャネル混合ウェイトにブロック対角構造を課す。
• トークン間で共有された二層MLPを通じた適応的ウェイト共有を導入する。
• 情報量の多い成分に焦点を当てるために周波数モードへソフトスレッディング（スパース性）を適用する。
• 非周期境界や局所特徴を扱うために適応的ウェイトマスキングと残差局所項を許容する。

実験結果

リサーチクエスチョン

• RQ1 AFNO は自己注意に対して競争力のある、またはそれを上回る精度を、視覚トランスフォーマーの計算コストを大幅に抑えた形で達成できるか。
• RQ2 ブロック対角のチャネル構造、ウェイト共有、およびスパース性を課すことが、高解像度画像タスクの一般化と効率性を向上させるか。
• RQ3 前学習タスク（インペインティング）、下流の few-shot segmentation、および Cityscapes の高解像度セグメンテーションにおいて、AFNO は既存のミキサーと比較してどのように性能を示すか。
• RQ4 AFNO は解像度不変性を持ち、非常に大きなトークン系列に対してメモリ消費を抑えつつスケールできるか。

主な発見

• AFNO は画像インペインティングにおいて自己注意と同等の PSNR/SSIM を達成しつつ、GFLOPs を約30%削減している。
• few-shot segmentation では 2/3 のデータセットで自己注意を上回り、他は同等または下回るが FLOPs は低い。
• SegFormer-B3 を用いた Cityscapes セグメンテーションで AFNO-100% は mIoU 80.9 を達成し、同程度の設定で Efficient Self-Attention、LS、GFN 系より優れている。
• ViT-S/4 による ImageNet-1K 分類では AFNO が top-1 accuracy 80.89%（自己注意は 80.87%）を達成し、パラメータ数と FLOPs は類似。
• アブレーション研究はスパース性（lambda ~0.01）、8 ブロック、および適応ウェイトが FNO および非適応変種より AFNO の性能向上に寄与することを示す。
• AFNO は解像度不変性と効率的なチャネル混合を示し、高解像度ビジョンタスクに対する伝統的な自己注意の原理的代替を提供する。