[논문 리뷰] FNet: Mixing Tokens with Fourier Transforms

James Lee-Thorp, Joshua Ainslie외 2명|arXiv (Cornell University)|2021. 05. 09.

Topic Modeling인용 수 14

한 줄 요약

FNet는 자기주의 주의를 파라미터가 없는 푸리에 변환으로 대체하여 효율적인 시퀀스 모델링을 구현한다. 이는 512 입력 길이에서 GPU에서 80% 빠르고 TPU에서 70% 빠르게 학습되며, BERT의 GLUE 정확도의 92–97%를 달성한다. 또한 장거리 시퀀스에서 뛰어난 확장성과 작은 모델 크기에서 뛰어난 성능을 보인다.

ABSTRACT

We show that Transformer encoder architectures can be sped up, with limited accuracy costs, by replacing the self-attention sublayers with simple linear transformations that "mix" input tokens. These linear mixers, along with standard nonlinearities in feed-forward layers, prove competent at modeling semantic relationships in several text classification tasks. Most surprisingly, we find that replacing the self-attention sublayer in a Transformer encoder with a standard, unparameterized Fourier Transform achieves 92-97% of the accuracy of BERT counterparts on the GLUE benchmark, but trains 80% faster on GPUs and 70% faster on TPUs at standard 512 input lengths. At longer input lengths, our FNet model is significantly faster: when compared to the "efficient" Transformers on the Long Range Arena benchmark, FNet matches the accuracy of the most accurate models, while outpacing the fastest models across all sequence lengths on GPUs (and across relatively shorter lengths on TPUs). Finally, FNet has a light memory footprint and is particularly efficient at smaller model sizes; for a fixed speed and accuracy budget, small FNet models outperform Transformer counterparts.

연구 동기 및 목표

• 자기주의 주의를 떠나 성능 손실 없이 단순하고 파라미터가 없는 선형 변환으로 대체할 수 있는지 조사하는 것.
• 자기주의 주의의 대체로써 푸리에 변환이 구조적이고 효율적인 토큰 혼합 메커니즘으로서의 효과성을 평가하는 것.
• 푸리에 기반 혼합이 특히 장거리 시퀀스에서 더 빠른 학습과 추론을 가능하게 함을 보여주는 것.
• 일부 자기주의 주의 레이어와 푸리에 혼합을 조합한 하이브리드 모델을 통해 정확도-속도 트레이드오프를 개선하는 것.
• 실제로 장거리 작업에서 자기주의 주의의 근사치보다도 다른 혼합 메커니즘이 승리할 수 있음을 입증하는 것.

제안 방법

• FNet는 각 Transformer 인코더 블록의 자기주의 주의 서브레이어를 빠른 푸리에 변환(FFT) 알고리즘을 사용한 파라미터가 없는 푸리에 변환으로 대체한다.
• 푸리에 변환은 시퀀스 차원에서 작동하며, 주파수 도메인 변환을 통해 위치 간 토큰을 혼합한다.
• 각 블록은 푸리에 혼합 후에 레이어 정규화를 적용한 다음, ReLU와 드롭아웃을 포함한 표준 피드포워드 네트워크를 거친다.
• 모델는 표준 위치 임베딩과 분류 작업을 위한 학습 가능한 풀러 헤드를 사용한다.
• 하이브리드 모델의 경우 특정 위치(예: 상단 레이어)에 소수의 자기주의 주의 레이어를 삽입하고, 나머지는 푸리에 혼합을 사용한다.
• 학습은 표준 최적화와 학습률 스케줄링을 사용하며, 추론은 GLUE 및 Long Range Arena 벤치마크에서 평가된다.

실험 결과

연구 질문

• RQ1 파라미터가 없는 푸리에 변환이 자기주의 주의를 효과적으로 대체할 수 있을까? 이는 경쟁 가능한 정확도를 유지할 수 있는가?
• RQ2 FNet의 학습 및 추론 속도는 표준 및 장거리 시퀀스에서 BERT 및 효율적인 Transformer 변종과 비교해 어떻게 되는가?
• RQ3 정확도와 속도를 균형 잡는 데 최적의 하이브리드 FNet 모델에서 자기주의 주의 레이어의 배치와 수는 어디인가?
• RQ4 FNet은 다양한 입력 길이에서 성능과 효율성 면에서 어떻게 확장되는가? 특히 장거리 작업에서 어떻게 되는가?
• RQ5 자기주의 주의를 구조적 선형 변환(예: FFT)으로 대체하는 것이 파라미터화되거나 희소화된 자기주의 주의 근사치보다 유리한가?

주요 결과

• GLUE 벤치마크에서 FNet는 BERT-Base 및 BERT-Large 성능의 각각 92–97%를 달성하며, 512 입력 길이에서 GPU에서 80% 빠르고 TPU에서 70% 빠르게 학습된다.
• Long Range Arena 벤치마크에서 FNet는 가장 정확한 효율적인 Transformer들과 동일한 정확도를 보이며, GPU에서 모든 시퀀스 길이에서 그들보다 빠른 속도를 기록한다.
• 장거리 시퀀스에서는 FNet이 GPU에서 평가된 모든 효율적인 Transformer 모델보다 빠르며, TPU에서는 짧은 시퀀스에서 가장 빠른 모델들을 초월한다.
• 상단에 두 개의 자기주의 주의 레이어만 포함한 하이브리드 FNet 모델은 BERT의 정확도의 97–99%를 달성하면서도 여전히 40–70% 빠르게 학습된다.
• FNet는 더 작은 메모리 프로파일을 가지며, 특히 작은 모델 크기에서 고정된 속도 및 정확도 예산 하에서 표준 Transformer 모델보다 뛰어난 성능을 보인다.
• 푸리에 변환의 O(N log N) 복잡도는 조절된 자기주의 주의조차도 O(N²)보다 실질적인 확장성 면에서 뛰어나다는 것을 보여준다.