기술 요약: STAR-KV

문제 정의

100K 이상의 토큰 컨텍스트 윈도우를 지원하기 위한 대규모 언어 모델(LLM)의 급격한 확장은 추론 과정에서 심각한 시스템 수준의 병목 현상을 초래했으며, 이는 주로 Key-Value(KV) 캐시에 의해 발생합니다. KV 캐시는 재계산을 피하기 위해 이전에 처리된 토큰의 텐서를 저장하며, 이로 인해 메모리 점유율이 컨텍스트 길이에 따라 선형적으로 증가합니다. LLaMA-3.1-8B와 같이 128K 컨텍스트를 가진 시나리오에서는 KV 캐시가 전체 GPU 메모리의 최대 81%를 차지할 수 있으며, 이는 종종 모델 가중치에 필요한 저장 용량을 초 exceed(초과)합니다.

저차원 분해(low-rank decomposition)는 은닉 차원(hidden dimension)의 중복성을 활용하여 KV 캐시를 압축하는 유망한 방법으로 등장했지만, 기존 방식들은 두 가지 근본적인 문제에 직면해 있습니다:

1. 고정되거나 휴리스틱한 랭크 선택: 기존 방법들은 균일한 랭크를 사용하거나 오프라인 휴리스틱(예: Fisher 정보량)에 의존하며, 이는 공격적인 압축 상황에서 모델 품질을 유지하는 데 어려움을 겪습니다. 압축률을 적당한 수준(30~50%) 이상으로 높이면 일반적으로 상당한 정확도 저하가 발생합니다.
2. 재구성 오버헤드: 압축된 키(key) 및 값(value) 텐서를 추론 시점에 재구성하는 것은 추가적인 연산을 도입하여 지연 시간(latency)을 증가시킵니다. 또한, 표준 저차원 방식은 키와 값 투영(projection)의 서로 다른 민감도나 결과적인 잠재 표현(latent representation)을 양자화할 때 발생하는 특정 과제를 고려하지 못하는 경우가 많습니다.

방법론

저자들은 미세 조정된 적응형 랭크 제어와 효율적인 양자화를 달und하도록 설계된 프레임워크인 STAR-KV(Soft Thresholding for Adaptive Rank KV Cache Compression)를 제안합니다. 이 방법론은 세 가지 핵심 구성 요소로 이루어져 있습니다:

1. 소프트 임계값 처리를 통한 적응형 랭크 선택

고정된 랭크를 사용하는 대신, STAR-KV는 키 및 값 투영 행렬의 특이값(singular values)에 적용되는 학습 가능한 임계값을 사용하여 랭크 선택을 매개변수화합니다.

• 미분 가능한 임계값 처리: 그래디언트 기반 최적화를 가능하게 하기 위해, 저자들은 미분 불가능한 하드 임계값(hard threshold)을 이동된 쌍곡 탄젠트 함수($\tanh$)에 기반한 소프트 임계값 연산자로 대체했습니다. 이를 통해 모델은 각 디코더 블록과 어텐션 헤드에 대해 독립적으로 최적의 임계값($\alpha$)을 학습할 수 있습니다.
• 학습 목적 함수: 프레임워크는 지식 증류 손실(Knowledge Distillation Loss, 교사 모델의 출력 분포를 일치시키기 위함)과 적응형 압축 손실(Adaptive Compression Loss)로 구성된 결합 손실 함수를 최적화합니다. 압축 손실은 $\sum e^{-\alpha_i}$로 공식화되어, 훈련 초기에는 공격적인 압축을 장려하면서 임계값이 증가함에 따라 모델 성능이 안정화되도록 자연스럽게 유도합니다.
• 추론: 소량의 캘리브레이션 세트(<6 GPU 시간)로 훈련되면, 소프트 임계값은 오프라인에서 하드 임계값으로 변환되며, 추론 전 투영 행렬이 절단(truncate)됩니다.

2. 하이브리드 분해 전략

본 논문은 결합 분해(Joint Decomposition, JD)(연결된 헤드를 인수 분해)와 헤드별 분해(Head-wise Decomposition, HD)(헤드를 독립적으로 인수 분해) 사이의 트레이드오프를 분석합니다.

