LogQuant: Log-Distributed 2-Bit Quantization of KV Cache with Superior Accuracy Preservation

🔼 Figure 2 visualizes the unpredictable nature of attention weights in LLMs over time. It displays the maximum attention score for each token position across four consecutive decoding steps, using the Llama3-8B-Instruct model on both GSM8K and OpenBookQA datasets. The unpredictable fluctuations highlight the challenges in accurately predicting important tokens for efficient memory management, especially when considering compression techniques.

read the caption

Figure 2: The maximum attention score of each token position across four consecutive decoding steps, marking the high attention positions for illustrating the unpredictable nature of attention scores. This analysis was conducted using Llama3-8B-Instruct (Dubey et al., 2024) on the GSM8K (Cobbe et al., 2021) and OpenBookQA (Mihaylov et al., 2018) datasets.

🔼 Figure 3 illustrates the distribution of attention weights across different token positions within the context window of a large language model. Boxplots summarize the attention scores across all attention heads, showing the median and interquartile range (25th and 75th percentiles) for each token position. The figure highlights that the attention scores for the first two tokens (referred to as ‘sink tokens’), which are typically the most recently processed tokens, exhibit a higher median and overall distribution of scores compared to the combined scores of the subsequent 128 tokens. This observation supports the notion that the most recent tokens carry more weight in the attention mechanism. The data presented in this graph is derived from experiments conducted using the Llama3-8B-Instruct model on the GSM8K dataset.

read the caption

Figure 3: Attention distribution across different token positions, represented as boxplots based on 25% quantiles across all attention heads. The median and overall distribution of attention scores for sink tokens (Xiao et al., 2023) (tokens 0 and 1) are greater than the sum of the most recent 128 tokens. The attention scores are derived from experiments using Llama3-8B-Instruct (Dubey et al., 2024) and the GSM8K (Cobbe et al., 2021) dataset.

🔼 Figure 4 compares the effectiveness of different token selection methods for compressing the key-value (KV) cache in large language models (LLMs). It shows the attention coverage achieved by four different methods: LogQuant, Kivi, Streaming, and H2O. The comparison is made across various LLMs (Llama3-8B-Instruct, Qwen-2-7B-Instruct, Phi-3-mini-128k-Instruct) and uses a subset of the GSM8K dataset. The x-axis represents the length of the reserved portion of the KV cache, while the y-axis shows the average attention scores captured by each selection method. The figure demonstrates that LogQuant achieves better attention coverage than the other methods, indicating its superior ability to select and retain important tokens while reducing memory usage. The first two sink tokens (tokens with consistently high attention scores) are excluded from the analysis to focus on the relative performance of the selection methods.

read the caption

Figure 4: The attention coverage without the first two sink tokens for different selection methods (Liu et al., 2024c; Xiao et al., 2023; Zhang et al., 2024) and different models (Dubey et al., 2024; Yang et al., 2024; Abdin et al., 2024), tested on a subset of the GSM8K (Cobbe et al., 2021) dataset. Details of LogQuant will be introduced in Section 2.5.

🔼 This figure compares the effects of two different KV Cache compression strategies: quantization and eviction. It demonstrates that using quantization to reduce the numerical precision of tokens instead of removing them entirely (eviction) leads to significantly less distortion of the attention distribution. The plot visualizes the L1 error between the original attention distribution and the distributions after compression using both methods.

read the caption

Figure 5: Eviction and Quantization Loss on Attention Distribution

🔼 LogQuant’s KV cache compression workflow is illustrated. Initially, 3W tokens are kept at full precision. A log-sparse filtering strategy is then applied to the first 2W tokens, resulting in half of them being quantized. This process reduces the number of full-precision tokens, ultimately compressing the reserved token length back down to 2W. This cyclical process ensures efficient memory management.

read the caption

Figure 6: LogQuant’s KV cache compression workflow. The number of reserved original-precision tokens increases from 2⁢W2𝑊2W2 italic_W to 3⁢W3𝑊3W3 italic_W. We then apply a log-sparse strategy to filter the first 2⁢W2𝑊2W2 italic_W tokens, quantize half of these tokens, and compress the reserved token length back to 2⁢W2𝑊2W2 italic_W.

🔼 This figure displays the accuracy (exact match) results on the GSM8K dataset for various language models using different compression ratios. It visualizes the performance trade-off between compression and accuracy for different models and compression strategies, allowing for a comparison of the effectiveness of LogQuant relative to other approaches.

