Quantization-aware training

How much memory the weights need

A trained network is mostly a long list of numbers, its weights. Large models are usually trained and served in bf16, a 16-bit floating-point format, so each weight takes 2 bytes. Stored as 4-bit integers, each weight takes half a byte, plus a small overhead for the scale factors that map the integers back to real values.With one bf16 scale for every group of 128 weights, the overhead is 16 bits ÷ 128 = 0.125 bits per weight, so “4-bit” storage costs about 4.13 bits per weight. Figure 1 multiplies this out for five model sizes.

Smaller weights also move faster. When a model generates text one token at a time, each step reads every weight from memory, so the time per token is often set by memory bandwidth rather than arithmetic. Packed 4-bit weights shrink that traffic, but the speed-up arrives only if the runtime has kernels that unpack the format efficiently.

The price is accuracy. Rounding every weight to one of 15 or 16 values perturbs the network, and the perturbation grows as bits are removed. Quantization-aware training pays less of that price by training the network with the rounding in place.

Rounding a number onto a grid

A uniform quantizer replaces a real number with the nearest of a fixed set of evenly spaced levels. This page uses the symmetric version that is common for weights. Choose a clipping range [−α, α]. With b bits, place the levels s = α / q_max apart, where q_max = 2^b−1 − 1. The spacing s is called the scale. Each value x is stored as an integer q and read back as x̂:At 4 bits q_max = 7, so q runs from −7 to 7: 15 levels. The sixteenth code, −8, goes unused so that the grid is symmetric about zero. Many libraries use the full range instead.

q = clamp(round(x / s), −q_max, q_max)
x̂ = s · q

The error x − x̂ comes in two kinds. Inside the range, rounding error is at most half a step, s/2. Outside it, the value is replaced by the edge level ±α, and this clipping error grows with the distance from the range.

For these weights the best α is 1.2 standard deviations at 2 bits, 2.5 at 4 bits and 3.9 at 8 bits. With only a few levels, it pays to spend them on the dense middle of the distribution and give up the rare large values. Real weight tensors are not exactly normal, so quantization libraries choose α from the data. Common rules take the largest absolute value, a high percentile, or the α with the lowest measured error.

Activations may use an asymmetric range with a zero-point offset, because their values need not be centered on zero. After a ReLU, for example, every value is zero or positive, and a symmetric grid would waste half its levels.

One scale for many weights

A tensor stores one scale for many weights, not one per weight. With a single scale for the whole tensor, called per-tensor scaling, α has to reach the largest weight. One outlier then stretches the grid for every other weight in the tensor. Giving each row its own scale, called per-row or per-channel scaling, confines the damage to the outlier's row.

Smaller groups isolate outliers better but store more scales. Four-bit formats for language models usually go further than per-row scales and use one scale per group of 32 to 128 consecutive weights.

Activations are harder to quantize than weights. They change with every input, and some channels carry large values on nearly every token, so low-bit activation quantization often needs finer scales, calibration data, clipping, or a transformation of the values first. This is one reason why quantizing only the weights to 4 bits is more common than quantizing weights and activations to 4 bits together, a combination written W4A4.

Putting the rounding inside training

There are two ways to arrive at a quantized model. Post-training quantization (PTQ) quantizes a finished model without further end-to-end training. Some PTQ methods use a small calibration set to choose scales; others work from the weights alone. At 8 bits PTQ often preserves quality. At 4 bits and below, the choice of quantizer and calibration matters more, and the accuracy loss can be large.

Quantization-aware training (QAT) continues training, but computes the loss with quantized weights, and sometimes quantized activations, in the forward pass. The loss then includes the damage done by rounding, and gradient descent reduces that damage along with the rest of the loss. Figure 4 shows one training step.

The rounding step is called fake quantization. It rounds a tensor to the grid and immediately multiplies by the scale again, so the result is still a floating-point tensor whose values lie on the grid, and training runs on ordinary floating-point kernels.Jacob et al., 2018 describe training with simulated quantization for integer-only inference. When training ends, the master weights are converted one last time, to the integer codes q and scales s that the inference runtime expects, and W is discarded.

Gradients through a staircase

Rounding breaks gradient descent. As a function of its input, round() is a staircase, flat between jumps. Its derivative is zero almost everywhere and undefined at the jumps, so backpropagation through it delivers no gradient to the master weights.

The straight-through estimator (STE) works around this.Bengio, Léonard and Courville, 2013 study the straight-through estimator for neurons with hard thresholds. In the backward pass it treats rounding as the identity, with derivative 1. The version on this page also keeps the derivative of clipping, which is 1 inside [−α, α] and 0 outside, so a clipped weight receives no gradient. Other QAT methods use other surrogate gradients, especially when they learn the scale.LSQ (Esser et al., 2020) defines a gradient for the step size s itself, so the clipping range is trained along with the weights.

