Linear Attention과 Performer

커널화의 핵심

어텐션 연산 softmax(Q·Kᵀ)·V가 O(n²)인 까닭은 가운데에 n×n 행렬이 생기기 때문이야. Performer(Choromanski et al., 2021)를 비롯한 linear attention 계열은 이런 질문을 던졌어. Softmax를 커널 함수 φ로 바꿔서 softmax(Q·Kᵀ) ≈ φ(Q)·φ(K)ᵀ이 되게 할 수 없을까?

그런 φ가 있으면 φ(Q)·φ(K)ᵀ·V를 φ(Q)·(φ(K)ᵀ·V)로 다시 쓸 수 있어. 괄호를 어디에 두느냐가 중요해. 오른쪽 식은 d×d 행렬을 먼저 계산하고(O(n·d²)), 다시 φ(Q)를 곱해 투영해(O(n·d²)). 전체 비용은 O(n·d²)이야. 시퀀스 길이에는 선형이고 모델 차원에는 제곱으로 늘어나므로, 모델 차원이 적당한 긴 시퀀스에서는 아주 큰 이득이 생겨.

왜 주류가 되지 못했을까?

수학은 성립했지만 실용적인 결과는 기대에 못 미쳤어. 무작위 특징 지도로 softmax를 근사하면 긴 문맥에서 오차가 더 커져. Softmax를 충실히 근사하는 데 필요한 φ의 차원은 어텐션 분포의 엔트로피에 따라 늘어나는데, 긴 문맥이 만드는 높은 엔트로피의 분포에서는 근사 비용이 폭발해 버려.

실제 벤치마크에서는 속도 향상이 가장 필요한 길이에서 Performer 계열 모델이 표준 어텐션보다 의미 있게 품질을 잃었어. 공동체가 얻은 결론은 단순해. 정확한 어텐션의 품질은 포기하기 어렵다. 근사가 아주 뛰어나지 않으면 얻는 절약보다 잃는 게 더 커.

남긴 유산

Performer가 주류가 되지는 못했지만, 나중에 다른 형태로 더 잘 작동한 개념적 틀을 세웠어. 바로 영리한 분해를 통한 구조화 어텐션이야. Kimi Linear, MHLA, gated linear attention은 모두 이 아이디어의 후손이야. 더 영리한 φ와 더 풍부한 상태 장치를 써서 원래 Performer의 품질 문제를 극복하려 해.

Code

Linear attention factorization — 핵심 트릭·python

# Standard attention: O(n^2 * d)
# attn = softmax(Q @ K.T) @ V              # 가운데 n x n matrix

# Linear attention with kernel phi: O(n * d^2)
# Numerator: phi(Q) @ (phi(K).T @ V)        # 가운데 d x d matrix
# Denominator: phi(Q) @ (phi(K).T @ ones)
# Output: numerator / denominator

# 트릭은 결합법칙:
# (phi(Q) @ phi(K).T) @ V   ==   phi(Q) @ (phi(K).T @ V)
# 좌변: O(n^2 * d). 우변: O(n * d^2).

Exercise

PyTorch로 무작위 특징 지도를 이용한 Performer식 linear attention을 구현해. 작은 (1, 8, 1024, 64) 입력에서 표준 scaled-dot-product attention과 출력을 비교하고, 속도와 두 출력 사이의 차이를 함께 측정해. 짧은 시퀀스에서는 차이도 작지만 속도 이득도 작을 거야. 긴 시퀀스에서는 속도 이득이 커지는 만큼 출력 차이도 커지는지 봐. 바로 그 맞바꿈 때문에 Performer는 승자가 되지 못했어.