RACE Attention: A Strictly Linear-Time Attention for Long-Sequence Training

이 논문은 시퀀스 길이에 대해 엄격하게 선형 시간 복잡도를 가지며 기존 소프트맥스 어텐션의 계산 병목 현상을 해결하여 긴 컨텍스트 학습을 가능하게 하는 'RACE 어텐션'을 제안하고, 다양한 작업에서 강력한 성능을 유지하면서 최대 7500 만 토큰까지 단일 패스로 처리할 수 있음을 입증합니다.

핵심

1. 문제: "모든 친구를 한 번에 만나야 하는 비효율"

기존의 AI 는 글을 읽을 때 **'Softmax Attention'**이라는 방식을 사용합니다. 이를 비유하자면 다음과 같습니다.

비유:
당신이 100 명 친구가 있는 파티에 갔다고 상상해 보세요.
AI 가 "지금 이 친구 (단어) 가 내 이야기와 얼마나 관련이 있을까?"를 판단할 때, 100 명 친구 모두를 한 명씩 찾아다니며 눈을 마주치고 대화를 해야 합니다.

• 친구가 100 명이면 100 번의 대화.
• 친구가 100 만 명이면 100 만 번의 대화.
• 친구가 1 억 명이면... 1 억 번의 대화를 해야 합니다.

이 방식은 친구 (데이터) 가 적을 때는 괜찮지만, 친구가 수천만 명 (긴 문서, 긴 영상) 으로 늘어나면 시간과 에너지 (컴퓨터 메모리) 가 너무 많이 소모되어 AI 가 멈추거나, 아주 비싼 슈퍼컴퓨터가 아니면 처리 자체가 불가능해집니다.

2. 해결책: "RACE Attention - '요약된 메모'로 빠르게 파악하기"

이 논문에서 제안한 RACE Attention은 이 비효율적인 방식을 완전히 바꿉니다. 모든 친구를 일일이 만나지 않고, 친구들을 몇 개의 '조 (Group)'로 묶어 대표자만 만나서 전체 상황을 파악하는 방식입니다.

핵심 아이디어 1: "친구들을 조로 묶기 (LSH)"

AI 는 긴 글을 읽을 때, 비슷한 내용을 가진 단어들을 비슷한 '통 (Bucket)'에 자동으로 분류합니다.

• 기존 방식: "A 라는 단어와 B 라는 단어가 얼마나 닮았는지?"를 1 대 1 로 계산.
• RACE 방식: "A 와 B 가 같은 '통'에 들어갈 확률이 높은가?"를 확인.
  • 만약 두 단어가 같은 통에 들어있다면, 서로 닮았다고 가정하고 **통의 대표자 (요약 정보)**와만 대화합니다.
  • 이렇게 하면 1 억 명의 친구를 만나야 할 일을, 수십 개의 통 대표자만 만나면 되는 것입니다.

핵심 아이디어 2: "부드러운 분류 (Soft Hashing)"

기존의 비슷한 기술들은 친구를 통에 넣을 때 "딱! A 통 아니면 B 통"이라고 딱딱하게 나누는 경우가 많았습니다. 하지만 RACE 는 **"A 통에 70%, B 통에 30% 씩 섞여 있을 수도 있다"**고 부드럽게 (Soft) 분류합니다.

• 비유: 친구를 분류할 때 "너는 축구부야!"라고 딱 잘라 말하는 대신, "너는 축구부에도 70% 어울리고, 농구부에도 30% 어울리네?"라고 정확한 비율을 고려합니다.
• 이 덕분에 AI 가 학습할 때 실수를 바로잡을 수 있어 정확도도 떨어지지 않습니다.

3. 왜 이것이 대단한가요? (실제 성과)

이 기술은 단순히 이론적인 이야기가 아니라, 실제 하드웨어에서 놀라운 속도를 보여줍니다.

• 기존 방식 (FlashAttention): 최신 GPU(슈퍼컴퓨터) 를 써도 약 400 만 단어를 넘으면 메모리가 터지거나 계산이 너무 느려져서 멈춥니다.
• RACE 방식:
  • **CPU(일반 컴퓨터)**에서도 7,500 만 단어를 한 번에 처리했습니다.
  • GPU에서도 1,200 만 단어를 처리했습니다.
  • 속도 비교: 400 만 단어를 처리할 때, 기존 방식은 550 초가 걸렸는데 RACE 는 0.1 초 만에 끝냈습니다. (약 5,500 배 빠름!)

4. 결론: "AI 의 시야를 무한히 넓히다"

