OCTOPUS: Optimized KV Cache for Transformers via Octahedral Parametrization Under optimal Squared error quantization

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

긴 컨텍스트에서 KV 캐시의 메모리 대역폭이 추론 속도와 배치 규모를 좌우한다. 기존 회전 기반 양자화가 좌표 단위로 근사하던 반면, OCTOPUS는 좌표 triplet를 공동 양자화하고 방향과 노름을 분리해 비트를 비대칭적으로 배치함으로써 동일 비트에서 재구성 오차를 낮춘다. 텍스트/비디오/오디오 전 영역에서 다른 회전 기반 코덱들을 상회하거나 동률을 달성하고, decode 시 풀 키를 구성하지 않고도 키를 재구성한다.

왜 중요한가

긴 컨텍스트에서 KV 캐시의 메모리 대역폭이 추론 속도와 배치 규모를 좌우한다. 기존 회전 기반 양자화가 좌표 단위로 근사하던 반면, OCTOPUS는 좌표 triplet를 공동 양자화하고 방향과 노름을 분리해 비트를 비대칭적으로 배치함으로써 동일 비트에서 재구성 오차를 낮춘다. 텍스트/비디오/오디오 전 영역에서 다른 회전 기반 코덱들을 상회하거나 동률을 달성하고, decode 시 풀 키를 구성하지 않고도 키를 재구성한다.

핵심 기여

Octahedral triplet direction quantizer with norm marginals

rotated 좌표를 3개 단위 트리플로 묶어 방향 ni를 S^2에서 ξ, η의 2D 좌표로 매핑하고, 각 트리플의 ρi 노름과 ni 방향을 함께 양자화하는 KV 캐시 프리미티브를 제안한다. 코드북은 marginals에 맞춰 설계되며, 디코더에서 K̂를 직접 materialize하지 않고 키를 레지스트리에서 재구성한다.

MSE-optimal non-uniform bit allocation via Lagrangian optimization

각 트리플에 총 예산 Btri = 2bdir + bnrm을 배분하고, 방향/노름 비트를 (b+1, b−1)로 구성하는 비대칭 분할이 최적임을 Finite-dimensional에서의 Z 차원에서 보이는 라그랑지안 해를 통해 입증한다.

Optional 1-bit QJL residual for unbiased dot-product

QJL 잔차를 추가로 스케치하여 rotated-frame 잔차의 점곱 추정치를 편향 없이 추정한다. 잔차의 부호를 independent한 R′로 회전시켜 σ를 저장하고, q̂^T k̂에 보정항을 더한다.

Generalization beyond LLMs with fused KV pipeline

이 코덱은 LLM에 한정되지 않고 임의의 attention을 갖는 autoregressive 트랜스포머의 KV 캐시에 적용 가능하며, 3×3 연산의 공동 양자화와 fused Triton kernels로 end-to-end KV 파이프라인에서 decode를 가속한다.

핵심 아이디어 이해하기

출발점: KV 캐시의 비트-대역폭 문제는 긴 컨텍스트에서 심각한 병목이다. 기존 rotation-based 코덱은 좌표별 Lloyd-Max 양자화를 통해 marginals를 근사하지만, 실제로는 좌표 간 상관관계가 남아 있다. OCTOPUS은 u를 triplets로 분할하고, 각 triplet의 방향을 octahedral map으로 2D로 축약하되 노름 ρi를 별도로 양자화한다. Marginals fξ(ξ)와 fρ(ρ)로 정의된 비트 배분 문제를 Lagrangian으로 풀고, (b+1, b−1) 비트 분할이 finite-dimensional에서 최적임을 보인다. 디코딩은 K̂를 구현 레지스트리에서 재구성하고, QJL 잔차를 선택적으로 추가해 편향 없이 dot-product를 추정할 수 있다. 이 접근은 텍스트/비디오/오디오에서 기존 rotation-based 코덱과 비교해 낮은 MSE를 달성하며, 비트가 더 떨어질수록 우수한 이점을 보인다.

방법론

Raw key k를 길이 d에서 γ = ∥k∥2와 ũ = k/γ로 분해한다. 2) sign-flipped Walsh‑Hadamard rotation R = H diag(s)로 precondition한다. 3) u를 ntri = ⌈d/3⌉ triplets ti로 분할하고 각 ti에서 ρi와 ni를 산출한다. 4) ni를 Oct(ni)로 맵핑해 ξ, η로 표현하고 ρi의 Marginal을 계산한다. 5) Cξ(bdir)와 Cρ(bnrm) 코드북으로 ξ, η, ρ를 양자화하고, 3×3 joint rounding으로 최적 근사치를 선택한다. 6) 압축 상태 S(k) = (γ, Idir, Inrm)을 저장한다. 7) Decode 시마다 ti를 다시 재구성하고 q = sum_i ρ̂i qrot, in̂i를 이용한 점곱을 계산한다. 8) Optional OCTOPUS-QJL로 잔차 잔여를 스케치하여 unbiased dot-product를 얻는다.

