AdapShot: 의미 정보 기반 KV 캐시 재사용을 통한 적응형 다중-샷 인-컨텍스트 학습

AdapShot: Adaptive Many-Shot In-Context Learning with Semantic-Aware KV Cache Reuse

다중-샷 인-컨텍스트 학습(ICL)은 광범위한 예시를 활용하여 대규모 언어 모델(LLM)의 추론 능력을 향상시키는 유망한 방법론으로 부상했습니다. 그러나 기존 방법은 일반적으로 미리 정해진 고정된 수의 샷에 의존합니다. 이러한 정적인 접근 방식은 다양한 쿼리의 난이도 변화에 적응하지 못하여, 때로는 충분한 컨텍스트를 제공하지 못하거나 노이즈로 인해 성능 저하를 일으킬 수 있습니다. 또한, 긴 컨텍스트의 막대한 계산 및 메모리 비용은 다중-샷 학습의 실용성을 심각하게 제한합니다. 이러한 제한 사항을 해결하기 위해, 우리는 샷 수를 동적으로 최적화하고 효율적인 추론을 위해 KV 캐시 재사용을 활용하는 AdapShot을 제안합니다. 구체적으로, 우리는 출력 엔트로피를 활용하여 최적의 샷 수를 결정하는 탐색 기반 평가 메커니즘을 설계했습니다. 탐색 및 추론 단계 모두에서 불필요한 사전 채우기 계산을 방지하기 위해, 의미 정보를 고려한 KV 캐시 재사용 전략을 통합했습니다. 이 재사용 전략 내에서, 위치 인코딩 불일치 문제를 해결하기 위해, 캐시된 키-값 쌍의 유연한 재정렬을 가능하게 하는 분리 및 재인코딩 방법을 도입했습니다. 광범위한 실험 결과, AdapShot은 최첨단 DBSA에 비해 평균적으로 약 10%의 성능 향상과 4.64배의 속도 향상을 달성하는 것으로 나타났습니다.

Many-Shot In-Context Learning (ICL) has emerged as a promising paradigm, leveraging extensive examples to unlock the reasoning potential of Large Language Models (LLMs). However, existing methods typically rely on a predetermined, fixed number of shots. This static approach often fails to adapt to the varying difficulty of different queries, leading to either insufficient context or interference from noise. Furthermore, the prohibitive computational and memory costs of long contexts severely limit Many-Shot's feasibility. To address the above limitations, we propose AdapShot, which dynamically optimizes shot counts and leverages KV cache reuse for efficient inference. Specifically, we design a probe-based evaluation mechanism that utilizes output entropy to determine the optimal number of shots. To bypass the redundant prefilling computation during both the probing and inference phases, we incorporate a semantics-aware KV cache reuse strategy. Within this reuse strategy, to address positional encoding incompatibilities, we introduce a decoupling and re-encoding method that enables the flexible reordering of cached key-value pairs. Extensive experiments demonstrate that AdapShot achieves an average performance gain of around 10% and a 4.64x speedup compared to state-of-the-art DBSA.