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

2026-05-28

AI Summary

This work addresses the performance degradation commonly observed in low-bit post-training quantization (PTQ) for generative tasks, particularly in long-text generation and complex reasoning scenarios where maintaining the output quality of full-precision models remains challenging. The authors propose a logit-aware quantization method specifically targeting the final Transformer layer, which, for the first time, incorporates the language model head into block-level optimization. By replacing the conventional mean squared error (MSE) objective with cross-entropy between logits, the approach aligns the token probability distributions of the quantized and full-precision models. Notably, this method introduces no additional inference overhead and achieves substantial accuracy improvements on complex generation tasks across multiple mainstream large language models, while preserving performance on language modeling and understanding benchmarks comparable to that of the original full-precision models.

Abstract

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 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 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.

Problem

• post-training quantization
• large language models
• text generation
• token probability distribution
• low-bit quantization

Innovation

• Logit-aware Quantization
• Final-block Quantization
• Low-bit PTQ
• Cross-entropy Optimization
• LLM Generation Quality