ThriftAttention keeps 90% quality with 5% compute

개요

ThriftAttention은 쿼리-키 행렬의 5%만 FP16으로 연산하고 나머지는 FP4로 유지하는 혼합 정밀도 기법으로, FP4 연산 대비 약 90%의 품질을 유지하면서 컴퓨팅 비용을 크게 절감합니다.

주요 내용

* 핵심 아이디어: 기존 저비트 어텐션 방식이 전체 어텐션 연산을 FP4로 양자화하여 품질 저하를 초래하는 반면, ThriftAttention은 FP4 양자화 시 품질 저하를 유발하는 소수의 영향력 있는 쿼리-키 쌍에 집중하여 이들만 FP16으로 연산합니다.
* 성능: ThriftAttention은 FP4와 FP16 어텐션 간 성능 격차의 평균 89.1%를 회복시키며, FP16 컴퓨팅을 5%만 사용하고도 FP16 품질의 94.2%를 달성합니다 (LongBench 데이터셋 기준).
* 효율성: FP16 연산을 5%로 제한함으로써 모델의 표현력을 유지하면서도 FP4의 대역폭 감소 이점을 활용하며, 메모리 사용량도 비례적으로 줄어듭니다.
* 적용 가능성: 긴 컨텍스트 처리가 필요한 검색 증강 생성(Retrieval-Augmented Generation), 장문 요약, 코드 완성 등에서 FP4 방식의 품질 저하 문제를 해결할 수 있습니다.
* 한계점: 선택된 쿼리-키 블록이 드물지만 결정적인 토큰 상호작용을 놓칠 수 있으며, 멀티 GPU 설정 시 동기화 오버헤드가 발생할 수 있습니다. 또한, 블록 선택 단계가 병목이 될 경우 성능 향상이 줄어들 수 있습니다.

시사점

ThriftAttention은 긴 컨텍스트 LLM에서 컴퓨팅 효율성을 크게 높이면서도 품질 저하를 최소화할 수 있는 실용적인 방안을 제시하며, 기존 FP16 기반의 고비용 연산이 필요했던 장문 처리 작업의 생산성 향상 가능성을 보여줍니다.