기술 요약: MOSAIC

문제 정의

Key-Value (KV) 캐시는 시퀀스 길이와 배치 크기에 따라 그 크기가 선형적으로 증가하기 때문에, 대규모 언어 모델(LLM)의 긴 컨텍스트 추론에서 주요한 메모리 병목 현상이 됩니다. 기존의 토큰 제거(eviction) 방식(예: H2O, SnapKV, L2Norm, AdaKV)은 토큰의 중요도에 따라 캐시 항목을 선택적으로 삭제하여 메모리를 줄이지만, 일반적으로 모든 트랜스포머 레이어에 대해 **균등한 메모리 예산(uniform memory budget)**을 적용합니다. 이러한 접근 방식은 특정 토큰에 대한 레이어별 요구 사항이 서로 다르다는 점을 간과합니다. 균등한 할당은 캐시가 거의 필요 없는 레이어에는 캐시를 낭비하고, 더 많이 필요한 레이어는 결핍시키는 결과를 초래합니다.

기존의 레이어별 할당 방식은 사전에 할당 형태를 고정하거나, 어텐션 통계로부터 이를 도출하거나, 단일 타겟 예산에 대해 그룹별로 최적화한 후 재조정(rescaling)하는 방식을 취합니다. 이러한 방법들은 전체 레이어 예산 공간을 공동으로 탐색하거나, 재탐색 없이 다른 메모리 제약 조건으로 최적의 할당을 전이하는 데 실패합니다.

방법론: MOSAIC

MOSIC(Multi-Objective Search for Allocating Inference Cache Across LLM Decoder Layers)은 타겟 메모리 제약 조건 하에서 각 레이어별 KV 캐시 예산을 자동으로 최적화하도록 설계된 다목적 신경 구조 탐색(NAS) 프레임워크입니다. MOSAIC는 KV 캐시 할당을 정적 양자화(static quantization)와 유사하게 취급합니다. 즉, 할당은 캘리브레이션 데이터 상에서 오프라인으로 고정되며, 추론 시 매 프롬프트마다 할당 단계를 거칠 필요가 없습니다.

MOSAIC 파이프라인은 네 단계로 구성됩니다:

1. 1단계: 무제약 다중 예산 탐색 (Unconstrained Multi-Budget Search)

   • 고정된 평균 예산 제약 없이 레이어별 예산에 대한 공동 탐색을 수행합니다.
   • 분류기 가이드 베이지안 최적화(Classifier-guided Bayesian Optimization, BO) 절차를 사용합니다. 각 평가 후, 다층 퍼셉트론(MLP) 분류기가 랭크-1(rank-1) 구성과 그 외의 구성을 분리합니다. 이후 차분 진화(Differential evolution)를 통해 특정 구성이 랭크-1일 확률을 극대화하는 다음 후보를 제안합니다.
   • 목적 함수는 다목적입니다: 레이어당 평균 예산을 최소화하는 동시에 캘리브레이션 서브셋에 대한 태스크 점수를 최대화하는 것입니다.
   • 이 단계는 방대한 탐색 공간(예: 32개 레이어 모델의 경우 약 $10^{27}$개의 구성)을 탐색하여 예산-점수 트레이드오프의 파레토 프런트(Pareto front)를 추적합니다.

2. 2단계: 승자 선택 (Winner Selection)

   • 1단계의 파레토 프런트에 있는 구성들을 전체 테스트 데이터(단순 캘리브레이션 서브셋이 아닌)를 사용하여 재평가합니다.
   • 보편적인 앵커 역할을 할 단 하나의 "승자(winner)" 구성을 선택합니다(가장 높은 전체 데이터 점수를 가진 형태).

3. 3단계: 승자 재조정 (Winner Rescaling)

4. 2단계의 승자를 가져와 임의의 타겟 평균 예산($B$)에 맞게 레이어별 예산 비율을 재조정합니다.

   • 이를 통해 하나의 예산에서 발견된 최적의 형태를 추가적인 탐색 없이 다른 모든 예산(예: 2400개 토큰에서 128개 토큰으로)으로 전이할 수 있습니다.

5. 4단계: 예산 제약 적응형 탐색 (Budget-Constrained Adaptive Search, 선택 사항)

   • 특정 배포 예산에 대해, 재조정된 승자를 정교화하는 선택적 제약 탐색을 수행합니다.
   • 이 단계는 고정된 평균 예산 제약 내에서 탐색하며, 재조정된 승자, 휴리스틱 형태, 그리고 라틴 하이퍼큐브 샘플을 시드로 사용하여 탐색을 시작합니다.
   • 이는 해당 특정 예산에서 단순 재조정보다 더 나은 성능을 보이는 구성을 찾는 것을 목표로 합니다.

