핵심 토큰을 통한 단순화된 희소 어텐션

초록

희소 어텐션은 긴 문맥 추론의 비용을 줄일 수 있지만, 대부분의 변형은 새로운 아키텍처 구성 요소를 도입합니다. 우리는 아키텍처 변경 없이 희소 어텐션을 구현하는 더 간단한 접근 방식인 간소화된 희소 어텐션(SSA)을 소개합니다. 구체적으로, 먼저 요지 토큰(gist tokens)이 삽입된 시퀀스에 대해 지속적 사전 학습을 수행합니다. 표준적인 다음 토큰 손실을 평소와 같이 최적화하지만, 요지 토큰은 어텐션 마스크를 사용하여 언어 모델이 문맥의 어떤 부분에 주목할 수 있는지 제한합니다; 이는 모델이 각 청크의 중요한 정보를 요지 토큰에 압축하도록 학습시킵니다. 추론 시, SSA는 현재 쿼리와 작은 요지 토큰 집합 간의 어텐션을 통해 청크를 점수화하고, 상위 k개 청크를 선택적으로 펼쳐 해당 원시 토큰을 다시 도입합니다. 쿼리가 요지 토큰에 대해서만 점수화되므로, 전체 KV 캐시에 대한 단순 점수화와 관련된 메모리 대역폭 비용을 피할 수 있으며, 희소 어텐션 방법에 사용되는 보조 KV 캐시 접근 방식이 필요하지 않습니다. LongBench에서 SSA는 동일한 압축 비율에서 압축 및 추론 시간 희소 어텐션 기준선을 일관되게 능가합니다. 더욱 놀랍게도, 검색 증강 생성에서 SSA는 지속적 사전 학습 후 전체 어텐션보다 5.7포인트 이상 뛰어난 성능을 보일 수 있습니다. 우리는 이를 SSA의 선택적 펼침 능력 덕분이라고 생각합니다. 이는 쿼리와 관련된 청크에 어텐션을 집중시키고 노이즈를 효과적으로 걸러냅니다. SSA는 계층적 요지-오브-요지 변형(H-SSA)으로 확장되며, 최대 32배의 높은 압축 비율에서 정확도를 유지하거나 개선하면서 로그-선형 디코딩 복잡도를 달성합니다. 코드는 https://github.com/yuzhenmao/simplified-sparse-attention/에서 확인할 수 있습니다.

English

Sparse attention can reduce the cost of long-context inference, but most variants introduce new architectural components. We introduce Simplified Sparse Attention (SSA), a simpler approach to sparse attention that requires no architectural changes. Concretely, we first perform continued pretraining on sequences interleaved with gist tokens. We optimize the standard next-token loss as usual, but the gist tokens use an attention mask to restrict what parts of the context the language model can attend to; this teaches the model to pack each chunk's important information into the gist tokens. At inference time, SSA scores chunks via attention between the current query and the small set of gist tokens, selectively unfolding the top-k chunks by reintroducing their corresponding raw tokens. Since the query is scored only against the gist tokens, we avoid the memory-bandwidth cost associated with naive scoring against the full KV cache, without requiring the auxiliary KV cache approach used by sparse attention methods. On LongBench, SSA consistently outperforms compression and inference-time sparse-attention baselines under the same compression ratio. More strikingly, in retrieval-augmented generation, SSA can even outperform full attention after continued pretraining by over 5.7 points. We attribute this to the ability of SSA's selective unfolding, which concentrates attention on the query-relevant chunks and effectively filters out noise. SSA further extends to a hierarchical gist-of-gist variant (H-SSA) that achieves log-linear decoding complexity while maintaining or improving accuracy at high compression ratios up to 32x. The code is available at https://github.com/yuzhenmao/simplified-sparse-attention/.