논문 정보

제목: FNet: Mixing Tokens with Fourier Transforms

arXiv: 2105.03824

Introduction

트랜스포머 아키텍처의 핵심은 attention 메커니즘이다. attention은 하나의 토큰을 표현할 때 입력 시퀀스의 다른 모든 토큰과의 연관성을 계산하고, 이를 가중치로 정보를 종합한다. 이는 “모든 토큰이 서로를 참조할 수 있어야 한다”는 귀납적 편향(inductive bias)이 포함된 것이다. 따라서 레이어를 거칠수록 입력 토큰들의 조합으로부터 점점 더 상위 수준의 표현이 형성되며, 이러한 구조 덕분에 트랜스포머가 다양한 문법적(syntactic), 의미적(semantic) 관계를 포착할 수 있었다.

이 논문에서는 트랜스포머 인코더의 비교적 복잡한 self-attention 레이어를, 토큰을 섞는(token mixing) 더 단순한 메커니즘으로 완전히 대체할 수 있는지를 검증한다.

먼저 attention 서브레이어를 학습 가능한 두 개의 행렬곱으로 교체하였다. 하나는 시퀀스 차원을, 다른 하나는 은닉(hidden) 차원을 섞는다. 입력 $X \in \mathbb{R}^{N \times d}$에 대해 이는 다음과 같이 표현된다.

\[Y = W_{\text{seq}} \, X \, W_{\text{h}}, \qquad W_{\text{seq}} \in \mathbb{R}^{N \times N}, \; W_{\text{h}} \in \mathbb{R}^{d \times d}\]

이러한 단순한 선형 혼합만으로도 유의미한 성능이 나오자, 저자들은 더 빠르게 계산할 수 있는 구조화된(structured) 선형 변환으로 범위를 넓혀 실험하였다.

그 결과, 학습 파라미터가 전혀 없는 푸리에 변환(Fourier Transform)이 학습 가능한 선형 혼합과 거의 같은 성능을 보이는 것을 확인했으며, $O(N \log N)$의 FFT알고리즘 덕분에 긴 입력에 대해서도 매우 효율적으로 동작하였다. 이 모델을 FNet이라 명명하여 소개한다.

이 연구는 attention을 효율적으로 개선하려는 “efficient Transformer” 연구들과도 맞닿아 있다. attention 행렬을 희소화하거나 선형화하여 몇몇 연구들은 $O(N)$ 수준의 복잡도를 달성한다. 그러나 이들 기법은 근사 과정에서 발생하는 큰 상수 비용을 숨기고 있어, 실제로는 FNet보다 확장성이 떨어지는 경우가 많다.

이 논문에서는 크게 다음 4가지 결과를 보여준다:

단순한 선형 변환만으로도 텍스트의 다양한 관계를 모델링할 수 있다. 파라미터가 없는 푸리에 변환에 feed-forward 레이어의 MLP를 결합하는 것만으로도 충분한 성능이 나온다. 이는 적어도 일부 NLP 과제에서 attention이 트랜스포머 성능의 핵심 요소가 아닐 수도 있음을 시사한다.
FNet은 속도, 메모리, 정확도 사이에서 좋은 절충점을 제공한다. GLUE 벤치마크에서 BERT-Base 정확도의 92%, BERT-Large 정확도의 97%를 달성하면서, 학습 속도는 GPU에서 80%, TPU에서 70% 빠르다.
소수의 attention 레이어만 추가해도 정확도를 대부분 회복할 수 있다. 마지막 두 레이어만 self-attention으로 되돌린 FNet-Hybrid는 BERT 정확도의 97~99%를 달성하면서도 40~70% 빠르다. 이는 attention이 정확도 향상에 도움은 되지만, 모든 레이어에 필요하지 않을 수 있음을 보여준다.
FNet은 긴 입력에서 efficient Transformer보다 더 나은 속도–정확도 향상을 보인다. Long Range Arena(LRA) 벤치마크에서 FNet은 가장 정확한 efficient Transformer들과 비슷한 정확도를 보이면서, GPU를 이용한 모든 시퀀스 길이에서의 학습, 추론이 더 빨랐다.

Model

Discrete Fourier Transform

푸리에 변환은 함수를 그것을 구성하는 주파수 성분들로 분해하여 표현한다. 길이 $N$인 시퀀스 ${x_n}$ ($n = 0, \dots, N-1$)이 주어졌을 때, 이산 푸리에 변환(DFT)은 다음과 같이 정의된다.

\[X_k = \sum_{n=0}^{N-1} x_n \, e^{-\frac{2\pi i}{N} nk}, \qquad 0 \le k \le N-1\]

