Transformer-FSQ: Discrete Latents for Extreme KV-Cache Compression

Andras Szecsenyi; Thomas L. Lee; Vaisakh Shaj; Elliot J. Crowley; Amos Storkey

cs.LGNov 6, 2026

Transformer-FSQ: Discrete Latents for Extreme KV-Cache Compression

Authors: Andras Szecsenyi, Thomas L. Lee, Vaisakh Shaj, Elliot J. Crowley, Amos Storkey

Organizations: School of Informatics The University of Edinburgh · School of Engineering The University of Edinburgh

Abstract

Private, personalised language models must run on local consumer devices, where large, frontier models cannot fit, and retaining long contexts is non-trivial. Sequential memory---the key-value (KV) cache---grows linearly with context length, quickly dwarfing the model itself: for Qwen3-1.7B-Base, 200,000 retained entries require 21 GiB, 6.6 times more than the model weights. Yet the KV-cache compression methods we measured fail beyond moderate compression ratios. We introduce Transformer-FSQ, addressing extreme KV-cache compression by replacing standard key and value projections with a finite scalar quantisation bottleneck. Our cache stores a packed vector of discrete integer codes per token, with no per-token quantisation metadata. Converting a pretrained transformer trains only these projections on 0.25 tokens per original parameter. We evaluate Transformer-FSQ against five baseline families on models from 0.6B to 8B parameters. On long-context retrieval (RULER), from 12 $\times$ to 97.5 $\times$ Transformer-FSQ is the most accurate method we evaluate, across scales. A fused Triton kernel decodes directly from the integer codes; with baselines given their authors' kernels where available, it decodes at $1.01\times$ (95% interval $0.98$ -- $1.05$ ) the throughput of native SDPA attention on the uncompressed Qwen3-8B-Base at $32\times$ and $1.54\times$ at $97.5\times$ . Extreme KV-cache compression is therefore a practical way to extend how much context a fixed device can hold.

Figures & tables

Panel of Figure 1 — Figure 1: KV-cache capacity on a 16 GiB GPU. With ${\approx}0.74$ GiB left for cache after Qwen3-8B model weights, extreme compression ( $256\times$ ) expands token capacity from 5.4k to 1.39M, enabling million-token contexts on consumer hardware.

Figure 2: RULER at

8{,}192

tokens against achieved compression on Qwen3-0.6B, 1.7B and 8B. Solid: finetuned; dashed: training-free; grey horizontal line: uncompressed. Points average up to three seeds (one on 8B), with 95% bootstrap intervals over seeds and then items stratified by task. Beyond about

12\times

, Transformer-FSQ scores highest on all three models.

Panel of Figure 3 — Figure 3: Discrete latent geometry: (a) VQ searches an explicit codebook, inducing irregular Voronoi cells ( Gray and Neuhoff, 1998 ) . (b) FSQ quantises coordinates independently onto a fixed Cartesian grid $\mathcal{L}_{1}\times\mathcal{L}_{2}$ , an implicit codebook that is never materialised or searched.

Panel of Figure 4 — Figure 4: One decoding step. Left: standard attention caches keys and values. Right: Transformer-FSQ caches one bit-packed code $z_{t}$ per token, shared by all KV heads. Keys are reconstructed and rotated at every cached position ( ${\times}T$ ); values are decoded once, after attention pools the codes, $\bar{z}=\sum_{t}a_{t}z_{t}$ . Orange: added by Transformer-FSQ; grey: frozen backbone. Details: Appendix F .

Method	Family	Cache contents per token	Trained	Deepest	Max tokens (16 GiB, $\geq 0.30$ )
PyramidKV	token eviction	full-precision KV of kept tokens	no	$174\times$	552k
KIVI	scalar quantisation	2-bit KV, full-precision window	no	$5.2\times$	684k
VecInfer	vector quantisation	codebook indices, window, sinks	no	$11.0\times$	592k
Palu	low-rank projection	quantised low-rank SVD latent	yes	$10.2\times$	1.42M
TransMLA	latent attention	latent, fp16 rotary key	yes	$14.2\times$	542k
Transformer- FSQ (ours)	finite scalar quantisation	bit-packed integer code	yes	$410\times$	2.21M

