[논문 리뷰] You Need Better Attention Priors: A Unifying Framework for Attention Priors

TL;DR

트랜스포머의 표준 어텐션은 모든 토큰 위치를 동일하게 취급하는 암묵적인 '균일 사전(uniform prior)'을 가정합니다. 이 논문은 엔트로피 최적 수송(Entropic Optimal Transport, EOT) 이론을 통해 어텐션 메커니즘을 일반화하고, 이 균일 사전을 데이터로부터 직접 학습하는 '학습 가능한 사전(trainable prior)'으로 대체하는 **GOAT(Generalized Optimal transport Attention with Trainable priors)**를 제안합니다. GOAT는 FlashAttention과 같은 최신 하드웨어 가속 커널과 완벽히 호환되면서도, 최근 이슈가 된 주의 싱크(Attention Sink) 문제를 자연스럽게 해결합니다. 실험 결과, GOAT는 언어 모델링, 장문맥 처리, 생물학적 서열 분석, 이미지 인식 등 광범위한 분야에서 기존 위치 인코딩 방식(RoPE, ALiBi 등)을 능가하는 뛰어난 성능과 일반화 능력을 보여주었습니다.

연구 배경 및 동기

트랜스포머의 핵심인 어텐션 메커니즘은 모델이 입력 시퀀스 내의 관련 정보에 집중하도록 돕습니다. 하지만 표준 어텐션, 즉 소프트맥스를 적용한 점곱 어텐션(scaled dot-product attention)은 모든 토큰 위치의 중요도가 동등하다는 암묵적인 가정을 내포하고 있습니다. 이는 특정 위치(예: 문장의 시작)나 토큰 간의 상대적 거리 정보가 중요한 작업에서 최적이 아닐 수 있습니다.

이러한 한계를 극복하기 위해 상대적 위치 인코딩(Relative Position Encoding), RoPE, ALiBi 등 다양한 방법이 제안되었지만, 이들은 특정 휴리스틱에 기반하거나 어텐션 계산 방식 자체를 수정해야 하는 경우가 많았습니다. 이 연구는 보다 근본적인 질문에서 출발합니다: "어텐션의 균일한 사전 가정을 버리고, 데이터에 가장 적합한 사전(prior)을 모델이 직접 학습하게 할 수는 없을까?"

이 질문에 답하기 위해, 연구진은 어텐션을 엔트로피 최적 수송(EOT)이라는 수학적 프레임워크로 재해석하여, 학습 가능한 사전을 도입하는 우아하고 강력한 방법론을 제시합니다.

연구	접근법	장점	한계
Vaswani et al.	절대 위치 인코딩	구현이 간단함	시퀀스 길이에 대한 일반화 성능이 떨어짐
Shaw et al.	상대적 위치 인코딩	상대적 위치 정보를 명시적으로 모델링	어텐션 로직이 복잡해지고 계산 비용 증가
ALiBi	선형 편향 추가	외삽(extrapolation) 성능이 우수하고 간단함	고정된 선형 편향으로, 표현력이 제한될 수 있음
RoPE	회전 위치 인코딩	상대 위치 정보를 쿼리/키에 직접 주입	위치와 콘텐츠 정보가 얽혀 분리가 어려움
본 논문 (GOAT)	EOT 기반 학습 가능한 사전	데이터 기반의 유연한 사전 학습, 높은 일반화 성능	사전을 위한 추가 파라미터 및 학습 필요

핵심 기여

어텐션의 일반화: 어텐션 메커니즘을 엔트로피 최적 수송(EOT) 문제로 재해석하여, 기존 어텐션이 EOT의 특수한 경우(균일 사전을 사용)임을 보였습니다. 이는 어텐션에 대한 통합적인 이론적 기반을 제공합니다.
학습 가능한 사전(Trainable Prior) 도입: 고정된 균일 사전 대신, 모델이 데이터로부터 직접 유용한 위치적/구조적 편향을 학습하는 log π 항을 도입하여 표현력과 일반화 성능을 극대화했습니다.
최신 커널과의 호환성: GOAT는 어텐션의 최종 로짓(logit)에 편향을 더하는 방식으로 구현되어, FlashAttention, FlashAttention-2와 같은 최적화된 I/O-aware 커널을 수정 없이 그대로 활용할 수 있어 실용성이 매우 높습니다.
주의 싱크(Attention Sink) 문제 해결: GOAT의 학습 가능한 사전은 모델이 초기 토큰에 집중하는 '주의 싱크' 현상을 명시적으로 학습하고 제어할 수 있게 하여, 장문맥 처리에서 안정성을 높입니다.

