Lexico: Extreme KV Cache Compression via Sparse Coding over Universal Dictionaries

We thank you for recognizing the novelty of applying sparse dictionary learning for KV cache compression and our extensive experimental evaluation—a strength also highlighted by reviewers Vzmg and EtKf. We address each concern below.

Q1: Task-specific dictionaries

Yes, using task-specific dictionaries may improve task accuracies, but would incur the additional training cost of custom fine-tuning. While we find this a promising direction to improve performance, our motivation to employ a universal dictionary was to allow it to be used off-the-shelf for various tasks without the need for additional training. That said, we can see that one may want to explore task-specific dictionaries to further enhance the performance of Lexico in specific domains.

As MMLU-Pro Law does not have a training split, we tested task-specific dictionaries for GSM8K. Specifically, we present empirical results with the following dictionaries:

1. Universal dictionary: Dictionary trained on WikiText-103.
2. Dictionary pre-trained on GSM8K: Dictionary trained from scratch on the GSM8K training set.
3. Universal dictionary fine-tuned on GSM8K: Dictionary trained on WikiText-103 and then fine-tuned on GSM8K.

From the experimental results, we observe that training a dictionary from scratch solely on GSM8K significantly underperforms compared to the universal dictionary, especially at low sparsities. This indicates overfitting to the small GSM8K training set, failing to generalize even to its test set. In contrast, fine-tuning a universal dictionary on GSM8K sometimes provides marginal improvements but adds the additional overhead of task-specific fine-tuning. Therefore, we may expect a slight increase in Lexico's performance on MMLU-Pro given a dataset with similar text distributions.

Q2: Table 10 computation of keys and values

Apologies for the confusion. To clarify, the reported latency for the Lexico forward pass in Table 10 represents the entire process of generating one token. The row stating "Lexico: forward pass using $q (K_\text{csr}D_k^\top ) ^\top$" should be corrected to "Lexico: forward pass using $q (K_\text{csr}D_k^\top ) ^\top$ and $V_\text{csr}D_v^\top$", as the overall latency measurement indeed includes the additional compute for value vector reconstruction. We will revise the manuscript to explicitly note that the Lexico forward pass timing encompasses both the key and value reconstruction computations.

W1: Do value vectors also lie in low-rank subspaces?

Yes, our findings show that value vectors also lie in a low-rank subspace. We will include an analogous figure as Figure 2 for value vectors. Below we computed the average relative reconstruction errors over 1k tokens from the WikiText test split using a universal dictionary of size $N=4096$.

Although the sparse approximation for values is less effective than for keys, we believe that our empirical results on real-world benchmarks (LongBench, GSM8K) clearly demonstrate the effectiveness of our dictionary-based sparse representation.

W2: Dictionary training time across different sparsity levels

Training dictionaries is cheap because they only need to be trained once and can be provided by a third-party. As we report in L203, a sparsity of s=32 and a dictionary size of N=1024 requires only 2 hours on an A100 GPU. We provide the training runtime for different sparsities and dictionary sizes for Llama-3.1-8B-Instruct below and will incorporate these results in a revised edition.

Thank you for your thoughtful review. We are glad that you appreciated the novelty of Lexico, our effort in optimizations, and the overall readability of the paper. We address the remaining weaknesses and questions below.

W1: Paper organization

We greatly appreciate the reviewer’s feedback regarding the organization of our paper. In our revision, we plan to:

• Reorganize the Layout: We will reposition diagrams and tables so that they enhance readability, moving tables to the top of pages rather than interspersed between paragraphs. This can be done without difficulty in the case of acceptance given the additional ninth page.
• Move Mistral results for readability: We will move the Mistral 7B results to the supplementary material, addressing concerns of redundancy while still providing experiments sufficient for our message. The original intent was to demonstrate that Lexico is not tailored to any particular model.
• OMP Algorithm and Latency Measurements: We agree that these are important and will adjust the layout to fit them in the main text.

W2: Tradeoff between encoding time and accuracy

