핵심 개념
장문 컨텍스트 언어 모델의 학습 및 추론 비용을 줄이기 위해 부분 컨텍스트를 활용하는 효율적인 모델 아키텍처를 제안한다.
초록
이 논문은 장문 컨텍스트 언어 모델의 학습 및 추론 비용을 줄이기 위한 방법을 제안한다.
-
모델 사전 학습 후 별도의 장문 컨텍스트 학습 단계를 거치는 기존 접근법의 문제점을 지적한다. 이는 학습 비용이 많이 들고 추론 시 메모리 사용량이 증가하는 단점이 있다.
-
이를 해결하기 위해 LONGGEN이라는 하이브리드 모델 아키텍처를 제안한다. LONGGEN은 모델의 하단과 상단 층에 GPU 친화적인 희소 어텐션 메커니즘을 적용하고, 중간 층에는 완전 어텐션을 유지한다. 이를 통해 학습 및 추론 효율성을 높이면서도 장문 컨텍스트 성능을 유지할 수 있다.
-
다양한 희소 어텐션 패턴을 실험하여 최적의 설계를 찾아내었다. 특히 중간 층의 완전 어텐션이 복잡한 장문 추론 작업에 필수적임을 밝혀냈다.
-
LONGGEN은 사전 학습된 모델을 5B 토큰의 장문 데이터로 추가 학습하여 컨텍스트 길이를 4K에서 128K로 확장할 수 있다. 이 과정에서 학습 시간을 36% 단축하고 추론 시 KV 캐시 메모리를 62% 줄일 수 있다.
-
장문 컨텍스트 벤치마크에서 LONGGEN은 완전 어텐션 모델과 유사한 성능을 보이면서도 훨씬 효율적인 것으로 나타났다.
통계
장문 컨텍스트 학습 시 LONGGEN은 완전 어텐션 모델 대비 36% 빠른 학습 속도를 보인다.
LONGGEN은 KV 캐시 메모리를 62% 줄일 수 있어, 추론 시 1.67배 빠른 prefilling 속도와 1.41배 빠른 decoding 속도를 달성한다.
인용구
"LONGGEN builds on three key insights: (1) Sparse attention patterns, such as window attention (attending to recent tokens), attention sink (initial ones), and blockwise sparse attention (strided token blocks) are well-suited for building efficient long-context models, primarily due to their GPU-friendly memory access patterns, enabling efficiency gains not just theoretically but in practice as well. (2) It is essential for the model to have direct access to all tokens. A hybrid architecture with 1/3 full attention layers and 2/3 efficient ones achieves a balanced trade-off between efficiency and long-context performance. (3) Lightweight training on 5B long-context data is sufficient to extend the hybrid model's context length from 4K to 128K."