Quantization-Aware Healing: a compressed, 4-bit model

What happened

Making a large language model smaller almost always comes with a cost. The now-standard recipe for efficient deployment is to compress the architecture first, cutting the parameter count by removing layers, heads, or neurons, and then quantize the remaining weights down to 4 bits to shrink memory and compute further.

Both steps save a lot, but together they systematically degrade the capabilities people actually care about: reasoning, mathematical problem-solving, and code generation. Because of this, serious deployment pipelines add a recovery step, usually called healing, before the model goes into production. Recent open-weight releases such as gpt-oss, NVIDIA's Nemotron family, and our own Hypernova 60B all rely on some version of this compress-then-heal approach.

Our latest paper, Quantization-Aware Healing: A Practical Recipe for Recovering Compressed, 4-Bit LLMs, asks a question that the field has mostly left open: once a model has already been through structural compression, not just quantization, how well does that recovery step actually work, and what is the right way to do it?

We introduce Quantization-Aware Healing (QAH), and applied to a GPT-OSS 120B model compressed to 60B parameters and quantized to MXFP4, it produces a model that beats its own full-precision (bfloat16) version on 7 of 9 benchmarks. The 4-bit model ends up smaller, cheaper to run, and more accurate than the checkpoint it was quantized from.

This inverts the usual relationship between a 4-bit model and the 16-bit model it came from. Most efficiency pipelines follow the same three steps: compress the architecture, quantize the compressed weights, then heal the damage. The difference between methods is entirely in that last step. The dominant healing recipe is quantization-aware training (QAT).

An alternative, quantization-aware distillation (QAD), avoids re-running that history. Instead of a task loss, it distills a frozen full-precision teacher directly into the quantized student through a KL-divergence loss on the output logits. This works well when the only change is quantization, because a genuine full-precision version of the exact same model exists to act as teacher.

But once a model has gone through structural compression, fewer layers, heads, or neurons, and not just fewer bits, that assumption breaks. There is no independently trained full-precision version of the smaller architecture. The only candidate teacher is the recovered bfloat16 checkpoint, which is itself a distilled approximation of the original model.

Distilling from it anchors the quantized student to a degraded target and caps its accuracy at that recovered checkpoint's own ceiling. So the question of how to heal a model that has been both structurally compressed and quantized was, until now, genuinely open. QAH removes that ceiling with one change: it distills directly from the original, pre-compression model rather than from the recovered one.

Teacher and student do not even share an architecture. The teacher is full-size and full-precision, the student is half the size and running in MXFP4. Because a teacher's output distribution is architecture-agnostic, nothing about the size or shape mismatch prevents the transfer. The student never sees hard labels, only the teacher's output distribution, matched through KL divergence on the logits. This reframes what the quantization stage is doing.

Under QAH it is no longer a lossy postprocessing step applied after healing is finished. It is a second, full pass of distillation against the original teacher, supervision that the bfloat16 checkpoint never received. The 4-bit student is not compensating for information lost to quantization; it is picking up information the earlier recovery stage did not have the time or data to transfer.

There is also a stability benefit that falls out of the loss itself. Because KL distillation ties the student to a fixed teacher distribution, once the student catches up there is no further pressure for it to drift. A cross-entropy task loss, by contrast, keeps pushing the student toward hard labels indefinitely. That difference turns out to matter for both accuracy and training stability, as the comparison below shows.

Why it matters

It inserts fake-quantization operators into the forward pass and keeps fine-tuning the model on a task loss, so the weights learn to tolerate the low-precision representation. In practice this means re-running an already expensive multi-stage post-training process, supervised fine-tuning, RLHF, agentic tuning, through a noisier, lower-precision forward pass. It is costly, and as our results show, it can also become unstable if training continues too long past its best point.

What to watch

To make QAH work at long context, where the healing corpus includes documents up to 32k tokens, we reuse the memory-efficient chunked KL-divergence loss from our companion paper on efficient distillation. That loss computes the KL one slice of the sequence at a time and never materializes the full vocabulary-by-sequence grid, which is what makes 32k-token healing fit inside a fixed GPU memory budget.

We covered the mechanics of that loss in a previous post. QAH overview. After structural compression and quantization, capabilities drop sharply. QAH distills from the original model, a frozen teacher whose logits are precomputed offline, rather than from the recovered checkpoint. Source: paper Figure 1.