Table 1: Baselines and their limits on a 16 GiB GPU. Deepest : the deepest ratio at which the original authors report downstream accuracy, priced as in § 4.4 (PyramidKV: needle retrieval); for Transformer-FSQ, the deepest we evaluate, where RULER falls to

0.003

on Qwen3-0.6B-Base. Max tokens : input tokens whose cache fits beside the converted model’s weights in 16 GiB, on Qwen3-0.6B-Base at RULER

\geq 0.30

; for PyramidKV, tokens represented after eviction.

Figure 5: LongBench gov_report ROUGE-L against achieved compression on Qwen3-0.6B, 1.7B and 8B. Solid: finetuned; dashed: training-free; grey horizontal line: uncompressed; bands: one standard deviation over seeds.

Method	Training	Cache	Width	B/tok	RULER
Nominal budget: $144$ B per token per layer ( $28.4\times$ compression)
Transformer-FSQ (ours)	Discrete (STE, $L{=}16$ )	4-bit implicit grid	288	144	$\mathbf{0.232\pm 0.022}$
Continuous latent	Continuous (bf16)	16-bit float	72	144	$0.009\pm 0.008$
Latent + post-hoc FSQ	Continuous (bf16)	4-bit post-hoc FSQ	288	144	$0.175\pm 0.030$
Latent + post-hoc KIVI	Continuous (bf16)	4-bit blocks + scales	288	180	$0.206\pm 0.021$
Palu	Low-rank SVD	4-bit SVD + scales	144	144	$0.000\pm 0.000$

Table 2: Discreteness at matched byte budgets: RULER at

8{,}192

tokens on Qwen3-0.6B-Base (three seeds,

0.25

tokens per parameter) at nominal

144

and

80

B per token per layer; mean

\pm

95% bootstrap half-width. The uncompressed model (

4{,}096

B) scores

0.718\pm 0.023

. Continuous latents are

4\times

narrower to fit the budget; post-hoc variants quantise a wide continuous latent at evaluation. KIVI’s per-block scales add

36

and

20

B; TransMLA spends

32

B on rotary keys. Training through the 4-bit grid (STE,

L{=}16

) scores highest at both budgets.

Figure 6: Decode throughput against context length on NVIDIA L40S in bfloat16; rows: model, columns: batch size, both axes logarithmic. Grey dashed: the uncompressed model with native SDPA, timed in the same process. Orange: Transformer-FSQ at

32\times

97.5\times

and

409.6\times

, darker for deeper compression; a line that ends early ran out of memory. Protocol: Appendix D.5 .

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

	Qwen3-0.6B-Base	Qwen3-1.7B-Base	Qwen3-8B-Base
parameters	$596{,}049{,}920$	$1{,}720{,}574{,}976$	$8{,}190{,}735{,}360$
layers	$28$	$28$	$36$
model width $C$	$1024$	$2048$	$4096$
query heads	$16$	$16$	$32$
key-value heads $n_{kv}$	$8$	$8$	$8$
head width $d_{h}$	$128$	$128$	$128$

Appendix

Table 3: Models and their memory on a 16 GiB GPU.

c_{\text{full}}

: the uncompressed half-precision cache per token per layer (Eq. 9 ), the reference for every compression ratio. The uncompressed context limits are the baseline points of Figure 11 . ‡ Qwen3-8B-Base runs one seed throughout (§ 5.1 ).

	Qwen3-0.6B / 1.7B-Base	Qwen3-8B-Base
finetuning budgets $\tau$	$0.025^{\ast},\ 0.075^{\ast},\ 0.25$	$0.25$
optimiser	AdamW	AdamW
peak learning rate	$10^{-5}$	$10^{-5}$
schedule	linear warm-up over 5% of steps, cosine decay to 5% of peak
weight decay	$0.28$ , constant	$0.28$ , constant
data	ClimbMix	ClimbMix

