장문맥 처리의 병목 풀다…로그 선형 복잡도 블록 스파스 어텐션
원제 Block Sparse Attention with Log-Linear Complexity
추천 14댓글 1
Key Point
장문맥 처리가 핵심인 LLM 개발에서 어텐션 계산량을 획기적으로 줄이는 알고리즘이 나왔기 때문에, 모델 확장성과 추론 속도에 영향을 받는 엔지니어와 연구자들이 눈여겨봐야 한다.
핵심 요약
- 언어모델을 긴 문맥에 맞추려면 셀프 어텐션의 이차 비용이 제약이 된다.
- 블록 스파스 어텐션은 효율적 대안이지만, 유지할 블록을 선택하는 과정이 여전히 이차 복잡도다.
- 논문은 PISA라는 피라미드 Top-K 선택 전략을 제안해 후보를 단계적으로 좁혀나간다.
- 거친 수준부터 시작해 각 단계에서 LogSumExp 스코링으로 다음 단계 후보를 선택한다.
- 풀링으로 O(log N)개 수준의 키를 구성해 전체 복잡도를 O(N log N)으로 줄인다.
- Triton 커널로 하드웨어 최적화한 학습·추론 구현을 개발했으며, 쿼리-키 스코어 행렬을 메모리에 구성하지 않는다.
- 상식 추론 벤치에서 기존 방식과 비슷한 성능을 보이면서 검색 작업에서는 더 나은 결과를 얻었다.