각 출력 $X_k$는 모든 입력 $x_n$에 회전 인자(twiddle factor) $e^{-\frac{2\pi i}{N}nk}$를 곱해 더한 값이다. 즉 출력의 각 성분이 시퀀스 전체의 정보를 담게 되며, FNet은 바로 이 성질을 토큰 혼합(token mixing)에 활용한다. 또한 회전 인자가 위치 인덱스 $n$에 의존하므로, 변환 결과에는 토큰의 위치 정보도 함께 반영된다.

즉 Fourier 서브레이어는 $[N, D]$ 입력을 고정된 무늬들로 측정하여, 출력 행렬의 각 원소가 모든 토큰과 모든 특징의 정보를 담도록 재표현하는 역할까지만 하는 것이다. 섞인 표현에서 의미 있는 패턴을 추출하는 일은 뒤따르는 feed-forward 서브레이어가 학습을 통해 수행한다.

예시를 통해 원리를 살펴보면 다음과 같다.

입력과 DFT 행렬

토큰 4개, 임베딩 차원 4인 입력 $X \in \mathbb{R}^{4 \times 4}$를 생각해 보자. 행은 토큰, 열은 임베딩 차원(특징)이다.

\[X = \begin{array}{c|cccc} & h_0 & h_1 & h_2 & h_3 \\ \hline \text{나는} & 1 & 0 & 2 & 1 \\ \text{밥을} & 0 & 1 & 1 & 0 \\ \text{정말} & 2 & 1 & 0 & 1 \\ \text{먹었다} & 1 & 0 & 1 & 2 \end{array}\]

$N = 4$일 때 회전 인자는 $e^{-\frac{2\pi i}{4}nk} = (-i)^{nk}$이므로, DFT 행렬은 다음과 같다. (정규화 계수는 생략)

\[W = \begin{bmatrix} 1 & 1 & 1 & 1 \\ 1 & -i & -1 & i \\ 1 & -1 & 1 & -1 \\ 1 & i & -1 & -i \end{bmatrix}\]

$W$의 각 행은 서로 다른 빠르기로 진동하는 기본 패턴이다.

$k$	패턴 (위치 0~3)	의미
0	$[1,\ 1,\ 1,\ 1]$	변화 없음 (전체 합)
1	$[1,\ -i,\ -1,\ i]$	길이 4 동안 한 번 진동
2	$[1,\ -1,\ 1,\ -1]$	한 칸마다 부호가 바뀜 (가장 빠른 진동)
3	$[1,\ i,\ -1,\ -i]$	$k=1$의 켤레 (반대 방향 회전)

$W$는 대칭 행렬이므로, 두 1D DFT는 행렬곱으로 다음과 같이 표현된다.

$\mathcal{F}_{\text{h}}(X) = XW$: 오른쪽에서 곱하며, 각 행(토큰) 안에서 특징들을 섞는다.
$\mathcal{F}_{\text{seq}}(X) = WX$: 왼쪽에서 곱하며, 각 열(특징) 안에서 토큰들을 섞는다.

따라서 Fourier 서브레이어는 다음과 같이 쓸 수 있다.

\[Y = \Re\big(\mathcal{F}_{\text{seq}}(\mathcal{F}_{\text{h}}(X))\big) = \Re(WXW)\]

이는 학습 가능한 선형 혼합 $W_{\text{seq}} X W_{\text{h}}$에서 두 학습되던 행렬을 고정된 DFT 행렬로 바꾼 형태이다.

1단계: 은닉 차원 DFT (각 토큰 안에서 섞기)

각 행에 DFT를 적용한다. 예를 들어 “나는” 행 $[1, 0, 2, 1]$은 다음과 같이 계산된다.

\[\begin{aligned} k=0&: \; 1 + 0 + 2 + 1 = 4 \\ k=1&: \; 1 + 0\cdot(-i) + 2\cdot(-1) + 1\cdot(i) = -1 + i \\ k=2&: \; 1 - 0 + 2 - 1 = 2 \\ k=3&: \; 1 + 0\cdot(i) + 2\cdot(-1) + 1\cdot(-i) = -1 - i \end{aligned}\]

네 행에 모두 적용하면 다음과 같다.

\[XW = \begin{array}{c|cccc} & k_h{=}0 & k_h{=}1 & k_h{=}2 & k_h{=}3 \\ \hline \text{나는} & 4 & -1+i & 2 & -1-i \\ \text{밥을} & 2 & -1-i & 0 & -1+i \\ \text{정말} & 4 & 2 & 0 & 2 \\ \text{먹었다} & 4 & 2i & 0 & -2i \end{array}\]

