OCTOPUS : Cache KV optimisé pour Transformers via une paramétrisation octaédrique sous quantification optimale de l'erreur quadratique

Résumé

La mémoire cache clé-valeur (KV) domine la largeur de bande et l'empreinte mémoire dans l'inférence autorégressive à contexte long. Les codecs récents préconditionnés par rotation (TurboQuant, PolarQuant) montrent qu'une rotation aléatoire structurée suivie d'un quantificateur scalaire par coordonnée adapté à une marginale analytiquement tractable constitue une recette quasi-optimale pour la compression de la cache KV. OCTOPUS fait progresser ce paradigme grâce à la quantification conjointe de triplets de coordonnées rotatées. La direction de chaque triplet est projetée sur un carré via une paramétrisation octaédrique, et les deux coordonnées résultantes ainsi que la norme du triplet sont quantifiées selon Lloyd-Max face à des marginales adaptées à l'implémentation. L'optimisation de l'erreur quadratique par triplet conduit à une allocation de bits strictement non uniforme, dépendant uniquement de la dimensionnalité totale des clés. Nous constatons, par balayages, que l'optimum de qualité en dimension finie est constant sur tous les décodeurs réels que nous testons. Le codec est indépendant des données, en ligne et déterministe étant donné une graine. Pour le texte, la vidéo et l'audio, OCTOPUS égale ou surpasse tous les codecs à rotation antérieurs à chaque largeur de bit et métrique rapportée, avec un écart qui croît lorsque le nombre de bits diminue pour une compression extrême. De plus, une implémentation Triton fusionnée reconstruit les clés à la volée sans matérialiser la clé non compressée, de sorte que le codec n'ajoute aucune bande passante ni latence de décodage par rapport à la déquantification existante. Page du projet : 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/