Technical Summary: Logit-aware Final-block Quantization (LFQ)

1. Problem Statement

As Large Language Models (LLMs) scale, low-bit weight-only post-training quantization (PTQ) has become a standard for memory-efficient deployment. While block-wise PTQ methods (e.g., FlexRound, OmniQuant, Block-AP) successfully recover performance on language modeling and understanding tasks by minimizing the Mean Squared Error (MSE) between full-precision (FP) and quantized block outputs, they exhibit a critical failure in generative tasks.

The paper identifies that standard block-wise PTQ degrades the quality of text generation, particularly in long-form responses and complex reasoning chains (e.g., AIME, MATH). This degradation stems from two specific limitations in the optimization objective:

1. Omission of the Unembedding Layer: Existing block-wise methods optimize transformer blocks in isolation, ignoring the final linear projection layer (the LM head) that maps hidden states to token probabilities.
2. Reliance on MSE: Minimizing MSE between block outputs does not guarantee the alignment of token probability distributions. A small MSE in the hidden state space can result in a significant shift in the final token sampling distribution, causing the model to select incorrect tokens (top-1 flipping) or lose confidence in critical "aha" moments during reasoning.

This misalignment is less observable in non-autoregressive tasks but compounds over long generation trajectories, steering the model away from the FP baseline's reasoning path.

2. Methodology: Logit-aware Final-block Quantization (LFQ)

To address the misalignment between quantized models and FP baselines in generation tasks, the authors propose Logit-aware Final-block Quantization (LFQ).

Core Mechanism

LFQ modifies the standard block-wise PTQ pipeline by treating the final Transformer block differently from the preceding blocks:

• Intermediate Blocks (1 to $N-1$): Quantized using standard block-wise PTQ methods by minimizing the MSE between the outputs of the FP and quantized blocks.
• Final Block (Block $N$): Quantized by minimizing the Cross-Entropy (CE) loss between the logits of the FP model and the quantized model.

Optimization Objective

The optimization for the final block is defined as:
$$\min_{W_q} \mathcal{L}_{CE}(X W_{FP} W_{Head}, X W_q W_{Head})$$
Where:

• $W_{FP}$ and $W_q$ are the full-precision and quantized weights of the final block.
• $W_{Head}$ is the unembedding layer (LM head).
• $X$ is the input to the final block.
• The loss is computed over the token probability distributions derived from the logits ($X W W_{Head}$).

By explicitly incorporating the LM head and using Cross-Entropy (which is equivalent to minimizing KL divergence), LFQ directly aligns the token-level probability distribution of the quantized model with that of the FP model. This ensures that the quantized model reproduces the FP model's top-1 token predictions and maintains the correct ordering of top-$k$ tokens.

Compatibility

LFQ is designed as a drop-in enhancement for existing block-wise PTQ frameworks (FlexRound, OmniQuant, Block-AP). It does not alter the underlying quantization scheme (e.g., group size, bit-width) or the packing/unpacking kernels, ensuring it remains memory-efficient and compatible with existing inference engines.

3. Key Contributions

The paper claims three primary contributions:

1. Diagnosis of Generation Degradation: The authors are the first to demonstrate that the conventional block-wise PTQ objective (minimizing MSE at intermediate outputs) is misaligned with reproducing FP token predictions, leading to non-negligible accuracy gaps in text generation tasks.
2. Proposal of LFQ: The introduction of LFQ, a method that quantizes the final Transformer block via logit-level cross-entropy minimization. This approach consistently improves generation quality across diverse block-wise PTQ methods while preserving language modeling and understanding performance.
3. Comprehensive Validation: Extensive experiments across diverse model families (Qwen2.5, Qwen3, Llama 3.1/3.2, DeepSeek-R1, L1-Qwen) and tasks (IFEval, GSM8K, MATH500, AIME). The results show LFQ narrows the gap between low-bit quantized models and FP baselines, often achieving near-parity on complex reasoning benchmarks.

4. Experimental Results

The authors evaluated LFQ on instruction-tuned, reasoning, and Mixture-of-Experts (MoE) models under 4-bit and 3-bit weight-only quantization.

• Generation Quality: LFQ consistently outperforms standard block-wise PTQ on text generation benchmarks.
  • On AIME 2024, LFQ nearly matched the FP baseline for L1-Qwen-7B-Max and DeepSeek-R1-Distill-Llama-8B, whereas standard PTQ suffered substantial degradation.
  • On MATH500 and GSM8K, LFQ improved pass@8 and greedy decoding accuracy, often reducing the performance gap to FP baselines to within 1–2 percentage points.
  • On IFEval, LFQ improved instruction-following accuracy across all tested quantization schemes.
• Preservation of Other Tasks: LFQ maintained performance comparable to existing block-wise PTQ methods on language modeling (WikiText2 perplexity) and understanding (MMLU), confirming that the modification to the final block does not degrade general capabilities.
• Ablation Studies:
  • LM Head & CE: Removing either the LM head or the Cross-Entropy objective resulted in lower generation accuracy, confirming both are essential.
  • Block Scope: Applying LFQ to only the final block was found to be sufficient; applying it to multiple top blocks yielded marginal gains in generation but sometimes degraded MMLU performance.
  • Comparison with LQEC: While LoRA-based error compensation (RILQ) performed better on language modeling/understanding, LFQ consistently outperformed it on text generation tasks. The authors note that LFQ and RILQ address orthogonal aspects and can be combined for complementary benefits.

5. Significance and Claims

The paper positions LFQ as a simple yet effective solution to a critical bottleneck in deploying low-bit LLMs: the degradation of generative capabilities.

• Practical Impact: LFQ enables the deployment of low-bit models (3-bit/4-bit) for complex reasoning and long-form generation tasks without the prohibitive computational cost of Quantization-Aware Training (QAT) or the performance loss of standard PTQ.
• Theoretical Insight: The work highlights that minimizing MSE in hidden states is insufficient for generative alignment; direct optimization of the token distribution via the LM head is necessary for high-fidelity generation.
• Modesty: The authors frame LFQ as an enhancement to existing block-wise methods rather than a replacement. They emphasize that it preserves the efficiency and compatibility of current PTQ pipelines while specifically targeting the generation quality gap. The method is shown to be robust across different learning rates and calibration datasets (C4 vs. WikiText2).