OCTOPUS: Оптимизированный KV-кэш для трансформеров посредством октаэдрической параметризации при оптимальном квантовании с квадратичной ошибкой

Аннотация

Кэш ключ-значение (KV) доминирует в пропускной способности и объеме памяти при длинноконтекстном авторегрессивном инференсе. Недавние кодеки с предварительным вращением (TurboQuant, PolarQuant) показывают, что структурированное случайное вращение с последующим покоординатным скалярным квантователем, согласованным с аналитически вычислимым маргинальным распределением, является почти оптимальным рецептом для сжатия KV. OCTOPUS развивает эту парадигму с помощью совместного квантования повернутых троек координат. Направление каждой тройки отображается на квадрат с помощью октаэдрической параметризации, а два полученных координата и норма тройки квантуются по Ллойду–Максу с учетом маргинальных распределений, согласованных с реализацией. Оптимизация среднеквадратичной ошибки на тройку дает строго неравномерное распределение битов, зависящее только от общей размерности ключей. Мы обнаруживаем, что оптимум качества в конечной размерности при переборе параметров является постоянным на каждом реальном декодере, который мы тестируем. Кодек является независимым от данных, онлайн и детерминированным при заданном начальном значении. Для текста, видео и аудио OCTOPUS соответствует или превосходит все предыдущие кодеки с вращением при каждой заявленной разрядности и метрике, причем отрыв увеличивается по мере уменьшения количества битов для экстремального сжатия. Кроме того, объединенная реализация Triton восстанавливает ключи на лету без материализации несжатого ключа, поэтому кодек не добавляет пропускной способности или задержки при декодировании по сравнению с существующим деквантованием. Страница проекта: 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/