Appendix

Table 4: Finetuning hyperparameters.

\tau

: finetuning tokens per parameter;

\tau=0

is conversion without finetuning. All methods at a model scale share data, optimiser and budget. The released models’ results (§ 5.2 ) use

\tau=0.25

; Figure 7 shows the effect of

\tau

on Qwen3-0.6B-Base. ∗ Qwen3-0.6B-Base only.

Figure 7: RULER against finetuning tokens on Qwen3-0.6B-Base (

0.025

0.075

and

0.25

per parameter) at nominal

512

and

1{,}024

B per token per layer. Labels give achieved ratios (Palu stores

400

and

784

B). Transformer-FSQ plateaus with fewer finetuning tokens than the baselines.

Model	$\tau=0.025$	$\tau=0.075$	$\tau=0.25$
Qwen3-0.6B-Base	0.19 $n{=}30$ ; 0.15–0.25	0.56 $n{=}27$ ; 0.44–0.78	2.59 $n{=}54$ ; 1.47–5.14
Qwen3-1.7B-Base	–	–	7.32 $n{=}63$ ; 5.57–9.19
Qwen3-8B-Base	–	–	138.34 (34.6 h) $n{=}9$ ; 108.88–158.72

Appendix

Table 5: Conversion and finetuning GPU-hours per run on NVIDIA GH200 (one GPU per run; four for Qwen3-8B). Each cell: the mean over

n

finished runs of the finetuned methods (Transformer-FSQ, Palu, TransMLA; all budgets and seeds), then the min–max range; –: not run; for Qwen3-8B, mean wall-clock hours in parentheses.

\tau

: finetuning tokens per parameter.

Figure 8: RULER relative to the uncompressed model: score at

8{,}192

tokens divided by the same model’s uncompressed score; bands as in Figure 2 . Below about

5\times

, VecInfer and Palu keep the most; from

32\times

97.5\times

, Transformer-FSQ does.

Figure 9: LongBench narrativeqa F1 (0–100) against achieved compression. Solid: finetuned; dashed: training-free; grey horizontal line: uncompressed. PyramidKV scores highest on 0.6B and 1.7B; on 8B, Transformer-FSQ scores above Palu.

Figure 10: LongBench trec accuracy (0–100) against achieved compression. Line styles as in Figure 9 .

Figure 11: Tokens held on a 16 GiB GPU against RULER at

8{,}192

tokens. Crosses: out of memory. The stacked line includes only stacks in which Transformer-FSQ supplies at least half the compression in log terms. At matched RULER thresholds, Transformer-FSQ and its eviction stacks hold up to 1.7

\times

more tokens than the baselines.

Figure 12: Mean RULER at $8{,}192$ tokens for each Transformer-FSQ $\,\times\,$ PyramidKV stack on Qwen3-0.6B-Base. Grey cells fail the rule, fixed before the runs, that Transformer-FSQ supplies at least half the compression in log terms; values in Table 6 . Retrieval degrades gradually up to $32\times$ total and falls off at extreme stacked ratios.

FSQ	PyramidKV	Total ratio	B/tok/layer	Score $\pm$ half-CI	Seeds
–	–	$1.0\times$	4096	$0.72\pm 0.02$	1
$8\times$	–	$8.0\times$	512	$0.56\pm 0.03$	3
$16\times$	–	$16.0\times$	256	$0.32\pm 0.03$	3
$32\times$	–	$32.0\times$	128	$0.22\pm 0.01$	3
$97.5\times$	–	$97.5\times$	42	$0.04\pm 0.02$	3
$8\times$	$2\times$	$16.0\times$	256	$0.39\pm 0.02$	1

Appendix

Table 6: Stack-grid results at

8{,}192

