[논문 리뷰] Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers

John Guibas, Morteza Mardani외 4명|arXiv (Cornell University)|2021. 11. 24.

CCD and CMOS Imaging Sensors인용 수 78

한 줄 요약

이 논문은 Adaptive Fourier Neural Operators (AFNO)를 비전 트랜스포머의 효율적인 토큰 믹서로 소개하여 Fourier 기반 전역 합성에 적응적이고 희소하며 블록-구조의 채널 믹싱을 결합해 준다고 하며 준 선형 복잡도와 다운스트림 성능 향상을 달성한다, 특히 고해상도 및 소샷 분할 작업에서.

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.

연구 동기 및 목표

• 고해상도 비전 트랜스포머에서 제곱형(Self-Attention) 복잡도 이상의 효율적인 토큰 믹싱 동기를 제시한다.
• 이미지의 불연속성과 고해상도를 다루도록 Fourier Neural Operators를 적응시켜 AFNO를 개발한다.
• 효율성과 일반화를 개선하기 위해 블록 대각 가중치, 적응적 가중치 공유, 희소성 등의 아키텍처 수정사항을 도입한다.
• 상향 샘플링/인페인팅 사전학습, 소샷 분할, Cityscapes/고해상도 분할에서 AFNO의 효과를 입증한다.
• 최신 매ixed와의 코드 및 실증 비교를 제공한다.

제안 방법

• 토큰 믹싱을 연산자 학습을 통한 연속적 글로벌 컨볼루션으로 프레이밍한다.
• 토큰 믹싱을 위한 FFT 기반 구현과 함께 Fourier Neural Operators(FNO)를 채택한다.
• 가중치 수를 줄이기 위해 채널-믹싱 가중치에 블록 대각 구조를 강제한다.
• 토큰 전체에 걸쳐 공통의 2-계층 MLP를 통해 적응적 가중치 공유를 도입한다.
• 정보성이 높은 구성요소에 집중하도록 주파수 모드에 소프트스쓰링(희소성)을 적용한다.
• 적응적 가중치 마스킹과 잔여 로컬 항을 허용하여 비주기적 경계와 로컬 특징을 다룬다.

실험 결과

연구 질문

• RQ1 AFNO가 비전 트랜스포머에서 자기 주의(self-attention)보다 경쟁력 있거나 우수한 정확도를 substantially reduced computational cost로 달성할 수 있는가?
• RQ2 블록 대각 채널 구조, 가중치 공유 및 희소성의 강제화가 고해상도 이미지 작업에서 일반화와 효율성을 향상시키는가?
• RQ3 사전학습 작업(inpainting), 다운스트림 소샷 분할, 고해상도 Cityscapes 분할에서 AFNO가 기존 믹서와 비교해 어떤 성능을 보이는가?
• RQ4 AFNO는 해상도에 비의존적이며 매우 큰 토큰 시퀀스에 대해 메모리 소모 없이 확장 가능한가?

주요 결과

• AFNO는 자기 주의 대비 약 30%의 GFLOPs를 사용하면서 이미지 인페인팅에서 경쟁력 있는 PSNR/SSIM을 달성한다.
• 소샷 분할에서 AFNO는 2/3 데이터셋에서 자기 주의보다 우월하며, 나머지 데이터셋에서 동일하거나 더 낮은 FLOPs로 비슷한 성능을 보인다.
• Cityscapes 분할에서 SegFormer-B3와 함께 AFNO-100%가 mIoU 80.9를 달성하며, 비교 가능한 설정에서 Efficient Self-Attention, LS, GFN 변형들보다 뛰어나다.
• ViT-S/4에서의 ImageNet-1K 분류에서 AFNO는 Top-1 정확도 80.89%를 달성(자기 주의 80.87%와 유사한 파라미터 수 및 FLOPs)한다.
• 절단 연구(ablation)에서 희소성(lambda ~0.01), 8블록, 적응적 가중치가 FNO 및 비적응 변형 대비 성능 향상에 기여한다.
• AFNO는 해상도 불변성과 효율적인 채널 믹싱을 보여 고해상도 비전 작업에 전통적 자기 주의에 대한 원칙적 대안을 제시한다.