📊 statistics

WildCat: Near-Linear Attention in Theory and Practice

WildCat는 무작위 피벗 숄레스키 서브샘플링(randomly pivoted Cholesky subsampling)을 통해 작은 코어셋을 선택하고 가중치를 부여함으로써 근선형 시간 복잡도와 초다항 오차 감소를 달면서, 이미지 및 언어 작업 전반에서 이론적 보장과 실질적 성능 모두에서 이전의 근사 기법들을 능가하는 고정밀, 저비용 어텐션 압축 방법이다.

원저자: Tobias Schröder, Lester Mackey

게시일 2026-06-02

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Tobias Schröder, Lester Mackey

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 친구와 몇 시간 동안 나눈 대화를 기억하려고 노력하고 있다고 상상해 보세요. 표준 AI 모델(오늘날의 챗봇을 구동하는 것과 같은 모델)에서 그 대화의 모든 단어를 한꺼번에 기억해 내려는 것은, 마치 새로운 책을 쓰면서 동시에 머릿속에 도서관 전체의 책들을 담아두려는 것과 같습니다. 단어가 추가될수록 그 과정은 더 어려워지며, 필요한 메모리는 **이차 함수적(quadratically)**으로 증가합니다. 즉, 대화의 길이가 두 배가 되면, 기억하는 데 드는 노력은 단순히 두 배가 되는 것이 아니라 네 배로 늘어납니다. 결국 컴퓨터는 메모리가 부족해지거나 생각하는 데 너무 오랜 시간이 걸리게 됩니다.

이 논문은 이 문제를 해결하기 위해 WILDCAT(Weighted Iterative Low-rank Decomposition for Coreset ATtention)이라는 새로운 방법을 소개합니다. 이 방식이 어떻게 작동하는지 쉬운 비유를 통해 설명하겠습니다.

문제점: "도서관" 병목 현상

AI의 메모리를 거대한 도서관이라고 생각해 보세요. AI가 질문에 답해야 할 때, 보통은 관련 정보를 찾기 위해 모든 책(대화 속의 모든 단어)을 일일이 훑어봐야 합니다.

기존 방식: 만약 도서관에 책이 1,000권 있다면, AI는 1,000권을 확인합니다. 만약 도서관이 10,000권으로 늘어나면, AI는 10,000권을 확인해야 하지만, 그 과정에서 발생하는 노력은 폭발적으로 증가합니다. 이는 마치 모든 짚단 하나하나를 다른 모든 짚단과 비교하며 건더미 속에서 특정 바늘을 찾으려는 것과 같습니다.

해결책: "전문가 패널" (WILDCAT)

WILDCAT은 전략을 바꿉니다. 모든 것을 완벽하게 기억하려고 하는 대신, 대화 속의 모든 단어가 똑같이 중요한 것은 아니라는 점을 깨닫습니다. 어떤 단어는 매우 중요하지만, 어떤 단어는 그저 채우기용일 뿐입니다.

WILDCAT은 크게 두 가지 일을 수행합니다.

"코어셋(Coreset)" 선정 (전문가 패널):
당신에게 10,000명의 거대한 군중이 있고, 그 방의 전반적인 분위기를 파악해야 한다고 가정해 봅시다. 모든 사람을 인터뷰하는 대신, WILDCAT은 영리하고 빠른 알고리즘("Randomly Pivoted Cholesky")을 사용하여, 예를 들어 50명 정도의 아주 작고 대표성 있는 그룹을 뽑아냅니다.

마법 같은 점: 단순히 무작위로 뽑는 것이 아닙니다. 전체 군중을 가장 잘 대변할 수 있는 가장 중요한 사람들을 선별합니다. 이는 마을의 가장 목소리가 크고 다양한 구성원들을 뽑아 전체를 대변할 수 있는 의회를 구성하는 것과 같습니다.

목소리에 가중치 부여하기:
이 작은 50명의 그룹을 뽑은 후, WILDCAT은 이들을 모두 똑같이 취급하지 않습니다. 각자에게 "가중치"를 부여합니다.

