LowRAM, Key-Value 예측 관계로 KV Cache 압축

배경

자기회귀 생성 모델은 디코딩 과정에서 이전 토큰의 Key와 Value를 반복적으로 참조합니다. 이 정보를 KV Cache에 보관하기 때문에 컨텍스트가 길어질수록 캐시 메모리도 선형적으로 증가합니다. 그 결과 긴 문서를 처리할 때 배치 크기와 처리량이 제한되고, 추론 서버의 메모리 비용도 커집니다. 기존의 사후 학습 압축 방법은 토큰 간, 레이어 간, 특징 차원 간 중복을 줄이거나 낮은 수치 정밀도를 사용하는 방식이 많았습니다.

LowRAM은 Key와 Value를 동일하게 압축할 필요가 없다는 관점에서 출발합니다. 논문은 같은 어텐션 헤드 안에서 특정 토큰의 Value가 대응하는 Key를 통해 상당 부분 예측될 수 있다고 관찰했습니다. 연구진은 이 관계를 이용해 테스트 시점에 맞출 수 있는 비대칭 압축 파이프라인을 설계했습니다.

핵심 구성

• 저랭크 Key 압축: 레이어를 가로지르는 상관관계를 저랭크 인수분해로 표현해 저장해야 하는 Key의 양을 줄입니다.
• Value 예측: 헤드마다 경량 예측기를 맞추고 Key를 입력으로 Value를 복원합니다. 예측이 잘 되는 토큰은 원본 Value 전체를 저장하지 않아도 됩니다.
• 어려운 토큰의 잔차 보존: 복원 오차가 큰 토큰에만 보정 잔차를 저장합니다. 따라서 제한된 저장 공간을 예측이 잘 되지 않는 사례에 집중할 수 있습니다. 논문은 거의 손실 없는 성능에서 Value Cache를 4배 압축했다고 보고합니다.
• 테스트 시점 적응: 예측기는 메타러닝으로 초기화해 prefill 단계의 추가 비용을 줄일 수 있습니다. 멀티쿼리 추론처럼 일회성 비용을 허용하는 환경에서는 처음부터 학습하는 방식도 가능합니다.
• 선택적 복원: 디코딩 때 모든 KV 쌍을 완전히 복원하지 않고 선택된 일부만 재구성합니다. 4배 KV 압축과 262K 컨텍스트 조건에서 피크 디코딩 처리량이 최대 5배 높아졌다고 보고되었습니다.

의미와 고려할 점

LowRAM의 핵심은 새로운 압축 비율 자체보다 Key와 Value의 구조적 차이를 별도의 압축 축으로 활용했다는 데 있습니다. 이 방식은 양자화나 캐시 제거와 다른 종류의 중복을 다루므로 함께 적용할 수 있습니다. 논문에서 세 기법을 결합한 결과는 26.9배 압축과 기준 성능의 94.5% 유지였습니다.

긴 문서 질의응답이나 동일한 컨텍스트를 공유하면서 여러 질의를 처리하는 서비스에서는 적응된 예측기를 재사용할 수 있어 특히 유용할 수 있습니다. 다만 예측기 적합, 잔차 선택, 선택적 복원에는 추가 계산과 구현 복잡성이 따릅니다. 실제 이득은 하드웨어와 컨텍스트 길이, 요청 패턴에 따라 달라질 수 있습니다. 그럼에도 LowRAM은 KV Cache를 단순히 저정밀화하거나 삭제하는 데서 나아가, 모델 내부 표현 사이의 예측 가능성을 활용하는 방향을 보여줍니다.

출처: Hugging Face Daily Papers