이 단계에서 행은 여전히 토큰이다. 각 토큰은 자기 자신의 특징들만 섞었을 뿐, 다른 토큰의 정보는 아직 들어오지 않았다. 반면 열은 더 이상 원래의 특징 $h_m$이 아니라 특징 축의 주파수 성분 $k_h$가 된다.

2단계: 시퀀스 DFT (토큰끼리 섞기)

이번에는 각 열에 DFT를 적용한다. 첫 번째 열 $[4, 2, 4, 4]^\top$은 다음과 같이 계산된다.

\[\begin{aligned} k=0&: \; 4 + 2 + 4 + 4 = 14 \\ k=1&: \; 4 + 2\cdot(-i) + 4\cdot(-1) + 4\cdot(i) = 2i \\ k=2&: \; 4 - 2 + 4 - 4 = 2 \\ k=3&: \; 4 + 2\cdot(i) + 4\cdot(-1) + 4\cdot(-i) = -2i \end{aligned}\]

네 열에 모두 적용하면 다음과 같다.

\[WXW = \begin{array}{c|cccc} & k_h{=}0 & k_h{=}1 & k_h{=}2 & k_h{=}3 \\ \hline k_{\text{seq}}{=}0 & 14 & 2i & 2 & -2i \\ k_{\text{seq}}{=}1 & 2i & -6+2i & 2 & 0 \\ k_{\text{seq}}{=}2 & 2 & 2 & 2 & 2 \\ k_{\text{seq}}{=}3 & -2i & 0 & 2 & -6-2i \end{array}\]

3단계: 실수부 추출

마지막으로 실수부만 취하면 Fourier 서브레이어의 출력이 된다.

\[Y = \Re(WXW) = \begin{bmatrix} 14 & 0 & 2 & 0 \\ 0 & -6 & 2 & 0 \\ 2 & 2 & 2 & 2 \\ 0 & 0 & 2 & -6 \end{bmatrix}\]

두 DFT의 적용 순서는 결과에 영향을 주지 않는다. 행렬곱의 결합법칙에 의해 $W(XW) = (WX)W$이기 때문이다.

출력의 의미

행렬곱 $WXW$의 $(a, b)$ 원소를 풀어 쓰면 다음과 같다.

\[Y_{ab} = \Re\Big(\sum_{n,m} W_{an}\, X_{nm}\, W_{mb}\Big) = \sum_{n,m} X_{nm}\, P^{(a,b)}_{nm}, \qquad P^{(a,b)}_{nm} = \cos\!\left(\tfrac{\pi}{2}(an + bm)\right)\]

즉 $Y_{ab}$는 입력 $X$와 고정된 2차원 물결 무늬 $P^{(a,b)}$를 원소별로 곱해 모두 더한 값이다. 다시 말해 $X$가 그 무늬와 얼마나 닮았는지를 측정한 값이다. 예를 들어 $P^{(0,0)}$은 모든 원소가 1인 무늬이므로, $Y_{00} = 14$는 모든 토큰과 모든 특징을 더한 문장 전체의 총량이 된다.

이 관점에서 $Y$의 행과 열은 다음과 같이 해석할 수 있다.

행은 특정 토큰이 아니라, 문장 전체를 바라보는 하나의 관점이다. 행 $a$는 네 토큰을 패턴 $a$로 묶은 결과이다. 예를 들어 행 0은 네 토큰을 그대로 더한 문장 요약이고, 행 2는 $[+, -, +, -]$ 패턴으로 인접 토큰 간의 차이를 묶은 것이다.
열은 원래 특징이 아니라, 특징들을 조합해 만든 파생 특징이다. 열 $b$는 각 토큰의 특징을 패턴 $b$로 요약한 값이 토큰을 따라 어떻게 변하는지를 나타낸다.

결과적으로 출력의 모든 원소는 입력의 16개 값 전부로부터 계산되며, 이것이 Fourier 서브레이어가 토큰을 섞는 방식이다.

해석 시 유의할 점

1. 의미적 해석이 아니다.
$Y$의 각 원소에 문법적 관계와 같은 언어적 의미가 있는 것은 아니다. 무늬 $P^{(a,b)}$는 입력과 무관하게 고정되어 있기 때문이다. Attention은 무엇을 얼마나 섞을지를 입력이 결정하는 반면, FNet은 섞는 방식이 고정되어 있고 섞인 결과를 어떻게 해석할지만 이후 레이어에서 학습한다.

