라운드 로빈 샘플링 기반 동적 블록 희소 어텐션

초록

RRAttention은 헤드별 라운드 로빈 방식으로 쿼리를 샘플링하고, 스트라이드 수준에서 중요도를 추정해 블록 단위로 Top‑τ 선택을 수행한다. 이를 통해 $O(L^2)$의 복잡도를 $O(L^2/S^2)$로 낮추면서도 전역 패턴을 포착하고, 쿼리 독립성을 유지한다. 실험 결과 HELMET과 Video‑MME에서 99% 이상의 정확도를 유지하면서 128K 토큰 길이에서 2.4배 속도 향상을 달성한다.

상세 분석

RRAttention은 기존 동적 희소 어텐션이 안고 있던 “전처리 필요·전역 평가 부족·쿼리 독립성 위배·패턴 제한·소프트맥스 granularity” 다섯 가지 핵심 트레이드오프를 동시에 해결한다는 점에서 혁신적이다. 핵심 아이디어는 헤드 라운드 로빈 샘플링이다. 전체 시퀀스를 stride $S$ 로 나눈 뒤, 헤드 $h$마다 $P(i,h)=iS+(S-1-(h\mod S))$ 위치의 쿼리를 선택한다. 이렇게 하면 $S$개의 헤드가 서로 다른 위치를 순차적으로 커버하므로, 고정된 샘플링 포인트가 놓칠 수 있는 중요한 수직/대각 패턴을 모두 탐지할 수 있다.

샘플링된 쿼리는 해당 스트라이드 전체 키를 평균(또는 합)한 벡터와 내적하여 스트라이드‑레벨 중요도 $I^{(h)}{i,j}$ 를 계산한다. 여기서 $i$는 쿼리 스트라이드, $j$는 키 스트라이드이며, 정규화 계수 $1/(S\sqrt{d})$ 를 적용해 스케일을 맞춘 뒤 행별 소프트맥스를 적용해 확률 형태의 중요도 $P^{(h)}{i,j}$ 를 얻는다. 이 과정은 $O(L^2/S^2)$ 연산만 필요하므로, 기존 $O(L^2)$ 대비 $S^2$ 배의 효율성을 제공한다.

다음 단계에서는 스트라이드‑레벨 중요도를 블록 단위로 집계한다. 블록 $m$(쿼리)과 블록 $n$(키)에 대해 $S^{(h)}{m,n}=\sum{i\in m}\sum_{j\in n}P^{(h)}{i,j}$ 를 구하고, 각 쿼리 블록마다 Top‑τ 임계값을 적용해 누적 중요도가 $\tau$ 를 초과하는 키 블록만을 선택한다. 이때 선택된 블록은 동적 블록 매트릭스 $B{\text{dynamic}}$ 로 표시되며, 마지막 쿼리 블록에 대한 완전한 커버를 보장하기 위해 정적 매트릭스 $B_{\text{static}}$ 와 OR 연산을 수행한다.

RRAttention의 설계는 쿼리 독립성을 보장한다. 각 헤드가 독립적인 쿼리 위치를 사용하므로, 다른 쿼리와의 교차 간섭이 없으며, 이는 어텐션 분포의 의미적 일관성을 유지한다. 동시에 전역 평가가 가능하다. 스트라이드‑레벨 집계는 전체 키 공간을 고려하므로, 초기 쿼리 블록이 놓치는 장거리 의존성도 포착한다.

복잡도 측면에서, 스트라이드 $S$ 를 적절히 조정하면 메모리 사용량과 연산량을 선형적으로 조절할 수 있다. 논문에서는 $S=4$ 혹은 $8$ 을 사용해 128K 토큰 길이에서 2.4배 속도 향상을 보고했으며, Top‑τ 선택은 입력에 따라 자동으로 sparsity 를 조절한다(Adaptive Top‑τ). 실험 결과는 HELMET의 다양한 언어 이해 태스크와 Video‑MME의 멀티모달 비디오 이해 태스크에서 99% 이상 전체 어텐션 성능을 유지하면서, 기존 FlexPreFill, XAttention 대비 0.51.0% 이상의 정확도 향상과 1030% 정도의 연산 감소를 달성했다.

또한, RRAttention은 전처리‑프리이다. 사전 학습된 distillation 모델이나 오프라인 패턴 검색이 필요 없으며, 기존 모델에 바로 적용 가능하다. 이는 배포 환경에서의 유연성을 크게 높인다. 마지막으로, stride‑level softmax를 사용함으로써 블록‑레벨 연산과 GPU 메모리 효율성을 동시에 확보한다.

요약하면, RRAttention은 헤드 라운드 로빈 샘플링, 스트라이드‑레벨 중요도 추정, Top‑τ 블록 선택이라는 세 가지 핵심 메커니즘을 결합해, 동적 희소 어텐션의 핵심 요구사항을 모두 만족시키는 동시에 실용적인 속도·메모리 이점을 제공한다.