기술 요약: 효율적인 자기회귀 이미지 생성을 위한 헤드 인식 키-값 압축

문제 제기

이미지를 시각 토큰 시퀀스로 표현하는 자기회귀 (AR) 시각 생성 모델은 모델 크기와 토큰 수를 확장함으로써 고해상도 및 사진 같은 현실적인 결과를 달성했습니다. 그러나 이러한 확장은 자기회귀 디코딩 중에 필요한 키 - 값 (KV) 캐시로 인해 상당한 메모리 오버헤드와 낮은 추론 처리량을 초래합니다. 자연어 처리 (NLP) 분야의 최근 연구들은 토큰 선택, 퇴출 (eviction), 양자화 등의 KV 캐시 관리를 탐구해 왔으나, 이러한 기법을 시각 도메인에 직접 적용하는 것은 이미지의 고유한 공간적 구조로 인해 어렵습니다.

기존의 시각 특화 방법인 LineAR 과 같은 접근법은 토큰이 주로 공간적으로 인접한 이웃에 주의를 기울인다는 관찰을 바탕으로 토큰의 하위 집합만 유지하여 캐시 사용량을 줄이려 시도합니다. 그러나 이러한 방법들은 모든 어텐션 헤드에 고정된 균일한 예산을 할당합니다. 이 전략은 어텐션 헤드 간의 이질성을 간과하여 비최적의 메모리 할당을 초래합니다. 구체적으로, 일부 헤드는 국소적 이웃에 집중하여 (작은 예산 필요) 다른 헤드는 객체 구조 및 배경 영역과 같은 장거리 의존성을 포착하는 데 (더 큰 예산 필요) 사용됩니다. 균일한 할당은 이러한 서로 다른 패턴에 적응하지 못해 메모리 낭비 또는 생성 품질 저하를 초래합니다.

방법론: HeadKV

저자들은 헤드 인식 KV 캐시 압축을 수행하는 새로운 훈련 불필요 (training-free) 프레임워크인 HeadKV를 제안합니다. 핵심 통찰은 AR 이미지 생성 모델의 어텐션 헤드가 **국소적 (Local)**과 **전역적 (Global)**이라는 두 가지 뚜렷한 패턴을 보이며, 주어진 입력에 대해 동일한 레이어 내의 토큰 위치 전반에 걸쳐 헤드의 행동이 일관된다는 것입니다.

HeadKV 파이프라인은 세 가지 주요 단계로 구성됩니다:

1. 캐시 프리필링 (Cache Prefilling): 모델은 조건부 입력을 인코딩하고 압축 없이 초기 시각 토큰 시퀀스를 생성합니다. 이 단계는 헤드별 어텐션 분포를 추정하기 위해 충분한 컨텍스트를 축적합니다.
2. 헤드 그룹화 (분류): 캐시가 미리 정의된 길이에 도달하면, 프레임워크는 임계값 기반 어텐션 누적 기준을 사용하여 각 어텐션 헤드를 "국소적" 또는 "전역적"으로 분류합니다.
   • 각 헤드의 경우, 현재 토큰 위치에서 역방향으로 어텐션 점수를 임계값 ($\tau = 0.9$) 에 도달할 때까지 누적합니다.
   • 임계값에 도달하는 데 필요한 토큰 수가 미리 정의된 국소 윈도우 크기보다 작으면 해당 헤드는 국소적으로 분류됩니다. 그렇지 않으면 전역적으로 분류됩니다.
   • 이 분류는 생성 초기 단계에서 한 번 수행되며 전체 디코딩 과정에서 재사용되므로 추가 훈련이나 데이터셋 수준의 통계를 필요로 하지 않습니다.
3. 헤드 인식 캐시 압축: 헤드 유형에 따라 다른 압축 전략이 적용됩니다:
   • 국소적 헤드: 이러한 헤드는 주로 인접한 컨텍스트에 주의를 기울입니다. 프레임워크는 슬라이딩 윈도우 메커니즘을 통해 가장 최근 토큰만 유지하는 컴팩트한 예산을 할당합니다. 과거 토큰은 완전히 생략됩니다.
   • 전역적 헤드: 이러한 헤드는 장거리 의존성을 포착합니다. 선형적인 메모리 증가 없이 정보성이 높은 원거리 토큰을 보존하기 위해 저자들은 계층적 토큰 퇴출 (Stratified Token Eviction, STE) 전략을 설계했습니다.
     • 계층적 토큰 퇴출: 후보 과거 토큰은 현재 위치와의 거리에 따라 근거리 및 원거리 토큰이라는 두 개의 불교합 부분집합으로 분할됩니다.
     • 중요도 점수는 과거 쿼리 특징을 사용하여 각 부분집합에 대해 독립적으로 계산됩니다.
     • 가장 낮은 점수를 가진 토큰이 각 부분집합에서 퇴출됩니다 (예: 근거리에서 $M/2$, 원거리에서 $M/2$). 이는 국소적으로 지배적인 토큰이 의미상 관련 있는 원거리 토큰을 가리는 것을 방지하여 국소적 세부 사항과 전역적 구조의 균형을 잡은 보존을 보장합니다.

