PaperlayerView exact revision on arXiv

Source-linked AI summary

OCTOPUS: Optimized KV Cache for Transformers via Octahedral Parametrization Under optimal Squared error quantization

Mark Boss, Vikram Voleti, Simon Donné, Shimon Vainer

arXiv:2605.21226v1cs.LGcs.AI

TL;DR

Long-context autoregressive inference is constrained by the memory bandwidth and footprint of KV caches. OCTOPUS jointly quantizes rotated coordinate triplets using octahedral direction mapping and non-uniform Lloyd-Max quantization, and is the best rotation-based codec across evaluated text, video, and audio tasks at matched bit widths.

  • Problem

    KV-cache reads dominate long-context autoregressive inference, motivating improved compression beyond existing eviction, scalar, and rotation-preconditioned codecs.

  • Method

    OCTOPUS maps rotated coordinate-triplet directions octahedrally and Lloyd-Max-quantizes direction coordinates and triplet norms with an MSE-optimal non-uniform bit split.

  • Results

    Across language, video, and audio tasks, OCTOPUS is the best rotation-based codec at matched bit widths K=V ∈{4, 3, 2}, with larger gaps at lower budgets.

  • Takeaways & Limitations

    The reported construction generalizes rotation-based KV compression beyond language models to autoregressive video and audio transformers.

  • Takeaways & Limitations

    OCTOPUS adds arithmetic overhead over scalar Lloyd-Max decoding and remains slower than bf16 SDPA, making it most attractive when KV bandwidth or capacity is limiting.

Abstract

from arXiv · show

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/

1 Introduction

OCTOPUS compresses rotated KV-cache triplets by separately quantizing their octahedrally parameterized directions and norms, using an MSE-optimal non-uniform bit split and fused on-the-fly reconstruction. It generalizes rotation-based KV compression beyond language models and is reported as the best rotation-based codec at matched widths K=V ∈ {4, 3, 2}.

  • Introduction: OCTOPUS splits pre-conditioned keys into triplets and Lloyd-Max-quantizes each triplet’s norm plus its octahedrally mapped direction.The octahedral map represents a unit 3-vector with two scalars in [−1, 1]2 using O(1) arithmetic and piecewise-linear encode/decode.
  • Introduction: Fused Triton kernels reconstruct keys on the fly from packed bit indices without materializing the full key tensor.This makes the compress-decode pipeline suitable for KV-cache inference.
  • Introduction: An MSE-optimal Lagrangian yields the implemented non-uniform (b+1, b−1) bit split at d=128.The split follows a finite-dimensional stationarity condition for per-triplet squared error.
  • Introduction: OCTOPUS-QJL optionally adds a 1-bit residual per rotated coordinate to drive seed-averaged dot-product bias to zero.The residual incurs one sign bit per rotated coordinate.
  • Introduction: OCTOPUS is reported as the best rotation-based codec at matched bit widths K=V ∈ {4, 3, 2} across long-context language, video, and audio generation.The reported evaluation includes Qwen2.5-7B-Instruct-1M, chunk-wise video diffusion, and frame-wise causal video forcing.

2 Related Work

Prior KV-cache compression methods include token eviction, scalar quantization with residuals, sparse coding, and rotation-preconditioned codecs. OCTOPUS builds on rotation-based quantization while introducing octahedral direction encoding for transformer decoding and composing with QJL.

  • KV-cache compression: KV-cache compression spans token eviction, per-channel scalar quantization with per-token residuals, sparse coding, and rotation-preconditioned codecs.These approaches respectively retain likely-useful tokens, model individual key-coordinate distributions, trade larger code tables for ultra-low rates, or apply data-oblivious random orthogonal operators.
  • Rotation-preconditioned quantization: TurboQuant uses random orthogonal rotations and 1-D Lloyd-Max codebooks whose design depends only on dimensionality and bit rate.It characterizes rotated unit-vector coordinates as symmetric-Beta marginals and uses a structured Walsh-Hadamard transform with random sign flips as the standard fast preconditioner.
  • Unit-direction encodings and unbiased estimators: OCTOPUS is presented as the first transformer-decoding use of octahedral direction quantization, extending compact direction encodings from real-time rendering.QJL instead provides an unbiased inner-product estimator using a 1-bit Johnson-Lindenstrauss sketch, which OCTOPUS composes with under a tag.

3 Method

