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

TL;DR

LogQuant introduces a log-based 2-bit quantization strategy for KV caches in LLM inference, driven by the observation that attention spike positions follow a log distribution. By combining a log-distributed token selection with a position-agnostic attention calculation, it preserves crucial tokens while reducing memory, outperforming prior 2-bit methods like KiVi in accuracy and delivering notable throughput gains under a fixed memory budget. Empirical results across Llama3, Qwen, and Phi models on GSM8K and LongBench demonstrate improved task accuracy, particularly on Math and Code Completion, alongside hardware-efficient benefits such as ~25% higher throughput and ~60% larger batch sizes on an H100 with HuggingFace. The approach integrates cleanly with the Transformer ecosystem, offering a practical path to efficient, long-context inference in resource-constrained environments.

Abstract

We introduce LogQuant, a groundbreaking 2-bit quantization technique for KV Cache in large language model (LLM) inference, delivering substantial memory savings while preserving superior performance. Previous methods either assume that later tokens are more important or attempt to predict important tokens based on earlier attention patterns. Both approaches, however, can result in performance bottlenecks or frequent mispredictions. LogQuant takes a different approach. By applying a log-based filtering mechanism, it selectively compresses the KV Cache across the entire context, achieving better performance with the same or even reduced memory footprint compared to existing methods. In benchmark tests, it enhances throughput by 25% and boosts batch size by 60% without increasing memory consumption. For challenging tasks such as Math and Code Completion, LogQuant improves accuracy by 40% to 200% at the same compression ratio, outperforming comparable techniques.LogQuant integrates effortlessly with popular inference frameworks like Python's transformers library. Implementation can be available in https://github.com/Concyclics/LogQuantKV.

Figures (8)

• Figure 1: The observed log-distribution pattern is evident not only in the magnitude of attention scores but also in the positions of attention spikes. These spikes become sparser as the model attends to tokens further from the most recent position, indicating that the model not only focuses on nearby tokens. This phenomenon, illustrated here with Llama3-8B-Instruct dubey2024llama3 on the GSM8K dataset GSM8K, is consistent across different tasks and models, as further detailed in Section \ref{['sec:Methodology']}.
• 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 dubey2024llama3 on the GSM8K GSM8K and OpenBookQA OpenBookQA2018 datasets.
• 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 streamingLLM (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 dubey2024llama3 and the GSM8K GSM8K dataset.
• Figure 4: The attention coverage without the first two sink tokens for different selection methods liu2024kivistreamingLLMzhang2024h2o and different models dubey2024llama3yang2024qwen2abdin2024phi, tested on a subset of the GSM8K GSM8K dataset. Details of LogQuant will be introduced in Section \ref{['sec:LogQuant']}.
• Figure 5: Eviction and Quantization Loss on Attention Distribution