통합 멀티모달 모델의 KV 캐시를 태스크·타입별로 나눠 압축하는 UniCache

무엇인가

통합 멀티모달 모델은 이해·생성·편집을 하나의 네트워크에서 처리한다. BAGEL로 대표되는 Mixture-of-Transformers(MoT) 계열은 모달리티별 Transformer 전문가를 공유 어텐션으로 묶는다. 문제는 문맥이 길어질수록 추론 중 KV 캐시를 저장하고 반복 접근하는 비용이 커진다는 점이다. 기존 KV 캐시 압축 기법은 대부분 고정된 추론 패러다임과 단일 모달리티 캐시를 전제로 설계됐다. 반면 통합 모델에서는 태스크마다 관여하는 KV 캐시 타입 자체가 다르고, 그 구성과 동역학도 달라서 하나의 전역 압축 정책으로는 태스크별 요구를 놓친다.

어떻게 동작하나

논문은 이 한계를 네 가지로 정리한다. 첫째, 태스크별 캐시 구성이 다르다. 둘째, 같은 태스크 안에서도 캐시 타입마다 토큰 제거(eviction)와 양자화(quantization)에 대한 민감도가 다르다. 셋째, 캐시 타입별로 어텐션 기여도가 달라 균일 예산 배분이 맞지 않는다. 넷째, 캐시 타입의 상대적 중요도와 전체 KV의 중요도가 태스크와 추론 단계에 따라 변한다. 이미지 편집이 대표적 예다. 편집 지시를 따르면서 요청과 무관한 원본 이미지 내용을 보존해야 하는데, 제거 기반 기법은 원본 세부를 잃고 양자화 기반 기법은 편집 자체를 수행하지 못하는 상반된 실패가 나타난다.

무엇과 다른가

UniCache의 첫 단계는 태스크 인지 캐시 분할이다. BAGEL 기준으로 이미지 이해는 instruction KV와 source-ViT KV를, 텍스트-이미지 생성은 instruction KV를, 이미지 편집은 instruction·source-ViT·source-VAE KV를 함께 쓴다(source-VAE가 원본 세부를 담당). 이 조건부 KV들은 프리필에서 생성돼 추론 내내 재사용된다. 여기에 어텐션 싱크를 이루는 boundary KV와 디노이징 스텝마다 다시 계산되는 current-latent KV는 압축 대상에서 제외되는 보호 세그먼트로 둔다. 다음으로 오프라인 캘리브레이션으로 각 태스크-타입 쌍에 압축 정책을 배정한다. 어텐션이 집중되는 instruction·source-ViT KV에는 중요한 토큰을 남기는 제거 정책을, 어텐션이 넓게 퍼지는 source-VAE KV에는 토큰 커버리지를 유지하는 양자화 정책을 붙인다. 이 배정은 추론 중 고정되고, 압축 강도만 온라인 예산 배분으로 조정된다.

어떻게 쓰나

예산 배분은 청크 단위 어텐션 통계로 이뤄진다. 레이어와 추론 스텝을 청크로 묶어, 각 후보 캐시 타입의 유지된 키에 대한 어텐션 확률을 헤드·쿼리·레이어-스텝에 걸쳐 평균 낸 값이 그 타입의 중요도가 된다. 이 값에 비례해 공유 저장 예산을 나누되, 직전 할당량을 상한으로 둔다. 어떤 타입이 상한에 도달하면 남은 예산은 어텐션 질량이 양수인 미포화 타입에 비례 재분배된다. 전체 예산은 유지 비율 ρ로 조절하는데, 이해 태스크는 조건부 문맥이 자기회귀 디코딩 내내 살아 있으므로 고정 비율을 쓴다. 생성·편집은 조건부 KV의 총 어텐션 질량이 디노이징이 진행될수록 줄어드는 것을 반영해, 초기 대비 현재 질량 비율에 따라 ρ를 하한까지 낮춘다. 마지막으로 타입별 압축은 서로 독립적으로 병렬 실행된다. 제거 정책은 남길 토큰 수를, 양자화 정책은 비트 정밀도와 그룹 단위를 결정하고, 압축된 후보 세그먼트와 보호 세그먼트는 원래 토큰 순서로 재조립돼 어텐션 계산에 들어간다.

