기술 요약: ProxyKV

문제 정의

대규모 언어 모델 (LLM) 에서의 효율적인 장문맥 추론은 현재 Key-Value(KV) 캐시의 메모리 요구 사항에 의해 병목 현상을 겪고 있습니다. 시퀀스 길이 ($N$) 가 증가함에 따라 KV 캐시의 선형적 성장은 종종 GPU 메모리 용량을 초과하여, 막대한 I/O 오버헤드와 추론 지연을 초래합니다.

기존 KV 캐시 가지치기 (pruning) 방법들은 점수 산출 효율성과 가지치기 정밀도 사이의 근본적인 트레이드오프에 직면해 있습니다:

• 휴리스틱 방법 (예: H2O, SnapKV): 빠르고 지연 시간이 짧지만, 지역적 관찰 창에 의존하여 복잡한 추론에 필요한 전역적 의미 의존성을 종종 놓칩니다.
• 재구성 기반 방법 (예: KVZip): 타겟 모델에서 2 차 교사가 지시하는 프리필링 (prefilling) 패스를 수행하여 전역 문맥을 재구성함으로써 높은 정밀도를 달성합니다. 그러나 이는 막대한 프리필링 오버헤드를 초래하여 가지치기의 지연 시간 이점을 무효화합니다.

핵심적인 과제는 타겟 모델에서 2 차 프리필링의 무거운 계산 비용을 치르지 않고도, 저지연 휴리스틱과 고정밀 재구성 사이의 간극을 메우는 것입니다.

방법론

저자들은 타겟 모델에서 중요도 점수 산출 작업을 경량 비동기 "소규모 모델 프록시 (Small-Model Proxy)"로 이전하도록 설계된 교차 모델 프록시 가지치기 프레임워크인 ProxyKV를 제안합니다.

시스템 아키텍처

이 시스템은 세 가지 구성 요소로 이루어진 파이프라인을 통해 작동합니다:

1. 대규모 모델 타겟 ($M_l$): 프리필 임계 경로와 전체 디코드 단계를 담당하는 주 모델입니다.
2. 소규모 모델 프록시 ($M_s$): 별도의 실행 스트림 (또는 CUDA 스트림) 에서 비동기적으로 실행되는 동족 내 (intra-family) 의 더 작은 모델 (예: 15 억 또는 40 억 파라미터) 입니다. 이는 입력 문맥에서 크로스-헤드 어텐션 특징 ($X$) 을 추출합니다.
3. HybridAxialMapper ($\Phi_\theta$): 프록시의 특징 ($X$) 을 타겟 정렬 중요도 점수 ($\hat{Y}$) 로 매핑하는 학습 가능한 모듈입니다.

작업 흐름:
추론 중 입력 문맥은 두 모델로 동시에 배분됩니다. 프록시는 어텐션 특징을 생성하며, 이는 HybridAxialMapper 를 통해 처리되어 Top-K 임계값 설정을 통해 이진 가지치기 마스크 ($M$) 를 생성합니다. 그런 다음 타겟 모델은 가지치기된 KV 캐시만을 사용하여 디코딩을 수행합니다. 중요한 점은, 프록시의 KV 캐시가 프리필 단계 직후에 즉시 해제되어 메모리 오버헤드가 일시적이라는 것입니다.

주요 구성 요소

1. HybridAxialMapper 아키텍처

이질적인 모델들 (서로 다른 헤드 수와 레이어 깊이) 간의 아키텍처 간극을 메우기 위해, 매핑기는 시간적 특징 추출을 크로스-헤드 정렬과 분리하는 3 단계 파이프라인을 사용합니다:

• 1 단계: 시간적 특징 추출: 1D 컨볼루션 스템 (Stem) 이 원시 프록시 특징을 모델 중립적인 잠재 공간으로 투영하여 헤드 구성을 정규화합니다.
• 2 단계: 시간 축 문맥 인코더: Transformer 인코더가 사인파 위치 인코딩을 사용하여 시퀀스를 처리함으로써 전역 문맥을 포착하며, 긴 시퀀스 (최대 17 만 토큰) 를 위해 슬라이딩 윈도우 방식을 활용합니다.
• 3 단계: 헤드 축 크로스-어텐션: 학습 가능한 타겟 쿼리가 프록시의 합성 헤드 토큰을 타겟의 특정 헤드 구성으로 집계합니다. 이 단계는 전체 네트워크를 재학습하지 않고도 임의의 타겟 헤드 수로 매핑기를 확장할 수 있게 합니다.

2. 다중 그레인 Hybrid Loss

