OCTOPUS: Caché KV optimizado para Transformers mediante Parametrización Octaédrica bajo cuantificación de error cuadrático óptimo

Resumen

La caché clave-valor (KV) domina el ancho de banda y la huella de memoria en la inferencia autorregresiva de contexto largo. Los códecs recientes con precondicionamiento por rotación (TurboQuant, PolarQuant) muestran que una rotación aleatoria estructurada seguida de un cuantizador escalar por coordenada ajustado a una marginal analíticamente tratable constituye una receta casi óptima para la compresión de KV. OCTOPUS avanza este paradigma mediante la cuantización conjunta de tripletes de coordenadas rotadas. La dirección de cada triplete se mapea a un cuadrado mediante una parametrización octaédrica, y las dos coordenadas resultantes junto con la norma del triplete se cuantizan según Lloyd-Max frente a marginales ajustadas a la implementación. La optimización del error cuadrático por triplete produce una asignación de bits estrictamente no uniforme que depende únicamente de la dimensionalidad total de las claves. Encontramos que el óptimo de calidad en dimensión finita, determinado mediante barridos, es constante en todos los decodificadores reales que evaluamos. El códec es ciego a los datos, en línea y determinista dada una semilla. En texto, video y audio, OCTOPUS iguala o supera a todos los códecs con rotación previos en cada ancho de bit y métrica reportados, con una ventaja que crece a medida que los bits disminuyen para una compresión extrema. Además, una implementación fusionada en Triton reconstruye las claves sobre la marcha sin materializar la clave sin comprimir, por lo que el códec no añade ancho de banda ni latencia en la decodificación por encima de la des cuantización existente. Página del proyecto: 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/