이 논문의 핵심 메시지는 다음과 같습니다.

"더 강력한 컴퓨터를 사서 해결할 문제가 아닙니다. 알고리즘을 똑똑하게 바꾸면, 일반 컴퓨터로도 거대한 데이터를 순식간에 처리할 수 있습니다."

RACE Attention은 AI 가 이제까지 불가능했던 수백만 단어의 책, 긴 영상, 복잡한 코드를 한 번에 읽고 이해할 수 있게 해주는 열쇠가 됩니다. 앞으로 우리가 AI 와 대화할 때, 훨씬 더 길고 복잡한 이야기를 나누게 될 날이 머지않았음을 의미합니다.

---

한 줄 요약:
기존 AI 가 긴 글을 읽을 때 '모든 단어를 일일이 비교'해서 느려진다면, RACE Attention은 '비슷한 단어들을 그룹화해서 대표자만 비교'함으로써 속도는 수천 배 빨라지고 정확도는 그대로 유지하는 획기적인 기술입니다.

기술 요약

1. 문제 정의 (Problem)

현대 트랜스포머 (Transformer) 모델의 핵심인 Softmax Attention은 시퀀스 길이 ($N$) 에 대해 **이차적 시간 복잡도 ($O(N^2)$)**를 가집니다. 이는 긴 문맥 (Long Context) 을 처리할 때 계산 비용과 메모리 사용량이 급증하여 실용적인 한계에 부딪히게 됩니다.

• 현실적 한계: 최신 GPU 최적화 커널인 FlashAttention-2/3 를 사용하더라도, NVIDIA GH200 (96GB) 같은 최상위 하드웨어에서도 단일 어텐션 레이어의 순전파/역전파를 수행할 때 시퀀스 길이가 약 400 만 토큰을 초과하면 실행이 불가능해집니다.
• 기존 대안의 부족: 선형 어텐션 (Linear Attention), Performer, Linformer 등의 기존 근사 방법들은 정확도 저하, 임베딩 차원에 대한 이차적 복잡도, 또는 추론/학습 시 불안정성 등의 문제를 겪고 있습니다. 특히 YOSO 와 같은 방법론은 미분 불가능성 (non-differentiable) 으로 인해 엔드 - 투 - 엔드 학습 시 성능이 떨어집니다.

2. 방법론 (Methodology)

저자들은 RACE (Repeated Arrays-of-Count Estimators) Attention을 제안합니다. 이는 Softmax Attention 을 대체할 수 있으며, 시퀀스 길이와 임베딩 차원에 대해 **엄격하게 선형 (Strictly Linear)**인 시간 복잡도를 가지는 새로운 어텐션 메커니즘입니다.

핵심 아이디어

1. 각도 커널 (Angular Kernel) 의 강화:

   • Softmax 의 지수 함수 대신, 코사인 유사도에 기반한 **각도 커널 (Angular Kernel)**을 사용합니다.
   • 단순한 각도 커널은 고차원 공간에서 구별력이 부족하므로, 이를 강화 파라미터 $\gamma$를 사용하여 고차 다항식 (Higher-degree monomial) 형태로 변환하여 Softmax 와 유사한 날카로운 분포를 만듭니다.
   • 유사도 함수: $sim(Q_i, K_j) = (1 - \frac{\cos^{-1}(Q_i^\top K_j)}{\pi})^\gamma$

2. RACE 스케치 (RACE Sketch) 를 통한 선형 근사:

   • 전체 $N \times N$ 어텐션 행렬을 계산하지 않고, 국소 민감 해싱 (LSH) 기반의 RACE 스케치를 사용하여 어텐션 출력을 근사합니다.
   • 소프트 해싱 (Soft Hashing): 기존 RACE 의 이산적 (discrete) 해싱을 미분 가능한 (differentiable) 부드러운 확률 할당으로 변경합니다. 쿼리와 키를 하이퍼큐브의 모서리 (Corners) 로 부드럽게 매핑하여 엔드 - 투 - 엔드 학습을 가능하게 합니다.
   • 통계량 집계: 각 해싱 테이블에서 키 (Key) 와 값 (Value) 의 가중 합을 버킷 단위로 집계한 후, 쿼리 (Query) 가 해당 버킷 통계량을 참조하여 출력을 계산합니다.