배포 지표가 점wise 점수 재구성이 아닌 Top-K 유지 품질임을 인식하여, 저자들은 다섯 가지 항으로 구성된 복합 손실 함수 ($L_{total}$) 를 설계합니다:

• $L_{bin}$ (다중 비율 이진): 여러 유지 비율에 걸쳐 이진 결정 경계를 감독하며, 공격적인 압축을 강조합니다.
• $L_{mse}$ (값 가중 MSE): 중요한 토큰에 대한 수치적 크기 정렬을 보장합니다.
• $L_{fine}$ (내부 Top-K 순위): 유지된 Top-K 세트 내의 상대적 순서를 강제합니다.
• $L_{global}$ (질량 인식 전역 순위): Top-K 토큰과 비 Top-K 토큰 간의 분리 (신호 대 잡음 경계) 를 최적화합니다.
• $L_{cos}$ (코사인 정렬): 예측된 점수 벡터의 전역 방향을 타겟과 정렬시킵니다.

이 손실 설계는 경직된 회귀에서 가지치기 품질과 더 직접적으로 정렬된 상대적 순위 일관성으로 목적을 전환시킵니다.

주요 기여

1. 비동기 프록시 프레임워크: 점수 산출 병목 현상을 타겟의 임계 경로에서 제거하고 비동기 소규모 모델로 이전함으로써, Llama-3.1-8B(듀얼 GPU) 에서 최대 3.21 배의 프리필링 속도 향상을 달성하고 17 만 토큰까지 속도 향상을 유지하는 새로운 가지치기 프레임워크입니다.
2. HybridAxialMapper 및 Hybrid Loss: 이질적인 모델들 간의 구조적 간극 (예: 80 억 타겟 대 15 억 프록시) 을 성공적으로 메우는 학습 가능한 아키텍처와 손실 함수로, 기존 연구에서 사용된 정적 헤드 대 헤드 정렬의 경직성을 극복합니다.
3. 실증적 검증: Llama-3.1, Qwen-2.5, Qwen-3 계열 (70 억~320 억 타겟) 을 아우르는 광범위한 실험을 통해 ProxyKV 가 KVZip 오라클의 평균 정확도 **약 98.7%**를 회복하면서 지연 시간을 크게 줄임을 입증했습니다.

결과

• 정확도: ProxyKV 는 LongBench, SCBench, RULER 벤치마크 전반에 걸쳐 KVZip 오라클 (고정밀 기준) 의 성능과 일치합니다. 공격적인 유지 비율 ($\rho \le 0.2$) 에서도 대부분의 작업에서 1~2 포인트 이내로 추적하며 KVZip 평균 정확도의 약 98.7% 를 회복합니다.
• 효율성:
  • Llama-3.1-8B 에서 ProxyKV 는 KVZip 대비 (듀얼 GPU) 프리필링 지연 시간을 3.21 배 단축하고, 단일 공유 GPU 에서 약 1.5 배 단축합니다.
  • HybridAxialMapper 자체는 프리필 월 (wall) 시간의 **4.7%~6.8%**만을 차지합니다.
  • 이 방법은 약 8 배의 타겟/프록시 크기 비율로 320 억 타겟 (Qwen-3) 까지 효과적으로 확장됩니다.
• 강건성: 이 프레임워크는 홀드아웃 데이터셋 (예: SCBench RepoQA) 에 대한 강력한 제로샷 일반화를 보여주며, 휴리스틱 방법 (SnapKV) 이 실패하는 복잡한 멀티홉 추론 작업에서도 성능을 유지합니다.

중요성 및 주장

이 논문은 ProxyKV 가 장문맥 LLM 추론을 제약해 온 "점수 산출 비용–정확도 간극"을 성공적으로 메웠다고 주장합니다. 동족 내 어텐션 상관관계와 학습 가능한 교차 축 매핑을 활용함으로써, 막대한 프리필링 오버헤드 없이 재구성 기반 방법의 고정밀도를 달성합니다.

저자들은 타겟 모델이 확장됨에 따라 (예: 320 억으로) 정적 레이어 대 레이어 정렬 (IAM 또는 KVZip+IAM 에서 사용됨) 이 구속적인 제약이 되지만, 학습 가능한 HybridAxialMapper 는 여전히 효과적임을 강조합니다. 이는 ProxyKV 가 프리필 단계 동안 지연 시간과 메모리 압력을 크게 줄이면서 오라클에 가까운 정확도를 회복하는 확장 가능한 효율적 장문맥 추론 경로를 제공함을 시사합니다. 이 방법은 전체 모델 재학습이나 복잡한 하드웨어 수정 없이 경량 동족 내 프록시와 작은 학습 가능한 매핑기만 필요로 하는 배포 시 솔루션으로 제시됩니다.