제안 방법론: GOAT

엔트로피 최적 수송(EOT)과 어텐션

**엔트로피 최적 수송(EOT)**은 한 확률 분포를 다른 확률 분포로 변환하는 데 필요한 '비용'을 최소화하는 방법을 찾는 이론입니다. 여기에 엔트로피 항을 추가하여, 변환이 너무 결정적이지 않고 부드럽게 이루어지도록 조절합니다.

논문은 표준 어텐션이 EOT 문제의 해와 수학적으로 동일한 형태를 가짐을 보입니다. 여기서 '비용'은 쿼리와 키의 유사도(-s)에 해당하고, '사전 분포'(π)는 균일 분포(uniform distribution)에 해당합니다.

GOAT의 공식

GOAT는 이 사전 분포 π를 학습 가능한 파라미터로 대체합니다. 최종적인 어텐션 확률 분포 $p^*$ 는 다음과 같이 계산됩니다.

p^* = \text{softmax}\left(\frac{s}{\tau} + \log \pi\right)

$s$ : 쿼리( $Q$ )와 키( $K$ )의 유사도 점수 행렬 ( $s = QK^T$ )로, 콘텐츠 기반의 유사도를 나타냅니다.
$\tau$ : 온도(temperature) 파라미터입니다.
$\log \pi$ : 모델이 학습하는 사전(prior) 분포의 로그값입니다. 이 항이 GOAT의 핵심으로, 콘텐츠와 무관하게 위치나 구조에 대한 편향을 학습합니다.

표준 어텐션은 $\log \pi$ 가 0 (즉, $\pi$ 가 균일 분포)인 특수한 경우와 같습니다. GOAT는 이 항을 학습함으로써 ALiBi의 선형 편향, 주의 싱크의 초기 토큰 편향 등 다양한 패턴을 데이터로부터 직접 배울 수 있습니다.

구현 방식

GOAT는 쿼리와 키 벡터를 콘텐츠(content) 부분과 위치(position) 부분으로 분리하여 $\log \pi$ 를 효율적으로 계산합니다.

쿼리: $q = [q_c, q_p]$ (콘텐츠 쿼리, 위치 쿼리)
키: $k = [k_c, k_p]$ (콘텐츠 키, 위치 키)

최종 로짓은 두 부분의 합으로 구성됩니다.

\text{Logits} = \underbrace{q_c k_c^T}_{\text{콘텐츠 점수}} + \underbrace{q_p k_p^T}_{\log \pi \text{ (학습된 사전)}}

이 구조 덕분에 콘텐츠 기반 상호작용과 위치 기반 편향이 분리되어 더 안정적이고 해석 가능한 학습이 가능해집니다.

# GOAT의 개념적 로짓 계산
# q, k는 각각 (batch, num_heads, seq_len, head_dim) 크기를 가짐
# head_dim은 content_dim과 position_dim으로 나뉨
q_c, q_p = torch.split(q, [content_dim, position_dim], dim=-1)
k_c, k_p = torch.split(k, [content_dim, position_dim], dim=-1)

# 1. 콘텐츠 기반 유사도 점수 계산
# 결과: (batch, num_heads, seq_len, seq_len)
content_logits = q_c @ k_c.transpose(-2, -1)

# 2. 학습 가능한 사전(위치 편향) 계산
# 이 부분이 log π 역할을 하도록 학습됨
position_logits = q_p @ k_p.transpose(-2, -1)

# 3. 최종 로짓 = 콘텐츠 점수 + 위치 편향
final_logits = content_logits + position_logits

# 이후 final_logits에 스케일링, 마스킹을 적용하고 softmax를 취함
attention_weights = torch.softmax(final_logits / temperature, dim=-1)

실험 결과 분석

GOAT는 언어, 생물학, 비전 등 다양한 분야에서 일관되게 강력한 성능을 보였습니다.