3. 알고리즘 흐름:

   • $L$개의 해싱 테이블을 사용합니다.
   • 각 테이블에서 $P$개의 랜덤 초평면 (Hyperplanes) 을 사용하여 쿼리와 키를 투영합니다.
   • 투영된 벡터를 $R=2^P$개의 버킷으로 부드럽게 할당 (Soft Bucketization) 합니다.
   • 버킷별 키의 질량 (Mass) 과 값 (Value) 의 합을 집계하여 어텐션 출력을 선형 시간으로 재구성합니다.

3. 주요 기여 (Key Contributions)

1. 엄격한 선형 시간 복잡도:

   • RACE Attention 은 시퀀스 길이 $N$과 임베딩 차원 $d$에 대해 $O(N \cdot d)$의 선형 복잡도를 달성합니다.
   • 기존 FlashAttention 이 $O(N^2)$의 메모리/시간 제약을 받는 반면, RACE 는 전체 어텐션 행렬을 생성하지 않아 메모리 사용량을 획기적으로 줄입니다.

2. 이론적 보장 및 미분 가능성:

   • YOSO 와 달리 미분 가능한 부드러운 LSH를 도입하여 엔드 - 투 - 엔드 학습을 지원합니다.
   • 근사 오차에 대한 이론적 증명: $L$ (테이블 수), $P$ (초평면 수), $\beta$ (온도 파라미터) 가 근사 오차 (Bias-Variance Trade-off) 에 어떻게 영향을 미치는지 수학적으로 증명했습니다.

3. 압도적인 확장성 (Scaling):

   • GPU: NVIDIA GH200 에서 1200 만 토큰까지 단일 어텐션 레이어의 순전파/역전파를 처리 가능 (FlashAttention 은 400 만 토큰에서 한계).
   • CPU: Intel Xeon Gold 5220R 에서 7500 만 토큰까지 처리 가능.
   • 이는 현재 최첨단 어텐션 구현체들의 능력을 훨씬 능가하는 수치입니다.

4. 하드웨어 가속기보다 알고리즘의 우위:

   • 최상위 GPU 에서 FlashAttention-2/3 를 사용하는 것보다, 일반 CPU 에서 RACE Attention 을 사용하는 것이 더 긴 시퀀스 (약 400 만 토큰 이상) 에서 더 빠른 속도를 기록했습니다. 이는 알고리즘의 복잡도 감소가 하드웨어 가속보다 더 결정적임을 보여줍니다.

4. 실험 결과 (Results)

• 정확도:

  • 언어 모델링 (WikiText-103, PTB), 마스킹 언어 모델링 (Tiny Stories), 텍스트/이미지 분류 (QNLI, CIFAR-10 등) 에서 기존 Strong Baselines (FlashAttention, Linformer, Performer 등) 과 동등하거나 더 나은 성능을 보였습니다.
  • 특히 64K 시퀀스 길이에서도 높은 정확도를 유지했습니다.

• 성능 및 효율성:

  • 속도: 400 만 토큰 시퀀스에서 FlashAttention-2 대비 약 5500 배, FlashAttention-3 대비 약 2600 배 빠릅니다.
  • 메모리: 긴 시퀀스 처리 시 OOM (Out Of Memory) 오류 없이 처리 가능하며, 활성화 (Activation) 메모리 사용량이 크게 감소합니다.
  • 비교: Linformer 나 Performer 와 같은 기존 선형 어텐션 방법들은 긴 시퀀스에서 정확도가 떨어지거나, 큰 상수 인자로 인해 실제로는 RACE 보다 느리고 메모리 효율이 낮았습니다.

5. 의의 및 결론 (Significance)

• 긴 문맥 학습의 실용화: RACE Attention 은 오늘날의 상용 하드웨어 (Commodity Hardware) 에서 수십만에서 수천만 토큰에 달하는 긴 문맥을 가진 모델을 실제적으로 훈련 (Training) 하고 추론 (Inference) 할 수 있는 길을 열었습니다.
• 이론과 실전의 결합: 단순한 휴리스틱이 아닌, RACE 스케치와 LSH 이론에 기반한 엄밀한 수학적 프레임워크를 제공하여 효율성과 정확도 간의 트레이드오프를 정량화할 수 있게 했습니다.
• 미래 지향성: 추론 시 키 - 값 캐싱 (KV Caching) 을 효율적으로 지원할 수 있으며, 커스텀 CUDA 커널을 통해 더 최적화된 성능을 기대할 수 있는 기반을 마련했습니다.

결론적으로, 이 논문은 Softmax Attention 의 이차적 병목 현상을 근본적으로 해결하여, 차세대 초장문맥 (Ultra-long Context) AI 모델 개발을 위한 새로운 표준을 제시합니다.