OCTOPUS rotates keys, jointly quantizes coordinate triplets by separating their norms from octahedrally parameterized directions, and reconstructs them inside attention without materializing the uncompressed key. Its analytically motivated codebooks, non-uniform bit allocation, and local joint-rounding search are data-oblivious and decoder-compatible.

  • Pipeline: OCTOPUS encodes each key as a global norm, packed octahedral-coordinate indices, and packed triplet-norm indices, while decoding the lossy key inside attention without materializing the key matrix.The method assumes the key dimension d is a power of two for the Walsh-Hadamard transform.
  • Triplet quantization: Each rotated key is partitioned into zero-padded triplets whose norms and directions are quantized separately, with directions mapped from S2 to [−1, 1]2 by an octahedral bijection.The octahedral fold has a constant Jacobian per octant, enabling scalar Lloyd-Max quantization of the two square coordinates.
  • Bit allocation: The triplet budget Btri=2bdir+bnrm uses a non-uniform allocation because the uniform split bdir=bnrm=b is sub-optimal in squared error.The method reports a finite-dimensional optimum at (b+1, b−1), with the exact sweep finding this minimum for b ∈ {2, 3, 4}.
  • Bit allocation: 31–41% MSE reduction relative to uniform (b, b) is achieved by the implemented (b+1, b−1) split on synthetic Gaussian keys at d=128.Every other tested diagonal step increased MSE, including +44 to +73% at δ=+2 and an order of magnitude or more at δ=−2.
  • Joint rounding: Joint rounding selects direction candidates using a local 3×3 neighborhood, which was byte-identical to full-grid search across 104 random rotated triplets in d=128 for bdir ∈ {2, . . . , 5}.The encoder changes while the bitstream layout, codebooks, and decoder remain untouched, so the same fused attention kernel decodes deployed states.

4 Experiments

Across synthetic, language, video, and audio evaluations, OCTOPUS matches or outperforms rotation-preconditioned baselines, especially at extreme low-bit rates, while adding arithmetic cost at decode time. Its strongest gains appear in reconstruction fidelity, retrieval, perplexity, and perceptual quality when competing codecs collapse below 4 bits.

  • Experimental setup: Experiments compare OCTOPUS and OCTOPUS-QJL with TurboQuant-MSE, TurboQuant-QJL, and PolarQuant using matched symmetric K=V bit widths and shared rotation, value codec, and residual window.The only varying component across comparison rows is the K codec.
  • Synthetic fidelity: At b=4, OCTOPUS MSE is 1.3× below TurboQuant-MSE; at b=2, it is 2.4× below PolarQuant, while OCTOPUS-QJL reduces IP error 3× below TurboQuant-QJL.On needle retrieval, OCTOPUS-QJL tracks fp32 within 0.001, and OCTOPUS preserves 0.92 softmax mass at b=2.
  • Long-context language modeling: On Qwen2.5-7B-Instruct-1M, OCTOPUS leads every rotation codec: the WikiText-2 gap is +2.7% versus +3.1/4.4/8.0% at b=4 and +34.7% versus +63/187/772% at b=2.All language-model rows use the same residual window, value group size, and boundary-block protection recipe.
  • Needle retrieval: For multi-key needle retrieval, OCTOPUS retains 1.00 recall at b=3, while at b=2 OCTOPUS and OCTOPUS-QJL retain 0.81 and 0.83 versus 0.04 and 0.01 for PolarQuant and TurboQuant-QJL.At b=4, all codecs reach 1.00 recall.
  • Video and audio: Across video and audio, low-bit degradation favors OCTOPUS: it remains at 6.75 dB LSD and +1.07 dB SNR at b=2, while competing codecs reach 12.6–13.2 dB LSD with negative mean SNR.PolarQuant degrades 1.4× faster than OCTOPUS on mean LPIPS as bits decrease, and OCTOPUS adds more arithmetic than scalar Lloyd-Max decoding.

5 Conclusion · A Encoder and decoder algorithms · B Mathematical details