Why the master weight stays in full precision

An update applied to Ŵ directly would be rounded away. A weight on level 1 that is nudged by −0.03 rounds straight back to level 1, however many times it is nudged. The master weight W keeps the nudges. Each one moves W a little, and Ŵ changes only when their sum carries W across a rounding boundary, halfway between two levels.

The same mechanism has a failure mode. A master weight that sits near a boundary can flip Ŵ back and forth between two levels as noisy updates push it across and back, as happens early in Figure 6. Nagel et al. report that these oscillations hurt low-bit QAT of efficient networks such as MobileNets, and propose ways to damp them or freeze the oscillating weights.Nagel, Fournarakis, Bondarenko and Blankevoort, “Overcoming Oscillations in Quantization-Aware Training”, ICML 2022.

PTQ and QAT on a small network

The lab below trains a network in your browser and quantizes it both ways. The network is small: one input, 32 tanh hidden units and one output, 97 parameters in all. Its task is to fit sin(3x) + 0.3 cos(9x) at 64 evenly spaced inputs in [−1, 1]. Each run has three stages:Details: full-batch Adam with a cosine learning-rate decay from 0.01 to 0. Weights are quantized per tensor with the symmetric scheme above, and the scale is recomputed from the largest master weight at every step, so no weight is ever clipped. Biases and activations are not quantized. The scale is treated as constant during backpropagation, even though it is recomputed for the next step. There is no distillation.

1. Train without quantization for 1,500 steps (the FP baseline).
2. PTQ: round that model's weights to the grid at 2, 3 and 4 bits, and measure the loss.
3. QAT: starting from the same unquantized model, train 2,000 more steps at each bit width with fake-quantized weights and the STE.

Across these eight seeds, removing bits increases PTQ loss. QAT often reduces that loss by adjusting the master weights and unquantized biases over another 2,000 updates. At 2 bits the grid has three levels, −s, 0 and +s, and the outcome swings from seed to seed. Figure 8 shows all eight seeds at once.

The lab is far simpler than QAT for a language model: weights only, per tensor, no activation quantization and no distillation. The losses are measured on the training points, so they do not measure generalization. QAT receives 2,000 extra training steps while the FP and PTQ references stay fixed; there is no equally long unquantized continuation for comparison. The experiment demonstrates the mechanism and its possible failures, but does not predict gains for a language model.

Which tensors a transformer quantizes

A transformer keeps most of its parameters in a few kinds of large matrix: the linear projections of the attention and feed-forward (MLP) blocks. In Llama 3 8B these projections hold 87% of the parameters, so they are the first thing to quantize.Computed from the published architecture: 32 layers, model dimension 4,096, feed-forward dimension 14,336 and 8 key/value heads (Llama 3 paper, Table 3), with a 128,256-token vocabulary and separate input and output embeddings (the released model configuration). Total: 8.03 billion parameters. The table shows one illustrative layout, W4A8: 4-bit weights and 8-bit activations. Published recipes differ in which tensors they quantize, at what precision, and how finely they scale them.

The choices in the example follow from the earlier sections. The projections are large, so 4 bits saves the most memory there, and per-group scales contain outliers. Activations change with every token, so each token gets its own scale. The KV cache grows with context length and can become a large share of memory in long conversations; this layout stores it in 8 bits, and some deployments use 4. Normalization and softmax are cheap and sensitive to rounding, and the residual path is where every layer writes its output, so this layout leaves them in bf16.

Common choices in LLM quantization-aware training

• Start from a pretrained checkpoint. Many methods fine-tune a higher-precision model for a short quantization-aware phase. BitNet b1.58 instead trains with quantization from the start.
• Distill. Some recipes train the quantized model to match the output probabilities of a higher-precision model. Others use ordinary next-token targets, or combine the two.
• Learn the scales. LSQ and its successors give the step size its own gradient. Other methods recompute or calibrate the scales.
• Handle outliers in the model's structure. Per-group weight scales, per-token activation scales, and rotations that spread outlier channels before quantizing, as in QuaRot and SpinQuant, complement QAT; they can be used with or without it.

Published recipes

These five recipes differ in bit width, training signal and starting point. Each row links to the paper or official announcement that supports it.

A practical order follows. PTQ comes first at 8 bits and for weight-only 4-bit deployment, because it needs no retraining. QAT becomes worth its training cost as precision drops, when activations are quantized too, or when PTQ misses the quality target. How much QAT recovers depends on the model, the task, the quantizer, the calibration data, the kernels and the training recipe.