tokens on Qwen3-0.6B-Base (RULER, 13 tasks of 50 items; Appendix D.3 ). Every stack that passes the rule fixed before the runs is listed, against Transformer-FSQ alone, the uncompressed model, and PyramidKV alone at the same achieved total budget (own block below). Ratios and B/token/layer are achieved, not nominal. Score: mean over complete seeds

\pm

item-bootstrap half-width (Appendix C.5 ); seeds : complete finetuning seeds.

Panel of Figure 13 — Figure 13: One decoding step of one layer: native SDPA (left) and the Transformer-FSQ kernel (right). Top band: GPU memory (HBM); everything below runs on chip. Native SDPA streams $4096$ B of 16-bit keys and values per cached token. At $32\times$ , the kernel streams $128$ B of packed codes per token, reconstructs keys on the tensor cores and accumulates attention-weighted codes; $W_{v}$ is applied once, after the blocks are merged.

Figure 14: Decode throughput against compression ratio, Qwen3-0.6B-Base: median tokens/s on NVIDIA L40S; rows: sequence length, columns: batch size, both axes logarithmic. Grey: the uncompressed model with native SDPA, timed in the same process; orange circles: Transformer-FSQ at every ratio timed; teal upward triangles: PyramidKV; green diamonds: Palu; blue downward triangles: KIVI; purple plus signs: VecInfer. KIVI and VecInfer run in their native float16, the rest in bfloat16. Missing markers: out of memory or outside the method’s grid.

Figure 15: Decode throughput against compression ratio, Qwen3-1.7B-Base; layout as in Figure 14 .

Figure 16: Decode throughput against compression ratio, Qwen3-8B-Base; layout as in Figure 14 .

Figure 17: Measured against analytical KV-cache memory on Qwen3-0.6B-Base in bfloat16 at

8{,}192

context (

7{,}855

prompt tokens). Grey: the analytical byte count (§ 4.4 ) at

8{,}192

tokens; coloured: the buffer allocated after the prompt, which equals the analytical bytes per token times the prompt’s tokens exactly.

Nominal budget: 144 B/token/layer (28.4 $\times$ )
Compression / quantisation	Width	Bytes	RULER
Continuous latent	72	144	$0.009\pm 0.008$
Latent + post-hoc FSQ	288	144	$0.175\pm 0.030$
Transformer-FSQ (ours)	288	144	$\mathbf{0.232\pm 0.022}$
Latent + trained min–max 4-bit	224	140	$0.187\pm 0.020$
Latent + post-hoc KIVI	288	180	$0.206\pm 0.021$

Appendix

Table 7: Cache width, precision and when quantisation is applied. RULER at

8{,}192

tokens on Qwen3-0.6B-Base (three seeds,

0.25

tokens per parameter); the uncompressed cache (

4{,}096

B per token per layer) scores

0.718\pm 0.023

. Scores: mean

\pm

95% bootstrap half-width; bold: best. Bytes: actual storage per token per layer, metadata included; Width: bottleneck dimension. Trained FSQ rounds during training (straight-through estimator); post-hoc FSQ and KIVI quantise a continuously trained latent only at evaluation. KIVI adds fp16 scales and offsets (

25\%

over budget); Palu stores 4-bit SVD codes with scales; TransMLA stores 4-bit content codes plus 16-bit rotary keys.

Cache	B/token/layer	Ratio	RULER ( $8{,}192$ )
FSQ, 4 levels	144	28 $\times$	$0.178\pm 0.013$
FSQ, 16 levels	144	28 $\times$	$0.232\pm 0.022$
FSQ, 256 levels	144	28 $\times$	$0.161\pm 0.011$
FSQ, 4 levels	80	51 $\times$	$0.104\pm 0.013$
FSQ, 16 levels	80	51 $\times$	$0.177\pm 0.017$
FSQ, 256 levels	80	51 $\times$	$0.035\pm 0.014$

Appendix

Table 8: Levels per coordinate: RULER at

8{,}192

tokens on Qwen3-0.6B-Base (

0.25

finetuning tokens per parameter) for FSQ with 4, 16 and 256 levels at

