기술 요약: HeatKV - 시각적 자기회귀 모델링을 위한 헤드 튜닝 KV 캐시 압축

문제 제기

시각적 자기회귀 (VAR) 모델은 이미지를 다중 스케일 토큰 맵의 시퀀스로 생성함으로써 빠른 추론 속도와 유리한 확장성을 제공하며, 이미지 생성을 위한 확산 모델의 경쟁력 있는 대안으로 부상했습니다. 그러나 KV 캐시의 특성상 이러한 모델은 심각한 메모리 제약에 직면합니다. 일반적으로 한 번에 하나의 토큰을 생성하는 대규모 언어 모델 (LLM) 과 달리, VAR 모델은 증가하는 해상도로 전체 토큰 맵을 생성합니다. 결과적으로 KV 캐시의 토큰 수는 생성 스케일에 따라 초선형적으로 증가합니다.

Infinity-2B 및 Infinity-8B 와 같은 최첨단 VAR 모델은 단일 1024 × 1024 이미지를 생성하는 데 최대 10GB 의 KV 캐시 메모리가 필요할 수 있습니다. 이러한 높은 메모리 사용량은 특수 하드웨어 가속기에서 지원 가능한 최대 배치 크기와 동시 사용자 수를 제한합니다. ScaleKV 및 HACK 과 같은 기존 VAR 압축 방법들은 종종 attention 헤드나 레이어를 이진 카테고리 (예: "고예산" 대 "저예산") 로 분류하는 것과 같은 거시적 전략에 의존합니다. 저자들은 이러한 이진 분류가 지나치게 제한적이라고 주장합니다. 동일한 레이어 내에서도 서로 다른 헤드 간에 attention 패턴이 크게 달라지기 때문입니다. 이로 인해 메모리가 낭비되거나 최적화되지 않은 가지치기가 발생할 수 있습니다.

방법론: HeatKV

저자들은 개별 attention 헤드와 특정 이미지 스케일의 세분화된 단위에서 캐시 할당을 조정하는 새로운 전략인 HeatKV(Head-tuned KV-cache Compression) 를 제안합니다. 이 방법은 세 가지 주요 구성 요소로 이루어집니다:

1. 헤드 - 스케일 중요도 (오프라인 보정)

분산을 고려한 지표를 사용하는 이전 방법들과 달리, HeatKV 는 **스케일 의존적 누적 Attention 점수 **(S-CAS)를 기반으로 attention 헤드를 순위 매깁니다.

• 메커니즘: 소규모 오프라인 보정 데이터셋을 사용하여, 미래 스케일의 토큰에서 이전 스케일의 토큰으로의 attention 확률을 집계합니다.
• 지표: 주어진 레이어 $\ell$, 헤드 $h$, 소스 스케일 $k$에 대해, S-CAS 는 모든 미래 스케일 ($\tau > k$) 에서 스케일 $k$의 토큰으로 향하는 평균 attention 질량을 측정합니다.
• 결과: 이는 각 소스 스케일에 대한 특정 "헤드 - 스케일" 쌍의 순위 목록을 생성하여, 특정 헤드의 특정 이전 스케일 캐시가 미래 토큰 생성에 얼마나 중요한지 나타냅니다. 이를 통해 시스템은 특정 헤드가 스케일 1 은 유지해야 하지만 스케일 2 는 폐기할 수 있음을 식별할 수 있으며, 이는 이진 방법으로는 달성할 수 없는 세분화된 수준입니다.

2. 예산 할당

HeatKV 는 각 스케일 $k$ 생성 후 총 캐시 크기가 사용자가 정의한 메모리 예산 $b$ 내에 유지되도록, 가지치기해야 하는 헤드의 수 ($N_k$) 를 계산합니다.

• 제약 조건: 이 방법은 "Streaming-LLM" 원칙을 준수하여, 첫 $s$개의 스케일을 "attention sink"(안정성에 필수적) 로 유지하고 현재 생성 중인 스케일은 절대 캐시하지 않습니다.
• 계산: 알고리즘은 예산 제약을 충족하기 위해 가지치기해야 하는 최소한의 헤드 수를 결정하여, 생성이 진행됨에 따라 캐시를 동적으로 축소할 수 있게 합니다.

3. 제거 순서 최적화 (탐욕적 조기 가지치기)

VAR 에서의 중요한 과제는 단일 생성 단계 (트랜스포머 레이어 간) 에서 캐시 크기가 초선형적으로 증가한다는 점입니다. 스케일 간에만 가지치기를 수행하는 표준 방법은 하드웨어 한계를 초과하는 일시적인 메모리 초과 할당을 초래할 수 있습니다.

