ThriftAttention: 5%의 연산량으로 90%의 품질 유지

쿼리-키 행렬 (query-key matrix)의 단 5%만을 전정밀도 (full precision)로 계산함으로써, 4비트 어텐션 (4-bit attention)에서 손실된 정확도의 거의 90%를 복구할 수 있습니다 [1]. 이 방법은 출력 분포를 지배하는 아주 작은 어텐션 블록 (attention blocks) 서브셋을 분리하여 이를 FP16으로 업그레이드하고, 나머지는 FP4로 남겨둡니다. 소프트맥스 (soft-max)가 공동으로 평가되기 때문에, 저정밀도 꼬리 (low-precision tail) 부분도 고정밀도 헤드 (high-precision head)의 이점을 얻게 되어, 아주 적은 연산 비용으로 전정밀도에 가까운 점수를 산출합니다.

기존의 저비트 어텐션 (low-bit attention) 파이프라인은 전체 어텐션 연산을 FP4로 양자화 (quantise)합니다. 이는 Blackwell급 GPU에서 추론 (inference) 속도를 높여주는 기술이지만, 시퀀스 (sequences)가 수천 토큰을 넘어서면 품질이 급격히 저하됩니다 [1]. 이러한 접근 방식은 모든 쿼리-키 (query-key) 쌍을 동일하게 중요하다고 취급하며, 양자화 오차 (quantisation error)가 소수의 영향력 있는 상호작용에 집중된다는 경험적 사실을 무시합니다. 그 결과 발생하는 성능 저하로 인해 FP4는 검색 증강 생성 (retrieval-augmented generation), 긴 문서 요약 (long-document summarisation), 대규모 코드 완성 (code-completion) 분야에서 매력적이지 않은 선택지가 되었습니다.

ThriftAttention의 휴리스틱 (heuristic)은 가장 영향력 있는 쿼리-키 블록을 선택하여 FP16으로 계산하며, 이를 통해 FP4에서 FP16으로 전환했을 때의 성능 격차 중 평균 89.1%를 회복합니다 [1]. “ThriftAttention은 FP4 추론의 효율성 이점을 유지하면서도 FP4와 FP16 어텐션 사이의 품질 격차 대부분을 복구합니다.” 블록 행렬의 5%로 전정밀도 작업을 제한함으로써, 시스템은 FP4의 감소된 대역폭 (bandwidth) 이점을 누리면서도 모델의 표현력 (representational power) 대부분을 보존합니다.

LongBench 벤치마크 세트에서 혼합 정밀도 (mixed-precision) 모델은 순수 FP16 모델의 점수인 0.469에 비해 0.452를 기록했습니다. 즉, FP16 연산량을 단 5%만 사용하면서도 베이스라인 품질의 94.2%를 유지한 것입니다 [1]. README 보고에 따르면 동일한 수준의 품질 회복이 확인되었으며, 이는 아주 적은 양의 고정밀도 예산만으로도 대부분의 다운스트림 지표 (downstream metrics)를 전체 정밀도 오라클 (full-precision oracle)의 좁은 오차 범위 내로 가져오기에 충분함을 입증합니다. 메모리 사용량 (memory footprints)은 비례하여 줄어들며, 추론 지연 시간 (inference latency) 또한 전체 FP16 대비 감소하지만, 논문에서 8k 토큰 프롬프트에 대한 정확한 감소 배수를 명시하지는 않았습니다.

이 접근 방식은 여전히 드물지만 결정적인 토큰 간 상호작용을 간과할 수 있는 휴리스틱 (heuristic)에 의존하고 있으며, 논문은 엔드 투 엔드 (end-to-end) 학습 파이프라인 대신 기존 LLM 제품군에 대한 추론만을 평가했습니다. 선택기 (selector)를 멀티 GPU 설정으로 확장할 경우 동기화 오버헤드 (synchronization overhead)가 발생할 수 있으며, 블록 선택 (block-selection) 단계가 병목 현상 (bottleneck)이 될 경우 보고된 이점은 감소할 수 있습니다. 결과적으로, 5%의 FP16만으로 충분하다는 주장은 워크로드의 토큰 분포와 모델의 어텐션 패턴 (attention patterns)에 따라 달라질 수 있습니다.

롱 컨텍스트 (Long-context) 서비스는 기존의 일반적인 어텐션 커널 (vanilla attention kernel)을 ThriftAttention으로 교체할 수 있으며, 이는 감소된 FP16 연산량에 비례하여 메모리 사용량을 줄여줍니다. 다만, 논문에서 20배의 대역폭 (bandwidth) 감소를 정량화하여 제공하거나 10%의 품질 손실 상한선을 구체적으로 명시하지는 않았습니다. 32k 토큰 문서 요약과 같이 기존에 전체 FP16이 필요했던 벤치마크들은, 속도 향상이 실제 프로덕션 규모의 처리량 (throughput)으로 이어지는지 확인하기 위해 혼합 정밀도 드롭인 (mixed-precision drop-in) 방식으로 재실행되어야 합니다.

참고 문헌

1. ThriftAttention: Selective Mixed Precision for Long-Context FP4 Attention