CoWindow: 주의 메커니즘의 중복 계산을 줄인 구조화된 어텐션
원제 CoWindow Attention: Full Causal Coverage Is a Collective Property
추천 59댓글 1
Key Point
효율적인 장문맥 처리가 필수인 생성형 AI에서 Full Attention의 중복 계산을 제거하면서도 성능을 유지하는 방법이 제시되었으므로, 스케일링 효율성을 개선하려는 팀들의 관심을 받을 만합니다.
핵심 요약
- 기존 Full Attention 방식은 모든 어텐션 헤드가 전체 인과 히스토리를 반복 처리하면서 중복 계산과 메모리 오버헤드가 발생한다.
- 연구팀이 제안한 CoWindow Attention(CoWA)은 KV 헤드 간에 인과 히스토리 접근을 분산시키는 구조화된 어텐션이다.
- 모든 헤드는 근거리 대각 윈도우와 접두사-싱크 윈도우를 공유하며, 나머지 먼 과거는 보완적 장거리 윈도우로 분할되어 전체 인과 커버리지를 달성한다.
- 학습된 라우터나 인덱서 없이 위치 기반으로 정의되어 학습·추론에서 일관되며 KV-헤드 텐서 병렬화와 정렬된다.
- 8K 토큰 환경에서 100% 커버리지를 유지한 CoWA는 89.73% 정확도로 FullAttn의 89.97%에 가까우며, 중복된 장거리 윈도우는 훨씬 낮은 성능을 보인다.
- 128K 토큰 벤치마크에서 CoWA는 훈련 시 정방향 지연시간 7.4배, 역방향 7.4배 감소, 추론 디코딩 지연시간 3.0배 감소를 달성한다.
- 훈련 중 피크 메모리는 FullAttn과 동일하지만 디코딩 중에는 7.6배 더 낮다.
- 0.6B~14B 파라미터 범위의 스케일링 법칙 훈련에서 CoWA는 FullAttn의 퍼플렉시티를 유지하면서 총 훈련 연산량을 줄인다.
- 14B와 32B 모델이 FullAttn과 유사한 지식, 추론, 장문맥 검색 성능을 보인다.
- 결과는 전체 인과 커버리지가 모든 헤드의 중복 속성이 아니라 헤드 앙상블의 집단적 속성이 될 수 있음을 보여준다.