하이스파스 오라클 토큰 선택과 KV 캐시 공유를 통한 하이브리드 희소 어텐션

초록

HySparse는 전체 어텐션 레이어와 희소 어텐션 레이어를 교차 배치하고, 앞선 전체 어텐션에서 얻은 토큰 중요도와 KV 캐시를 그대로 활용해 토큰 선택과 연산을 최적화한다. 이를 통해 대규모 모델에서 메모리와 연산 비용을 크게 절감하면서도 성능을 향상시킨다.

상세 분석

HySparse는 기존 희소 어텐션 방식이 갖는 두 가지 근본적인 한계를 동시에 해결한다. 첫 번째는 토큰 중요도를 예측하기 위해 별도의 프록시 모델이나 추가 파라미터를 도입하는 복잡성이다. HySparse는 바로 직전의 전체 어텐션 레이어에서 생성된 Q·K·V 행렬을 이용해 각 토큰의 어텐션 스코어를 계산하고, 이를 토대로 토큰의 중요도를 정량화한다. 즉, 전체 어텐션이 제공하는 ‘오라클’ 역할을 활용해 토큰 선택을 수행함으로써 별도 예측 모듈이 필요 없으며, 선택 과정이 정확하고 일관된다. 두 번째는 희소 어텐션이 연산량은 줄이지만 KV 캐시를 재사용하지 못해 메모리 사용량이 크게 감소하지 못한다는 점이다. HySparse는 전체 어텐션 레이어에서 이미 저장된 KV 캐시를 그대로 재활용한다. 구체적으로, 전체 어텐션이 생성한 K와 V를 블록 단위로 분할하고, 각 희소 어텐션 레이어는 동일한 블록 인덱스를 사용해 필요한 부분만 선택한다. 이렇게 하면 희소 어텐션 단계에서 추가적인 KV 저장이 필요 없으며, 전체 메모리 사용량이 10배 가까이 절감된다.

아키텍처는 ‘Full‑Attention → Sparse‑Attention → Full‑Attention → …’ 형태로 레이어를 교차 배치한다. 전체 어텐션 레이어는 기존 트랜스포머와 동일하게 Q·K·V를 계산하고, 토큰별 어텐션 스코어 행렬 S를 구한다. 이후 S를 블록 단위로 집계해 Top‑K 토큰을 선택하고, 선택된 토큰에 해당하는 KV 인덱스를 저장한다. 선택된 토큰은 다음 희소 어텐션 레이어에서만 사용되며, 이때는 기존 KV 캐시를 그대로 읽어와 연산한다. 희소 어텐션 자체는 블록‑스파스 어텐션 방식을 채택해, 각 블록 내에서만 완전 어텐션을 수행하고 블록 간에는 토큰 선택을 통해 연산을 제한한다.

실험에서는 7B 파라미터를 가진 Dense 모델과 80B 파라미터 MoE 모델에 HySparse를 적용하였다. 80B MoE 모델에서는 전체 49 레이어 중 5개만 전체 어텐션을 사용했음에도 불구하고, 기존 전체 어텐션 대비 퍼플렉시티와 정확도에서 모두 우수한 결과를 얻었다. 특히 KV 캐시 저장량이 거의 10배 감소하면서도 추론 속도는 크게 저하되지 않았다. 이는 대규모 언어 모델에서 메모리 병목을 해소하고, 비용 효율적인 서비스 배포가 가능함을 의미한다.

전체적으로 HySparse는 ‘오라클 기반 토큰 선택’과 ‘KV 캐시 공유’라는 두 축을 통해 희소 어텐션의 효율성을 극대화하고, 기존 전체 어텐션과 비교해 연산·메모리·성능 측면에서 모두 실질적인 이점을 제공한다.