OCTOPUS jointly quantizes contiguous rotated coordinate triplets by mapping directions to octahedral square coordinates and quantizing them with norm information under an asymmetric bit split. Its encoder runs in one pass with register-resident intermediates, while fused split-K decoding reconstructs keys on the fly; the mathematical details establish rotation-based inner-product preservation and coordinate-marginal structure.

  • 5 Conclusion: OCTOPUS maps each contiguous 3-coordinate rotated direction block to two scalars on [−1, 1]^2 and combines them with the triplet norm into three integers.Lloyd-Max quantizers are matched to the oct-coordinate and norm marginals under the asymmetric (b+1, b−1) bit split.
  • 5 Conclusion: The codec retains TurboQuant’s data-oblivious, online guarantees and combines without modification with 1-bit QJL.These properties are stated as inherited or directly compatible codec guarantees.
  • A Encoder and decoder algorithms: The encoder processes each key in one pass while keeping the rotated vector, triplet norms, octahedral coordinates, and integer indices in registers.Algorithm 1 is described as the implemented encoder.
  • A Encoder and decoder algorithms: The decoder fuses bit unpacking, octahedral decoding, centroid lookup, value dequantization, and online softmax in a single split-K flash-decoding kernel.Algorithm 2 follows the split-K flash-decoding style of Dao et al. and Shah et al.
  • A Encoder and decoder algorithms: The encoder’s 3×3 optimal-rounding refinement is disabled by setting ∆ to {(0, 0)}, recovering the legacy scalar-rounding baseline.The default candidate set is ∆ = {−1, 0, 1}^2.
  • B Mathematical details: For any q, k ∈ R^d and sign vector s ∈ {±1}^d, the rotated representation preserves the inner product as q⊤k = (Rq)⊤(Rk) = γ (Rq)⊤u.This is Proposition 1’s inner-product invariance statement.
  • B Mathematical details: Quantizing u after rotating the query preserves an unbiased dot product in expectation, while each rotated coordinate follows the symmetric-Beta marginal under the stated spherical model.The marginal result assumes uniformly distributed ˜u on S^{d−1} and an independent uniform random sign vector.

C Derivations … D Bit-allocation sweep

The derivations establish OCTOPUS’s sphere-preserving rotation, triplet-wise octahedral representation, implementation-matched Lloyd-Max marginals, and joint reconstruction rules. The bit-allocation sweep supports the implemented (b+1, b−1) split as the unique diagonal choice improving MSE for b∈{2,3,4}.

  • C.1 Magnitude-Direction Split; C.2 Sign-Flipped WHT Rotation: OCTOPUS separates each key into its norm and unit direction, then applies a sign-flipped normalized Hadamard rotation that preserves the sphere and has inverse R−1 = R⊤ = DsH.The implementation stores the original norm and uses an ϵ-safe denominator for zero or tiny keys.
  • C.4 Triplet Split (Section 3.2): Rotated directions are zero-padded, split into contiguous triplets, and represented as ti = ρini with ni∈S2.Zero-radius triplets receive a finite placeholder through an ϵ-safe divisor.
  • C.6 Octahedral Encode; C.7 Octahedral Decode: Octahedral encoding folds the L1-normalized direction onto a square, and decoding inverts the fold before final normalization back to S2.The implementation uses sign(0)=+1, clamps denominators, and reconstructs from dequantized oct-coordinate centroids.
  • C.9 Triplet MSE Bound; C.10 Expected MSE Budget: The triplet MSE analysis decomposes reconstruction error into radius and direction terms, while the expected-MSE budget uses independence of Gaussian block radius and direction with high-rate Lloyd-Max scaling.The quantized variable is the scalar norm ρi itself.
  • C.11 Lagrangian Bit Allocation; C.12 Bit-Gap Scaling: The high-rate allocation minimizes f(bnrm,bdir)=A4^-bnrm+D4^-bdir subject to bnrm+2bdir=Btri, with the factor of two reflecting two octahedral coordinates.The resulting finite-dimensional allocation is supported empirically rather than by an asymptotically growing direct-ρ bit gap.
  • C.13 Joint Rounding; C.14 Score Factorization: Joint rounding selects the radius centroid nearest the projected dot product si, not the centroid nearest the true radius ρi; reconstructed rotated triplets are ˆui=ˆρiˆni.The score factorization then applies the usual attention scale after forming the raw score.

E Rounding ablation