• 해결책: HeatKV 는 스케일 간뿐만 아니라 각 트랜스포머 레이어 이후에도 예산 제약을 강제하는 **탐욕적 조기 가지치기 **(Greedy Early Pruning) 알고리즘 (알고리즘 1) 을 도입합니다.
• 전략: 특정 레이어 이후 캐시 크기가 예산을 초과하면, 알고리즘은 해당 레이어 또는 이전 레이어에서 추가로 가지치기할 헤드 - 스케일을 탐욕적으로 선택합니다. 이는 S-CAS 순위가 낮은 헤드를 우선적으로 가지치기하며, 동점인 경우 메모리 해소를 극대화하기 위해 더 나중의 레이어나 더 나중의 소스 스케일에서 가지치기를 우선시합니다.
• 최적성: 저자들은 (명제 1) 이 탐욕적 접근 방식이 최소 카디널리티의 조기 가지치기 헤드 집합을 산출하여 예산 위반을 방지하면서도 사용 가능한 메모리를 최대화한다고 증명합니다.

주요 기여

1. 세분화된 관찰: 저자들은 VAR 모델에서 attention 패턴이 동일한 attention 헤드 내에서도 서로 다른 스케일 간에 크게 변하지만, 서로 다른 입력 프롬프트 간에는 상대적으로 안정적임을 입증했습니다. 이는 이진 헤드 분류에서 연속적이고 스케일 의존적인 예산 할당으로의 전환을 동기화합니다.
2. S-CAS 지표: 스케일 의존적 누적 Attention 점수의 도입은 특정 메모리 예산에 맞춰진 정적 가지치기 일정을 구성할 수 있게 하여, 각 헤드가 필수적인 스케일은 유지하면서 비필수적인 스케일을 폐기할 수 있게 합니다.
3. 최적화된 가지치기 알고리즘: 스케일 수준이 아닌 레이어 수준에서 예산 제약을 강제하는 탐욕적 제거 알고리즘의 개발은 VAR 의 초선형 성장 특성상 중요한 일시적인 메모리 오버플로우를 방지합니다.
4. 최첨단 성능: 이 방법은 VAR 모델 KV 캐시 압축 분야에서 새로운 최첨단 성능을 달성하여, 세분화된 할당이 거시적 접근법보다 훨씬 우수함을 입증했습니다.

실험 결과

저자들은 MS-COCO 2017 검증 세트를 사용하여 Infinity-2B 및 Infinity-8B모델에서 HeatKV 를 평가했습니다. 그들은 4%, 10%, 20% 의 메모리 예산에 걸쳐 HeatKV 를 StreamingLLM, SnapKV, HACK, ScaleKV 와 같은 베이스라인과 비교했습니다.

• 정확도: HeatKV 는 모든 베이스라인을 일관되게 능가했습니다. 특히 Infinity-2B 모델에서 HeatKV 는 10% 예산에서 ScaleKV 가 20% 예산에서 달성한 것보다 더 나은 FID, LPIPS, PSNR 점수를 기록했습니다. 4% 예산에서 HeatKV 는 Infinity-2B 에서 2.70 의 FID 를 유지하여, 다음으로 가장 좋은 방법 (ScaleKV 의 4.92) 보다 훨씬 낮았습니다.
• 프롬프트 정렬: GenEval 을 사용하여 HeatKV 는 낮은 예산에서도 강력한 프롬프트 정렬을 유지했으며, 점수는 전체 캐시 베이스라인과 비교할 만했습니다.
• 지각적 품질: 인간 선호도 점수 (HPSv2.1) 결과는 HeatKV 가 높은 지각적 품질을 유지했으며, Infinity-2B 에서 4% 예산 시 다른 방법들보다 우수함을 나타냈습니다.
• 효율성: HeatKV 는 기존 방법 대비 메모리 할당에서 2 배 더 높은 압축률을 달성하면서도 유사하거나 더 나은 이미지 정확도를 유지했습니다. 이 방법은 또한 더 큰 배치 크기를 가능하게 하여 전체 캐시 추론에 비해 생성 지연을 효과적으로 줄였습니다.

중요성 및 주장

이 논문은 HeatKV 가 제한된 메모리 용량을 가진 하드웨어 (예: 소비자용 GPU) 에 VAR 모델을 배포하거나 기존 인프라에서 더 많은 동시 사용자를 서비스하는 데 있어 중요한 진전을 나타낸다고 주장합니다.

저자들은 이 방법이 모델 아키텍처나 학습 과정의 변경을 요구하지 않으며, 정적 가지치기 일정을 생성하기 위해 소수의 오프라인 보정 세트 (최소 1~10 개의 프롬프트) 만 의존한다고 강조합니다. 결과는 세분화된, 헤드별 캐시 할당이 공격적인 메모리 제약 하에서 VAR 모델의 의미 구조와 생성 품질을 유지하는 데 필수적이며, attention 헤드나 레이어의 거시적 이진 분류에 의존하는 방법들보다 우수함을 시사합니다.

논문은 HeatKV 가 메모리를 최적화하지만, 향후 연구는 효율성을 더욱 향상시키기 위해 attention 희소성 활용이나 모델 양자화와 같은 다른 최적화 기술과 결합하는 것을 탐구할 수 있다고 결론지었습니다.