Improving Quantized FP4 Weight Quality via Logit Distillation

Table of Contents

Introduction

Quantization has become one of the most impactful techniques for deploying Large Language Models (LLMs) efficiently. By reducing the precision of weights and activations, quantization significantly lowers GPU memory (VRAM) requirements and accelerates inference—offering speedups in both memory-bound and compute-bound scenarios. However, the process is inherently lossy, and quantization-induced degradation in model quality remains a critical challenge. Ensuring that quantized models remain accurate and reliable is essential for their real-world usability.

The new FP4 data type has recently been introduced, with hardware support in both NVIDIA’s Blackwell GPUs and AMD’s upcoming MI355X accelerators. There are two distinct FP4 formats: MXFP4 and NVFP4.

MXFP4 is the more flexible format, designed to support mixed-precision matrix multiplication. It represents weights in 4-bit E2M1 format, combined with group-wise scaling factors shared over vectors of size 32. These scales are stored in 8-bit E8M0 format, which restricts them to powers of two. While this choice simplifies hardware implementation, it severely limits representational precision. In practice, replacing E4M3 scales with E8M0 results in a noticeable drop in model quality. In contrast, the NVFP4 format uses smaller group sizes (16) and E4M3 8-bit scales, leading to much better quality retention after quantization. However, NVFP4 is limited to homogeneous FP4 × FP4 tensor core operations, which restricts it in scenarios requiring mixed-precision accumulation such as memory-bound workloads with higher precision activations.

The table below summarizes the difference between the two FP4 dtypes:

To address the quantization error, particularly in the MXFP4 format, we explore a lightweight approach to recover model accuracy. Unfortunately, due to the symmetric nature of the quantization, calibration-free techniques like HQQ cannot be applied effectively in this context.

While more advanced methods such as HQQ+ can successfully restore accuracy for low-bit quantization by training low-rank adapters (e.g., LoRA), these approaches come with non-negligible overhead, which we aim to avoid in latency and memory-sensitive deployments.

Instead, we freeze both the 4-bit weights Wq and the block-scales scales and focus on training a simple post-hoc correction consisting of a post-scaling factor and a shift applied to FP4 matrix multiplication output: matmul_fp4(x, Wq, scales) * post_scale + shift, with both the post-scales and the shift are 16-bit vectors, keeping the method lightweight.

Unlike the two-level scaling mechanism used in NVFP4 to boost accuracy, we find that the majority of quality recovery comes from the bias correction, not the post-scaling. This insight is especially useful in practice: for popular models that already include a bias term (such as Qwen), the learned shift can be fused directly into the pre-trained bias, incurring virtually zero additional runtime overhead.

We delve into this approach in detail in the next section.

Approach

Weight Quantization

The first step is quantizing the model to FP4 weights. This is a one-shot process that involves first calculating the block-wise scales scales via absmax after reshaping to match the target block size: weights.abs().amax(dim=1, keepdim=True).
We then cast these scales to the target dtype (E8M0 for MXFP4, E4M3 for NVFP4), and the weights are divided by their corresponding block scales, then matched to the closest bin in the FP4 E2M1 format. This results in a tuple: Wq referring to the 4-bit quantized weights, and scales referring to the 8-bit block-wise scales.

This operation can be efficiently implemented using torch.searchsorted, or further accelerated with a custom Triton kernel — which would be especially useful for dynamic FP4 quantization. In our case, since the weights are static, this step is only performed once.

We also explored an exhaustive search strategy to optimize the block scales, which slightly improves the quantization accuracy — particularly for the NVFP4 format.

The code for quantizing to MXFP4 and NVFP4 is available here.

Distillation

The second step is the training phase. We introduce channel-wise trainable parameters post_scale and shift, such that the final output is computed as matmul_fp4(x, Wq, scales) * post_scale + shift. The post_scale and shift vectors are initialized to ones and zeros, respectively.

For training, we use a simple KL divergence loss (without temperature scaling) on the logits only, matching the outputs between the teacher model (unquantized bfloat16) and the student model (with frozen FP4 weights and trainable post-scale/shift). While it's common to apply a temperature during distillation, we found that setting the temperature parameter to 1 gave the best results. Including intermediate-layer losses did not yield further improvements.

We use a custom linear learning rate scheduler with 3% warm-up and a final learning rate of 1e-6. The MXFP4-calibrated model is trained for 2 epochs, while the NVFP4 model is trained for only 1 epoch, as it reaches our target accuracy of over 99% after the first epoch. Regarding the data, we train on a mixture of 50K examples randomly sampled from a variety of open-source datasets across domains such as math, knowledge, and code. These include prompts from MetaMathQA, Orca Math and Evol-Instruct.

We also experimented with making Wq learnable using techniques such as the Straight-Through Estimator (STE) and soft rounding. However, these approaches were unnecessary in practice and sometimes even degraded performance.

After examining the trained post_scale and shift parameters, we observe that most of the post_scales remain close to their initial values, while the shift terms consistently adapt across all layers. This indicates that a simple bias correction is sufficient for most of the quality recovery, contrary to the post-scaling approach used in NVFP4. Furthermore, if post_scale is used, the entire operation can be fused into a single FMA (fused multiply-add) instruction, incurring no additional compute overhead.

Results

We run experiments on the Llama-3.1-8B-Instruct model as a reference benchmark. Larger models are expected to achieve better quality post-quantization, as they tend to be less sensitive to quantization error.

Uncalibrated Results

We begin by reporting the performance of the uncalibrated models before the distillation phase. For reference, we also include results from HQQ using a group size of 64. As shown in Table 2, MXFP4 weights lead to significant degradation in quality — particularly on math-heavy tasks such as GSM8K. In contrast, the NVFP4 format preserves quality much better, with performance closer to that of HQQ.

Calibrated Results

We now report results after the distillation phase for both MXFP4 and NVFP4. For comparison, we also include results from INT4-calibrated methods such as GPTQ and AWQ in Table 3.

As shown, our simple distillation procedure successfully recovers accuracy. The relative performance with respect to the unquantized weights improved significantly, rising from 92.86% to 99.14% for MXFP4, and from 97.57% to 99.17% for NVFP4. This highlights that a straightforward bias correction, similar to what's used in the calibrated version of HQQ, is an effective strategy for restoring accuracy in both INT4 and FP4 weight quantization formats.

Conclusion

In this blog post, we presented a lightweight method for recovering accuracy in FP4 weight quantization using logits distillation. We demonstrated that a simple bias correction is sufficient to improve accuracy— even with the more constrained MXFP4 format — and can be seamlessly integrated into the FP4 inference pipeline with minimal to zero overhead.

While this work focused on static, pre-quantized weights, the same technique can be extended to handle activation quantization errors. In particular, combining it with post-post scaling (i.e., two-level scaling) offers a promising path for addressing quantization challenges in FP4-based pipelines.

We made both the MXFP4 and MVFP4 publicly available on Hugging Face 🤗!

Citation

@misc{badri2025fp4,
title = {Improving Quantized FP4 Weight Quality via Logit Distillation,
url = {https://mobiusml.github.io/fp4_blogpost/},
author = {Hicham Badri},
month = {June},
year = {2025}
}