2. 실수부만 취하면 정보가 일부 중복된다.
$k = 1$과 $k = 3$은 서로 켤레 관계이므로 $Y_{ab} = Y_{(4-a)\bmod 4,\,(4-b)\bmod 4}$가 성립한다. 실제로 $Y_{11} = Y_{33} = -6$이다. 따라서 16개의 값이 모두 독립적이지는 않으며, 이 때문에 실수부만 취한 변환은 역변환이 불가능하다.

3. residual 연결이 토큰 정체성을 되살린다.
Fourier 서브레이어 뒤에는 $\text{LayerNorm}(X + Y)$가 이어진다. 따라서 $n$번째 행은 다음 두 정보를 함께 가진다.

\[\underbrace{X_{n,:}}_{\text{토큰 } n \text{ 자신의 특징}} + \underbrace{Y_{n,:}}_{\text{문장 전체를 } n\text{번 패턴으로 본 요약}}\]

$Y$ 단독으로는 행이 토큰에 대응하지 않지만, residual 덕분에 각 행은 “자기 토큰의 정보 + 전역 정보”를 갖게 된다. 이후 feed-forward 서브레이어는 각 행을 독립적으로 처리하는 연산임에도, 이미 모든 토큰의 정보가 섞인 입력을 받으므로 사실상 문장 전체를 보게 된다.

4. 비선형성은 feed-forward 레이어에만 있다.
Fourier 서브레이어 또한 선형 변환이다.

토큰 혼합이 없다면

토큰 혼합의 중요성은 Fourier 서브레이어를 완전히 제거한 FF-only 모델과 비교하면 분명해진다. 논문에서 FF-only 모델은 이진 분류 과제에서 약 50%, 즉 무작위 추측 수준의 정확도에 그쳤다. 또한 시퀀스 방향 DFT만 남기고 은닉 방향 DFT를 제거한 모델은 정확도가 다소 떨어졌지만 FF-only보다는 훨씬 좋은 성능을 보였다. 은닉 방향 혼합은 feed-forward 레이어도 수행하는 일이므로, Fourier 서브레이어에서 핵심적인 역할은 토큰 방향의 혼합임을 알 수 있다.

FNet architecture

그림과 같이 FNet은 attention이 없는(attention-free) 트랜스포머 인코더 아키텍처이다. 각 인코더 블록은 Fourier 서브레이어 → Add & Normalize → Feed Forward → Add & Normalize 순으로 구성되며, 이 블록을 $N$개 쌓은 뒤 dense 레이어로 최종 출력을 만든다. 즉 기존 트랜스포머 인코더 블록에서 self-attention 서브레이어만 Fourier 서브레이어로 교체한 구조이다.

Fourier 서브레이어는 (시퀀스 길이, 은닉 차원) 크기의 임베딩 입력에 2D DFT를 적용한다. 2D DFT는 시퀀스 차원을 따라 수행하는 1D DFT $\mathcal{F}{\text{seq}}$와 은닉 차원을 따라 수행하는 1D DFT $\mathcal{F}{\text{h}}$를 차례로 적용한 것과 같으며, 출력에서는 실수부만 취한다.

\[y = \Re\big(\mathcal{F}_{\text{seq}}(\mathcal{F}_{\text{h}}(x))\big)\]

두 1D DFT는 교환 가능하므로 적용 순서는 결과에 영향을 주지 않는다. 또한 실수부만 남기기 때문에, 뒤따르는 feed-forward 서브레이어나 출력 레이어가 복소수를 다룰 수 있도록 수정할 필요가 없다. Fourier 서브레이어 자체에는 학습 파라미터가 없다.

이 변환의 효과는 앞의 예시에서 보았듯이, 출력의 각 원소가 모든 토큰의 정보를 담게 된다는 것이다. 따라서 각 위치를 독립적으로 처리하는 feed-forward 서브레이어가 시퀀스 전체의 정보를 충분히 얻을 수 있다.

시간 영역과 주파수 영역의 교대

푸리에 변환에는 쌍대성(duality)이 있어, DFT를 두 번 적용하면 원래 신호가 (순서만 뒤집힌 채) 상수배되어 돌아온다.

\[\mathcal{F}(\mathcal{F}(x))_n = N \cdot x_{(-n) \bmod N}\]

즉 정방향 DFT를 한 번 더 적용하는 것은 역변환과 사실상 같은 역할을 한다. 따라서 FNet의 블록을 차례로 통과하는 과정은, 입력을 “시간” 영역(토큰 위치)과 주파수 영역 사이에서 번갈아 오가게 하는 것으로 볼 수 있다.

이때 합성곱 정리(convolution theorem)에 의해, 주파수 영역에서 계수를 곱하는 것은 시간 영역에서 합성곱을 하는 것과 같다.

