기술 요약: S4R - 압축된 롱 컨텍스트 KV 캐싱을 위한 선택적 샘플링, 부공간 및 희소 재구성

1. 문제 정의

수십만 또는 수백만 토큰에 달하는 대규모 언어 모델(LLM)의 컨텍스트 윈도우 확장은 심각한 메모리 병목 현상을 야기했습니다. 자기회귀적 생성(autoregressive generation)을 위해 중간 계산을 저장하는 데 필요한 키-값(Key-Value, KV) 캐시는 컨텍스트 길이에 따라 선형적으로 증가하며, 종종 모델 파라미터 자체의 메모리 점유율을 초과하기도 합니다.

기존의 KV 캐시 압축 솔루션은 효율성과 적응성 사이에서 근본적인 트레이드오프에 직면해 있습니다:

• 오프라인/고정 방식 (Offline/Fixed Methods): Palu, LoRC, Eigen Attention과 같은 방식은 추론 전 캘리브레이션 데이터나 가중치에 대해 저차원 분해(예: SVD)를 수행합니다. 디코딩 시에는 효율적이지만, 분포 변화(distribution shifts)에 민감하며 특정 프롬프트와 정렬되지 않는 고정된 압축 설계를 사용합니다.
• 온라인/프롬프트 의존 방식 (Online/Prompt-Dependent Methods): ShadowKV와 xKV는 추론 중 현재 프롬프트로부터 저차원 표현을 추정합니다. 컨텍스트 변화에는 더 강건하지만, 전체 프롬프트를 분해하는 과정이 매우 비싼 프리필링(prefilling) 단계에서 과도한 계산 비용을 발생시키며, 디코딩 중에 전체 캐시를 재구성해야 하는 경우가 많아 처리량(throughput)을 저해합니다.

본 논문은 다음과 같은 간극을 식별합니다: 고정 압축의 낮은 프리필링 오버헤드와 프롬프트 의존 방식의 컨텍스트 적응성을 결합하되, 전체 프롬프트 분해나 전체 캐시 재구성 비용을 발생시키지 않는 방법론입니다.

2. 방법론: S4R

S4R(Selective Sampling, Subspaces, and Sparse Reconstruction)은 **프리필링(Prefilling)**과 **디코딩(Decoding)**이라는 두 단계 과정을 통해 이 딜레마를 해결합니다.

2.1 프리필링 단계: 선택적 부공간 구축

S4R은 전체 롱 컨텍스트 시퀀스에 대해 SVD(Singular Value Decomposition)를 적용하는 대신, 선택적으로 샘플링된 토큰의 부분 집합으로부터 저차원 부공간을 구축합니다.

• 싱크 토큰 보존 (Sink Token Preservation): 글로벌 피벗 역할을 하는 "싱크 토큰(sink tokens)"으로 식별된 처음 $s$개의 토큰은 압축에서 제외되어, 장거리 의동성(long-range dependencies)을 안정적으로 유지하기 위해 풀 정밀도(full precision)로 저장됩니다.
• 하이브리드 샘플링 (Hybrid Sampling): 나머지 비-싱크 토큰들에 대해, S4R은 다음의 하이브리드 전략을 사용하여 대표성 있는 부분 집합 $X[I]$를 샘플링합니다:
  1. 초기 위치에서의 균등 샘플링(Uniform sampling).
  2. 가장 최근 토큰들의 연속된 블록(Contiguous block).
• 정규화 및 SVD (Normalization & SVD): 고진폭 벡터가 분해를 지배하는 것을 방지하기 위해 샘플링된 키와 값을 특징 차원(feature dimension)을 따라 정규화합니다. 그 후, 이 부분 집합에 대해 절단된 SVD(Truncated SVD)를 수행하여 저차원 기저($B_K, B_V$)를 도출합니다.
• 계수 저장 (Coefficient Storage): 비-싱크 KV 상태는 학습된 기저 위로 투영되며, 전체 차원의 벡터 대신 저차원 계수($C_K, C_V$)만을 저장합니다.
• 쿼리 투영 (Query Projection): 로컬 쿼리 또한 동일한 저차원 부공간으로 투영되어, 이후 효율적인 관련성 추정을 가능하게 합니다.

2.2 디코딩 단계: 희소 재구성

생성 과정에서 S4R은 전체 캐시를 재구성하는 것을 피합니다. 대신, 관련성이 있다고 판단되는 특정 토큰들만을 희소하게 재구성합니다.

• 잠재 공간 관련성 추정 (Latent-Space Relevance Estimation): 토큰의 관련성은 잠재 계수 공간에서 직접 추정됩니다. 투영된 쿼리 계수는 캐시된 키 계수와 곱해져, 전체 재구성 없이도 어텐션 로짓(attention logits)을 계산합니다.
• 희소 선택 (Sparse Selection): 다음 기준에 따라 재구성할 토큰의 부분 집합을 선택합니다:
  1. 로컬 윈도우 (Local Window): 가장 최근의 $W$개 토큰은 항상 유지됩니다.
  2. 글로벌 중요도 (Global Importance): 잠재 공간에서 도출된 평균 관련성 점수를 기반으로, 나머지 히스토리 중 상위 $k$개의 토큰을 선택합니다.
