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

John Guibas, Morteza Mardani+4 more|CaltechAUTHORS (California Institute of Technology)|Nov 24, 2021

CCD and CMOS Imaging Sensors参考文献 43被引用数 25

ひとこと要約

AFNOは、視覚トランスフォーマー用の効率的なフーリエベースのトークンミキサーで、準線形の計算量、適応チャネルミキシング、スパ spars性を実現し、few-shot segmentationと高解像度Cityscapes segmentationでいくつかのミキサーを上回る。

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.

研究の動機と目的

• 自己注意が二次スケールする高解像度ビジョンに対して、効率的なトークンミキシングの動機づけ。
• 効率性のためのアーキテクチャ変更を組み込みつつ、フーリエニューラルオペレーターを画像トークンに適用してAFNOを開発する。
• AFNOが事前学習、少数ショット分割、および高解像度分割で、精度と効率の有利なトレードオフを提供することを示す。
• ImageNet、Cityscapes、CelebA-Faces、ADE-Cars、LSUN-Cats などのデータセットにおいて、最先端ミキサー（GFN、LS、自己注意）に対してスケーラビリティと競争力のある性能を示す。

提案手法

• トークンミキシングをフーリエ領域を介した連続的なグローバル畳み込みとして捉える。
• パラメータ数を削減し並列性を可能にするため、チャネルミキシングウェイトにブロック対角構造を課す。
• トークン全体で共有される二層MLPを介した適応型ウェイト共有を導入。
• ソフト閾値処理と縮小（LASSO風）によって周波数モードのスパース性を取り入れ、コンパクトな表現を促進。
• 離散フーリエ変換を用いたFFTベースのグローバル畳み込みを離散的な画像グリッドに適用し、非周期的境界に対処するための局所残差項を追加する。

実験結果

リサーチクエスチョン

• RQ1 AFNOは自己注意や他の高速ミキサーと比較して、視覚タスクで精度を維持または向上させつつ、線形から準線形のトークンミキシング時間を達成できるか？
• RQ2 ブロック対角ウェイト、適応ウェイト共有、スパース性といったアーキテクチャ変更は、高解像度画像の一般化と効率性を改善するか？
• RQ3 少数ショット分割と高解像度セマンティック分割の事前学習-転移パイプラインにおけるAFNOの性能は？
• RQ4 ImageNet-1K、Cityscapes、CelebA-Faces、ADE-Cars、LSUN-Cats などのデータセットにおいて、AFNOは最先端ミキサー（GFN、LS、効率的な自己注意）と競合できるか？

主な発見

Backbone	Mixer	Params	GFLOPs	Latency (sec)	SSIM	PSNR(dB)
ViT-B/4	Self-Attention	87M	357.2	1.2	0.931	27.06

• AFNOは、ブロック対角チャネルウェイトによりパラメータ数を削減しつつ、系列長に対して準線形の複雑さを達成する。
• 適応ウェイトとスパース性はFNOより性能を向上させ、アブレーションにおいて適応ウェイトと8ブロックが好ましい設定であることを示す。
• ImageNet-1Kインペインティングでは、AFNOは自己注意と同等のSSIM/PSNRを達成しつつGFLOPsを少なくする。
• 少数ショット分割では、三データセットのうち2つで自己注意を上回り、全体的には競争力がある。
• Cityscapes 1024x1024 と SegFormer-B3 で、AFNO-100% は報告されたミキサーの中で最良の mIoU を達成し、AFNO-25% は FLoPsを大幅に低く抑えつつ性能をほぼ一致させる。
• ImageNet-1K分類では、AFNO (ViT-S/4) が 80.89% トップ1、GFNが 78.77%、LSと競合。