주요 결과

Synthetic 데이터(d=128)에서 OCTOPUS은 2-bit에서 MSE 0.0897, cos 0.9547, IP 2.682를 달성하고, 3-bit에서 cos 0.9871, MSE 0.0260, IP 1.444, 4-bit에서 cos 0.9965, MSE 0.0071, IP 0.753를 보인다. OCTOPUS-QJL은 2-bit에서 IP 오차 2.015로 TurboQuant-QJL 대비 크게 개선된다. WikiText-2/C4 perplexity(Table 2)에서 b=4일 때 OCTOPUS은 WikiText-2 10.306 (+2.7%), C4 12.896 (+1.5%)를 기록하며, b=2일 때는 13.517 (+34.7%), 17.976 (+41.5%)로 비슷한 추세를 보인다. Needle-in-a-haystack 실험에서 b=4에서 모든 코덱이 recall 1.00, b=3에서 OCTOPUS 1.00, PolarQuant 0.86, b=2에서 OCTOPUS 0.81, OCTOPUS-QJL 0.83, PolarQuant 0.04, TurboQuant-QJL 0.01로 나타난다. Autoregressive 비디오(video)와 오디오(audio)에서는 b=4에서 LPIPS, PSNR, SSIM, CLIP 등의 지표에서 OCTOPUS이 대체로 baseline과 근접하거나 개선되며, b=2에서는 TurboQuant-MSE/QJL 및 PolarQuant 대비 OCTOPUS의 품질 저하가 억제된다. encode/decode 시간은 비트 폭이 낮을수록 더 큰 절감 효과를 보이며, OCTOPUS의 decode 커널은 0.49–0.66 ms 수준으로 측정된다. KV 캐시의 메모리-대-성능 Pareto_frontier에서 OCTOPUS은 Pareto 경계의 상단에 위치하며, OCTOPUS-QJL은 residual 비용으로 약 60 MB의 추가 비용이 든다.

기술 상세

Algorithm 1(인코더)과 Algorithm 2(디코더)로 구성되며, k를 γ, ũ로 분해하고 R = H diag(s)로 회전한 뒤, d 차원의 벡터를 3차원 triplet로 묶어 ρi와 ni로 분해한다. ni는 Oct−1으로 역변환되어 ξ, η로 맵핑되며, ρi는 코드북 Cρ으로 양자화된다. 각 triplet에 대해 3×3 조합으로 조정하는 joint_rounding으로 ðξi, ηi, ρi를 선택한다. 이때 최적의 bdir, bnrm 분할은 Eq. 10의 Lagrangian 해에 의해 (b+1, b−1)로 결정된다. QJL 잔차를 활용하면 q̂의 편향을 제거하지만, 1비트 잔차 비용이 추가된다.

실무 활용

Long-context KV 캐시를 사용하는 LLM, 멀티모달 생성 모델, 비디오/오디오 생성 파이프라인에서 KV 대역폭과 메모리 footprint를 감소시키고 품질을 유지한다.

LLM의 수십~수천 토큰 컨텍스트에서 KV 캐시의 대역폭 절감
비디오 디퓨전/실시간 영상 생성에서 KV 재구성 대역폭 감소
다중 모달 모델의 KV 캐시 공유 및 가속
끝단 엔진에서 QJL 잔차를 활용한 점곱 추정

코드 공개 여부: 미확인

키워드

KV-cacherotation-preconditioned codecsTurboQuantPolarQuantoctahedral mapTritonQJLLloyd-Max

용어 해설

Walsh-Hadamard transform: — 고속 직교 전처리로, 고차원 벡터를 빠르게 곱하는 구조적 변환이다. OCTOPUS에서 방향 벡터 u의 marginals를 거의 등방성으로 만들고 구현 효율을 높이며, O(d log d) 연산으로 수행된다.
octahedral map: — 단위 벡터를 equal-area 맵으로 S^2를 [−1,1]^2로 매핑하는 기하적 파라메트라이제이션이다. 3좌표 벡터의 방향을 두 scalars로 표현하고, 역방향은 재구성에 사용된다.
Lloyd-Max quantization: — 주어진 1-D 확률분포에 대해 평균제곱오차를 최소화하도록 코드북 centroids를 교대적으로 업데이트하는 최적 양자화 알고리즘이다. 각 좌표의 양자화에 사용된다.
Walsh-Hadamard rotation: — 무작위 부호 반전이 곁들여진 Walsh-Hadamard 행렬로, 방향 벡터의 marginals를 거의 등방성으로 만들기 위한 회전 기법이다. OCTOPUS에서 방향 벡터의 좌표 marginals를 안정적으로 양자화하기 위해 사용된다.