144

and

80

B per token per layer; mean

\pm

95% bootstrap half-width over three seeds and items. 16 levels score highest at both sizes.

Panel of Figure 18 — Figure 18: Alternative decoding topologies: FSQ-pre (left) and MLA-FSQ (right). FSQ-pre applies RoPE before the encoder, so codes carry position and keys need no per-step reconstruction. MLA-FSQ caches a half-precision rotary key $k^{r}$ beside the content codes and adds its query dot product to the logits. Shaded plates ( ${\times}T$ ) scale with context length $T$ . Appendix E gives why we keep post-projection Transformer-FSQ.

Method	seeds	niah_ single_1	niah_ single_2	niah_ single_3	niah_ multikey_1	niah_ multikey_2	niah_ multikey_3	niah_ multivalue	niah_ multiquery	vt	cwe	fwe	qa_ 1	qa_ 2	Mean
Uncompressed	1	1.00	0.96	0.98	0.90	0.82	0.42	0.83	0.88	0.80	0.30	0.51	0.62	0.30	0.72
Transformer-FSQ (2.0 $\times$ )	3	1.00	0.92	0.97	0.66	0.41	0.03	0.67	0.53	0.71	0.34	0.42	0.47	0.23	0.57
Transformer-FSQ (4.0 $\times$ )	3	1.00	0.89	0.97	0.69	0.43	0.03	0.68	0.54	0.61	0.31	0.42	0.53	0.25	0.57
Transformer-FSQ (8.0 $\times$ )	3	1.00	0.87	0.97	0.68	0.33	0.04	0.69	0.52	0.59	0.31	0.45	0.53	0.27	0.56
Transformer-FSQ (12.0 $\times$ )	3	0.99	0.85	0.66	0.40	0.11	0.01	0.33	0.41	0.18	0.31	0.24	0.43	0.17	0.39
Transformer-FSQ (16.0 $\times$ )	3	0.99	0.69	0.61	0.26	0.03	0.00	0.29	0.27	0.18	0.22	0.21	0.31	0.15	0.32

Appendix

Table 9: Per-task RULER on Qwen3-0.6B-Base (13 tasks of 50 items,

8{,}192

tokens, bfloat16). Finetuned methods use

0.25

tokens per parameter; KIVI, PyramidKV and VecInfer are training-free. Rows give achieved compression ratios; scores average complete seeds. Shading encodes the task score.

Method	seeds	niah_ single_1	niah_ single_2	niah_ single_3	niah_ multikey_1	niah_ multikey_2	niah_ multikey_3	niah_ multivalue	niah_ multiquery	vt	cwe	fwe	qa_ 1	qa_ 2	Mean
Uncompressed	1	1.00	1.00	1.00	0.98	0.90	0.62	0.97	0.98	0.62	0.36	0.68	0.60	0.32	0.77
Transformer-FSQ (2.0 $\times$ )	2	1.00	1.00	1.00	0.85	0.29	0.04	0.84	0.86	0.72	0.33	0.51	0.63	0.32	0.64
Transformer-FSQ (4.0 $\times$ )	3	1.00	0.99	0.99	0.86	0.43	0.04	0.85	0.86	0.79	0.33	0.47	0.62	0.33	0.66
Transformer-FSQ (8.0 $\times$ )	3	0.99	0.84	0.97	0.64	0.09	0.01	0.72	0.63	0.49	0.40	0.33	0.43	0.21	0.52
Transformer-FSQ (12.0 $\times$ )	3	0.99	0.63	0.93	0.50	0.05	0.00	0.56	0.43	0.48	0.34	0.18	0.32	0.19	0.43
Transformer-FSQ (16.0 $\times$ )	3	1.00	0.62	0.87	0.38	0.04	0.00	0.43	0.36	0.25	0.27	0.17	0.27	0.18	0.37

Appendix

Table 10: Per-task RULER on Qwen3-1.7B-Base (13 tasks of 50 items,

8{,}192

tokens, bfloat16). Finetuned methods use