The rounding ablation compares four encoder modes at matched bits while holding the bitstream and decoder fixed. The implemented 3×3 local search matches full direction search and reduces MSE by 6–14% at matched bit rate, with gains largest at tight budgets.

  • Experimental setup: Four rounding modes are compared at matched bits as a pure encoder ablation, using the same bitstream and decoder.Metrics are averaged over five seeds on n=4096 Gaussian keys with d=128 and nquery=64.
  • Search variants: The 3×3 local search is byte-identical to the full direction search at every tested bit width.Full-codebook direction search is the joint-optimum upper bound for the given (Cξ, Cρ) pair, while local_3x3 is Algorithm 1’s default.
  • Results: 6–14% MSE reduction is achieved by the implemented 3×3 search at matched bit rate, without changing the bitstream format or decoder.The ablation indicates that optimal rounding primarily shifts the encoder rather than the codebook.
  • Results: The rounding gain is largest at the tightest bit budgets, where one misrounding consumes a larger fraction of the remaining precision.The results support using a small local neighborhood for practical rounding.

F QJL effective-rate accounting

Table 6 reports the effective-rate cost of adding the one-bit residual side-car described in Section 3.6.

  • F QJL effective-rate accounting: Table 6 details the effective-rate cost of the added residual side-car.The accounting concerns the side-car introduced in Section 3.6.
  • F QJL effective-rate accounting: The residual side-car uses one bit.This is the specific overhead quantified in the table.
  • F QJL effective-rate accounting: The side-car is described in Section 3.6.Table 6 provides its corresponding rate accounting.

G Kernel speed and KV compression

OCTOPUS trades decode-time computation for KV memory savings by fusing decompression into the attention loop, while its encoding remains lightweight enough for long sequences. OCTOPUS-QJL increases the effective rate by exactly 0.5 bits per scalar to provide additional reconstruction headroom.

  • Decode and encode performance: 5–11× fused decode overhead versus cuDNN SDPA bf16 decreases at lower bit widths as packed data shrinks.Each decode step fuses centroid lookup for TQ-MSE or octahedral reconstruction for OCTOPUS into the attention loop, trading compute for KV memory savings.
  • Decode and encode performance: ≤2 ms TQ-MSE encode time is reported even at 65k tokens.OCTOPUS encoding instead uses a Kronecker-factored WHT and direct triad indexing.
  • Effective-rate accounting: 0.5 bits per scalar is the exact effective-rate increase from adding the 1-bit JL residual in OCTOPUS-QJL.The added residual provides roughly that much reconstruction headroom inside the standard dequantize-then-dot attention path.

H Long-context needle-in-a-haystack sweep · I Memory-budget Pareto · J Full per-modality tables

The section expands long-context retrieval results across context lengths and bit widths, visualizes deployment memory trade-offs, and supplies full per-codec metrics for each modality.

  • H Long-context needle-in-a-haystack sweep: Table 8 expands the long-context retrieval summary across the full context-length and bit-width grid.It extends the summary from Sec. 4.2.
  • H Long-context needle-in-a-haystack sweep: The long-context sweep covers both context length and bit width.Table 8 reports the expanded retrieval summary across their full grid.
  • I Memory-budget Pareto: Figure 4 recasts the LLM results as a deployment memory trade-off.The figure presents the results in a form intended to expose deployment memory choices.
  • I Memory-budget Pareto: The memory-budget analysis makes the Pareto frontier visible at fixed context length.This reframing highlights the frontier under a fixed context-length constraint.
  • J Full per-modality tables: Tables 9, 10, and 11 provide the per-codec metrics behind the modality summary.The tables supply the detailed measurements underlying Sec. 4.3’s modality-level summary.
  • J Full per-modality tables: The full per-modality tables support the modality summary in Sec. 4.3.Together, Tables 9, 10, and 11 provide the detailed per-codec results.

K Stills

The section presents representative worst-case video frames and identifies OCTOPUS as visually faithful across bit widths, while competing codecs collapse at b≤3. These qualitative examples complement aggregate LPIPS, PSNR, and SSIM results.

  • Worst-case frames: Figure 5 provides representative worst-case frames that complement the aggregate video LPIPS, PSNR, and SSIM results reported in Section 4.3.The frames are selected to illustrate codec divergence qualitatively rather than replace the aggregate metrics.
  • Worst-case frames: Figure 5 shows the single frame with the highest combined cross-codec L1 divergence from the fp16 baseline at each bit depth.Rows correspond to b=4, 3, and 2; columns show the baseline and each codec.
  • Worst-case frames: OCTOPUS remains visually faithful at every tested bit width, whereas competing codecs collapse at b≤3.The comparison uses the same frame index for both pipelines.
Loading 2605.21226v1…