\[\mathcal{F}^{-1}\big(\mathcal{F}(x) \odot \mathcal{F}(g)\big) = x \circledast g\]

따라서 주파수 영역에 있는 표현에 feed-forward 레이어의 가중치를 곱하는 것은, 원래 영역에서 보면 그와 관련된 커널로 합성곱을 하는 것에 해당한다. 이런 관점에서 저자들은 FNet을 곱셈과 합성곱을 번갈아 수행하는 모델로 해석할 수 있다고 설명한다. 다만 이는 직관적인 해석이며, 실제로는 실수부 추출, residual 연결, 비선형 활성화 함수가 끼어 있으므로 엄밀한 등가 관계는 아니다.

임베딩

DFT의 회전 인자 $e^{-\frac{2\pi i}{N}nk}$가 위치 인덱스 $n$에 의존하므로, Fourier 서브레이어는 그 자체로 위치 정보를 반영한다. 따라서 FNet은 원리적으로 positional embedding이 필요 없다. 그럼에도 논문에서는 BERT와 깔끔한 비교를 위해, BERT와 동일하게 단어 임베딩, 절대 위치 임베딩, 문장 구분(type) 임베딩을 합한 입력 임베딩을 사용하였다.

Implementation

실험 결과, GPU에서는 실험한 모든 시퀀스 길이(512–8192 토큰)에서 FFT가 행렬곱보다 빨랐다. 반면 TPU에서는 비교적 짧은 시퀀스(4096 토큰 이하)의 경우, DFT 행렬을 미리 계산해 캐싱해 두고 행렬곱으로 DFT를 계산하는 편이 FFT보다 빨랐다. 그보다 긴 시퀀스에서는 TPU에서도 FFT가 더 빨랐다.

이에 따라 GPU 구현에서는 항상 FFT를 사용하고, TPU 구현에서는 길이 4096까지는 행렬곱으로, 그보다 긴 시퀀스에서는 FFT로 2D DFT를 계산한다.

저자들은 이러한 차이의 원인으로 두 가지를 든다.

TPU는 GPU보다도 행렬곱에 더 고도로 최적화된 하드웨어이다.
GPU의 FFT 구현이 TPU의 FFT 구현보다 효율적이다.

따라서 TPU의 FFT 구현이 개선된다면 TPU에서의 FNet 성능도 더 향상될 것으로 예상된다. 모델은 JAX와 Flax 프레임워크로 구현되었다.

Results

1. Transfer learning (GLUE)

BERT와 동일한 Base/Large 설정으로 C4 데이터셋에서 사전학습한 뒤, GLUE에서 fine-tuning하여 비교하였다.

FNet은 BERT-Base 정확도의 92%, BERT-Large 정확도의 97%를 달성하면서 GPU에서 약 80%, TPU에서 약 70% 빠르게 학습된다. Large에서 격차가 줄어든 것은 FNet-Large가 BERT-Large보다 학습이 안정적이었기 때문으로 보인다.
FF-only는 이진 분류 과제에서 50% 수준에 그쳐 학습에 실패했다. 즉 토큰 혼합은 필수적이다.
Random의 낮은 성능은 아무렇게나 섞는 것이 아니라 구조화된 혼합이 필요함을 보여준다.
FNet-Hybrid는 마지막 두 Fourier 서브레이어만 self-attention으로 교체한 모델로, 적은 속도 손실로 BERT 정확도의 97~99%를 회복한다.
Linear는 Base에서 FNet보다 약간 정확하지만, 메모리를 더 많이 쓰고 긴 시퀀스로의 확장성이 떨어지며, Large에서는 학습이 불안정했다.

2. 속도–정확도 비교

모델 크기를 다양하게 바꿔 가며 비교하면, 큰 모델에서는 BERT와 FNet-Hybrid가, 작은 모델에서는 FNet과 Linear가 가장 좋은 속도–정확도 절충(Pareto frontier)을 보인다. 즉 FNet은 자원이 제한된 환경의 경량 모델에서 특히 유리하다.

3. Long Range Arena (LRA)

정확도: FNet의 평균 정확도(55.30)는 가장 정확한 모델인 Transformer(55.83)와 비슷하며, 대부분의 efficient Transformer보다 높다.
속도: GPU에서는 모든 시퀀스 길이(512–8192)에서 비교 모델 중 가장 빠르다. 길이 2048에서 Transformer 대비 학습은 3.2배, 추론은 3.4배 빠르다.
메모리: 모든 길이에서 메모리 사용량이 가장 적으며, Transformer가 메모리 부족(OOM)으로 학습하지 못하는 길이 8192에서도 7.4GB만 사용한다.

hemisus

Contact