Full Attention의 귀환: 단 100번의 학습 단계로 Full Attention을 Sparse로 전환하기
Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps
핵심 요약
Full Attention 모델을 최소한의 학습으로 Sparse 모델로 변환하여 성능 저하 없이 추론 속도를 획기적으로 높이는 RTPurbo 기술 소개.
- RTPurbo 기술 — Full Attention의 KV 캐시를 효율적으로 관리하여 추론 속도를 대폭 개선함.
- Sparse 전환 — 단 몇 백 번의 학습 단계만으로 모델을 Sparse하게 최적화함.
- 성능 향상 — 1M 컨텍스트에서 최대 9.36배의 프리필 속도와 2배 이상의 디코드 속도 향상을 달성함.
- 효율적 인덱싱 — 16차원 인덱서를 사용하여 중요한 토큰만 효율적으로 검색함.
대규모 언어 모델(LLM)의 긴 컨텍스트 추론은 Full Attention의 이차 비용(quadratic cost)으로 인해 병목 현상을 겪습니다. 기존의 효율적인 대안들은 네이티브 Sparse 학습이나 휴리스틱한 토큰 제거 방식에 의존하는 경우가 많아, 효율성, 학습 비용, 정확도 사이에서 바람직하지 않은 트레이드오프를 발생시킵니다.
본 연구에서는 Full-attention LLM이 이미 본질적으로 Sparse하며, 최소한의 적응만으로도 고도로 Sparse한 모델로 변환될 수 있음을 보여줍니다. 우리의 접근 방식은 다음 세 가지 관찰에 기반합니다:
(1) 전체 어텐션 헤드 중 일부만이 긴 컨텍스트 처리를 위해 Full Attention을 필요로 함
(2) 장거리 검색은 주로 저차원 부분 공간에 의해 결정되므로, 16차원 인덱서를 사용하여 관련 토큰을 효율적으로 검색할 수 있음
(3) 유용한 토큰 예산은 쿼리에 따라 크게 달라지므로, 고정된 top-p 희소화보다 동적인 top-p 선택이 더 적합함
이러한 통찰을 바탕으로, 우리는 검색 헤드에 대해서만 전체 KV 캐시를 유지하고 Sparse Attention을 위한 경량 토큰 인덱서를 도입한 RTPurbo를 제안합니다. 모델의 본질적인 희소성을 활용함으로써, RTPurbo는 단 몇 백 번의 학습 단계만으로 희소화를 달성합니다. 긴 컨텍스트 벤치마크 및 추론 작업에 대한 실험 결과, RTPurbo는 거의 손실 없는 정확도를 유지하면서도 상당한 효율성 향상을 제공하며, 1M 컨텍스트에서 최대 9.36배의 프리필 속도 향상과 약 2.01배의 디코드 속도 향상을 달성했습니다. 이러한 결과는 비용이 많이 드는 네이티브 Sparse 사전 학습 없이도 표준 Full-attention 학습으로부터 강력한 Sparse 추론을 얻을 수 있음을 시사합니다.
