LFQ: Logit-aware Final-block Quantization for Boosting the Generation Quality of Low-Bit Quantized LLMs

As large language models continue to scale, low-bit weight-only post-training quantization (PTQ) offers a practical solution to their memory-efficient deployment. Although block-wise PTQ is capable of matching the full-precision (FP) baseline on basic language modeling and understanding, its quality is degraded for \textit{generative} tasks---especially at longer responses and extended chains of thought, which is critical in boosting task accuracy. We attribute this shortfall to two factors: (i) the omission of the unembedding layer (the LM head) in block-wise optimization and (ii) the reliance on the mean squared error (MSE) objective. Both factors cause the token probability distribution of the quantized model to misalign with that of the FP model, yielding notable accuracy drops on text generation benchmarks. To rectify the discrepancy, we introduce \emph{Logit-aware Final-block Quantization (LFQ)}, a simple yet effective enhancement to block-wise PTQ that quantizes the final Transformer block by minimizing the cross-entropy between the logits of the FP model and those of its quantized counterpart. By aligning token probabilities at the logit level in the final block, LFQ consistently improves the accuracy of complex generation tasks over state-of-the-art block-wise PTQ across diverse model families, while maintaining parity with FP baselines on language modeling and understanding.

Large language models are powerful, but they require a large amount of memory, which makes them expensive and difficult to run in practice. One common way to reduce this cost is to compress the model after training, so that it can be stored and used more efficiently. However, we find that existing compression methods can hurt the model’s ability to generate high-quality answers. In this paper, we identify an important reason for this problem: current methods do not sufficiently preserve the model’s final word-prediction behavior. To address this, we propose LFQ, a simple improvement that makes the compressed model’s final predictions more closely match those of the original model. This helps the compressed model keep more of its generation ability while still reducing memory use. Across several model families and tasks, LFQ improves performance on challenging generation benchmarks while maintaining strong performance on simpler language tasks.