Self-Attention 이후: 트랜스포머 다음에 올 것들

트랜스포머 아키텍처가 AI를 움직여 온 지도 어느덧 8년이 되었습니다. 주요 언어 모델 대부분, 이미지 생성 시스템 대부분, 그리고 점점 늘고 있는 오디오·비디오 모델이 'Attention Is All You Need' 논문에서 소개된 셀프 어텐션 메커니즘 위에 만들어져 있습니다. 하지만 셀프 어텐션에는 근본적인 문제가 있습니다. 연산량과 메모리 비용이 시퀀스 길이의 제곱에 비례해서 늘어난다는 점입니다. 입력 길이가 두 배가 되면 비용은 네 배가 됩니다.

짧은 시퀀스에서는 이 문제가 크지 않습니다. 하지만 우리가 지향하는 128K 토큰 컨텍스트 윈도우, 그리고 사람들이 원하는 백만 토큰 윈도우에서는 심각한 병목이 됩니다. 여러 레이어에 걸쳐 연산을 재사용하는 어텐션 잔차, 제곱 비용을 없앤 선형 어텐션 변형, 어텐션과 저비용 메커니즘을 섞은 하이브리드 아키텍처 등 대안을 찾는 연구가 쏟아지고 있습니다. 트랜스포머가 사라지는 것은 아니지만, 분명히 다시 빚어지고 있습니다.

셀프 어텐션이 비싼 이유

대안들을 이해하려면 셀프 어텐션이 실제로 무엇을 계산하는지 알아야 합니다. N개의 토큰이 주어지면 셀프 어텐션은 모든 토큰 쌍 사이의 연관성 점수를 계산합니다. 토큰 1과 토큰 2, 토큰 1과 토큰 3, ..., 토큰 1과 토큰 N, 그다음 토큰 2와 나머지 모든 토큰, 이런 식입니다. 즉 N²개의 쌍입니다.

import torch
import torch.nn.functional as F
def self_attention(Q, K, V):
"""
Standard self-attention.
Q, K, V: (batch, seq_len, d_model)
The attention matrix is seq_len × seq_len.
For seq_len = 1024:   ~1M entries   (manageable)
For seq_len = 32768:  ~1B entries   (expensive)
For seq_len = 131072: ~17B entries  (very expensive)
"""
d_k = Q.size(-1)
# This matmul creates the N×N attention matrix
scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
weights = F.softmax(scores, dim=-1)
return torch.matmul(weights, V)

4K 토큰에서는 어텐션 행렬이 1,600만 개의 원소를 가집니다. 최신 GPU에서는 문제없는 크기입니다. 128K 토큰에서는 160억 개가 됩니다. 100만 토큰에서는 1조 개를 넘어섭니다. Flash Attention이 있어도(이것은 연산량을 줄이지는 않지만 메모리 접근 패턴을 크게 개선합니다) 결국 제곱 스케일링이 이기게 됩니다.

초기 트랜스포머가 512나 1024 토큰에 묶여 있었던 이유가 바로 이것입니다. 하드웨어와 최적화가 발전할 때마다 한계는 올라갔지만, 우리는 수학적인 벽과 싸우고 있는 셈입니다. 선형 스케일링(O(N))이 제곱 스케일링(O(N²))보다 근본적으로 낫다는 것은 명확하고, 대부분의 대안 아키텍처가 노리는 방향도 바로 그것입니다.

어텐션 잔차: 이미 계산한 것을 재사용하기

어텐션 비용을 줄이는 가장 실용적인 방법 중 하나는 어텐션을 대체하지 않습니다. 대신 이전 레이어에서 이미 한 계산을 재사용해서 각 어텐션 레이어를 더 싸게 만듭니다.

관찰은 이렇습니다. 깊은 트랜스포머(예를 들어 32개 레이어)에서는 인접한 레이어의 어텐션 패턴이 놀랄 만큼 비슷한 경우가 많습니다. 15번 레이어와 16번 레이어는 약간의 조정만 있을 뿐 비슷한 위치에 주목하는 경향이 있습니다. 모든 레이어마다 N² 어텐션 행렬을 처음부터 계산하는 것은 중복입니다. 작업의 상당 부분이 이미 한 레이어 전에 끝나 있기 때문입니다.

어텐션 잔차는 이 점을 활용해서 '잔차(residual)' 어텐션 패턴을 계산합니다. 즉 현재 레이어가 주목하고 싶은 대상과 이전 레이어가 계산한 결과의 차이입니다. 차이가 작다면(중간 레이어에서는 보통 그렇습니다) 계산이 더 저렴해집니다. 전체 어텐션 패턴은 이전 레이어의 패턴에 현재 레이어의 잔차를 더한 것입니다.