• 민감도 분석: 특이값 분석 결과, 동일한 랭크 제약 조건 하에서 값(value) 투영이 키(key) 투영보다 일반적으로 더 높은 근사 오차를 보이는 것으로 나타났습니다. 반대로, JD는 HD에 비해 낮은 근사 오차를 제공하지만, $h \times$ (헤드 수)만큼 높은 재구성 계산 비용을 발생시킵니다.
• 하이브리드 체계: STAR-KV는 다음과 같은 하이브리드 접근 방식을 채택합니다:
  • 값 투영(Value Projections): 값은 압축 오차에 더 민감하므로 충실도를 보존하기 위해 **결합 분해(JD)**를 사용합니다.
  • 키 투영(Key Projections): 키는 상대적으로 덜 민감하므로 재구성 오버헤드를 최소화하기 위해 **헤드별 분해(HD)**를 사용합니다.
  • 최적화: 값에 대한 JD의 오버헤드를 보상하기 위해, 프레임워크는 재구성 행렬을 출력 투영에 융합하도록 연산 순서를 재정렬하는 가중치 흡수(weight absorption) 기법을 활용하여 FLOPs를 크게 줄입니다.

3. 저차원 인지 혼합 정밀도 양자화

잠재 표현을 더욱 압축하기 위해, STAR-KV는 저차원 인수의 정렬된 특이값들로부터 발생하는 아웃라이어 채널 문제를 해결합니다.

• 블록 단위 하다마르(Block-wise Hadamard): 전역적인 하다마르 변환을 적용하는 대신, 이 방법은 잠재 차원을 "아웃라이어"(높은 크기를 가진 상위 20% 채널)와 "인라이어"(나머지 80%) 블록으로 분할합니다.
• 혼합 정밀도: 각 블록에 독립적인 하다마르 변환을 적용한 후 혼합 정밀도 양자화를 수행합니다. 아웃라이어 블록은 높은 정밀도(예: 4-bit)로 양자화하고, 인라이어 블크는 낮은 정밀도(예: 3-bit)로 양자화합니다.
• 제로 오버헤드: 하다마르 인수는 분해된 투영 행렬에 융합되어 추가적인 런타임 비용을 발생시키지 않습니다.

주요 결과

STAR-KV는 여러 LLM 제품군(LongChat, LLaMA-2, LLaMA-3)과 벤치마크(상식 추론, LongBench, RULER)를 통해 평가되었습니다.

• 압축 대비 정확도: STAR-KV는 압축되지 않은 모델과 비교하여 무시할 만한 수준의 정확도 저하만으로 최대 75%의 KV 캐시 압축을 달성합니다. LongChat-7B-v1.5에서 60% 압축 시 베이스라인과 일치하는 정확도를 달성하여, Palu 및 ReCalKV와 같은 최첨단 방식들을 능가했습니다.
• 결합 압축: 제안된 3.2-bit 혼합 정밀도 양자화와 결합했을 때, STAR-KV는 최대 20배의 전체 KV 캐시 감소를 달와합니다.
• Long-Context 성능: LongBench 및 RULER(4K 시퀀스 길이)에서 STAR-KV는 50%에서 크게 저하되고 경쟁력을 유지하기 위해 더 낮은 압축률(30%)이 필요한 Palu와 비교하여, 높은 압축률(60%)에서도 우수한 정확도를 유지합니다.
• 속도 향상: 커스텀 Triton 기반 GPU 커널로 구현된 STAR-KV는 다음을 제공합니다:
  • PyTorch SDPA 대비 어텐션 모듈에서 최대 6.9배의 속도 향상.
  • 롱 컨텍스트(128K) 또는 고배치 시나리오에서 최대 3.1배 높은 엔드 투 엔드 생성 처리량.

의의 및 주장

본 논문은 STAR-KV가 KV 캐시 관리의 정확도-압축 트레이드오프에 있어 새로운 파레토 최적선(Pareto front)을 구축한다고 주장합니다. 그 의의는 다음과 같습니다:

1. 적응형 입도(Granularity): 정적 또는 휴리스틱한 랭크 선택을 넘어, 특정 레이어와 헤드의 민감도에 적응하는 학습 가능한 미세 제어 메커니즘으로 전환했습니다.
2. 효율성: 하이브리드 분해 전략과 가중치 흡수를 통해 근사 오차와 재구성 오버헤드 사이의 균형을 성공적으로 맞추었습니다.
3. 시너지: 저차원 분해와 양자화가 블록 단위 하다마르 변환을 통해 시너지 효과를 낼 수 있음을 입증하여, 성능 저하 없이 극단적인 압축(20배)을 달성했습니다.
4. 실용성: 구조적 변경이나 광범한 재학습이 필요 없는 사후 학습(post-training) 솔루션을 제공하여, 최소한의 캘리브레이션 비용으로 기존 LLM에 적용 가능하게 만들었습니다.

저자들은 자신들의 접근 방식이 (최대 75% 저차원, 20배 총 압축) 공격적인 압축을 가능하게 하는 동시에, 압축되지 않은 베이스라인 및 기존 압축 방법들의 정확도를 유지하거나 오히려 개선함으로써 롱 컨텍스트 추론의 메모리 및 대역폭 병목 현상을 효과적으로 해결한다고 강조합니다.