언어 모델링 (C4 데이터셋): 125M 모델 학습 시, GOAT는 ALiBi나 RoPE보다 낮은 Perplexity를 달성하여 더 나은 언어 모델링 능력을 입증했습니다.
장문맥 처리 (Passkey Retrieval): 시퀀스 길이가 길어질수록 기존 모델들의 성능이 급격히 하락하는 반면, GOAT는 학습된 사전 덕분에 매우 긴 시퀀스에서도 안정적으로 높은 정확도를 유지했습니다. 이는 GOAT가 주의 싱크 문제를 효과적으로 해결하고 길이 일반화에 강건함을 보여줍니다.
생물학적 서열 모델링 (인간 게놈): DNA 서열 예측에서 GOAT는 RoPE 기반 모델보다 낮은 검증 NLL(Negative Log-Likelihood)을 기록하여, 복잡하고 긴 의존성을 가진 데이터에서도 뛰어난 성능을 보였습니다.
이미지 인식 (Vision Transformer): 표준 ViT는 학습 시 사용된 해상도와 다른 해상도의 이미지에 대한 제로샷 추론 성능이 떨어집니다. GOAT를 적용한 ViT는 위치 사전이 이미지 크기 변화에 더 유연하게 적응하여, 다양한 해상도에서 훨씬 뛰어난 성능을 보였습니다.

비판적 평가

GOAT는 강력한 이론적 배경과 실용성을 겸비한 인상적인 방법론이지만, 몇 가지 고려할 점이 있습니다.

파라미터 증가: 학습 가능한 사전을 위해 쿼리와 키 벡터에 위치 부분을 추가해야 하므로, 모델의 파라미터가 소폭 증가합니다.
사전의 효과: 학습된 사전의 효과는 태스크와 데이터의 특성에 따라 달라질 수 있습니다. 특정 데이터셋에서는 단순한 편향(e.g., ALiBi)만으로도 충분할 수 있으며, 이 경우 GOAT의 유연성이 과도할 수 있습니다.
해석의 어려움: log π가 무엇을 학습했는지 시각화하고 분석할 수 있지만, 그 내부 동작을 완전히 이해하는 것은 여전히 어려운 과제일 수 있습니다.

향후 연구 방향

GOAT는 어텐션 연구에 새로운 가능성을 열었습니다. 향후 연구는 다음과 같은 방향으로 확장될 수 있습니다.

사전 파라미터화: $q_p k_p^T$ 외에 더 효율적이거나 표현력이 풍부한 사전 파라미터화 방식 탐구
멀티모달 적용: 텍스트, 이미지, 오디오 등 여러 양식이 결합된 모델에서 각 양식의 특성에 맞는 사전을 학습하는 연구
동적 사전: 입력에 따라 동적으로 변하는 사전을 생성하여 모델의 유연성을 더욱 높이는 연구

실무 적용 가이드

GOAT를 실무에 적용할 때 다음 팁을 고려할 수 있습니다.

장문맥 모델에 우선 적용: 특히 시퀀스 길이가 수만 토큰 이상으로 길어지는 RAG, 문서 요약, 코드 생성 등의 태스크에서 큰 성능 향상을 기대할 수 있습니다.
사전 초기화: 장문맥 모델의 안정적인 학습을 위해, 학습 초기에 사전( $\log \pi$ )이 주의 싱크(초기 토큰에 높은 가중치)나 ALiBi와 유사한 형태가 되도록 초기화하는 것이 수렴에 도움이 될 수 있습니다.
FlashAttention 활용: GOAT는 FlashAttention과 호환되므로, 훈련 및 추론 시 반드시 최적화된 커널을 사용하여 메모리 및 속도 이점을 극대화해야 합니다.

결론

GOAT는 기존 어텐션 메커니즘의 암묵적인 '균일 사전' 가정을 명시적인 '학습 가능한 사전'으로 대체함으로써, 어텐션 연구의 새로운 지평을 열었습니다. 엔트로피 최적 수송이라는 탄탄한 이론적 프레임워크 위에서 어텐션을 재정의함으로써, 모델이 데이터에 가장 적합한 귀납적 편향(inductive bias)을 스스로 학습하게 만들었습니다. 그 결과, 길이 일반화와 장문맥 처리 능력에서 놀라운 성능 향상을 보였으며, 하드웨어 가속 커널과의 호환성을 통해 실용성까지 확보했습니다. 이 연구는 향후 트랜스포머 아키텍처의 발전에 중요한 이정표가 될 것입니다.

참고 자료

논문: You Need Better Attention Priors (arXiv:2401.15380)
코드 저장소: https://github.com/HazyResearch/goat

[논문 리뷰] You Need Better Attention Priors