[Deep Learning] FNet: Mixing Tokens with Fourier Transforms
FNet
트랜스포머의 self-attention은 CNN 또는 RNN 대비 장기 의존성 (long-term dependency)를 해결한다. 하지만 연산 complexity가 O(N²)으로 N이 길어질수록 bottlenect으로 작용한다.
FNet은 푸리에 변환 (fourier transform)을 통해 self-attention 대신 token 간의 정보를 학습하는 mixing mechanism으로 simple linear transformation 연산만 문장 내의 다양한 의미 관계를 모델링할 수 있다. 또한, 속도 메모리, 정확도 측면에서 효율성을 보인다.
그리고 FNet hybrid model은 단 두 개의 self-attention sublayer만을 사용하여 BERT에 준하는 성능을 보였다. 그리고 Long-Range Arena에서 다른 efficient Transformer 대비 속도와 정확도 측면에서 향상을 이뤘다.
아키텍처
Transformer encoder에서 self-attention sublayer를 fourier sublayer로 완전히 대체한다. fourier sublayer는 2D DFD를 수행한다. 원래 Transformer encoder와 동일하게 Feed Forward network를 사용하고 있으며, 이는 정보를 한번 더 섞어주는 역할을 하는 것으로 볼 수 있다. 푸리에 변환에서 positional information이 반영되므로 position embedding 없이 성을 잘 내지만, BERT 모델과 Fourier sublayer의 자체 효과를 비교하기 위해 position embedding을 사용했다.
Fourier sublayer (Mixing sublayer)
embedding input은 hidden dimension (1D), sequence length (1D)이며 output y는 Fourier transform 결과의 실수 부분만을 사용했을 때 가장 성능이 좋았다. MLP-Mixer에서 MLP 기반의 token/channel mixing으로 정보를 교환하듯이 여기서 mixing은 주파수 영역에서 각 성분은 모든 token들의 조합으로 표현된다. mixing sublayer는 learnable parameter를 사용하지 않는다.
이렇게 간단한 푸리에 변환 기반의 linear mixing transformation만으로 텍스트의 다양한 semantic relationship을 모델링할 수 있음을 보였다.