LinearKV: One Cached State Suffices for Position-Independent Caching in Hybrid LLMs

기술 요약: LINEARKV

문제 정의
대규모 언어 모델(LLM) 서빙은 긴 컨텍스트의 프리필(prefilling) 비용이 증가함에 따라 심각한 지연 시간(latency)과 처리량(throughput) 병목 현상에 직면해 있습니다. 위치 독립적 캐싱(Position-Independent Caching, PIC)은 프롬프트 내의 위치와 관계없이 독립적으로 프리필된 토큰 청크를 저장하고 재사용함으로써 이를 완화합니다. 그러나 기존의 PIC 방법들은 토큰 인덱스 기반의 Key-Value(KV) 캐시를 결합하고 선택적으로 복구할 수 있는 풀 어텐션(full-attention) 아키텍처를 위해 설계되었습니다. 풀 어텐션 레이어와 순환(recurrent) 레이어(예: Mamba-2, Gated DeltaNet)를 교차 배치하는 하이브리드 LLM은 이러한 기본 원칙을 깨뜨립니다. 순환 레이어는 토큰 인덱스 기반의 캐시를 유지하는 대신 컨텍스트를 고정된 크기의 상태(state)로 압축합니다. 이는 두 가지 과제를 발생시킵니다: (1) 토큰 인덱스 항목을 결합할 수 없는 상황에서 어떻게 분리된 캐시 청크로부터 순환 상태를 초기화하고 복구할 것인가 하는 점, (2) 통합된 PIC 전략이 서로 다른 하이브리드 아키텍처(예: 스칼라 감쇠 방식의 Mamba-2 vs 밀집 전이 방식의 Gated DeltaNet)에 걸쳐 일반화될 수 있는지 여부입니다.

방법론: LinearKV
저자들은 하이브리드 LLM을 위한 위치 독립적 캐싱을 위한 훈련 불필요(training-free) 프레임워크인 LinearKV를 제안합니다. 핵심 혁신은 매칭된 캐시 청크 적용 후 수행되는 디커플링된 초기화(decoupled initialization) 전략입니다:

1. 디커플링된 초기화:

   • 풀 어텐션 레이어: 이 레이어들은 표준 PIC와 동일하게 진행되며, 매칭된 청크로부터 추출된 KV 엔트리를 토큰 축을 따라 결합합니다.
   • 순환(Linear) 레이어: 상태를 결합하려고 시도하는 대신(고정된 크기의 요약본이므로 결합이 불가능함), 이 프레임워크는 함수 $f$를 통해 $K$개의 매칭된 로컬 상태 $\{S^{\text{local}}_1, \dots, S^{\text{local}}_K\}$를 하나의 초기 상태 $S_{\text{init}}$으로 매핑합니다.
   • 셀렉터 재사용: 이러한 디커플링을 통해 LinearKV는 기존의 PIC 토큰 선택 알고리즘(예: CacheBlend, EPIC, ProphetKV)을 변경 없이 재사용할 수 있습니다. 셀렉터는 재계산할 토큰 위치 세트 $U$를 결정합니다. 시스템은 $U$에 있는 KV 엔트리를 덮어쓰고, 컨텍스트 순서에 따라 $U$ 내의 토큰들을 통과하며 순환 상태를 $S_{\text{init}}$으로부터 전진시켜 캐시를 복구합니다.

2. 선형 상태 초기화기 (Linear-State Initializer, $f$):
   본 논문은 $f$에 대한 두 가지 선택지를 조사합니다:

   • 정확한 합성(Exact Composition): 대수적으로 원칙적인 접근 방식으로(동시 연구인 HYPIC에서 채택됨), 청크들의 전이 연산자를 사용하여 $K$개의 캐시된 상태를 합성함으로써 정확한 전체 접두사(full-prefix) 상태를 재구성합니다.
   • 단일 블록 초기화(Single-Block Initialization, LinearKV의 기본값): 더 단순한 접근 방식으로, $f$는 단일 캐시된 상태(일반적으로 마지막 매칭된 청크인 $S^{\text{local}}_K$)를 초기화제로 선택하고 나머지는 폐기합니다.