주요 기여

1. MOSAIC 프레임워크: 태스크 및 제거 방식당 단 한 번의 공동 탐색을 수행하여 승자 할당을 생성하는 4단계 워크플로우를 제공합니다. 이 승자는 재조정하여 어떤 타겟 예산에도 적용할 수 있어, 예산별로 매번 탐색할 필요가 없습니다.
2. 포괄적 평가: 네 가지 제거 방식(SnapKV, H2O, AdaKV, L2Norm), 두 개의 모델(Llama-3-8B-Instruct, Mistral-7B-Instruct-v0.2), 그리고 두 개의 벤치마크(LongBench, RULER)에 걸쳐 평가를 수행했습니다.
3. 학습된 할당 분석: 최적의 할당이 엄격한 피라미드 분포나 균등한 분포를 따르기보다, 일관되게 특정 중간 레이어(예: 32개 레이어 모델에서 10, 14–16, 20번 레이어)를 선호한다는 사실을 발견했습니다.

실험 결과

MOSIC은 다양한 설정에서 균등 할당 대비 유의미한 성능 향상을 입증했습니다:

• 성능 향상:

  • 154개의 설정(태스크, 방법, 예산 조합) 중 146개에서 균등 할당보다 우수한 성능을 보였습니다.
  • RULER(특히 어려운 니들 리트리벌 서브태스크)에서 최대 +48 포인트, LongBench(L2Norm 기반 Code 카테고리)에서 +11.9 포인트를 달성했습니다.
  • 균등 할당과 동일한 성능을 내면서도 최대 4배 적은(75% 감소) KV 캐시를 사용합니다.
  • 17개의 Llama 설정에서 MOSIC은 절반 이하의 캐시만으로 균등 할고 수준의 점수에 도달했습니다.

• 전이 가능성:

  • 3단계의 "승자"(추가 탐색 없이 재조정된 것)만으로도 87개 Llama 설정 중 75개에서 균등 할당을 능가했습니다.
  • 단 한 번의 탐색(9–40 GPU-시간)이 모든 타겟 예산에 대응합니다.

• EvolKV와의 비교:

  • 고정된 예산에서 그룹별로 탐색하는 EvolKV와 비교했을 때, MOSIC은 한 번의 탐색당 약 $10^{40}$배 더 큰 탐색 공간을 탐색합니다.
  • MOSIC은 이를 4배 적은 평가 횟수와 18% 적은 GPU 시간으로 달성했습니다.
  • EvolKV가 각 예산에 대해 전용 탐색을 수행함에도 불구하고, MOSIC은 6개 설정 중 5개에서 더 높은 점수를 기록했습니다.

• 태스크 vs. 방법 민감도:

  • 학습된 형태는 태스크 의존성이 매우 높습니다(Code를 위한 형태는 QA로 잘 전이되지 않음). 그러나 동일한 태스크 내에서는 서로 다른 제거 방식 간에는 상당히 잘 전이됩니다.

의의 및 주장

본 논문은 KV 캐시 관리를 균등 또는 휴리스틱 기반 할당에서 데이터 기반의 다목적 최적화로 전환함으로써 새로운 패러다임을 구축한다고 주장합니다. 주요 의의는 다음과 같습니다:

1. 공동 최적화의 우월성: 모든 레이어를 공동으로 탐색하는 것이 레이어를 순차적으로 최적화하거나 형태를 미리 고정하는 것보다 더 나은 할당을 생성합니다.
2. 한 번의 탐색, 다양한 예산: 단 한 번의 오프라인 탐색이 "승자" 형태를 생성하며, 이는 다양한 하드웨어 제한에 맞춰 재조정될 수 있어 매우 효율적입니다.
3. 레이어 이질성: 실험 결과는 트랜스포머 레이어가 캐시 요구 사항 측면에서 동질적이지 않음을 확인해 줍니다. 특히 소수의 중간 레이어들이 일관되게 불균형적으로 많은 캐시 예산을 필요로 한다는 점을 입증했습니다.

저자들은 두 개의 밀집 디코더(dense decoder) 모델만 평가했다는 점과, 현재 버전의 Mistral-7B에 대해 특정 모델-방법 조합에서 일부 단계를 제외했다는 한계점을 언급했습니다. 향후 연구로는 이를 Mixture-of-Experts(MoE) 모델 및 최신 제거 방식까지 확장하는 것을 제안합니다.