전제와 한계

실험은 MoT 기반 통합 모델 두 개(BAGEL, SenseNova-U1)에서 수행했다. 이해는 MME·MMBench·MMMU, 텍스트-이미지 생성은 GenEval, 이미지 편집은 PIE-Bench와 GEdit-Bench(영어 지시, Qwen2.5-VL-72B-Instruct를 판정자로 사용)를 쓴다. 비교 대상은 full-KV 모델과 H2O, KIVI, StreamingLLM, SnapKV, PyramidKV다. 공정 비교를 위해 모든 기법을 이해·편집 약 80%, 텍스트-이미지 생성 약 60%의 KV 메모리 절감으로 맞췄다. 결과적으로 UniCache는 이해·편집에서 5배, 생성에서 2.5배 KV 캐시 압축을 품질 손실 거의 없이 달성한다고 보고한다. 이해 태스크에서 KIVI는 MME 인지(Perception) 항목이 가장 높지만 인지(Cognition) 항목은 가장 낮은 반면, UniCache는 두 축 모두 원본 모델과 비슷한 수준을 유지한다. 생성에서는 거의 모든 하위 범주에서 기존 기법을 앞선다.

어블레이션은 구성 요소별 기여를 분리한다. boundary KV와 current-latent KV 보호를 제거하면 텍스트-이미지 생성 성능이 떨어지고, 어텐션 기반 할당은 KV 캐시 타입이 여러 개 얽히는 이미지 이해와 편집에서 특히 결정적이다. 태스크 인지 시간 스케줄을 빼면 생성에서 약간의 성능 저하가 생기는데, 스텝별 균일 예산이 디노이징에 따른 조건부 KV 어텐션 감소를 따라가지 못하기 때문이다. 스케줄 대안 비교에서는 어텐션 기반 감쇠가 GenEval에서 선형 감쇠와 동일하고 PIE-Bench 전 지표에서 가장 좋았다. 압축률을 높여도 UniCache가 더 강건하며, PIE-Bench에서 H2O와 KIVI는 두 지표 중 최소 하나에서 성능 저하를 피하지 못한다. 효율 측정은 NVIDIA A100 40GB 단일 GPU에서 이뤄졌고, 어텐션 브랜치당 약 20K KV 토큰 규모에서 처리량이 최대 1.78배 향상되고 관리 대상 KV 캐시 메모리는 약 80% 줄었다.

개발자 관점에서 이 논문은 멀티모달 서빙 스택에서 KV 캐시를 단일 전역 정책으로 다루면 태스크별로 다른 실패 모드가 생긴다는 점을 보여준다. 특히 편집 파이프라인에서 원본 이미지 세부 보존과 지시 수행이 서로 다른 압축 민감도를 가진다는 관찰은, 캐시 타입을 구분해 정책을 나누는 구현이 실무에서 유효할 수 있음을 시사한다. 다만 적용 전 확인할 것이 있다. UniCache는 학습이 필요 없는 대신 오프라인 캘리브레이션으로 정책 배정을 고정하므로 모델·태스크 조합마다 캘리브레이션이 필요하고, 보호 세그먼트와 후보 세그먼트를 구분하는 분할 규칙이 대상 아키텍처에 의존한다. 또한 되돌릴 수 없는 제거 때문에 상한에 걸린 뒤 남는 예산 일부가 쓰이지 않을 수 있다. 원문에 제시된 범위에서 실험은 BAGEL과 SenseNova-U1 두 MoT 모델에 한정되며, 별도의 한계 절은 마련돼 있지 않다.