비유: 만약 이 작은 그룹 중 한 사람이 매우 크고 주관이 뚜렷한 리더라고 가정해 봅시다. 그렇다면 그 사람의 목소리는 더 큰 비중을 차지합니다. 반대로 어떤 사람이 조용하다면, 그 사람의 목소리는 비중이 적어집니다. WILDCAT은 이 50명의 목소리 각각에 완벽한 "볼륨"을 계산하여, 오직 이들만 듣더라도 10,000명 전체의 목소리를 듣는 것과 똑같이 들리도록 만듭니다.

이것이 왜 중요한가

논문은 WILDCAT이 세 가지 주요 돌파구를 달성했다고 주장합니다.

속도 (선형에 가까운 속도): 10,000명의 군중 전체를 듣는 대신 50명의 작은 그룹에만 집중하기 때문에, 생각하는 데 걸리는 시간이 매우 느리게 증가합니다. 대화의 길이가 두 배가 되어도, 시간은 엄청나게 늘어나는 것이 아니라 아주 조금씩만 증가합니다. 이는 책의 모든 페이지를 읽는 것에서 핵심 요약본만 읽는 것으로 전환하는 것과 같습니다.
정확도 (초다항식적 정확도): 보통 무언가를 요약하면 세부 사항을 놓치기 마련입니다. 하지만 WILDCAT은 특별합니다. 중요한 세부 사항을 거의 놓치지 않는다고 주장하기 때문입니다. 논문은 수학적으로 대화가 길어질수록 오차(누락된 세부 사항)가 다른 거의 모든 기존 방식보다 훨씬 빠르게 줄어든다는 것을 증명합니다.
실용성: 저자들은 단순히 수학적 계산에 그치지 않고, 강력한 컴퓨터 칩(GPU)에서 작동하는 버전을 실제로 구축했습니다. 그들은 다음 작업들을 통해 테스트했습니다:
- 이미지 생성: 이미지가 흐릿하거나 이상해지지 않으면서도 다른 방식보다 더 빠르게 고품질 이미지를 생성했습니다.
- 이미지 분류: 기존의 느린 방식만큼 정확하게 사진 속 사물을 인식하면서도 훨씬 빠르게 수행했습니다.
- 긴 대화: 언어 모델이 메모리 부족 없이 훨씬 더 긴 대화 기록을 기억할 수 있게 해주었으며, 다른 "메모리 압축" 기술들보다 뛰어난 성능을 보였습니다.

핵심 요약

WILDCAT은 10시간짜리 회의를 듣고, 즉시 가장 중요한 50개 순간을 골라내어, 각각의 중요도를 적절히 배정한 뒤, 상사가 마치 모든 내용을 다 들은 것처럼 느끼도록 매우 정확하게 요약해 주는 초효율적인 비서와 같습니다. 다만 비서는 단 몇 페이지의 기록만 작성했을 뿐입니다.

이를 통해 AI 모델은 슈퍼컴퓨터 없이도 훨씬 더 긴 대화와 더 큰 작업들을 처리할 수 있으며, 이를 통해 더 빠르고, 저렴하며, 더 많은 맥락을 기억할 수 있게 됩니다.

기술 요약: WILDCAT: 이론과 실제에서의 근선형 어텐션 (Near-Linear Attention in Theory and Practice)

문제 정의

어텐션 메커니즘은 현대 트랜스포머 기반 아키텍처의 초석으로서, 시퀀스의 장거리 의존성을 모델링할 수 있게 해준다. 그러나 입력 시퀀스 길이 $n$ 에 따라 자원 요구량이 이차식( $\Theta(n^2d)$ )으로 증가하기 때문에 배포 시 비용이 매우 높다는 문제가 있다. 다양한 근사 어텐션 방법들이 존재하지만, 이론적 보장과 실제 성능 사이에는 상당한 격차가 존재한다:

실제적인 근사법들 (예: Reformer, Performer, ScatterBrain)은 엄격한 오차 보장이 부족하거나 높은 충실도를 보장하기 위해 이차 시간 복잡도를 요구한다.
이론적 진전 (예: Alman & Song, 2023)은 유계된 입력(bounded inputs)에 대해 근선형 시간( $O(n^{1+o(1)})$ )과 초다항식 오차 감소(super-polynomial error decay)가 가능하다는 것을 보여주었으나, 기존의 실제 구현체들은 근선형 시간 내에 이러한 감소를 달ach하지 못하거나 이를 보장하기 위해 이차 시간을 요구한다.