• 희소 재구성 (Sparse Reconstruction): 선택된 로컬 및 글로벌 토큰들만이 학습된 기저를 사용하여 계수로부터 재구성됩니다. 이들은 최종 어텐션 계산을 수행하기 위해 풀 정밀도 싱크 토کن들과 결합됩니다.

3. 주요 기여

1. 프롬프트 인식 저차원 압축 (Prompt-Aware Low-Rank Compression): S4R은 외부 캘리브레이션 데이터가 아닌, 선택적으로 샘플링된 프롬프트 토큰으로부터 키/값 기저를 구축함으로써, 약간의 프리필링 적응성을 대가로 더 나은 컨텍스트 정렬을 실현합니다.
2. 희소 재구성 전략 (Sparse Reconstruction Strategy): 디코딩 중에 정보가 풍부한 위치(로컬 윈도우 + 글로벌 중요 토큰)만을 유지함으로써, S4R은 전체 캐시 재구성의 계산 병목 현상을 피하고 처리량을 크게 향상시킵니다.
3. 싱크 토큰 보존 (Sink Token Preservation): 결정적인 글로벌 어텐션 피벗인 초기 몇 개의 "싱크" 토큰을 풀 정밀도로 명시적으로 보존하여, 핵심적인 정보를 압축할 때 발생하는 불안정성 문제를 해결합니다.
4. 효율성-정확도 트레이드오프 (Efficiency-Accuracy Trade-off): S4R 설계는 고정 방식의 낮은 프리필링 오버헤드와 프롬프트 의존 방식의 적응성을 성공적으로 결합합니다.

4. 실험 결과

저자들은 Llama (1B, 8B) 및 Qwen (4B, 14B) 모델 제품군을 사용하여 LongBench와 RULER 벤치마크에서 S4R을 평가했습니다.

• 정확도 (Accuracy):
  • LongBench: S4R은 풀 캐시 베이스라인에 근접한 정확도를 유지하면서 최대 5배의 KV 캐시 압축을 달 achieves 합니다. Qwen3-4B의 경우, 평균 점수 47.79를 기록하여(풀 캐시의 47.79와 동일), Palu와 같은 고정 방식이나 ShadowKV와 같은 경쟁력 있는 프롬프트 의존 방식보다 우수한 성능을 보였습니다.
  • RULER: 128k 컨텍스트에서 2.5배 압축 비율을 적용했을 때, S4R은 평균 점수 93.61을 달성하여 풀 KV 캐시(93.97)에 거의 근접했으며, xKV(93.47)를 약간 상회했습니다. 특히 멀티 키(multi-key) 및 멀티 쿼리(multi-query) 검색 작업에서 강력한 성능을 보였습니다.
• 효율성 (Efficiency):
  • S4R은 프롬프트 의존 방식의 오버헤드를 실질적으로 줄여줍니다. xKV와 비교했을 때, S4R은 첫 토큰 생성 시간(TTFT)을 79.90초에서 27.65초로, 512-토큰 생성 지연 시간을 260.76초에서 61.54초로 단축했습니다.
  • 최적화되지 않은 구현에서는 Palu와 같은 고정 방식보다 약간 느릴 수 있으나, 공격적인 압축 비율 하에서 S4R은 더 우수한 정확도-효율성 트레이드오프를 제공합니다.

5. 의의 및 주장

본 논문은 S4R이 롱 컨텍스트 LLM 추론을 위한 적응형 KV 캐시 압축의 실용적인 경로를 제공한다고 주장합니다. 주요 의의는 기존 저차원 방식들의 "딜레마"를 해결했다는 점에 있습니다:

• 오프라인 방식의 캘리브레이션 데이터 의존성을 피합니다.
• 온라인 방식의 과도한 프리필링 및 재구성 비용을 피합니다.

신중하게 선택된 토큰의 부분 집합으로부터 지배적인 KV 부공간을 근사할 수 있으며, 매 디코딩 단계마다 전체 재구성이 필요하지 않다는 것을 입증함으로써, S4R은 현저히 감소된 지속 메모리와 관리 가능한 계산 오버헤드로 고충실도(high-fidelity) 롱 컨텍스트 추론을 가능하게 합니다. 저자들은 현재 구현이 시스템 최적화가 이루어지지 않았음을 언급하며, 보고된 효율성 지표가 알고리즘 오버헤로를 반영하고 있음을 시사하였고, 커널 수준의 최적화를 통해 추가적인 이득이 가능함을 밝혔습니다.