FreqKV: Frequency Domain Key-Value Compression for Efficient Context Window Extension

Dear reviewer, thanks for taking the time to review our paper and for your insightful comments. We will discuss your concerns and refine our paper accordingly.

Comparison with other compression methods (Q1, Q3)

Our work is an efficient context window extension method that implements an iterative KV compression manner. Our intuition is to extend the context window of a pre-trained LLM effectively and efficiently. The baseline methods in Tab. 1, LongLoRA and LoCoCo, are the most relevant to our work. Our method compresses KV states iteratively to extend the context window and optimizes both fine-tuning and inference efficiency.

From the perspective of KV compression, We compare our method with different strategies like SnapKV and PyramidKV on LongBench. LongBench is the most commonly used for evaluation by these methods. Since our method is capable of compressing KV cache during inference, we compare it with these KV compression methods on LongBench to evaluate language understanding performance and generation quality. However, SnapKV and PyramidKV only compress the KV cache and conduct experiments within the original context window. They are not evaluated on contexts out of the window in their paper. We have tested SnapKV and PyramidKV on PG19 and Proof-pile and got ppl > 100 when extending contexts to 8K on llama2.

As for LoCoCo, they only report an average score on LongBench in their paper and the evaluation script is not given in their official repository. We have implemented it by ourselves and got a much worse performance.

Comparison with local attention (Q2)

We have evaluated the performance (PPL) of the Local Attention baseline on Proof-plie with llama2-7b as in the following table. "Local Attention" stands for keeping sink tokens and the latest tokens in the cache. It shares the same sink size and retaining size as FreqKV. The results show that the model benefits from our FreqKV when extending context length.

Further Interpretation of our work

As introduced in Sec. 1, we are inspired by the well-established technique to compress images in the frequency domain in CV. It has been observed that low-frequency channels are more important for CNNs. It occurs to us whether the hidden states of texts can be compressed in the frequency domain. We transform key states and value states in LLaMA-2-7b from the time domain to the frequency domain for power spectrum analysis. As shown in Figure 1, the energy distribution increasingly concentrates on low-frequency components as the computation process progresses inside the model. This suggests that high-frequency components, which contribute less to the overall information, can be discarded without a significant impact on performance, thereby enhancing computational efficiency.

Previous work in NLP has studied to train Transformer encoder architectures in the frequency domain like FNet and Fourier Transformer. Our work aims to extend the context window efficiently by compressing key-value states in the frequency domain for decoder-only LLMs. For contexts out of the original window size, the model only uses the compressed KV. Additional training adapts the model to leverage the new pattern of KV when extending to longer contexts.

Regarding all the typos you pointed out

Thank you for your suggestions. We will revise them in the paper.

We hope that the provided analysis and discussion have addressed your concerns.

Dear reviewer, thanks for taking the time to review our paper and for your insightful comments. We will discuss your concerns and refine our paper accordingly.

Choice of the model (Q1)

FreqKV is an efficient context window extension method that implements an iterative KV compression manner. Our intuition is to extend the context window of a pre-trained LLM effectively and efficiently. Most related works (LongLoRA, LoCoCo, and Activation Beacon) use llama2 as the base model to extend its context window. Following these studies, we use llama2 to evaluate the performance of context extension and long-context understanding. Models after llama2 have already been trained to support longer contexts (32k for Mistral and 128k for LLaMA3). Extending their context length requires more resources than we currently have, so we are not able to conduct experiments on these models.

Experiments on other benchmarks (Q2)

We find Ruler is a very recent benchmark and is not used for evaluation by the baseline methods in our paper. Due to time and resource constraints, we choose the Needle-in-a-Haystack task, aligning with most prior methods. FreqKV achieves an accuracy of 86.8%, significantly outperforming SnapKV's 49.6%. Unlike FreqKV, which extends the context window while compressing the KV cache, SnapKV focuses solely on KV compression, leading to an accuracy drop to 0 when the test context exceeds the original window size. Detailed results are provided in the appendix of the revised version.

We hope that the provided analysis and discussion have addressed your concerns.

Dear reviewer, thanks for taking the time to review our paper and for your insightful comments. We will discuss your concerns and refine our paper accordingly.

Compression overhead (Q1)