We compared OMP with alternative approaches—Iterative Hard Thresholding (IHT) and Lasso—on the WikiText test set using our universal dictionary (size N=4096). While IHT and Lasso are iterative methods that can achieve better accuracy with additional iterations (trading off computation time for improved performance), OMP’s greedy nature does not allow for such tuning. In our experiments, even after increasing the number of iterations (and thus compute time) for IHT and Lasso, they still produced higher reconstruction errors than OMP across various sparsity levels. Moreover, OMP was the fastest since its complexity scales with the sparsity s, which in our case is much smaller than N.

If the reviewer is aware of alternative approaches that may offer better tradeoffs, we would be eager to explore them.

Q1: Estimating the dictionary via least squares

Indeed a closed-form solution via least squares exists in each update. However, this approach proved prohibitively slow for large training sets. The solution requires forming and inverting an $N \times N$ matrix, an operation that scales on the order of $\mathcal{O}\left ( N^3 \right )$. With $N = 4096$ in our experiments, such step quickly dominates overall runtime and overshadows any advantage of jumping straight to the global minimizer.

Related work: Comparing sparse coding to product quantization

Thank you for the suggestion. We will add the reference and comparison to our Related Work. Product quantization and sparse coding both use a form of "codebook" for vector compression. Sparse coding represents each vector as a sparse linear combination of learned dictionary atoms, while product quantization partitions vectors (e.g., the first k coordinates) and clusters each sub-vector, concatenating the nearest centroids. PQCache (https://arxiv.org/abs/2407.12820) learns its centroids from input tokens; our method trains the dictionary offline. Despite these differences, both exploit the key vectors’ clustering structure for significant compression.

L149: overcomplete definition

It means that the dictionary has more atoms (number of vectors) than the dimensionality of the vectors. This allows multiple ways to represent the same signal, making it easier to find a good sparse representation.

Clarification on head-level compression

The compression is performed on key vectors of size $m$ (not $d$). The "one per layer" phrase in the paper refers to the dictionary rather than the KV cache. For each token, we have $L \times H$ key vectors of size $m$, while we have $L$ key dictionaries of shape $m \times N$. Thus, all $H$ heads within a layer share the same dictionary. An analogous setup applies to the value cache and dictionary.

Figure 2 axes

In Figure 2, each axis represents 7 tokens across all 8 heads of layer 10, yielding 7×8=56 elements (with some elements cropped for clarity). The plotted values are from key vectors for each head (with dimensionality m). In the left panel, both axes correspond to the first 7 tokens from the same sentence, whereas in the right panel each axis represents the first 7 tokens from two different sentences. The elements are sorted by similarity to highlight clustering.

Table 1 error

We measure the relative reconstruction error (RRE) for each key or value vector as $\text{RRE}=\frac{\|\mathbf{k}-\hat{\mathbf{k}}\|_2}{\|\mathbf{k}\|_2}$, where $\mathbf{k}$ is the original vector and $\hat{\mathbf{k}}$ its reconstruction. In Table 1, each reported value is the average RRE across 8k key/value vectors—corresponding to 1k tokens, each producing 8 head vectors—obtained from forward passes on the specified datasets. The standard deviation is computed over those 8k samples.

We thank the reviewer for the thoughtful and constructive feedback. We are pleased that you appreciated Lexico’s flexible design, near-lossless performance on benchmarks, and the advantage of using a dictionary for KV cache compression. It is also encouraging that reviewers Vzmg and EhUD similarly highlighted our method's strong performance, particularly regarding its improved memory vs. performance trade-offs compared to other compression methods. We address the weaknesses and questions below.

Hardware measurements

Section E of the paper already presents a latency analysis and, hence, claiming we "lack real hardware measurements" is misleading. In light of the reviewer's emphasis on this point, we will move that section into the main body in the revised manuscript. We have also extended the same setup described in Section E to include varying sparsity levels for a more fine-grained analysis; these results will be incorporated into the revised edition.

Although Lexico may incur higher latency at larger sparsities, it is designed to address memory-constrained scenarios where a modest latency trade-off is acceptable, and hardware or library optimizations (e.g., support for FP8 multiplications) can further narrow this gap while preserving substantial memory savings.

Lexico for CoT/reasoning tasks

We are glad the reviewer noted the potential for Lexico in long-form reasoning tasks. To investigate this, we tested a reasoning model, DeepSeek-R1-Distill-Llama-8B, on the AIME 2024 exam. We evaluated the performance trade-off when using Lexico to compress the KV cache memory of long reasoning traces:

At a moderate compression level ($s=24$), Lexico retains full-precision performance (30%). When the compression becomes more aggressive, as with Lexico ($s=14$) and KIVI (2 bits), the model's score declines, though Lexico still scores higher than KIVI at similar compression levels.

One approach to improve performance is to train task-specific or reasoning-specific dictionaries. Recent reasoning models generate traces that extend well beyond the prompt, so much of the KV cache is used for these model-generated traces. Because these traces may inhabit a different subspace than the prompt KV cache, training the Lexico dictionary specifically on output generations could be a promising direction for improving the performance-memory trade-off in long-form reasoning tasks. We demonstrate this potential for task-specific dictionaries and relevant ablations in the "Dictionary sources and adaptability across domains" section in our response to reviewer Vzmg.

We are pleased that you appreciated the clarity of our experiments, Lexico's effectiveness, and our use of sparse coding. EtKf and EhUD similarly recognized the soundness of our experimental design. We address your concerns below, and would be grateful if you reevaluate our work in light of them.

LongBench

We do not agree that our evaluation is limited. We have covered a comprehensive set of task categories such as document QA, summarization, few-shot tasks, and code completion; this aligns with prior work (e.g., KIVI) that evaluated on the same subset. For completeness, we have Lexico on all LongBench tasks with Llama-3.1-8B-Instruct below.

The newly added tasks demonstrate a similar trend as the other tasks. The performance remains nearly lossless when $s=16$ (21% of the original KV cache size). Our results remain consistent with our subset of tasks in Table 2. Even at $s=8$ (12% of the original), our approach maintains a strong performance-memory tradeoff; note that 12% is a regime 2-bit quantization cannot achieve. These results show Lexico's strength in handling long contexts.

Needle In A Haystack (NIAH)

We show that Lexico is resilient to such stress-testing. Using Llama-3.1-8B-Instruct with an 8k context size, we followed the settings in PyramidKV. We compared against token eviction methods—which are well-suited for tasks that require only a few tokens from the context—and report ROUGE-1 multiplied by 10. Note that the columns represent the needle depth ratio.

We find that Lexico outperforms competitors on a 20% KV cache budget, while it delivers similar performance on a 10% budget. Note that the 10% budget is not feasible for 2-bit quantization methods. While we believe that NIAH is a valuable synthetic benchmark, performance on the task does not strongly reflect that of practical performance, as exemplified in the underperformance of eviction-based methods on GSM8K (Figure 4).

Dictionary training cost

Training dictionaries is cheap because they only need to be trained once and can be provided by a third-party. As we report in L203, sparsity of s=32 and a dictionary size of N=1024 require only 2 hours on an A100 GPU. We provide the training runtime for different sparsities and sizes for Llama-3.1-8B-Instruct.

Latency

We kindly direct the reviewer to the Hardware measurements section in our response to reviewer EtKf (R2).

Dictionary sources and adaptability across domains

Per your suggestion, we investigated whether training a dictionary on a specific domain enhances the performance. In our ablation below, we trained a dictionary on different sources:

1. Dictionary trained on WikiText-103.
2. Dictionary trained on GSM8K training set.
3. Dictionary trained on WikiText-103 then finetuned on GSM8K.

We observe that training a dictionary from scratch on GSM8K significantly underperforms, especially at low sparsities. Overfitting to GSM8K fails to generalize even to its test set. In contrast, finetuning the universal dictionary on GSM8K provides marginal improvement but adds the additional overhead of custom fine-tuning.

Missing reference

We appreciate the reviewer’s suggestion and will incorporate QJL into our related work. While QJL introduces a 1-bit quantization approach for KV cache compression, its experiments are confined to 3- and 5-bit schemes. In contrast, our work benchmarks Lexico against 2- and 4-bit state-of-the-art methods, and further pushes the efficiency by exploring memory regimes lower than 2 bits. Whereas QJL applies its transformation solely to keys (with per-token quantization for values), our approach compresses both keys and values, offering a more comprehensive and efficient solution.