0.25

tokens per parameter; KIVI, PyramidKV and VecInfer are training-free. Rows give achieved compression ratios; scores average complete seeds. Shading encodes the task score.

Method	seeds	niah_ single_1	niah_ single_2	niah_ single_3	niah_ multikey_1	niah_ multikey_2	niah_ multikey_3	niah_ multivalue	niah_ multiquery	vt	cwe	fwe	qa_ 1	qa_ 2	Mean
Uncompressed	1	1.00	0.98	1.00	1.00	0.96	1.00	0.99	0.99	0.90	0.92	0.90	0.90	0.44	0.92
Transformer-FSQ (8.0 $\times$ )	1	1.00	0.96	0.96	0.82	0.24	0.04	0.84	0.78	0.60	0.70	0.85	0.62	0.22	0.66
Transformer-FSQ (12.0 $\times$ )	1	1.00	1.00	1.00	0.66	0.14	0.02	0.79	0.77	0.58	0.58	0.73	0.52	0.20	0.61
Transformer-FSQ (32.0 $\times$ )	1	1.00	1.00	0.90	0.62	0.04	0.00	0.56	0.52	0.40	0.46	0.38	0.36	0.18	0.49
Transformer-FSQ (97.5 $\times$ )	1	0.96	0.74	0.30	0.20	0.00	0.00	0.17	0.14	0.00	0.07	0.15	0.14	0.14	0.23
Transformer-FSQ (409.6 $\times$ )	1	0.00	0.00	0.00	0.00	0.00	0.00	0.00	0.00	0.00	0.00	0.00	0.06	0.00	0.00

Appendix

Table 11: Per-task RULER on Qwen3-8B-Base (13 tasks of 50 items,

8{,}192

tokens, bfloat16). Finetuned methods use

0.25

tokens per parameter; KIVI, PyramidKV and VecInfer are training-free. Rows give achieved compression ratios; scores average complete seeds. Shading encodes the task score.

Compression 1.4 $\times$ –2.8 $\times$
Method	B/tok	Ratio	seeds	$n$	skip	gov_ report	narrative- qa	trec
Uncompressed	–	–	1	60	0	18.2	13.0	10.0
Transformer-FSQ	2048	2.0	3	180	0	17.1 ${\pm}$ 1.3	5.3 ${\pm}$ 1.2	25.0 ${\pm}$ 5.0
TransMLA	2048	2.0	3	180	0	12.7 ${\pm}$ 0.0	2.7 ${\pm}$ 0.6	16.7 ${\pm}$ 5.8
VecInfer	1601 §	2.6	1	60	0	18.1	8.0	10.0
Compression 2.8 $\times$ –5.7 $\times$

Appendix

Table 12: Per-task LongBench on Qwen3-0.6B-Base, each task in its own metric (over-length prompts truncated), over the 20 documents per task that every configuration scored. Finetuned methods use

0.25

tokens per parameter; KIVI, PyramidKV and VecInfer are training-free. Transformer-FSQ and uncompressed rows run in bfloat16, all other rows in float32.

Compression 1.4 $\times$ –2.8 $\times$
Method	B/tok	Ratio	seeds	$n$	skip	gov_ report	narrative- qa	trec
Uncompressed	–	–	1	60	0	17.9	12.0	40.0
Transformer-FSQ	2048	2.0	3	180	0	15.9 ${\pm}$ 0.6	8.9 ${\pm}$ 0.9	30.0 ${\pm}$ 13.2
TransMLA	2048	2.0	3	180	0	14.6 ${\pm}$ 0.1	4.1 ${\pm}$ 1.5	20.0 ${\pm}$ 5.0
VecInfer	1601 §	2.6	3	180	0	17.1 ${\pm}$ 0.5	10.1 ${\pm}$ 1.6	38.3 ${\pm}$ 2.9
Compression 2.8 $\times$ –5.7 $\times$

Appendix