본 논문은 계산 효율성(근선형 시간), 스펙트럼 정확도(초다항식 오차 감소), 그리고 현대 하드웨어에서의 실질적인 구현 가능성을 동시에 갖춘 접근 방식의 필요성을 제기한다.

방법론

저자들은 입력 키의 작고 가중치가 부여된 코어셋(coreset)에 대해서만 어텐션을 수행함으로써 소프트맥스 어텐션 메커니즘을 근사하는 방법인 WILDCAT(Weighted Iterative Low-rank Decomposition for Coreset ATtention)을 소개한다. 이 방법론은 세 가지 핵심 구성 요소로 이루어져 있다:

1. 무작위 피벗 촐레스키(Randomly Pivoted Cholesky, RPC)를 통한 가중 코어셋 선택

모든 $n$ 개의 키에 어텐션을 수행하는 대신, WILDCAT은 $r$ 개의 키( $r \ll n$ ) 부분집합을 선택하여 코어셋 $K_S$ 를 형성한다.

알고리즘: 선택 과정은 커널 행렬 $H = \exp(\beta K K^\top)$ 의 부분 촐레스키 분해를 구축하는 병렬화된 무작위 피벗 촐레스키(RPC) 알고리즘을 사용한다(Chen et al., 2022a에서 차용). 이 알고리즘은 현재 잔여 커널 행렬의 대각선에서 피벗 인덱스를 샘플링한다. 이는 선택된 지점들이 전체 데이터셋을 스펙트럼 측면에서 대표하도록 보장한다.
효율성: 선택 과정은 $O(nr^2)$ 시간에 실행된다. GPU 처리량을 극대화하기 위해 저자들은 입력을 $B$ 개의 빈(bin)으로 분할하여 병렬로 처리하는 분할 정복 전략을 채결하여, 실행 시간을 $O(nr^2/B^3 + nrd/B^2)$ 로 단축했다.

2. 님스트롬 근사(Nyström Approximation)를 통한 최적 재가중치 부여

코어셋 $K_S$ 가 선택되면, 어텐션 행렬의 재구성 오차를 최소화하기 위해 최적의 가중치를 계산한다.

님스트롬 가중치: 가중치 $W$ 는 님스트롬 방법을 사용하여 $W = h(K_S, K_S)^+ h(K_S, K)$ 로 도출된다. 이는 전체 커널 특징을 근사하기 위한 코어셋 특징들의 최적 선형 결합을 제공한다.
값 압축(Value Compression): 값 행렬 $V$ 는 $V_S = WV$ 로 압축되어, 최종 어텐션 출력을 오직 $r$ 개의 코어셋 요소만을 사용하여 계산할 수 있게 한다.

3. 불변성 및 재스케일링(Invariance and Rescaling)

저자들은 키의 전역적 재중심화(global recentering)와 키 및 쿼리의 재스케일링에 따른 소프트맥스 어텐션의 불변성을 활용한다.

재중심화: 키는 평균을 빼는 방식으로 재중심화된다 ( $K \leftarrow K - \bar{k}$ ).
온도 스케일링(Temperature Scaling): 커널 행렬을 저계수(low-rank) 근사가 가능하게 만드는 것(온도 $\tau$ 를 높임)과 쿼리 기반 오차 팽창 계수를 최소화하는 것 사이의 트레이드오프를 조절하기 위해 람베르트 W 함수(Lambert-W function)와 쿼리 및 키의 노름(norm)에 의존하는 폐쇄형 재스케일링 파라미터 $\tau$ 를 도출한다.

주요 기여

본 논문은 WILDCAT에 대해 다음과 같은 바람직한 속성들을 확립한다:

근선형 실행 시간: WILDCAT은 정확한 어텐션의 이차 비용을 피한다. $r \in n^{o(1)}$ 크기의 코어셋에 대해 어텐션을 수행함으로써, 알고리즘은 $O(n^{1+o(1)})$ 의 근선형 시간 내에 실행된다.
초다항식 오차 감소: 유계된 입력 가정 하에, WILDCAT은 $O(n^{-\sqrt{\log(\log(n))}})$ 의 초다ynomial 오차 감소를 가지며 정확한 어텐션을 근사한다. 이는 이전의 실용적인 근사법들이 이차 시간 없이 달성할 수 없었던 것을 근선형 시간 내에 달성한 것이다.
성장에 대한 강건성: 이론적 분석(Corollary 2)에 따르면, WILDCAT은 입력 항목이나 차원이 $n$ 에 대해 초로그(super-logarithmically)로 성장하더라도 초다항식 오차 감소를 유지하며, 이는 기존의 어텐션 계산 복잡도 한계를 뛰어넘는다.
KV 캐시 압축: 이 방법은 KV 캐시의 메모리 점유율을 선형 $\Theta(nd)$ 에서 근상수(near-constant) $O(rd)$로 줄이는 COMPRESSKV 모듈을 포함한다. 이는 장문맥 추론을 용이하게 한다.
실제적 구현: 저자들은 이론적 스펙트럼 보장과 실제 성능 사이의 간극을 메우기 위해 GPU에 최적화된 PyTorch 구현체를 제공한다.

실험 결과

저자들은 세 가지 도메인에서 WILDCAT을 평가하였다:

이미지 생성 (BigGAN): WILDCAT은 정확한 어텐션 대비 4.33배의 속도 향상을 달了했다. Inception Score(IS)의 저하가 가장 작았으며(1.22%), ImageNet 검증 세트에 대해 Fréchet Inception Distance(FID) 저하가 전혀 없었으며, 이는 Reformer, Performer, KDEformer, Thinformer보다 우수한 성능을 보였다.
이미지 분류 (T2T-ViT): ImageNet 분류 벤치마크에서 WILDCAT은 모든 근사법 중 가장 높은 평균 Top-1 정확도(82.19%, 정확한 어텐션의 82.55% 대비)를 기록했다. 또한 계산량이 가장 많은 첫 번째 레이어에서 11.59배의 속도 향상을 제공했다.
장문맥 언어 모델 (KV 캐시 압축): Qwen2.5-7B-Instruct 모델을 사용하여 13개의 LongBench-E 태스크에 대해 테스트한 결과, WILDCAT(COMPRESSKV를 통해)은 다양한 압축 수준(75%, 87.5%, 93.75%)에서 다섯 가지 주요 KV 캐시 압축 방법(StreamingLLM, PyramidKV, SnapKV 포함)보다 일관되게 우수한 평균 압축 품질을 보였다.
FlashAttention 2와의 비교: FlashAttention 2와의 벤치마크 결과, 시퀀스 길이가 증가함에 따라 ( $2^{13}$ 에서 $2^{18}$ 까지) WILDCAT의 속도 향상은 1.1배에서 68배까지 성장하는 동시에 근사 오차는 꾸준히 감소하였다.

의의 및 주장

본 논문은 어텐션 근사에서의 이론-실제 간극을 성공적으로 메웠다고 주장한다.

이론적 돌파구: 유계된 실제 입력에 대해 근선형 시간 내에 초다항식 오차 감소가 가능하다는 것을 입증하였으며, 이는 이전에는 이론적으로만 가능하거나 이차 시간을 요구했던 보장이다.
실용적 유용성: 이 방법은 단순한 이론적 구조물이 아니라, 생성, 분류, 언어 이해를 위한 최첨단 모델에서 효율적으로 구현되고 GPU 상에서 검증되었다.
확장성: KV 캐치 메모리 요구량을 근상수 크기로 줄임으로써, WILDCAT은 자원이 제한된 하드웨어에서도 장문맥 언어 모델을 배포할 수 있게 한다.

저자들은 한계점도 언급하였다. 구체적으로, 현재의 작업은 인과적 마스킹(causal masking)이 있는 스트리밍 생성 문제를 다루지 않았으나(prefix sum이나 divide-and-conquer를 통한 확장이 가능함을 시사함), RPC의 순차적인 피벗 선택 특성이 향후 오버샘플링 메커니즘을 통해 개선될 수 있음을 밝혔다.