As introduced in Sec. 4.2, KV is compressed in the frequency domain. The whole compression process consists of DCT, filtration of high-frequency components, and IDCT. KV is transferred to the frequency domain by DCT, and then compressed. IDCT is used to convert the frequency components back to the time domain. (Line 202) It does not increase the size of the compressed KV.

Our compression occurs on-the-fly during inference. As illustrated in Sec. 4.3, when the cache is not filled, subsequent tokens could get in and fill the cache. When the cache is filled, it will be compressed so that more tokens could get in until the cache is filled again.

We maintain S = 4 sink tokens uncompressed. The retaining ratio γ in compression is set to 0.5. Therefore, the retaining size during each compression is L = γ · (N − S) = 2046. As long as the cache size reaches its capacity of N = 4096, the 4092 states since the 5-th state in the cache will be compressed into 2046 states. The compression is performed every N − L − S tokens with the complexity of O(NlogN).

To quantify compression overhead, we have measured FLOPs (TFLOPs) with different context lengths on llama2-7b. "Sink+Recent" stands for keeping sink tokens and the latest tokens in the cache. It shares the same sink size and retaining size as FreqKV. The difference in FLOPs between the two methods shows the overhead of compression. The statistics are given in the following table. It shows that the computation overhead of our compression process grows less than 0.5% even with a length of 16K and could be negligible.

Study of attention sinks (Q2)

We have evaluated the performance (PPL) of the attention sink baseline on Proof-plie as in the following table. "Sink+Recent" stands for keeping sink tokens and the latest tokens in the cache. It shares the same sink size and retaining size as FreqKV. The results show that the model benefits from our FreqKV when extending context length.

More long-context benchmarks (Q3)

We have evaluated FreqKV on the Needle-in-a-Haystack task and compared it with SnapKV. FreqKV achieves an accuracy of 86.8%, significantly outperforming SnapKV's 49.6%. Unlike FreqKV, SnapKV is solely a KV compression method and cannot extend the LLMs' context window, causing its accuracy to drop to 0 when the context exceeds the original window. Detailed results are provided in the appendix of the revised version.

Since GSM8k is a math reasoning benchmark with question lengths under 1k tokens, it is not considered a long-context benchmark, and we do not evaluate our method on it.

Why llama2 but not llama3 (Q4)

FreqKV is an efficient context window extension method that implements an iterative KV compression manner. Our intuition is to extend the context window of a pre-trained LLM effectively and efficiently. Most related works (LongLoRA, LoCoCo, and Activation Beacon) use llama2 as the base model to extend its context window. Following these studies, we use llama2 to evaluate the performance of context extension and long-context understanding. Models after llama2 have already been trained to support longer contexts (32k for Mistral and 128k for LLaMA3). Extending their context length requires more powerful GPUs and is more difficult to evaluate.

We hope that the provided analysis and discussion have addressed your concerns.

Dear reviewer,

As you are interested in the performance of FreqKV on llama3, we have conducted experiments on Llama3-8B-Instruct to validate its effectiveness. Llama3.1/3.2 have already been trained to support longer contexts (128K). Extending their context length requires more resources than we currently have, so we are not able to conduct experiments on these models.

The original context window length of Llama3-8B-Instruct is 8K. It is equipped with GQA (Grouped-Query Attention), which means it has a lower proportion of parameters for attention modules than Llama2-7B-Base/Chat (MHA). We use FreqKV to extend the context window from 8K to 16K and evaluate the performance on six Single-Doc QA and Summarization tasks from LongBench. The results of the vanilla model and the SOTA KV compression method SnapKV are also reported in the following table.

Llama3 and SnapKV are evaluated with a context length of 8K. Their performance drops significantly when setting the context length to 16K. FreqKV is evaluated with a context length of 16K. While SnapKV maintains the performance of the vanilla model, FreqKV enables the model to handle longer contexts and achieves improvements on most tasks. It demonstrates the effectiveness of FreqKV when applied to Llama3-8B-Instruct.

Moreover, we evaluate FreqKV and SnapKV on Needle-in-a-Haystack with the document length ranging from 1K to 16K. FreqKV performs well in extending the context window and achieves an average accuracy of 95.2%, significantly outperforming SnapKV which achieves only 51.5%.

Detailed results are provided in the Appendix. E of the revised version. These experimental results demonstrate the effectiveness of FreqKV when applied to Llama3-8B-Instruct.

