PN-Mix, 쿼리-키 페어 노름으로 저비트 KV 캐시 양자화 개선

긴 출력 추론에서 커지는 KV 캐시 부담

대규모 언어 모델이 긴 답변을 생성하면 과거 토큰의 키와 밸류를 저장하는 KV 캐시가 계속 커집니다. 디코딩 단계에서는 이 캐시를 반복해서 읽어야 하므로 GPU 메모리 용량뿐 아니라 메모리 대역폭도 병목이 됩니다. 양자화는 저장 공간과 전송량을 줄일 수 있지만, 2비트와 같은 낮은 정밀도를 모든 값에 적용하면 어텐션 계산 오차가 커져 모델 정확도가 떨어질 수 있습니다.

키 자체가 아니라 어텐션 관계를 보는 방식

기존의 혼합 정밀도 키 캐시 양자화 방법은 주로 키 캐시의 통계량을 이용해 중요한 채널을 찾고, 해당 채널에 더 많은 비트를 배정합니다. PN-Mix는 중요도 판단을 실제 어텐션 계산에 더 가깝게 가져갑니다. RoPE는 채널을 쌍으로 다루기 때문에, 이 연구는 RoPE 채널 쌍의 양자화 오차가 어텐션 로짓을 어떻게 흔드는지 분석했습니다.

논문의 핵심 관찰은 쿼리와 대응하는 키의 노름 곱이 큰 채널 쌍일수록 양자화 오차가 로짓 변화에 크게 기여할 수 있다는 점입니다. PN-Mix는 이 페어 노름을 중요도 신호로 사용해 해당 채널 쌍에 더 높은 정밀도를 할당합니다. 모든 채널을 동일하게 압축하거나 키 측 통계만으로 순위를 정하는 대신, 현재 입력에서 어텐션에 중요한 상호작용에 정밀도 예산을 집중하는 방식입니다.

중요도와 비트 배분은 프리필 단계에서 계산됩니다. 이후 생성 과정에서는 이 결과를 재사용하므로, 새 토큰이 생성될 때마다 중요도를 다시 추정할 필요가 없습니다. 긴 출력 서비스에서 반복 계산을 줄이면서 입력 적응형 양자화를 구현하려는 설계입니다.

보고된 실험 결과

표준 추론, 코드 생성, AIME 벤치마크에서 PN-Mix는 평균 2.26비트 설정으로, 평균 2.32~2.39비트를 사용하는 기존 혼합 정밀도 방법과 비슷하거나 더 나은 평균 정확도를 보였다고 보고됩니다. Pareto 분석에서는 평균 정밀도를 1.5비트까지 낮춰도 성능 저하가 상대적으로 완만했습니다. 반면 비교 대상 방법은 2.25비트 아래에서 성능이 더 급격히 하락했습니다.

이는 KV 캐시 압축에서 비트 수만큼이나 비트를 어디에 배정하는지가 중요하다는 점을 보여줍니다. 저장된 키의 크기나 분포만 보는 것보다, 실제 어텐션 로짓을 만드는 쿼리-키 상호작용을 활용하는 편이 저비트 구간의 성능을 더 잘 보존할 수 있다는 의미입니다.

의미와 한계

PN-Mix는 긴 컨텍스트와 긴 출력 추론에서 캐시 저장량 및 데이터 이동을 줄일 수 있는 방향을 제시합니다. 특히 RoPE 기반 decoder-only LLM에서는 채널 단위가 아니라 채널 쌍 단위로 중요도를 분석함으로써 모델의 계산 구조에 맞춘 혼합 정밀도 배분을 시도합니다.

다만 제시된 연구의 적용 범위는 RoPE 기반 decoder-only LLM과 키 캐시에 맞춰져 있습니다. 다른 위치 인코딩, 다른 어텐션 구조, 키와 밸류를 함께 양자화하는 환경에서도 같은 효과가 나타날지는 추가 검증이 필요합니다. 그럼에도 어텐션을 실제로 좌우하는 상호작용을 기준으로 양자화 예산을 배분한다는 원칙은 장문 추론 시스템 설계에 중요한 시사점을 제공합니다.

출처: Hugging Face Daily Papers