OpenAttention, 1%만 읽고 정확도 99% 유지하는 희소 어텐션 제안

희소 어텐션에서 읽지 않고 남겨 둔 꼬리 부분을 대리 토큰 하나로 보정하는 기법 OpenAttention이 논문으로 공개됐다. 연구진은 RULER-32K HARD 벤치마크에서 Llama-3.1-8B이 키-값 캐시의 1%만 읽고도 밀집 어텐션 정확도의 99%를 유지했다고 밝혔다.

연구진이 직접 만든 어텐션 커널은 1%만 읽는 조건에서 디코드 어텐션 단계를 밀집 어텐션보다 10배 빠르게 실행했다고 한다. 논문은 모델, 문맥 길이, 데이터셋 전반에서 최고 수준(state-of-the-art)의 정확도를 달성했다고도 밝혔다.

읽지 않은 꼬리를 어떻게 다루나

희소 어텐션은 키-값 캐시 중 일부만 읽는다. 읽지 않은 꼬리를 버리면 출력에 편향이 생긴다. 기존 방식처럼 꼬리에서 표본을 뽑는 방법은 읽는 비율이 5% 이하일 때 실패하는데, 이 구간에서는 선택기가 놓친 핵심 토큰(heavy hitter)이 꼬리에 있을 가능성이 높다고 논문은 설명한다.

연구진은 꼬리 전체를 개수를 반영한 평균 하나로 접은 대리 토큰이, 실용적인 선택기를 쓸 때 모든 읽기 비율에서 표본 추출보다 출력 오차가 낮다는 점을 관찰했다. 또 읽는 비율이 낮으면 표본 추출의 가중치 오차가 꼬리를 버리는 방식의 거의 두 배가 되는 반면, 대리 토큰 하나의 가중치 오차는 꼬리를 버리는 방식보다 낮다는 것을 수식으로 보였다고 한다.

파이프라인 구성

OpenAttention은 군집화, ErrorGate, Top-k 선택을 거친 뒤 개수를 반영한 평균(count-aware mean) 대리 토큰을 더하는 순서로 이뤄진다. 군집화와 선택에는 추가 학습이 필요 없는 Query-Geometry Clustering(QGC)과 Query-Geometry Selector(QGS)를 설계했고, 질의 방향을 고려해 키를 압축한다.

기존 방법이 군집 중심값만으로 순위를 매기는 것과 달리, ErrorGate는 군집별로 추정한 가중치 오차에 따라 실제로 읽을 양을 배분한다.