주요 결과 및 기여
본 논문은 세 가지 하이브리드 모델(Granite/Mamba-2, OLMo/GDN, Qwen/GDN)과 세 가지 PIC 셀렉터를 통한 실증적 근거를 바탕으로 세 가지 주요 기여를 합니다:

1. 하이브리드 PIC의 정형화: 저자들은 하이브리드 PIC를 유일한 아키텍처 특정 연산이 선형 상태 초기화인 디커플링된 프레임워크로 정형화했습니다. 이를 통해 기존의 풀 어텐션 PIC 셀렉터를 수정 없이 플러그인할 수 있습니다.
2. 정확한 합성의 취약성: 대수적 직관과는 반대로, 본 논문은 정확한 합성이 불필요하며 아키텍처에 따라 취약함을 입증합니다.
   • GDN 모델(OLMo, Qwen)의 경우, 정확한 합성 방식과 단일 블록 초기화 방식 모두 동일하게 작동하며, 둘 다 최대 92%의 전체 품질을 회복합니다.
   • Mamba-2 모델(Granite)의 경우, 정확한 합성은 붕괴됩니다. Mamba-2의 스칼라 감쇠 특성 때문에, 독립적으로 프리필된 청크들로부터 발생하는 오차가 깊이(depth)를 통해 누적되어, 구성된 초기 상태가 높은 상대 오차를 갖게 됩니다. 결과적으로, 정확한 합성은 EPIC 셀렉터 하에서 약 46.6%의 전체 품질만을 회복하는 반면, 단일 블록 초기화 방식은 **86.8%**를 회복합니다.
   • 저자들은 이러한 실패가 아키텍처 자체 때문이 아니라, $K$개의 독립적으로 조건화된 딥 연산자(deep operators)의 집합체에서 기인함을 밝혀냈으며, 단일 블록 방식은 이러한 오차 누적을 피합니다.
3. 효율성 이점: 단일 블록 초기화기는 계산 비용이 더 저렴합니다. 이는 정확한 합성에 필요한 밀집 전이 행렬(dense transition matrices)을 온라인으로 합성하는 비용을 피할 수 있게 해줍니다. 이를 통해 정확도를 유지하거나 향상시키면서도 정확한 합성 방식 대비 TTFT(Time-to-First-Token)를 추가로 5~17% 단축합니다.

실험 결과

• 벤치마크: LongBench QA, RULER (8K–32K), LongBench Summarization에서 평가되었습니다.
• 성능:
  • Granite (Mamba-2): LinearKV(단일 블록)는 모든 셀렉터에서 정확한 합성 방식보다 성능이 월등히 뛰어납니다. 예를 들어, 20% 재계산 예산 하의 EPIC에서 LinearKV는 정확한 합성의 46.6% 대비 86.8%의 전체 품질을 달야냅니다.
  • OLMo & Qwen (GDN): 두 방식 모두 유사한 성능을 보이며, LinearKV는 정확한 합성의 성능과 일치합니다.
  • 강건성: Mamba-2에서 단일 블록 초기화제의 우수성은 다양한 재계산 비율(0.03 ~ 0.4)과 컨텍스트 길이(8K ~ 32K)에 걸쳐 유지됩니다.
• 지연 시간: LinearKV는 TTFT를 전체 프리필 시간의 0.46배~0.62배로 줄여, 측정된 모든 시나리오에서 정확한 합성 방식보다 우수한 성능을 보였습니다.

의의
본 논문은 하이브리드 LLM의 경우, 캐시된 청크로부터 정확한 전체 접두사 상태를 재구성하는 "대수적으로 원칙적인" 접근 방식이 계산 비용이 높을 뿐만 아니라 Mamba-2와 같은 특정 아키텍처에서 모델 성능에 해로울 수 있음을 확립합니다. 단일 캐시 상태가 견고한 초기자로서 충분하다는 것을 입증함으로써, LinearKV는 위치 독립적 캐싱의 이점을 극대화하면서도 지연 시간을 최소화하고 다중 소스 상태 합성에서 발생하는 오차 누적을 피할 수 있는 간단하고 훈련이 필요 없는 아키텍처 불가지론적(architecture-agnostic) 솔루션을 제공합니다. 이러한 발견은 순환 아키텍처에서 효과적인 캐싱을 위해 더 복잡한 상태 합성이 필수적이라는 가설에 도전합니다.