OCTOPUS: 최적의 제곱 오차 양자화 하에서 팔면체 매개변수화를 통한 Transformer용 최적화된 KV Cache
요약
OCTOPUS는 Transformer의 KV Cache 메모리 대역폭과 점유 공간 문제를 해결하기 위해 팔면체 매개변수화를 활용한 새로운 양자화 코덱을 제안합니다. 회전된 좌표 삼중항을 공동 양자화하여 제곱 오차를 최적화함으로써, 기존의 회전 기반 코덱보다 뛰어난 압축 성능을 보여줍니다. 특히 Triton을 이용한 구현을 통해 추가적인 지연 시간 없이 실시간으로 키를 재구성할 수 있는 효율성을 갖추었습니다.
핵심 포인트
- 팔면체 매개변수화를 통해 회전된 좌표 삼중항을 정사각형으로 매핑하여 공동 양자화 수행
- Lloyd-Max 양자화를 적용하여 구현 환경에 최적화된 주변 분포 대응
- 데이터에 무관하고(data-oblivious) 온라인 방식이며 결정론적인 특성 보유
- 극한의 저비트 압축 환경에서 기존 회전 코덱 대비 우수한 성능 입증
- Triton 기반 구현으로 디코딩 시간의 지연(latency) 없이 즉석 키 재구성 가능
Key-Value (KV) 캐시는 긴 문맥의 자기회귀 추론 (autoregressive inference)에서 메모리 대역폭과 점유 공간 (footprint)을 지배합니다. 최근의 회전 전처리 코덱 (rotation-preconditioned codecs) (TurboQuant, PolarQuant)은 구조화된 무작위 회전 (structured random rotation) 후에 분석적으로 다룰 수 있는 주변 분포 (marginal)에 맞춘 좌표별 스칼라 양자화기 (per-coordinate scalar quantizer)를 사용하는 것이 KV 압축을 위한 최적에 가까운 방식임을 보여주었습니다. OCTOPUS는 회전된 좌표 삼중항 (coordinate triplets)의 공동 양자화 (joint quantization)를 통해 이 패러다임을 발전시킵니다. 각 삼중항의 방향은 팔면체 매개변수화 (octahedral parameterization)를 통해 정사각형으로 매핑되며, 결과로 나온 두 좌표와 삼중항 노름 (triplet norm)은 구현에 맞춘 주변 분포에 대해 Lloyd-Max 양자화됩니다. 삼중항별 제곱 오차 (squared error)를 최적화하면 키 (keys)의 총 차원 수에만 의존하는 엄격한 비균등 비트 할당 (non-uniform bit allocation)을 얻을 수 있습니다. 스윕 (sweeps)을 통해 유한 차원 품질 최적값을 찾아낸 결과, 우리가 테스트한 모든 실제 디코더에서 일정하게 나타남을 확인했습니다. 이 코덱은 데이터에 무관하며 (data-oblivious), 온라인 방식이고, 시드 (seed)가 주어지면 결정론적 (deterministic)입니다. 텍스트, 비디오, 오디오 전반에 걸쳐 OCTOPUS는 보고된 모든 비트 너비와 지표에서 이전의 모든 회전 코덱과 대등하거나 이를 능가하며, 극한의 압축을 위해 비트 수가 낮아질수록 그 격차는 더욱 커집니다. 또한, 융합된 Triton 구현은 압축되지 않은 키를 실체화하지 않고 즉석에서 키를 재구성하므로, 이 코덱은 기존의 역양자화 (dequantization) 대비 디코딩 시간의 대역폭이나 지연 시간 (latency)을 추가하지 않습니다. 프로젝트 페이지: https://octopus-quant.github.io/
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.AI의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기