We hope that the provided results on Llama3 have addressed your concerns.

Dear reviewer, thanks for taking the time to review our paper and for your insightful comments. We will discuss your concerns and refine our paper accordingly.

Comparison with LongLoRA (W1)

Our intuition is to extend the context window of a pre-trained LLM effectively and efficiently. LongLoRA uses sparse attention to reduce the training FLOPs when extending to a longer context. However, it still requires the original attention on the full sequence during inference. The results of LongLoRA reported in Tab. 1 are using full KV of the sequence during evaluation. If they still use the sparse attention during inference, their performance drops significantly. On the contrary, our method compresses KV states iteratively to extend the context window and optimizes both fine-tuning and inference efficiency. We use the compressed KV during evaluation.

Compression overhead (W3)

Our compression occurs on-the-fly during inference. As illustrated in Sec. 4.3, when the cache is not filled, subsequent tokens could get in and fill the cache. When the cache is filled, it will be compressed so that more tokens could get in until the cache is filled again.

We maintain S = 4 sink tokens uncompressed. The retaining ratio γ in compression is set to 0.5. Therefore, the retaining size during each compression is L = γ · (N − S) = 2046. As long as the cache size reaches its capacity of N = 4096, the 4092 states since the 5-th state in the cache will be compressed into 2046 states. The compression is performed every N − L − S tokens with the complexity of O(NlogN).

To quantify compression overhead, we have measured FLOPs (TFLOPs) with different context lengths. "Sink+Recent" stands for keeping sink tokens and the latest tokens in the cache. It shares the same sink size and retaining size as FreqKV. The difference in FLOPs between the two methods shows the overhead of compression. The statistics are given in the following table. It shows that the computation overhead of our compression process grows less than 0.5% even with a length of 16K and could be negligible.

Extension to longer contexts (Q1)

To extend to longer contexts (like > 32K) without additional modifications, the only necessity might be GPUs with larger memory. As shown in Tab. 2, the memory usage for each GPU is 44.90GB. We are not able to extend our method to longer contexts directly since we only have 48GB-ADA6000 and 24GB-RTX4090.

However, one of the possible solutions is model quantization, which employs low-bit integer (INT) and low-bit floating-point (FP) formats [1,2]. Another direction is model pruning, which removes part of the model parameters [3,4]. These methods are orthogonal to context extension or compression and can be combined together for specific requirements.

[1] Integer or Floating Point? New Outlooks for Low-Bit Quantization on Large Language Models

[2] AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration

[3] LLM-Pruner: On the Structural Pruning of Large Language Models

[4] SliceGPT: Compress Large Language Models by Deleting Rows and Columns

Choice of the model (Q2)

FreqKV is an efficient context window extension method that implements an iterative KV compression manner. Our intuition is to extend the context window of a pre-trained LLM effectively and efficiently. Most related works (LongLoRA, LoCoCo, and Activation Beacon) use llama2 as the base model to extend its context window. Following these studies, we use llama2 to evaluate the performance of context extension and long-context understanding. Models after llama2 have already been trained to support longer contexts (32k for Mistral and 128k for LLaMA3). Extending their context length requires more resources than we currently have, so we are not able to conduct experiments on these models.

Experiments on Needle-in-a-Haystack (Q2)

We have evaluated FreqKV on the Needle-in-a-Haystack task and compared it with SnapKV. FreqKV achieves an accuracy of 86.8%, significantly outperforming SnapKV's 49.6%. Unlike FreqKV, SnapKV is solely a KV compression method and cannot extend the LLMs' context window, causing its accuracy to drop to 0 when the context exceeds the original window. Detailed results are provided in the appendix of the revised version.

Adaptive retaining ratios across layers (Q4)

We agree with you that different layers utilize KV in different patterns. As shown in Fig. 1, the energy of key states and value states is increasingly concentrated in the low-frequency components as the layer gets deeper. It might further improve FreqKV if we retain more cache in shallow layers and less cache in deep layers. Moreover, we can even use different retaining ratios in different heads as mentioned in Appendix A.

For the clarity of our method, we do not introduce an adaptive method for determining retaining ratios across layers. Most of the long-context compression methods can be equipped with an adaptive method to determine hyper-parameters. It could be promising to study specific differences and associations across different layers, heads, or other modules in future work.

