OCTOPUS:基于八面体参数化与最优平方误差量化的Transformer优化KV缓存
阅读原文· arxiv.org本文提出了OCTOPUS,一种用于Transformer长上下文推理的新型KV缓存压缩方法。它通过联合量化旋转后的坐标三元组来优化现有旋转编码器的性能。其核心在于将每个三元组的方向通过八面体参数化映射至正方形,并对量化坐标与三元组范数实施最优平方误差量化,从而实现了一种仅依赖于键总维度的严格非均匀比特分配策略。实验表明,OCTOPUS在文本、视频和音频任务上,在所有测试的比特宽度和指标上均达到或超越了此前的旋转编码器,并在极端压缩时优势尤为显著。此外,其融合Triton实现能在线重建键值,不引入额外的解码带宽或延迟。
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/