이 논문은 Transformer를 "커널 함수로 근사한 어텐션"의 관점에서 바라본다 — Softmax를 커널 특징 맵(kernel feature map)으로 대체하고 행렬 곱의 결합법칙을 활용하면, Transformer는 O(N²)에서 O(N)으로 선형화될 뿐 아니라 수학적으로 RNN과 동치임이 밝혀진다.
arXiv →Transformer(Vaswani et al., 2017)는 자연어 처리, 오디오, 이미지 생성 등 광범위한 분야에서 놀라운 성능을 보였습니다. 그러나 이 강력함에는 숨겨진 대가가 있었습니다 — 입력 시퀀스 길이 N에 대해 O(N²)의 시간 및 메모리 복잡도입니다. 긴 시퀀스를 처리할수록 비용이 기하급수적으로 늘어나는 구조적 문제였습니다.
이 병목의 근원은 self-attention입니다. N개 입력에 대해 N×N 크기의 어텐션 행렬을 계산해야 하고, 이를 저장·처리하는 비용이 O(N²)가 됩니다.
"Transformers achieve remarkable performance in several tasks but due to their quadratic complexity, with respect to the input's length, they are prohibitively slow for very long sequences."번역: Transformer는 여러 과제에서 탁월한 성능을 보이지만, 입력 길이에 대한 이차적 복잡도로 인해 매우 긴 시퀀스에서는 현실적으로 사용하기 어렵다.
핵심 문제는 이렇습니다: "왜 지금까지 이 한계가 해소되지 않았는가?" 기존 연구들은 근사 또는 구조적 제약으로 복잡도를 줄였지만, Softmax 어텐션의 수학적 구조 자체를 바꾸지 않았습니다. Softmax 함수가 만들어내는 전역 정규화(global normalization)가 O(N²) 계산을 필수적으로 만들기 때문입니다.
이 논문은 묻습니다: "Softmax를 대체하는 수식이 있다면, 행렬 곱의 순서를 바꿔 복잡도를 근본적으로 줄일 수 있지 않을까?"
이 논문의 핵심 통찰은 Softmax 어텐션의 유사도 함수를 커널(kernel)로 해석하는 것입니다. 일반화된 어텐션은 다음과 같이 쓸 수 있습니다:
"The definition of attention in equation 2 is generic and can be used to define several other attention implementations... The only constraint we need to impose to sim(·), in order for equation 3 to define an attention function, is to be non-negative. This includes all kernels k(x,y): R²ˣᶠ → R⁺."번역: 어텐션의 정의는 일반적이며, 유사도 함수 sim(·)에 필요한 유일한 제약은 비음수성(non-negativity)이다. 이는 모든 양의 정수값 커널을 포함한다.
즉, Softmax를 쓰지 않아도 됩니다! 유사도 함수가 비음수만 보장된다면 어떤 함수든 어텐션을 정의할 수 있습니다. 이 통찰이 모든 것의 출발점입니다.
커널을 특징 맵(feature map) φ로 분해할 수 있다면, 즉 sim(q,k) = φ(q)ᵀφ(k)라면:
보충 이 아이디어는 커널 기반 어텐션 해석(Tsai et al., 2019)에서 영감을 받았습니다. 그러나 Tsai et al.이 어텐션을 이론적으로 이해하기 위해 커널 관점을 사용했다면, 이 논문은 그것을 실제 계산 속도 향상에 직접 활용합니다.
논문은 먼저 일반적인 Transformer 레이어를 다음과 같이 정의합니다:
| 변수 | 의미 | 비고 |
|---|---|---|
x ∈ ℝᴺˣᶠ | N개의 F차원 특징 벡터로 이루어진 입력 시퀀스 | 시퀀스 길이 N |
f_l(·) | 각 위치를 독립적으로 변환하는 함수 (보통 2층 FFN) | 위치 간 교류 없음 |
A_l(·) | self-attention 함수 — 위치 간 정보를 교환하는 유일한 부분 | O(N²) 병목 지점 |
직관적 해설
수학적 유도
Softmax 어텐션의 정의 (기존 방식):
| 변수 | 의미 | 비고 |
|---|---|---|
Qᵢ | i번째 위치의 쿼리 벡터 | xᵢWQ로 계산 |
Kⱼ | j번째 위치의 키 벡터 | xⱼWK로 계산 |
Vⱼ | j번째 위치의 값 벡터 | xⱼWV로 계산 |
sim(q,k) | 쿼리와 키 사이의 유사도 함수 | Softmax의 경우 exp(qᵀk/√D) |
직관적 해설
수학적 유도
커널 특징 맵 φ: ℝᴰ → ℝᶜ을 도입해 sim(q,k) = φ(q)ᵀφ(k)로 분해하면:
| 변수 | 의미 | 비고 |
|---|---|---|
φ(·) | 커널 특징 맵 (kernel feature map) | 비음수 값을 출력해야 함 |
C | 특징 맵의 출력 차원 | 유한 차원이어야 선형화 가능 |
Σφ(Kⱼ)VⱼT | 모든 키-값 쌍의 누적 합 행렬 [C×M] | 한 번만 계산하고 재사용 |
직관적 해설
수학적 유도
논문은 실험에서 다음 특징 맵을 사용합니다:
| 변수 | 의미 | 비고 |
|---|---|---|
elu(x) | Exponential Linear Unit: x≥0이면 x, x<0이면 exp(x)-1 | x≥0일 때 항등 함수 |
+1 | 비음수성 보장을 위한 오프셋 | elu(x) ∈ (-1, ∞)이므로 +1 추가 |
직관적 해설
수학적 유도
주황색: φ(x) = elu(x)+1 (항상 양수). 회색 점선: y=0 기준선. x < 0에서도 그래디언트가 살아있음을 확인할 수 있습니다.
자기회귀 생성을 위해서는 위치 i가 이전 위치 j ≤ i만 참조해야 합니다. 이를 선형 어텐션에 적용하면:
| 변수 | 의미 | 비고 |
|---|---|---|
Sᵢ [C×M] | i번째까지의 누적 키-값 행렬 (attention memory) | Sᵢ = Sᵢ₋₁ + φ(Kᵢ)VᵢT |
zᵢ [C] | i번째까지의 정규화 벡터 (normalizer memory) | zᵢ = zᵢ₋₁ + φ(Kᵢ) |
직관적 해설
수학적 유도
인과 선형 어텐션을 RNN 형태로 재작성하면:
| 변수 | 의미 | 비고 |
|---|---|---|
s₀ = 0, z₀ = 0 | 초기 은닉 상태 (hidden states) | RNN의 초기 상태와 동일 |
sᵢ [C×M] | 어텐션 메모리 (attention memory) | 매 타임스텝 O(1) 업데이트 |
zᵢ [C] | 정규화 메모리 (normalizer memory) | 매 타임스텝 O(1) 업데이트 |
직관적 해설
수학적 유도
인과 마스킹이 있는 선형 Transformer의 순전파를 단계별로 살펴봅니다:
그래디언트 계산 — 선형 시간·상수 메모리
단순히 수식 (12)를 deep learning 프레임워크에서 구현하면, 역전파를 위해 모든 중간 상태 Sᵢ를 저장해야 합니다. 이는 메모리를 max(D,M)배 증가시켜 긴 시퀀스에서 실용성을 잃습니다. 논문은 이를 해결하기 위해 그래디언트를 누적합(cumulative sum)으로 유도합니다:
직관적 해설
수학적 유도
| 설정 | MNIST 이미지 생성 | CIFAR-10 이미지 생성 | 음성 인식 (WSJ) |
|---|---|---|---|
| 레이어 수 | 8 | 16 | 9 |
| 어텐션 헤드 | 8 | 8 | 6 |
| 임베딩 차원 | 256 (헤드당 32) | 256 (헤드당 32) | 256 (헤드당 32) |
| FFN 차원 | 1024 (임베딩×4) | 1024 (임베딩×4) | 1024 (임베딩×4) |
| 출력 모델 | 10-logistic mixture | 10-logistic mixture | CTC loss (음소) |
| 옵티마이저 | RAdam (lr=1e-4) | RAdam (lr=1e-4) | RAdam (lr=1e-4→감소) |
| 배치 크기 | 10 | 4 (Linear/Reformer) | - |
| 훈련 기간 | 250 에포크 | 7일 | 수렴까지 |
| GPU | NVidia GTX 1080 Ti (11GB) | NVidia P40 (24GB) | - |
| 구현 | PyTorch (Paszke et al., 2019). CUDA 그래디언트 커널 ~200줄. 코드: linear-transformers.com | ||
복사 태스크(sequence duplication)에서 4층 8헤드 Transformer를 최대 128 길이 시퀀스로 훈련. 결과: 선형 어텐션은 Softmax와 동일한 최종 손실에 수렴하며, Reformer보다 낮은 손실을 보입니다.
"linear converges smoothly and reaches a lower loss than lsh due to the lack of noise introduced by hashing. In particular, it reaches the same loss as softmax."번역: 선형 어텐션은 해싱에 의한 노이즈 없이 부드럽게 수렴하며, LSH보다 낮은 손실을 달성한다. 특히 Softmax와 동일한 손실에 도달한다.
속도 및 메모리: NVidia GTX 1080 Ti (11GB)에서 다양한 시퀀스 길이 N ∈ {2⁹, 2¹⁰, ..., 2¹⁶}로 비교. Softmax는 N²로 스케일 — N=4,096이 최대. 선형 어텐션은 N=2¹⁶(65,536)까지 가능하며 메모리와 시간 모두 선형으로 스케일합니다.
| 방법 | Bits/dim ↓ | 처리량 (이미지/초) ↑ | 상대 속도 |
|---|---|---|---|
| Baseline Softmax | 0.621 | 0.45 | 1× |
| LSH-1 (Reformer) | 0.745 | 0.68 | 1.5× |
| LSH-4 (Reformer) | 0.676 | 0.27 | 0.6× |
| Ours Linear | 0.644 | 142.8 | 317× |
보충 Bits/dim(비트/픽셀)은 낮을수록 좋습니다. 선형 어텐션은 Softmax보다 0.023 bits/dim 차이(약 4% 열세)밖에 없으면서 317배 빠릅니다. 이 속도 향상이 가능한 이유: 추론 시 메모리가 상수여서 단일 GPU로 10,000개 MNIST 이미지를 동시 생성할 수 있기 때문입니다.
시퀀스 길이 3,072 (32×32×3). Softmax는 배치 크기 1로도 겨우 P40 24GB에서 실행. 7일 동안 모든 모델 훈련.
| 방법 | Bits/dim ↓ | 처리량 (이미지/초) ↑ | 상대 속도 |
|---|---|---|---|
| Baseline Softmax | 3.47 | 0.004 | 1× |
| LSH-1 | 3.39 | 0.015 | 3.75× |
| LSH-4 | 3.51 | 0.005 | 1.25× |
| Ours Linear | 3.40 | 17.85 | 4,462× |
7일 동안 선형 Transformer는 Softmax보다 3배 많은 에포크를 완료 — 결과적으로 더 나은 퍼플렉시티를 달성합니다. 이 실험은 시퀀스가 길수록 선형 어텐션의 이점이 더 커짐을 보여줍니다 (MNIST 317× → CIFAR-10 4,462×).
비자기회귀(non-autoregressive) 태스크에서도 선형 어텐션이 작동하는지 확인. CTC 손실을 사용하는 end-to-end 음성 인식. WSJ 80시간 데이터셋, 40차원 mel-scale filterbank 특징.
| 방법 | 검증 PER ↓ | 에포크당 시간 (초) ↓ |
|---|---|---|
| Baseline Bi-LSTM | 10.94 | 1047 |
| Softmax (표준 Transformer) | 5.12 | 2711 |
| LSH-4 (Reformer) | 9.33 | 2250 |
| Ours Linear | 8.08 | 824 |
PER(음소 오류율)에서 Softmax가 가장 우수합니다. 그러나 선형 Transformer는 에포크당 훈련 시간이 Softmax의 1/3 수준 — 같은 시간에 약 4배 더 많은 에포크를 완료합니다. 이 실험은 선형 어텐션이 자기회귀 태스크뿐 아니라 일반 순방향 태스크에도 적용 가능함을 보여줍니다.
평가 음성 인식에서 선형 어텐션은 Softmax보다 2.96 PER 열세입니다. 이는 음성 인식이 어텐션의 정밀한 전역 정보 통합이 중요한 태스크이기 때문일 수 있습니다. 긴 시퀀스 자기회귀 생성(이미지, 텍스트)에서의 이점이 분류형 태스크에서는 다소 줄어들 수 있습니다.
보충 [부록 데이터] 논문의 보완 실험에서는 키와 값을 저장하고 재사용하는 "stateful-softmax" 베이스라인과도 비교합니다. MNIST에서 stateful-softmax는 16.8×, CIFAR-10에서 80×로 크게 개선되지만, 여전히 선형 어텐션(317×, 4,462×)에 비해 훨씬 느립니다. Stateful-softmax는 상태 크기가 시퀀스 길이에 비례하여 증가하기 때문입니다.
| 얻은 것 | 잃은 것 |
|---|---|
| O(N) 시간·메모리 복잡도 | Softmax의 정확한 지수 커널 |
| 자기회귀 추론 시 상수 메모리 | 일부 정밀한 주의 집중 능력 |
| 수천~수만 배 빠른 추론 속도 | 기존 사전훈련 모델과의 호환성 |
| Transformer-RNN 동치 관계 발견 | 최적 특징 맵 탐색 필요 |
이 논문은 두 가지 독립적인 공헌을 담고 있습니다: 첫째 선형 어텐션이라는 실용적 도구, 둘째 Transformer–RNN 동치 관계라는 이론적 통찰. 이 두 공헌은 서로 다른 방향의 연구를 촉발했습니다.
긴 시퀀스에서 자기회귀 생성을 수행해야 하는 모든 응용 — 픽셀 단위 이미지 생성, 오디오 생성, 긴 문서 요약, DNA/단백질 시퀀스 모델링 — 에서 추론 비용을 획기적으로 줄이는 데 직접 활용할 수 있습니다.
논문은 결론에서 두 가지 명시적 후속 연구 방향을 제시합니다:
① RNN과 Transformer 사이의 정보 저장·검색 메커니즘 비교 탐구. ② 더 나은 특징 맵 선택 — 특히 랜덤 푸리에 특징(random Fourier features)으로 RBF 커널을 근사해 Softmax 사전훈련 모델과의 호환성 확보.
평가 이 논문은 이후 Performer(Choromanski et al., 2021), RWKV, Mamba 등의 선형 또는 선형화 시퀀스 모델 연구에 직접적인 영감을 주었습니다. 특히 "Transformer가 RNN과 수학적으로 동치"라는 관점은 이후 여러 모델이 "훈련은 병렬 Transformer처럼, 추론은 RNN처럼"이라는 아이디어를 추구하게 만들었습니다. 2020년대 중반의 SSM(State Space Model) 붐의 지적 선조라 볼 수 있습니다.
핵심은 어떤 두 행렬을 먼저 곱하느냐에 따라 중간 결과물의 크기가 달라진다는 것입니다. 기존 방식 (φ(Q)φ(K)ᵀ)V: φ(Q)φ(K)ᵀ = [N×C]·[C×N] = [N×N] — N이 크면 이 행렬이 거대해집니다. 제안 방식 φ(Q)(φ(K)ᵀV): φ(K)ᵀV = [C×N]·[N×M] = [C×M] — N에 무관한 고정 크기입니다. 행렬 곱의 결합법칙은 수학적으로 (AB)C = A(BC)이므로 결과는 동일하지만, 중간 단계에서 다루는 데이터 크기가 완전히 다릅니다.
Softmax는 exp(qᵀk/√D)인데, 이것은 특정 커널의 특징 맵이 무한 차원인 경우에 해당합니다. 논문은 다항식(polynomial) 커널처럼 유한 차원 특징 맵으로 표현되는 커널이 지수 커널과 "동등하게 잘 작동한다"는 선행 연구(Tsai et al., 2019)를 인용합니다. φ(x) = elu(x)+1은 완벽한 대체가 아니지만, 대부분의 태스크에서 Transformer가 학습할 수 있는 패턴을 포착할 수 있을 만큼 충분히 표현력이 있습니다.
"We observe that if a kernel with positive similarity scores is applied on the queries and keys, linear attention converges normally."번역: 쿼리와 키에 양의 유사도 점수를 가진 커널을 적용하면 선형 어텐션이 정상적으로 수렴한다.
훈련 시에도 O(N)입니다! 이것이 Sparse Transformer나 Reformer 등 다른 효율적 Transformer와의 핵심 차이입니다. 비인과(non-causal) 선형 어텐션에서는 Σⱼφ(Kⱼ)VⱼT를 한 번만 계산하고 모든 쿼리에 재사용하여 O(NCM)이 됩니다. 인과 마스킹에서는 누적합 트릭과 특수 CUDA 그래디언트 커널을 사용하여 순전파와 역전파 모두 O(N)·상수 메모리를 달성합니다.
세 가지 핵심 차이가 있습니다. ① 제약 없음: Reformer는 LSH를 위해 Q=K를 강제합니다. 이로 인해 Q≠K인 디코딩 태스크(예: seq2seq)에 사용할 수 없습니다. 선형 어텐션은 Q와 K가 달라도 됩니다. ② 자기회귀 추론 속도: Reformer는 추론 시 여전히 시퀀스 의존적 비용이 있지만, 선형 어텐션은 진정한 상수 시간·메모리 추론을 달성합니다. ③ 점근적 복잡도: Reformer O(N log N) vs 선형 어텐션 O(N).
이 동치 관계는 Transformer 아키텍처를 두 모드로 실행할 수 있음을 의미합니다: 훈련 시는 병렬 Transformer 모드 — 전체 시퀀스가 주어지면 GPU를 최대로 활용. 추론 시는 RNN 모드 — 이전 은닉 상태 (S, z)만 유지하고 토큰 하나씩 처리. 이것은 이후 RWKV, RetNet, Mamba 등이 채택한 "훈련은 병렬, 추론은 순차"라는 패러다임의 선구입니다.
논문에 따르면 PyTorch 구현과 문서, 예제가 linear-transformers.com에 공개되어 있습니다. CUDA 그래디언트 커널은 약 200줄로 구현됩니다.
Softmax 어텐션과 선형 어텐션이 생성하는 어텐션 가중치 행렬을 비교합니다. 인과 마스킹 적용 시 선형 어텐션의 작동 방식을 확인하세요.
각 행: 해당 위치의 어텐션 분포. 상삼각 = 0 (미래 마스킹).
누적 상태 Sᵢ의 크기 (||Sᵢ||). 시퀀스가 진행될수록 증가.
논문의 주요 Figure를 원본 이미지로 보여드립니다. 각 Figure의 논문 내 위치와 주목할 포인트를 함께 확인하세요.
