Transformers are RNNs: 선형 어텐션으로 만드는 빠른 자기회귀 트랜스포머
목차
개요
Transformers are RNNs는 자기 어텐션을 커널 특징 맵의 선형 내적으로 표현하여 시퀀스 길이에 대한 복잡도를 제곱에서 선형으로 낮춘 연구다. Katharopoulos, Vyas, Pappas, Fleuret가 ICML 2020에서 발표했고 arXiv 식별자는 2006.16236이다.
트랜스포머의 병목은 자기 어텐션의 전역 수용 영역에 있다. 길이 N의 입력을 처리할 때 시간과 메모리가 모두 N의 제곱에 비례하므로 매우 긴 시퀀스에서는 학습과 추론이 사실상 불가능해진다. 저자들은 유사도 함수를 커널로 일반화한 뒤 행렬 곱의 결합법칙을 적용해 이 복잡도를 O(N)으로 줄인다.
이 정식화의 부수적인 결과가 논문 제목이다. 인과 마스킹을 적용한 선형 어텐션은 내부 상태를 갱신하며 출력을 내는 반복 구조, 즉 RNN과 정확히 같은 형태로 다시 쓸 수 있다. 그 결과 자기회귀 추론에서 시퀀스 길이와 무관한 상수 시간·상수 메모리 예측이 가능해지며, 매우 긴 시퀀스에서는 최대 4000배 빠른 생성 속도를 보인다.
배경과 관련 연구
효율적인 트랜스포머
기존 연구는 가중치 가지치기, 가중치 분해, 가중치 양자화, 지식 증류 등으로 메모리 효율을 개선해 왔다. Clark et al.의 ELECTRA는 replaced token detection이라는 사전학습 목적함수로 샘플 효율을 높였고, Lample et al.은 product-key attention으로 계층 용량을 늘렸다. 그러나 이런 방법들은 학습·추론 속도를 개선할 뿐 시퀀스 길이에 대한 제곱 복잡도 자체를 바꾸지 못한다.
문맥 길이를 늘리려는 계열도 있다. Transformer-XL은 고정 길이 문맥을 넘어서는 의존성을 학습하지만 이전 문맥을 메모리에 유지하는 비용이 크다. Sukhbaatar et al.은 어텐션 헤드별 최적 span을 학습해 문맥을 확장했으나, 두 접근 모두 점근적 복잡도는 원래 모델과 동일하다.
가장 밀접한 선행 연구는 Sparse Transformer와 Reformer다. Child et al.은 어텐션 행렬의 희소 분해로 복잡도를 O(N√N)로 낮췄고, Kitaev et al.의 Reformer는 locality-sensitive hashing(LSH)으로 O(N log N)까지 줄였다. 다만 Reformer는 LSH를 쓰기 위해 키를 쿼리와 동일하게 제약하므로 키와 쿼리가 달라야 하는 디코딩 과제에는 쓸 수 없고, 자기회귀 추론 자체를 가속하지도 못한다. 선형 트랜스포머는 쿼리와 키에 아무 제약을 두지 않으면서 시퀀스 길이에 선형으로 확장된다.
자기 어텐션에 대한 이론적 이해
Tsai et al.은 어텐션을 입력에 커널 스무더를 적용하는 연산으로 보는 커널 기반 정식화를 제안했다. 이 논문은 같은 커널 관점을 이해가 아니라 연산 가속의 도구로 사용한다. 또한 쿼리와 키에 양의 유사도 점수를 내는 커널을 적용하면 선형 어텐션이 정상적으로 수렴한다는 점을 관찰한다.
Cordonnier et al.은 충분한 수의 헤드를 가진 멀티헤드 자기 어텐션이 어떤 합성곱 계층도 표현할 수 있음을 보였다. 이 논문은 대신 자기회귀 목적함수로 학습된 자기 어텐션 계층이 순환 신경망으로 볼 수 있음을 보인다.
선형화된 소프트맥스
범주 수가 큰 분류 모델에서 소프트맥스는 오랫동안 병목이었다. Blanc & Rendle, Rawat et al. 등은 소프트맥스를 특징 맵의 선형 내적으로 근사해 샘플링 기반 학습을 가속했다. 같은 시기 Shen et al.은 객체 탐지 과제에 선형화된 어텐션을 적용했다. 이 논문은 어텐션 연산의 선형화에 그치지 않고 추론과 학습 모두에서 선형 복잡도와 상수 메모리를 갖는 자기회귀 모델까지 전개한다는 점에서 구별된다.
방법론
트랜스포머의 일반화된 어텐션 정식화
길이 N, 차원 F의 입력 시퀀스 x에 대해 트랜스포머는 L개의 계층 합성으로 정의된다.
1
T_l(x) = f_l(A_l(x) + x)
f_l은 각 위치의 특징을 독립적으로 변환하는 2계층 피드포워드 네트워크이고, A_l이 시퀀스 방향으로 작동하는 유일한 부분인 자기 어텐션이다. 입력 x는 세 개의 투영 행렬 W_Q, W_K, W_V를 통해 각각 Q, K, V로 사상된다.
1
2
Q = x W_Q, K = x W_K, V = x W_V
A_l(x) = softmax(Q K^T / sqrt(D)) V
이를 임의의 유사도 함수로 일반화하면 다음과 같다.
1
V'_i = sum_{j=1..N} sim(Q_i, K_j) V_j / sum_{j=1..N} sim(Q_i, K_j)
여기서 sim(q, k) = exp(q^T k / sqrt(D))로 두면 원래의 소프트맥스 어텐션이 그대로 복원된다. 즉 소프트맥스 어텐션은 지수 커널을 유사도로 쓰는 특수한 경우일 뿐이다.
선형 어텐션과 결합법칙
일반화된 식이 어텐션으로 성립하기 위해 sim에 필요한 유일한 제약은 비음수성이다. 이 조건은 모든 커널 k(x, y)가 만족하므로, 커널의 특징 표현 φ(x)를 사용해 다음처럼 다시 쓸 수 있다.
1
V'_i = sum_{j=1..N} φ(Q_i)^T φ(K_j) V_j / sum_{j=1..N} φ(Q_i)^T φ(K_j)
핵심은 여기서 행렬 곱의 결합법칙을 적용하는 것이다.
1
V'_i = φ(Q_i)^T ( sum_{j=1..N} φ(K_j) V_j^T ) / ( φ(Q_i)^T sum_{j=1..N} φ(K_j) )
벡터화된 형태로 보면 분자의 변형이 더 분명해진다.
1
( φ(Q) φ(K)^T ) V = φ(Q) ( φ(K)^T V )
왼쪽 괄호 순서는 N×N 크기의 어텐션 행렬을 명시적으로 만들어야 하므로 O(N^2)이 된다. 오른쪽 괄호 순서는 φ(K)^T V와 φ(K)의 합을 한 번만 계산해 모든 쿼리에 재사용하므로 시간과 메모리가 모두 O(N)이 된다. 소프트맥스 어텐션이 O(N^2)인 이유는 기울기 계산을 위해 전체 어텐션 행렬을 저장해야 하기 때문인데, 이 정식화에서는 그 행렬 자체가 사라진다.
특징 맵 선택과 계산 비용
소프트맥스 어텐션의 총 연산량은 쿼리·키 차원 D와 값 차원 M에 대해 O(N^2 max(D, M))로 늘어난다. 선형 어텐션은 차원 C의 특징 맵을 먼저 계산한 뒤 O(N C M) 연산으로 새 값을 얻는다.
지수 커널에 대응하는 특징 함수는 무한 차원이므로 정확한 소프트맥스 어텐션의 선형화는 불가능하다. 반면 다항식 커널은 유한 차원의 정확한 특징 맵을 가지며, 지수 커널이나 RBF 커널과 동등한 성능을 낸다는 점이 선행 연구에서 확인되었다. 차수 2의 다항식 선형 트랜스포머는 O(N D^2 M) 비용을 가지므로, N이 D의 제곱보다 커지는 구간에서 유리해진다. 수만 개 원소의 시퀀스를 다루려는 상황에서는 이 조건이 실제로 성립한다.
논문의 실험은 상대적으로 짧은 시퀀스를 다루므로 양의 유사도를 보장하는 다음 특징 맵을 사용한다.
1
φ(x) = elu(x) + 1
elu는 exponential linear unit 활성 함수이며, x가 음수일 때 기울기가 0이 되는 문제를 피하기 위해 relu 대신 선택되었다. 이 특징 맵은 어텐션 연산을 O(N D M)으로 만든다.
인과 마스킹과 누적합
자기회귀 학습을 위해서는 i번째 위치가 j ≤ i인 위치에만 영향을 받도록 마스킹해야 한다.
1
V'_i = sum_{j=1..i} sim(Q_i, K_j) V_j / sum_{j=1..i} sim(Q_i, K_j)
같은 선형화를 적용하면 합의 범위만 1부터 i까지로 바뀐다. 여기서 두 개의 누적 상태를 도입한다.
1
2
3
S_i = sum_{j=1..i} φ(K_j) V_j^T
Z_i = sum_{j=1..i} φ(K_j)
V'_i = φ(Q_i)^T S_i / ( φ(Q_i)^T Z_i )
S_i와 Z_i는 각각 S_{i-1}, Z_{i-1}로부터 상수 시간에 계산된다. 따라서 인과 마스킹이 적용된 선형 트랜스포머도 시퀀스 길이에 대해 선형 복잡도를 유지한다.
상수 메모리 기울기 계산
이 식을 딥러닝 프레임워크에 그대로 구현하면 기울기 계산을 위해 모든 중간 S_i를 저장해야 한다. 이 경우 메모리 소비가 max(D, M)배 증가하여 긴 시퀀스나 깊은 모델에 적용하기 어려워진다. 저자들은 분자에 대한 기울기를 누적합 형태로 직접 유도하여 순전파와 역전파를 모두 선형 시간·상수 메모리로 처리한다.
분자를 V’_i, 손실의 분자에 대한 기울기를 ∇L로 표기하면 다음이 성립한다.
1
2
3
∇_{φ(Q_i)} L = ∇_{V'_i} L ( sum_{j=1..i} φ(K_j) V_j^T )^T
∇_{φ(K_i)} L = ( sum_{j=i..N} φ(Q_j) ∇_{V'_j} L )^T V_i
∇_{V_i} L = ( sum_{j=i..N} φ(Q_j) ∇_{V'_j} L )^T φ(K_i)
Q에 대한 누적합은 1에서 N 방향으로, K와 V에 대한 누적합은 N에서 1 방향으로 계산된다. 후자는 RNN의 시간축 역전파와 같은 구조다. 최종 알고리즘은 차원 C의 특징 맵에 대해 계산 복잡도 O(N C M), 메모리 O(N max(C, M))를 갖는다.
논문의 의사코드는 다음과 같은 형태다.
1
2
3
4
5
6
def forward(phi_Q, phi_K, V):
V_out, S = 0, 0
for i in range(1, N + 1):
S = S + phi_K[i] @ V[i].T
V_out[i] = phi_Q[i] @ S
return V_out
역전파도 같은 방식으로 S를 순방향과 역방향으로 한 번씩 누적하며 세 기울기를 계산한다. 이 상수 메모리 기울기 계산은 약 200줄의 CUDA 코드로 구현되었다.
트랜스포머는 RNN이다
인과 마스킹 정식화를 시간축 재귀로 다시 쓰면 트랜스포머 계층은 두 개의 은닉 상태를 갖는 RNN이 된다. 어텐션 메모리 s와 정규화 메모리 z가 그 상태다.
1
2
3
4
5
s_0 = 0
z_0 = 0
s_i = s_{i-1} + φ(x_i W_K) (x_i W_V)^T
z_i = z_{i-1} + φ(x_i W_K)
y_i = f_l( φ(x_i W_Q)^T s_i / ( φ(x_i W_Q)^T z_i ) + x_i )
Universal Transformer가 깊이 방향 재귀를 다루는 것과 달리 여기서의 재귀는 시간 방향이다. 이 정식화는 특징 함수에 아무 제약을 두지 않으므로 이론적으로는 소프트맥스 어텐션을 쓰는 모델까지 포함해 어떤 트랜스포머도 표현할 수 있다.
실용적 함의는 학습과 추론의 비대칭을 해소한다는 점이다. 학습 시에는 정답 시퀀스 전체가 주어지므로 계층별 병렬화가 가능하다. 추론 시에는 i번째 출력이 i+1번째 입력이 되어 병렬화가 불가능한데, 기존 트랜스포머는 여기에 더해 매 시점 이전 전체 시퀀스에 대해 어텐션을 다시 계산해야 하므로 시점당 비용이 현재 길이의 제곱으로 증가한다. 선형 트랜스포머는 φ(K_j) V_j^T 행렬을 내부 상태로 저장하고 매 시점 갱신하기만 하면 되므로 시점당 시간과 메모리가 상수다.
실험 설정
모든 실험에서 비교 대상은 소프트맥스 어텐션을 쓰는 full transformer와 Reformer다. Reformer는 공개 코드의 PyTorch 재구현을 사용했고 reversible layer는 쓰지 않았는데, 메모리 측정이 자기 어텐션 계층에 한정되므로 결과에 영향을 주지 않는다. 표기는 softmax가 표준 트랜스포머, linear가 제안 모델, lsh-X가 해싱 라운드 X회의 Reformer다. 선형 트랜스포머 학습에는 elu 기반 특징 맵을 사용했다.
과제별 설정은 다음과 같다.
| 과제 | 데이터셋 | 모델 구성 | 최적화 설정 |
|---|---|---|---|
| 복사 합성 과제 | 최대 길이 128, 심볼 10종과 구분자 | 4계층, 헤드 8개, 배치 64 | RAdam, lr 1e-3에서 3000 스텝 후 1e-4 |
| 이미지 생성 | MNIST 784픽셀 | 8계층, 헤드 8개, 임베딩 256, 헤드당 32차원, FF 차원 4배, 로지스틱 혼합 10개 출력, 배치 10 | RAdam, lr 1e-4, 250 에폭 |
| 이미지 생성 | CIFAR-10 | 16계층, 나머지 구성 동일, softmax는 배치 1, linear와 lsh는 배치 4 | RAdam, 7일 고정 학습 |
| 음성 인식 | WSJ 80시간, 40차원 mel filterbank | 9계층, 헤드 6개, 동일 임베딩 차원 | RAdam, lr 1e-4, 검증 오차 정체 시 절반으로 감소 |
Reformer 설정은 논문 권고에 따라 버킷 64개를 사용했다. MNIST에서는 783 길이 시퀀스를 29개 원소짜리 27개 청크로, CIFAR-10에서는 37개 원소짜리 83개 청크로 나눴다. 음성 인식에서는 은닉 크기 320의 3계층 양방향 LSTM도 비교 대상에 포함했고 Adam 최적화기에 lr 1e-3을 사용했다. 연산량 측정은 11GB 메모리의 NVidia GTX 1080 Ti에서, CIFAR-10 학습은 24GB의 NVidia P40에서 수행했다.
주요 결과
합성 과제 수렴성과 자원 사용량
인과 마스킹이 적용된 시퀀스 복제 과제에서 linear는 안정적으로 수렴하며 softmax와 동일한 최종 손실에 도달한다. lsh는 해싱이 도입하는 노이즈 때문에 더 높은 손실에 머문다.
메모리와 연산량 측정은 길이 2^9부터 2^16까지 시퀀스에 대해 어텐션과 기울기를 계산하며 수행했다. 배치 크기는 시퀀스 길이에 반비례하게 조정하고 샘플당 시간과 메모리를 보고한다.
| 방법 | 최대 처리 길이 | 확장 양상 |
|---|---|---|
| softmax | 4096 | 시퀀스 길이의 제곱에 비례 |
| lsh-4, lsh-8 | 16384 | 선형에 가까움 |
| linear | 16384 이상 | 선형 |
softmax는 11GB GPU에서 4096 길이가 한계였고 lsh-4와 lsh-8은 16384까지 처리했다. linear는 모든 구성에서 두 기준 모델보다 빠르고 메모리를 적게 썼다. Reformer의 점근 복잡도는 O(N log N)이지만 log N 항이 충분히 작아 실측에서는 선형에 가깝게 나타난다.
이미지 생성 MNIST와 CIFAR-10
MNIST에서 250 에폭 학습 결과다.
| 방법 | Bits/dim | Images/sec |
|---|---|---|
| Softmax | 0.621 | 0.45 (1배) |
| LSH-1 | 0.745 | 0.68 (1.5배) |
| LSH-4 | 0.676 | 0.27 (0.6배) |
| Linear | 0.644 | 142.8 (317배) |
linear는 softmax와 거의 같은 perplexity를 유지하면서 생성 처리량이 300배 이상 높다. 상수 메모리 덕분에 단일 GPU로 10,000장의 MNIST 이미지를 동시에 생성할 수 있었다. 픽셀 사이에 저장해야 하는 것이 s_i와 z_i뿐이므로 메모리가 시퀀스 길이와 무관한 반면, softmax와 Reformer는 길이에 비례해 메모리가 증가한다.
CIFAR-10에서 7일 고정 학습 결과다.
| 방법 | Bits/dim | Images/sec |
|---|---|---|
| Softmax | 3.47 | 0.004 (1배) |
| LSH-1 | 3.39 | 0.015 (3.75배) |
| LSH-4 | 3.51 | 0.005 (1.25배) |
| Linear | 3.40 | 17.85 (4,462배) |
시퀀스 길이가 MNIST의 약 4배로 늘어나자 격차가 더 벌어진다. linear는 같은 시간 안에 softmax보다 3배 많은 에폭을 완료했고 그 결과 더 나은 perplexity를 얻었다. softmax가 이미지 한 장을 생성하는 동안 linear는 4,460장을 생성한다.
보조 실험으로 키와 값을 캐싱해 순환 모델처럼 동작하는 stateful-softmax 기준선도 측정했다. 이 방식의 상태 크기는 시퀀스 길이에 비례하므로 복잡도는 여전히 제곱이다.
| 방법 | MNIST Images/sec | CIFAR-10 Images/sec |
|---|---|---|
| Softmax | 0.45 (1배) | 0.004 (1배) |
| Stateful-softmax | 7.56 (16.8배) | 0.32 (80배) |
| LSH-1 | 0.68 (1.5배) | 0.015 (3.75배) |
| LSH-4 | 0.27 (0.6배) | 0.005 (1.25배) |
| Linear | 142.8 (317배) | 17.85 (4,462배) |
stateful-softmax는 기본 구현보다 크게 빠르지만 CIFAR-10에서 linear가 여전히 50배 이상 앞선다. Reformer에 동일한 상태 기반 추론을 구현하는 것은 새 입력마다 정렬과 청킹을 다시 수행해야 하므로 간단하지 않다.
단일 이미지 생성 지연 시간
처리량과 별개로 배치 크기 1에서 이미지 한 장을 만드는 데 걸리는 총 시간도 측정했다.
| 방법 | MNIST CPU 초 | MNIST GPU 초 | CIFAR-10 CPU 초 | CIFAR-10 GPU 초 |
|---|---|---|---|---|
| Softmax | 72.6 | 10.2 | 8651.4 | 300.1 |
| Stateful-softmax | 7.4 | 10.4 | 71.9 | 70.4 |
| LSH-1 | 46.0 | 19.2 | 2318.9 | 221.6 |
| LSH-4 | 112.0 | 55.8 | 5263.7 | 683.9 |
| Linear | 5.5 | 7.3 | 45.1 | 61.3 |
모든 방법이 GPU를 충분히 활용하지 못해 표 4의 처리량보다 낮은 성능을 보인다. linear는 CIFAR-10 기준 softmax보다 약 6.6배 빠르다. 주목할 점은 linear가 모든 경우에서 GPU보다 CPU에서 더 빠른 유일한 방법이라는 것이다. 어텐션을 RNN으로 계산하는 비용이 워낙 작아 시퀀스에 대한 외부 루프 자체가 주된 병목이 되기 때문이다.
자동 음성 인식
비자기회귀 과제에서의 성능을 확인하기 위해 CTC 손실 기반 end-to-end 음성 인식을 평가했다. WSJ 데이터셋의 시퀀스는 평균 800프레임, 최대 2,400프레임이다.
| 방법 | Validation PER | Time/epoch (s) |
|---|---|---|
| Bi-LSTM | 10.94 | 1047 |
| Softmax | 5.12 | 2711 |
| LSH-4 | 9.33 | 2250 |
| Linear | 8.08 | 824 |
linear는 순환 신경망 기준선과 Reformer를 성능과 속도 양쪽에서 앞선다. 에폭당 학습 시간이 softmax의 3배 이상 빠르다. 다만 이 과제에서는 softmax 트랜스포머가 가장 낮은 음소 오류율을 기록한다.
학습 곡선 관점에서 보면 MNIST에서는 linear가 softmax와 대등하게 수렴하고 두 Reformer 변형을 크게 앞선다. CIFAR-10에서는 lsh-1과 linear가 softmax와 lsh-4보다 훨씬 많은 에폭을 완료하며 더 좋은 값에 도달한다. 음성 인식에서는 linear가 에폭당 3배 빨라 softmax보다 약 4배 많은 에폭을 소화했음에도 수렴 성능에서 뒤진다.
한계와 주의사항
첫째, 정확한 소프트맥스 어텐션의 선형화는 불가능하다. 지수 커널에 대응하는 특징 함수가 무한 차원이기 때문이며, 따라서 이 방법은 근사가 아니라 다른 유사도 함수를 쓰는 별개의 어텐션이다. 소프트맥스로 사전학습된 모델을 그대로 가져다 쓸 수 없다는 뜻이기도 하다.
둘째, 성능이 모든 과제에서 동등하지는 않다. WSJ 음성 인식에서 linear의 검증 PER은 8.08로 softmax의 5.12보다 뚜렷하게 높다. 논문 스스로 이 과제에서는 소프트맥스 어텐션이 더 낫다고 인정한다.
셋째, 이득의 크기가 시퀀스 길이에 강하게 의존한다. 차수 2의 다항식 특징 맵을 쓰는 경우 N이 D의 제곱보다 커야 계산상 유리해진다. 짧은 시퀀스에서는 선형화의 이점이 크지 않다.
넷째, 상수 메모리 기울기 계산은 프레임워크의 기본 autograd로는 얻어지지 않는다. 누적합 형태의 기울기를 직접 유도하고 별도 CUDA 커널로 구현해야 하며, 그렇지 않으면 메모리 소비가 max(D, M)배 늘어난다.
다섯째, 특징 맵 선택은 아직 열린 문제다. 논문이 사용한 elu(x)+1은 실험 규모에서 잘 동작하는 선택일 뿐이며, 저자들은 랜덤 푸리에 특징으로 RBF 커널을 근사하는 방향을 후속 연구 과제로 남겨두었다.
여섯째, 정성적 결과에서는 모델 간 차이가 크지 않다. perplexity가 대체로 비슷하므로 생성 샘플의 품질 차이도 두드러지지 않으며, 다만 Reformer 계열은 무조건부 샘플의 다양성이 눈에 띄게 적다는 관찰이 있다.
결론
이 논문은 행렬 곱의 결합법칙을 활용해 자기 어텐션의 시간·메모리 복잡도를 시퀀스 길이에 선형이 되도록 낮춘다. 인과 마스킹을 적용해도 누적 상태 S_i와 Z_i를 통해 선형 점근 복잡도가 유지되며, 기울기까지 누적합으로 유도해 상수 메모리 학습이 가능하다.
가장 중요한 관찰은 인과 마스킹된 트랜스포머 계층이 두 개의 은닉 상태를 갖는 RNN과 동일한 형태로 표현된다는 점이다. 이 등가성 덕분에 자기회귀 추론이 수천 배 빨라지며, 실제로 MNIST에서 317배, CIFAR-10에서 4,462배의 생성 처리량 향상을 얻었다.
저자들은 이 결과가 RNN과 트랜스포머 양쪽에서 정보의 저장·검색 방식을 연구하는 여러 방향을 연다고 본다. 또 다른 방향은 선형 어텐션을 위한 특징 맵 선택이며, RBF 커널을 랜덤 푸리에 특징으로 근사하면 소프트맥스 어텐션으로 사전학습된 모델을 활용할 가능성도 있다고 언급한다.