Xinghao Wang Pengyu Wang Dong Zhang Chenkun Tan Shaojun Zhou Zhaoxiang Liu Shiguo Lian Fangxu Liu Kai Song Xipeng Qiu

초록

대규모 언어 모델(Large Language Models, LLMs)의 컨텍스트 길이를 확장하는 것은 큰 이점을 제공하지만, 계산 비용이 매우 높다. 이 비용의 주요 원인은 시퀀스 길이에 대해 O(N²)의 복잡도를 가지는 자기주의(self-attention) 메커니즘으로, 메모리 사용과 지연 시간 측면에서 주요 성능 저하 요인이다. 다행히도, 주로 긴 시퀀스에서 주목할 만큼 주의 행렬(attention matrix)은 희소(sparse)한 경향이 있으며, 이는 최적화의 기회를 시사한다. 블록-희소 주의(Block-sparse attention)는 시퀀스를 블록 단위로 분할하고 일부 블록에 대한 계산을 건너뛰는 방식으로 등장한 유망한 해결책이다. 그러나 이 방법의 효율성은 기반이 되는 주의 패턴에 크게 의존하며, 이로 인해 블록 수준의 희소성은 최적화되지 않을 수 있다. 예를 들어, 한 블록 내의 쿼리에 중요한 키 토큰(key tokens)이 수많은 다른 블록에 분산되어 있을 경우, 계산의 중복이 발생하게 된다. 본 연구에서는 주의 메커니즘의 순열(permutation) 성질을 활용하여 블록 수준의 희소성을 증가시키고, LLM의 프리필링(prefilling) 단계에서의 계산 효율성을 향상시키는 즉시 적용 가능한 방법인 순열 기반 블록-희소 주의(Permuted Block-Sparse Attention, PBS-Attn)를 제안한다. 우리는 실제 세계에서 가장 도전적인 긴 컨텍스트 데이터셋을 대상으로 포괄적인 실험을 수행하였으며, PBS-Attn이 기존의 블록-희소 주의 방법보다 모델 정확도에서 일관되게 우수함을 입증하였다. 또한, 전체 주의(full attention) 기준선과 거의 동등한 성능을 달성하였다. 자체 개발한 순열 기반 FlashAttention 커널을 기반으로, PBS-Attn은 긴 컨텍스트 프리필링에서 엔드투엔드 성능을 최대 2.75배까지 향상시켰으며, 실용적 타당성을 입증하였다. 코드는 다음 링크에서 확인 가능하다: https://github.com/xinghaow99/pbs-attn