이것은 영상 압축의 원리와 비슷합니다. 프레임을 각각 독립적으로 저장하는 대신, 키프레임 하나와 그 키프레임으로부터의 차이(잔차)만 연속으로 저장합니다. 차이는 보통 전체 프레임보다 훨씬 작기 때문에 압축 효율이 극적으로 좋아집니다.

실제로 어텐션 잔차는 품질 저하를 최소화하면서 깊은 모델 중간 레이어의 어텐션 연산 비용을 30~50% 줄입니다. 처음과 마지막 몇 개 레이어는 패턴이 더 뚜렷하기 때문에 여전히 전체 어텐션 계산이 필요하지만, 대부분을 차지하는 중간 레이어들은 상당한 속도 향상을 얻습니다.

선형 어텐션: 제곱 비용 버리기

선형 어텐션 변형들은 어텐션을 O(N²) 대신 O(N)으로 확장되도록 재구성하려고 합니다. 기본 접근은 N×N 어텐션 행렬을 명시적으로 계산하지 않고, 선형 연산만으로 같은(또는 거의 같은) 출력을 얻는 방법을 찾는 것입니다.

핵심 트릭은 softmax의 커널 분해에 있습니다. 표준 어텐션은 softmax(QK^T)V를 계산합니다. softmax를 φ(Q) · φ(K)^T 형태로 분해 가능한 다른 커널 함수로 바꾸면, 계산 순서를 재배열할 수 있습니다. (φ(Q) · φ(K)^T) · V(N×N 중간 결과가 생김) 대신 φ(Q) · (φ(K)^T · V)(d×d 중간 결과가 생기며, 여기서 d는 모델 차원)를 계산하는 것입니다. 긴 시퀀스에서는 d << N이므로 비용이 극적으로 줄어듭니다.

def linear_attention(Q, K, V, feature_map=None):
"""
Linear attention via kernel feature maps.
Cost: O(N * d^2) instead of O(N^2 * d)
"""
if feature_map is None:
# ELU+1 is a common choice (from Katharopoulos et al.)
feature_map = lambda x: F.elu(x) + 1
Q = feature_map(Q)  # (batch, seq_len, d)
K = feature_map(K)  # (batch, seq_len, d)
# Key insight: compute K^T @ V first (d × d matrix)
# instead of Q @ K^T first (N × N matrix)
KV = torch.einsum('bnd,bnm->bdm', K, V)  # (batch, d, d)
# Then multiply by Q
output = torch.einsum('bnd,bdm->bnm', Q, KV)  # (batch, N, d)
# Normalize
Z = torch.einsum('bnd,bd->bn', Q, K.sum(dim=1))  # normalization
output = output / Z.unsqueeze(-1)
return output

문제는 softmax를 다른 커널 함수로 바꾸면 어텐션 분포가 달라진다는 점입니다. softmax 어텐션으로 학습한 모델이 선형 어텐션에 잘 옮겨가지 않을 수도 있습니다. 품질 격차는 크게 좁혀졌습니다. 최근 선형 어텐션 변형은 softmax 어텐션 품질의 95~98%를 달성합니다. 하지만 특히 긴 범위에서 정확한 검색이 필요한 작업에서는 격차가 여전히 남아 있습니다.

상태 공간 모델: 다른 패러다임

Mamba 같은 상태 공간 모델(SSM)은 근본적으로 다른 접근을 택합니다. 토큰 쌍 사이의 관계를 계산하는 대신, 토큰마다 갱신되는 고정 크기 은닉 상태를 유지하면서 순환(recurrence) 방식으로 시퀀스를 처리합니다. 이것은 본질적으로 O(N)입니다. 토큰이 두 배가 되면 시간도 네 배가 아니라 두 배가 걸립니다.

최신 SSM의 핵심 혁신은 순환 파라미터를 입력에 의존하게 만드는 것입니다(선택적 상태 공간, selective state spaces). 이를 통해 모델은 제곱 비용 없이도 콘텐츠 기반 어텐션에 가까운 기능을 얻습니다. 어떤 정보를 기억하고 어떤 정보를 잊을지 '선택'할 수 있습니다. Mamba 계열 모델은 많은 벤치마크에서 트랜스포머와 대등한 품질을 보이면서, 긴 시퀀스에서는 훨씬 빠릅니다.

트레이드오프도 있습니다. SSM은 토큰을 순차적으로 처리하기 때문에, 모든 토큰을 동시에 처리할 수 있는 트랜스포머에 비해 학습 시 병렬화가 어렵습니다. 학습 효율은 중요합니다. 추론이 2배 빠르지만 학습이 3배 느린 모델은 꼭 이득이라고 할 수 없습니다. 전체 연산의 대부분이 학습에 들어가기 때문입니다.

하이브리드 아키텍처: 실용적인 길

