TACT-KV: Tri-Axis Cosine Transform for Compressing Volumetric KV Caches in Medical VLMs

AI summary

Overview

Research area: Computer Vision / efficient inference for vision-language models (VLMs), specifically key-value (KV) cache compression for 3D medical imaging and video.

Technical level: Advanced. The paper assumes familiarity with transformer KV caching, attention, the discrete cosine transform, rate–distortion theory, and scalar quantization.

Scope (one sentence): The paper proposes TACT-KV, a calibration-free method that reshapes the visual KV cache of a VLM into a three-axis (slice, height, width) tensor, applies a separable 3D DCT-II, allocates bits per frequency bin by reverse water-filling, and quantizes with a fixed scalar codebook — evaluated on six medical VLMs over two 3D CT benchmarks and on four general-purpose VLMs over Video-MME.

What This Paper Is About

A single CT or MRI volume produces tens of thousands of vision tokens after the vision encoder, and the resulting KV cache grows linearly with both batch size and context, dominating GPU memory during inference at clinical sites that typically run on single-card or small-cluster GPUs. Existing KV compression methods work on token or channel statistics, or apply 1D/2D transforms, and therefore leave the three-axis spatial structure of a 3D scan unexploited. The paper's goal is to compress the visual KV cache aggressively — down to 1-bit nominal budgets — while keeping model predictions close to the uncompressed bf16 reference, so that larger decoding batches fit on one GPU.

Key Contributions

1. A volumetric view of the visual KV cache. TACT-KV is presented as the first method to treat the visual KV cache as a volumetric tensor and exploit three-axis spatial correlation from 3D scans, reshaping the cache into the (S, H, W) grid inherited from the vision encoder and applying a separable 3D DCT-II across all three axes.

2. A closed-form distortion result. Theorem 1 gives the ratio of uniform-allocation distortion to reverse water-filling distortion as exactly 1/γ, where γ is the spectral flatness of the per-bin variances — i.e. the gain from adaptive bit allocation is governed by how concentrated the cache energy is in the frequency domain.

3. A calibration-free, parameter-free pipeline. The method needs no calibration data and adds no learnable parameters; text KV stays at bf16, and all spectral statistics are estimated on the fly per layer per sequence in a single forward pass.

4. A broad empirical evaluation. Six medical VLMs (4B–32B: MedGemma 1.5/4B/27B and Hulu-Med 4B/7B/32B), two 3D CT benchmarks (RSNA-ICH, CT-RATE), four general-purpose VLMs on Video-MME, quantizer and transform ablations, a theorem-closure check, and single-GPU deployment profiling.

Main Findings

• Lowest 1-bit flip rate on eleven of twelve model–dataset pairs. Across the eleven leading pairs, the next-best method's average flip rate is 2.60× that of TACT-KV. Examples at 1-bit on RSNA-ICH: MedGemma 1.5 4B 3.78% versus TurboQuant 6.27%, KIVI 10.74%, KVTC 16.55%, AKVQ-VL 28.54%; Hulu-Med 7B 6.08% versus KIVI 25.38%, KVTC 60.38%.

• Baselines collapse at 1-bit. KIVI's per-channel grouped-uniform primitive reaches 77% flip on Hulu-Med 4B RSNA-ICH and 65% on Hulu-Med 32B CT-RATE; TurboQuant exceeds 35% flip on two of six VLMs; VidKV and AKVQ-VL reach worst-case 1-bit flip rates of 45.5% and 75.7%. The authors attribute this to fixed quantization grids lacking dynamic range for the heavy-tailed visual KV distribution.

• Predictions stay within the bf16 confidence band. At 1-bit on RSNA-ICH, macro-averaged ROC curves for TACT-KV track the bf16 reference within its 95% bootstrap CI on all six VLMs, while baseline curves drift outside the band on the Hulu-Med variants.

