OCTOPUS: Cache KV Otimizado para Transformers via Parametrização Octaédrica sob Quantização de Erro Quadrático Ótimo

Resumo

A cache chave-valor (KV) domina a largura de banda e o footprint de memória na inferência autoregressiva de contexto longo. Codecs recentes pré-condicionados por rotação (TurboQuant, PolarQuant) mostram que uma rotação aleatória estruturada seguida por um quantizador escalar por coordenada compatível com uma marginal analiticamente tratável é uma receita quase ótima para compressão KV. O OCTOPUS avança esse paradigma por meio da quantização conjunta de tripletos de coordenadas rotacionadas. A direção de cada tripleto é mapeada para um quadrado via uma parametrização octaédrica, e as duas coordenadas resultantes e a norma do tripleto são quantizadas por Lloyd-Max contra marginais compatíveis com a implementação. A otimização do erro quadrático por tripleto fornece uma alocação de bits estritamente não uniforme que depende apenas da dimensionalidade total das chaves. Descobrimos que o ótimo de qualidade em dimensão finita com varreduras é constante em todos os decodificadores reais que testamos. O codec é independente dos dados, online e determinístico dada uma semente. Em texto, vídeo e áudio, o OCTOPUS iguala ou supera todos os codecs de rotação anteriores em todas as larguras de bits e métricas relatadas, com uma liderança que cresce à medida que os bits diminuem para compressão extrema. Além disso, uma implementação Triton fundida reconstrói chaves em tempo real sem materializar a chave descomprimida, de modo que o codec não adiciona largura de banda ou latência no momento da decodificação além da dequantização existente. Página do projeto: https://octopus-quant.github.io/

English

The key-value (KV) cache dominates memory bandwidth and footprint in long-context autoregressive inference. Recent rotation-preconditioned codecs (TurboQuant, PolarQuant) show that a structured random rotation followed by a per-coordinate scalar quantizer matched to an analytically tractable marginal is a near-optimal recipe for KV compression. OCTOPUS advances this paradigm through joint quantization of rotated coordinate triplets. Each triplet's direction is mapped to a square via an octahedral parameterization, and the two resulting coordinates and the triplet norm are Lloyd-Max quantized against implementation-matched marginals. Optimizing the per-triplet squared error gives a strictly non-uniform bit allocation depending only on the total dimensionality of the keys. We find the finite-dimensional quality optimum with sweeps to be constant on every real decoder we test. The codec is data-oblivious, online, and deterministic given a seed. Across text, video, and audio, OCTOPUS matches or beats every prior rotation codec at every reported bit width and metric, with a lead that grows as bits drop for extreme compression. Furthermore, a fused Triton implementation reconstructs keys on the fly without materializing the uncompressed key, so the codec adds no decode-time bandwidth or latency over the existing dequantization. Project Page: https://octopus-quant.github.io/