주요 기여

1. 어텐션 이질성 분석: 이 논문은 AR 이미지 생성 모델에 대한 심층 분석을 제공하여 두 가지 뚜렷한 어텐션 패턴 (국소적 및 전역적) 을 식별하고, 고정된 입력에 대해 레이어 내 토큰 위치 전반에 걸쳐 헤드 행동이 불변임을 규명했습니다.
2. HeadKV 프레임워크: 어텐션 헤드 간 비대칭 예산을 할당하는 새로운 훈련 불필요 압축 프레임워크입니다. 초기 단계 분류 메커니즘에 따라 국소적 헤드에는 컴팩트한 예산을, 전역적 헤드에는 더 큰 예산을 할당합니다.
3. 계층적 토큰 퇴출 (STE): 토큰을 거리별로 분할하여 정보성이 높은 장거리 의존성을 유지하도록 보장하는 전역적 헤드를 위한 구체적인 전략으로, 종종 국소적 토큰을 선호하는 표준 중요도 점수화의 한계를 해결합니다.
4. 실증적 검증: 광범위한 실험을 통해 HeadKV 가 Janus-Pro, Lumina-mGPT, LlamaGen-XL 등 여러 모델에서 KV 캐시 메모리 사용량을 크게 줄이고 추론 처리량을 향상시키면서도 고품질 이미지를 유지함을 입증했습니다.

결과

저자들은 Janus-Pro-1B, Janus-Pro-7B, Lumina-mGPT-768, LlamaGen-XL 네 가지 자기회귀 이미지 생성 백본에서 HeadKV 를 평가했습니다.

• 이미지 품질: HeadKV 는 GenEval 및 DPG 와 같은 벤치마크에서 풀 캐시 베이스라인과 비교 가능한 성능을 달성합니다. 예를 들어, 압축 비율 ($\rho$) 이 1/8 인 Janus-Pro-7B 에서 HeadKV 는 GenEval 점수 0.99 (풀 캐시 대비 0.99) 와 DPG 점수 83.78 (풀 캐시 대비 84.20) 을 유지합니다.
• LineAR 와의 비교: 균일한 할당 방법인 LineAR 과 비교할 때 HeadKV 는 압축 비율이 증가함에 따라 성능 저하가 적습니다 (예: $\rho=1/8$). LlamaGen-XL 에서 극단적인 압축 하에 HeadKV 는 2.75 의 더 나은 FID-50K 점수를 달성하여 LineAR(2.95) 보다 우수합니다.
• 효율성: HeadKV 는 메모리 사용량을 크게 줄이고 처리량을 향상시킵니다.
  • LlamaGen-XL($\rho=1/6$) 에서 HeadKV 는 메모리를 63.60% 줄이고 풀 설정 대비 3.00 배의 속도 향상을 달성하여 LineAR(2.93 배) 을 능가합니다.
  • Janus-Pro-7B($\rho=1/6$) 에서 52.08% 의 메모리 절감 효과와 함께 3.68 배의 속도 향상을 달성합니다.
  • Lumina-mGPT-768에서는 그룹화에 따른 구현 오버헤드로 인해 지연 시간이 LineAR 보다 약간 길지만, 여전히 상당한 메모리 절감과 함께 1.35 배의 속도 향상을 달성합니다.

중요성 및 주장

이 논문은 HeadKV가 추가 훈련이나 외부 감독 없이 자기회귀 이미지 생성의 메모리 병목 현상에 대한 실용적인 해결책을 제공한다고 주장합니다. 어텐션 헤드의 고유한 이질성을 인식하고 활용함으로써 이 방법은 이전 연구들의 "일률적 (one-size-fits-all)" 캐시 할당을 넘어섭니다.

저자들은 초기 생성 단계에서 관찰된 내재적 어텐션 패턴에 의존하기 때문에 그들의 접근 방식이 다양한 입력과 모델 아키텍처에 걸쳐 일반화 가능하다고 강조합니다. 계층적 토큰 퇴출 전략은 표준 퇴출 방법에서 종종 손실되는 장거리 구조적 정보를 보존함으로써 엄격한 메모리 예산 하에서 생성된 이미지의 품질을 유지하는 핵심 구성 요소로 강조됩니다. 이 연구는 균일한 토큰 유지보다 적응형 헤드별 자원 할당이 시각 생성을 위한 더 효율적인 패러다임임을 시사합니다.