• Theorem 1 partly closes empirically. In Figure 3 (RSNA-ICH, N=3000 studies, six VLMs), three of six cells fall within a 30% deviation band, one sits at 40%, and two (Hulu-Med 4B and 7B) lie above the identity line at low rate. Extended to all twelve pairs at 1-bit, eight of twelve fall within 50% of the closed-form 1/γ_K; the four exceptions are Hulu-Med 4B and 7B on both datasets, where the paper states the 1-bit setting falls outside the high-rate regime.

• Both the transform and the adaptive allocation matter. On RSNA-ICH at 1-bit (MedGemma 1.5 4B / Hulu-Med 7B): 3-axis DCT + water-filling 3.78 / 6.08; 2-axis 5.00 / 6.21; 1-axis 6.00 / 9.65; no DCT 6.61 / 42.20; 3-axis DCT with uniform allocation 9.97 / 25.76.

• Gain-shape scalar quantization is the best backend. At 1-bit (MedGemma 1.5 4B / Hulu-Med 7B): gain-shape SQ 3.78 / 6.08, symmetric uniform 24.04 / 36.25, NormalFloat 23.26 / 37.05. The post-DCT coefficient distribution is heavy-tailed, with median absmax-to-std of 4.81 versus the Gaussian reference 3.01 at d=256.

• Benefits transfer to video. On Video-MME (no-subtitle, 2,700 four-way MCQs), TACT-KV has the highest accuracy on every VLM at 1-bit, with a worst-case drop from bf16 of 4.77 pp on Qwen3-VL-8B (72.33% to 67.56%); KIVI falls below 13% on every VLM at 1-bit.

• Large batch and throughput gains at equal HBM. On one 96 GiB RTX PRO 6000 at S=60 slices, per-request cache shrinks 6.9–7.6× and maximum decoding batch grows 4.00× (Hulu-Med 4B, 32 to 128), 6.25× (Hulu-Med 7B, 64 to 400), and 8.00× (Hulu-Med 32B, 8 to 64), yielding 1.17–2.02× higher aggregate decoding throughput. At matched batch sizes, step latency drops 1.45×, 1.14×, and 1.10× respectively.

Methodology in Plain English

The starting observation is that in a 3D scan, neighbouring slices look very much like each other, just as neighbouring pixels within a slice do. The authors show this similarity carries into the decoder's visual KV cache: tokens at the same in-plane position across adjacent slices are highly similar, and the cache is strongly autocorrelated along all three spatial axes. When data is this smooth, most of its "energy" sits in a few low-frequency components — so converting it into a frequency representation should make it easy to compress.

The pipeline has four stages, run per layer per sequence:

1. Reshape. The visual key and value tensors are laid out as a (slice, height, width) grid, matching the grid the vision encoder already produces. Text keys and values are left untouched at bf16.

2. Transform. A separable 3D DCT-II is applied along slice, height, and width. The transform is orthonormal and lossless on its own; it merely concentrates energy. Separability keeps the cost at O((S+H+W)·n_v·d) rather than O(n_v²·d) and cuts parameters from (SHW)² to S²+H²+W². The feature axis is left untransformed.

3. Bit allocation. For a given per-layer bit budget, each frequency bin gets bits by reverse water-filling: bins with more energy (higher variance) get more bits, low-energy bins get zero and are dropped, and a water level is found by bisection with a ±1 greedy correction to hit the budget exactly. Bits are capped at 8. The paper proves the gain over uniform allocation equals 1/γ.

4. Quantize. Surviving bins are gain-shape encoded: the vector's ℓ2 norm is factored out, the residual is rotated by a fixed shared orthogonal rotation, and each coordinate is looked up in a precomputed Lloyd–Max table at that bin's bit width. Per-layer occupancy is Σ(b_i·d + 32) bits plus a small index overhead.

Decoding reverses these steps and can be fused into the attention kernel, so the full reconstructed cache is never materialized. Evaluation uses macro prediction-flip rate (the fraction of per-study, per-label binary predictions that change relative to bf16) as the primary metric, with macro AUROC and 95% bootstrap CIs, and per-question accuracy for Video-MME.

Why This Matters