read the caption

Figure 7: Accuracy(EM) with different compression ratio in GSM8K tasks for different models.

🔼 Figure 8 illustrates a comparison of memory usage and throughput between LogQuant with 2-bit quantization and a 16-bit baseline. The experiment used the Hugging Face generation pipeline, the Llama 3.1-8B model, and an NVIDIA H100 GPU. The graph shows how throughput and memory consumption change as the batch size increases for both methods. This helps to demonstrate the memory efficiency and performance gains achieved by LogQuant.

read the caption

Figure 8: memory usage and throughput comparison between 2bit LogQuant and 16bit baseline under huggingface generation pipeline with llama3.1-8B and H100.

🔼 This table compares the L1 error, a measure of difference between the original attention scores and those obtained after applying either eviction or quantization techniques. It shows how much the attention distribution is altered by each method, indicating the potential impact on model accuracy. Lower L1 error values suggest better preservation of the original attention distribution.

read the caption

Table 2: Comparison of L1 error with original attention for eviction and quantization.

🔼 This table presents a comparison of the accuracy achieved using different bit precisions (2-bit and 4-bit) for both KiVi and LogQuant quantization methods, in relation to the baseline accuracy (using original precision) on Llama3.1-8B model for various tasks. The Delta (Δ) column indicates the difference in accuracy percentage between each method and the baseline. For detailed per-task accuracy scores, please refer to Table C6.

read the caption

Table 3: Accuracy of Different Precision on Llama3.1-8B. Refer to the Table C6 for the scores of each specific task. The ΔΔ\Deltaroman_Δ shows the difference to baseline.

🔼 This table presents the average performance across seven task groups (Math, Code, Few-shot, Multi-QA, Single-QA, Summarization, and Synthetic) for three different LLMs (Llama-3.1-8B, Qwen1.5-7B-Chat, and Phi-3-mini-128k) using 2-bit quantization for the KV cache. It compares the accuracy of LogQuant against Kivi and a 16-bit baseline. The best result for each model and task using 2-bit quantization is highlighted in bold. Detailed results for each individual task within each group can be found in Table D7.

read the caption

Table 4: Task Group Average Score for Different Models with 2-bit KV Cache Quantization. (The best result of 2-bit quantization is in bold. Refer to Table LABEL:tab:longbench_all for the scores of each specific task in LongBench.)

🔼 Table B5 presents a comprehensive overview of the datasets used for evaluating the performance of the proposed model. For each dataset, the table provides the task group it belongs to (e.g., Math, Single-Document QA, etc.), the dataset name, the average length of the data points (calculated as the number of words for English datasets and the number of characters for Chinese datasets), the evaluation metric used (e.g., Accuracy (EM) for Exact Match accuracy, Accuracy (CLS) for classification accuracy, F1 score, Rouge-L score), the language of the dataset (English or Chinese), and the total number of data samples.

read the caption

Table B5: Overview of all test datasets. ‘Avg len’ (average length) is computed using the number of words for the English (code) datasets and the number of characters for the Chinese datasets. ‘Accuracy (CLS)’ refers to classification accuracy, while ‘Accuracy (EM)’ refers to exact match accuracy

🔼 This table presents a comparison of the accuracy achieved by different quantization methods (KiVi and LogQuant) on the Llama3.1-8B-Instruct model across various datasets. It shows the performance of both methods using 2-bit and 4-bit quantization, comparing them against a baseline of original precision (16-bit). The results highlight the trade-off between compression ratio and accuracy, indicating how much accuracy is lost when using lower precision quantization. The datasets included are diverse, encompassing various natural language tasks.

read the caption

Table C6: Comparison on Llama3.1-8B-Instruct of different quantization precisions

🔼 This table presents the performance of different models (Llama-3.8B-Instruct, Llama-3.1-8B-Instruct, Qwen1.5-7B-Chat-AWQ, Qwen1.5-14B-Chat-AWQ, Qwen2-7B-Instruct) on various subtasks within the LongBench benchmark, categorized into groups like Math, Single/Multi-Document QA, Summarization, Few-shot Learning, Synthetic Tasks, and Code Completion. It shows the accuracy (score) achieved using the 16-bit baseline, KiVi’s 2-bit quantization, and LogQuant’s 2-bit quantization method for each model on each subtask. This allows for a comparison of the accuracy loss introduced by each quantization technique against the full precision baseline.

read the caption

Table D7: LongBench score of each dataset