Automatic method for determining the retaining ratio (Q3)

We find this question quite similar to the Question 4. We have not introduced an automatic method for determining the retaining ratio for the clarity of our method. It deserves further exploration in future work.

We hope that the provided analysis and discussion have addressed your concerns.

The authors provided helpful clarifications to some of my questions (specifically about the NIAH task), and the proposed method is interesting. However, the results are not particularly strong, and no results are shown for context lengths exceeding 32k or on larger models. For these reasons, I am keeping my rating marginally below acceptance.

For these reasons, I am keeping my rating marginally above acceptance.

We appreciate your response and acknowledgement. But we are a little confused about the rating. In ICLR, 6 is marginally above acceptance while your initial rating 5 is marginally below the acceptance (unlike NeurIPS). If you consider our work marginally above acceptance, the rating should be raised.

Dear reviewer, thanks for taking the time to review our paper and for your insightful comments. We will discuss your concerns and refine our paper accordingly.

Misunderstanding of our compression process (Weaknesses and Q3)

Most of the weaknesses and questions you mentioned are due to the misunderstanding of our compression process. Our work does not decompress the compressed KV for attention computation with IDCT.

As introduced in Sec. 4.2, KV is compressed in the frequency domain. The whole compression process consists of DCT, filtration of high-frequency components, and IDCT. KV is transferred to the frequency domain by DCT, and then compressed. IDCT is used to convert the frequency components back to the time domain (Line 202). It does not increase the size of KV which has been compressed. Therefore, the memory and computation will be saved for the future decoding procedure.

Compression overhead (Q1)

Our compression occurs on-the-fly during inference. As illustrated in Sec. 4.3, when the cache is not filled, subsequent tokens could get in and fill the cache. When the cache is filled, it will be compressed so that more tokens could get in until the cache is filled again.

We maintain S = 4 sink tokens uncompressed. The retaining ratio γ in compression is set to 0.5. Therefore, the retaining size during each compression is L = γ · (N − S) = 2046. As long as the cache size reaches its capacity of N = 4096, the 4092 states since the 5-th state in the cache will be compressed into 2046 states. The compression is performed every N − L − S tokens with the complexity of O(NlogN).

To quantify compression overhead, we have measured FLOPs (TFLOPs) with different context lengths. "Sink+Recent" stands for keeping sink tokens and the latest tokens in the cache. It shares the same sink size and retaining size as FreqKV. The difference in FLOPs between the two methods shows the overhead of compression. The statistics are given in the following table. It shows that the computation overhead of our compression process grows less than 0.5% even with a length of 16K and could be negligible.

Justification of training requirement (Q2)

As shown in the formula of Eq.1 and 2, DCT and IDCT can be regarded as operations of mixing tokens. We only retain low-frequency components and remove high-frequency components for compression. The elements in the compressed KV are not the initial tokens, but the mixture of tokens. (If we retain all the components in the frequency domain and conduct IDCT, the tokens are the same as the initial tokens.) Moreover, our work aims to extend the context window for a pre-trained LLM. For contexts out of the original window size, the model only uses the compressed KV. Therefore, additional training would be necessary for models to learn to leverage KV states that are compressed in the frequency domain when extending the context window.

We have also measured the performance (PPL) of FreqKV without training on Proof-pile with llama2-7b. The scores are reported in the following table. It shows that the model benefits from the training procedure and learns the new pattern introduced in KV.

Impact of different retaining ratios (Q4)

As given in Sec 6.1, we have conducted detailed studies into the impact of different retaining ratios including computation cost, inference time, and perplexity evaluation. It should be noted that although the FLOPs of different retaining ratios presented in Table 4 are very close, larger retaining ratios lead to smaller chunk sizes, which determines how many attention scores will be masked. As a result, the attention matrix becomes denser and costs an increase in inference overhead with a minimal performance improvement as the retaining ratio increases from 0.5 to 0.75. This study justifies our choice of setting a 0.5 retaining ratio as the default for effectiveness and efficiency.

We hope that the provided analysis and discussion have addressed your concerns.

I appreciate the authors' efforts in preparing the responses. Although the author's responses partially addressed my concerns, I still concur with the other reviewers' comments about the limitation in evaluation, especially in demonstrating the effectiveness of the proposed algorithm in the other models. Therefore, I will maintain my original score.