Research impact. The paper reframes KV compression for volumetric inputs as a transform-coding problem over an anatomical grid, and supplies a rate–distortion statement (1/γ) linking compression gain to spectral concentration. It also draws an explicit distinction between the temporal axis in video, where attention concentrates on a few salient frames, and the slice axis in medical scans, where diagnostic signal is spread across many slices — a distinction that argues against directly transplanting video-domain compression to medical imaging.

Real-world applications:

• CT report generation and triage on single GPUs. Serving multiple concurrent studies per card at clinical sites that operate on single-card or small-cluster hardware.
• Multi-abnormality screening. CT-RATE's 18-abnormality multi-label protocol is the kind of workload where the 8× batch expansion on the 32B model matters most.
• Long video question answering. The Video-MME results indicate the same compression benefits general-purpose VLMs on video tasks.
• Memory-constrained deployment of large medical VLMs. The batch scaling grows monotonically as weights consume more of HBM headroom (66.4 GB of 102.6 GB for Hulu-Med 32B), which is exactly the regime where cache memory would otherwise be the binding constraint.

Industry relevance. The method is calibration-free and adds no learned parameters, and the text KV path is untouched, which keeps it compatible with existing serving stacks. On-device profiling shows both throughput gains at larger batch and latency reductions at matched batch — the two levers serving operators actually tune.

Future Directions

• Fix the low-rate regime. The four pairs outside the closure check (Hulu-Med 4B and 7B on both datasets, with closure errors of 401% to 1524%) violate the high-rate Gaussian assumption underlying Theorem 1. A distortion model that remains valid near zero bits per coordinate is an open problem the paper itself flags.

• Handle backbones without clean (S, H, W) grids. The pipeline reads grid parameters directly from the vision encoder; whether the approach extends to architectures where that structure is obscured, or where sliding-window attention layers (left at bf16 here) could also be compressed, is not established.

• Push adaptation beyond a fixed codebook. The quantizer is a fixed Lloyd–Max table with a shared rotation and no calibration data. Learning or adapting the rotation and codebooks to per-layer spectral shape is an obvious extension the current design deliberately avoids.

• Stress-test under degraded spatial correlation. The paper notes that noise or motion flattens the spectrum and pushes γ toward one, shrinking the predicted allocation gain. Quantifying how much acquisition noise, motion artifact, or thin-slice protocols erode the advantage in clinical data is not reported.

Target Audience

Researchers and engineers working on efficient inference for large vision-language models, particularly those deploying medical or video VLMs under tight GPU memory budgets. It will also interest readers in transform coding and rate–distortion analysis applied to learned representations, and practitioners building serving infrastructure who need a calibration-free, drop-in cache compression scheme. Readers without background in KV caching, the DCT, or quantization theory will find the theoretical section demanding.

Authors’ abstract

Vision-language models (VLMs) applied to 3D medical scans face a structural memory wall. A single CT or MRI volume produces tens of thousands of vision tokens, and the resulting key-value (KV) cache scales linearly with both batch and context. Existing KV compression methods operate on the cache's feature or token statistics, or apply 1D and 2D transforms, leaving the three-axis spatial structure in the cache unexploited. We propose TACT-KV (Tri-Axis Cosine Transform), a calibration-free KV compression method that exploits strong spatial correlations in visual KV states along all three anatomical axes of 3D scans. This correlation concentrates the KV cache energy under a frequency-domain transform. TACT-KV applies a separable 3D discrete cosine transform to the KV cache, allocates a per-layer bit budget across frequency bins via reverse water-filling (more bits to high-energy components), and quantizes each bin with a fixed scalar codebook. We derive a close-form relation: the distortion reduction from adaptive bit allocation is dominated by the spectral concentration of the transformed cache. Across six medical VLMs and two 3D CT benchmarks, TACT-KV achieves the lowest 1-bit prediction-flip rate on eleven of twelve model and dataset pairs in our main comparison. The same advantage reproduces on four general-purpose VLMs evaluated on video tasks, indicating that the benefit of exploiting three-axis structure generalizes beyond medical imaging. At equal high-bandwidth memory (HBM), TACT-KV scales the decoding batch size on a single GPU by up to $8\times$.