Table 13: Per-task LongBench on Qwen3-1.7B-Base, each task in its own metric (over-length prompts truncated), over the 20 documents per task that every configuration scored. Finetuned methods use

0.25

tokens per parameter; KIVI, PyramidKV and VecInfer are training-free. Transformer-FSQ and uncompressed rows run in bfloat16, all other rows in float32.

Method	B/tok	Ratio	seeds	$n$	skip	gov_ report	narrative- qa	trec
Uncompressed	–	–	1	462	138	16.9	17.4	63.0
Compression 2.8 $\times$ –5.7 $\times$
KIVI	847 §	4.8	1	450	150	15.4	23.5	58.0
Compression 5.7 $\times$ –11.3 $\times$
Transformer-FSQ	512	8.0	1	600	0	15.0	10.8	43.5
VecInfer	721 §	5.7	1	450	150	15.6	16.3	59.0

Appendix

Table 14: Per-task LongBench on Qwen3-8B-Base, each task in its own metric (over-length prompts truncated), over the 180 gov_report, 70 narrativeqa and 200 trec documents that every configuration scored. Finetuned methods use

0.25

tokens per parameter; KIVI, PyramidKV and VecInfer are training-free. Transformer-FSQ and uncompressed rows run in bfloat16, all other rows in float32.

KV-Cache Compression

Key-value cache compression reduces the memory and bandwidth needed to retain attention states during long-context inference.

See all 285 papers

Training Transformers for KV Cache Compressibility
May 7, 2026Yoav Gelberg, Yam Eitan, Michael Bronstein +2Long-Context Language Modeling KV-Cache Compression
University of Oxford · 2Technion – Israel Institute of Technology · 3AITHYRA +1
Sequential KV Cache Compression via Probabilistic Language Tries: Beyond the Per-Vector Shannon Limit
Apr 10, 2026Gregory MagarshakKV-Cache Compression Memory Compression
Thin Keys, Full Values: Reducing KV Cache via Low-Dimensional Attention Selection
Feb 16, 2026Hengshuai Yao, Xing Chen, Ahmed Murtadha +1Memory-Efficient Attention KV-Cache Compression
Sapient Intelligence · Department of Computing Science, University of Alberta

May 7, 2026J/K move · Enter open · S save

Training Transformers for KV Cache Compressibility

Yoav Gelberg, Yam Eitan, Michael Bronstein, Yarin Gal+1

University of Oxford · 2Technion – Israel Institute of Technology · 3AITHYRA+1

Long-Context Language Modeling KV-Cache Compression

Open paper AskPDF arXiv

Long-context language modeling is increasingly constrained by the Key-Value (KV) cache, whose memory and decode-time access costs scale linearly with the prefix length. This bottleneck has motivated a range of context-compression methods, from token-level summarization to recent optimization-based KV compression methods. These post-hoc methods operate on the KV cache of a fixed pretrained model, so their effectiveness is fundamentally limited by how well the model's internal representations can be compressed. In this work, we formalize the notion of KV compressibility and show that it is a property of the learned representations, rather than of the context alone. We prove that almost any sequence-to-vector function admits both highly compressible and inherently non-compressible transformer implementations, highlighting the need to guide transformers toward compressible representations during training. Motivated by this, we propose KV-Compression Aware Training (KV-CAT), a continued pretraining procedure that incentivizes the emergence of compressible representations. We introduce a train-time KV sparsification policy that masks KV slots during training. This forces the model to use fewer KV slots and encourages it to learn representations amenable to post-hoc compression. Empirically, we show that KV-CAT improves the quality-budget tradeoff of downstream compression methods across retrieval, long-context question answering, and perplexity-based evaluation of compressed-prefix continuation.

Abstract

KV-Cache Compression

Explore similar work

Training Transformers for KV Cache Compressibility

Sequential KV Cache Compression via Probabilistic Language Tries: Beyond the Per-Vector Shannon Limit

Thin Keys, Full Values: Reducing KV Cache via Low-Dimensional Attention Selection