멀티그라뉼러 압축으로 구현한 통합 스파스 어텐션

초록

UniSparse는 복합 토큰이라는 압축된 표현을 이용해 다중‑그라뉼러 컨텍스트를 요약하고, 이를 기반으로 동적 블록‑스파스 마스크를 생성한다. 압축‑기반 프록시 연산은 GPU 친화적이며, 사전 학습 없이도 텍스트·영상·오디오 등 다양한 모달리티에서 99% 이상의 정확도를 유지하면서 최대 2.61배 빠른 어텐션을 제공한다.

상세 분석

UniSparse의 핵심 아이디어는 “복합 토큰”이라는 개념이다. 기존 스파스 어텐션은 전체 토큰 수준에서 중요도를 추정하거나, 단순히 고정된 히스토리 패턴을 적용한다. 그러나 토큰 간 의미적 연관성은 지역적 집합에서 강하게 나타나며, 이러한 지역을 평균 풀링 등 간단한 공간 집계 연산으로 요약하면 원본 토큰들의 상대적 중요도 순위가 크게 변하지 않는다. 논문은 HELMET 데이터셋을 이용해 압축 비율 c=8일 때 블록 중요도 순위의 스피어만 상관계수가 0.98 이상임을 실험적으로 입증한다. 이는 절대 점수보다는 순위가 스파스 마스크 결정에 더 중요하다는 점을 강조한다.

압축 단계는 두 단계로 나뉜다. 첫째, 시퀀스 차원에서 c_q, c_k 라는 정수 압축 계수를 적용해 쿼리와 키 행렬을 각각 S′=S/c 로 다운샘플링한다. 여기서 평균 풀링이 기본이지만, 헤드 차원까지 압축하는 옵션을 제공해 메모리 대역폭을 추가로 절감한다. 둘째, 압축된 토큰들 간에 S′×S′ 규모의 어텐션 스코어 행렬을 계산하고, 블록 수준에서 평균 혹은 최대값을 집계한다. 이렇게 얻어진 블록 중요도는 Top‑P 방식으로 상위 블록을 선택해 최종 마스크 M을 만든다. 마스크는 기존 FlashAttention 같은 블록‑와이즈 커널에 그대로 적용되므로 구현 복잡도가 낮고, GPU 캐시 활용도가 높다.

연산 복잡도 측면에서 UniSparse는 C_select를 O(L·S/c) 수준으로 낮춘다. 기존 동적 스파스 방법들은 종종 전체 토큰 간 유사도 계산이나 복잡한 라우터 네트워크를 도입해 C_select가 O(L²)와 근접하거나, 높은 상수 비용을 가진다. 반면 UniSparse는 단순 평균 풀링과 작은 행렬 곱셈만으로 마스크를 생성하므로, 전체 파이프라인에서 선택 단계가 차지하는 비중이 미미하다. 또한 하드웨어 친화적인 블록‑스파스 구조를 유지하면서도, 압축된 스코어가 원본 스코어와 높은 순위 상관을 보이기 때문에 정확도 손실이 1% 이하에 머문다.

다중 모달리티 적용 가능성도 주목할 만하다. 텍스트 토큰, 비디오 프레임 패치, 오디오 스펙트로그램 모두 1‑D 혹은 2‑D 시퀀스 형태로 표현될 수 있으며, 평균 풀링은 이러한 데이터에 공통적으로 적용 가능한 연산이다. 실험에서는 Llama‑3.1‑8B‑Instruct, Qwen‑2.5‑7B‑Instruct 등 다양한 LLM과, 이미지‑텍스트 멀티모달 모델에 UniSparse를 적용했을 때, 기존 정적 스파스 방법보다 0.51.2% 높은 정확도와 1.82.6배 빠른 처리 속도를 기록했다.

한계점으로는 압축 비율 선택이 모델·데이터에 따라 민감하게 작용한다는 점이다. 너무 높은 c값은 중요한 미세 정보를 손실시켜 순위 상관이 급격히 떨어질 수 있다. 또한 현재 구현은 평균 풀링 기반이므로, 토큰 간 비선형 관계를 포착하는 데는 한계가 있다. 향후 연구에서는 학습 가능한 압축 모듈이나 어텐션‑전용 클러스터링 기법을 도입해 압축 단계 자체를 데이터에 맞게 최적화할 여지가 있다.

요약하면, UniSparse는 복합 토큰이라는 압축된 표현을 통해 동적 블록‑스파스 마스크를 효율적으로 생성하고, 기존 고비용 프록시 연산을 대체함으로써 LLM의 장기 컨텍스트 처리에서 정확도와 속도 사이의 트레이드오프를 크게 완화한다.