요즘 프로덕션 모델의 흐름은 서로 다른 어텐션 메커니즘을 결합한 하이브리드 아키텍처입니다. 이유는 단순합니다. 모델의 부분마다 유리한 연산 방식이 다르기 때문입니다.

• 전역 추론을 위한 전체 어텐션. 일부 레이어는 전체 시퀀스에 걸쳐 주목해야 합니다. 수천 토큰 떨어진 관련 문맥을 찾아야 하기 때문입니다. 이 레이어들은 (Flash 최적화가 가능한) 표준 셀프 어텐션을 사용합니다.
• 인근 문맥을 위한 로컬 어텐션. 많은 레이어는 주로 가까운 토큰에 주목합니다(슬라이딩 윈도우 어텐션). 256~1024 토큰의 고정 윈도우를 사용하면 비용이 O(N·W)로 줄어듭니다. 여기서 W는 윈도우 크기입니다.
• 넓은 문맥을 위한 선형 어텐션. 일부 레이어는 시퀀스 전체의 정보를 모아야 하지만 정밀한 어텐션 가중치까지는 필요하지 않습니다. 선형 어텐션은 O(N) 비용으로 이를 제공합니다.
• 순차 처리를 위한 SSM 레이어. Mamba 계열 레이어는 어텐션 계산 없이도 순차적 의존성을 효율적으로 처리할 수 있습니다.

Jamba(AI21)와 여러 연구용 아키텍처는 레이어의 역할에 따라 이 메커니즘들을 번갈아 사용합니다. 초기 레이어는 로컬 어텐션을 씁니다(문법과 지역 패턴 처리). 중간 레이어는 선형 어텐션이나 SSM을 써서 더 넓은 표현을 만듭니다. 몇몇 전략적 레이어만 전체 어텐션을 써서 전역 추론과 검색을 담당합니다. 이렇게 하면 전체 어텐션이 필요한 모델 품질은 유지하면서 거의 선형에 가까운 전체 스케일링을 얻을 수 있습니다.

개발자가 주목해야 할 것들

언어 모델 위에 애플리케이션을 만든다면, 그 아래에서 일어나는 아키텍처 변화는 실제 작업에 구체적인 영향을 줍니다.

• 컨텍스트 윈도우는 계속 커질 것입니다. 어텐션 비용이 낮아지면서 컨텍스트 윈도우도 넓어집니다. 그러면 애플리케이션 구조도 바뀝니다. 관련 문맥을 4K 윈도우에 맞추려고 복잡한 RAG 파이프라인을 만드는 대신, 그냥 모든 것을 100만 토큰 프롬프트에 넣을 수도 있습니다. 단순함은 매력적이지만, 지연 시간과 비용은 아키텍처마다 다릅니다.
• 지연 시간 프로필이 바뀝니다. 트랜스포머는 어느 지점까지는 지연 시간이 비교적 일정하다가 그 이후 제곱으로 늘어납니다. 선형 어텐션과 SSM 모델은 더 완만하게, 선형으로 늘어납니다. 응답 시간이 중요한 애플리케이션이라면 모델의 스케일링 특성을 이해하는 것이 중요합니다.
• 품질 차이는 작업에 따라 다릅니다. 선형 어텐션 모델은 긴 문맥의 특정 위치에서 정확한 정보를 찾아야 하는 작업('47페이지 목록의 세 번째 항목이 뭐였지?' 같은)에서 약간 떨어질 수 있습니다. 반면 일반적인 이해가 필요한 작업에서는 동등하게 잘 동작합니다. 사용 사례를 잘 파악하세요.
• 추론 최적화가 더 중요해집니다. 모델이 서로 다른 어텐션 방식을 섞으면서 아키텍처가 복잡해질수록, 추론 엔진은 이질적인 연산을 효율적으로 처리해야 합니다. vLLM, TensorRT-LLM 같은 프레임워크가 적응하고 있지만, 커스텀 아키텍처는 바로 지원되지 않을 수도 있습니다.

트랜스포머는 대체되고 있는 것이 아니라 진화하고 있습니다. 셀프 어텐션은 여전히 토큰 간 관계를 모델링하는 가장 표현력 높은 메커니즘입니다. 하지만 모든 레이어, 모든 곳에서 N² 전체 비용으로 쓸 필요는 없습니다. 앞으로 몇 년간의 모델들은 어텐션을 외과적으로 사용할 것입니다. 가장 중요한 곳에는 전체 정밀도를, 나머지 모든 곳에는 더 저렴한 대안을 씁니다. 그 결과 지금의 품질에 맞먹거나 넘어서면서도 더 빠르고, 더 긴 문맥을 다루고, 운영 비용은 더 낮은 모델이 나올 것입니다. 주목할 만한 변화입니다.