Why is quantization done in neural networks? Why is it necessary? –– FP32 operations are 16x more energy consuming, 20x slower, and 4x space consuming than INT8.
If we achieve moderate loss in accuracy with quantized version of the models, it’s a win-win scenario. Moreover, floating-point operations are not supported for some environments like early microcontrollers.
We’re mainly talking about inference here because quantized model runs for eternity, and speeds up both prefill which is a compute-bound phase and decode which is memory-bound.
How to achieve quantization? –– Quantize weights and activations. Common formats are W4A16 or generally W?A?. W4A16 means weights are quantized to 4 bits and activations at 16 bit.
Why is activation quantization harder than weights? –– Typical quantization rule involves FP16 to INT4 which means representing a 16-bit number (~6.10e-5 for FP16 or ~1.18e-38 for BF16) to 4 bit integers (-16 to 15). Activations are dynamic and depend on the inputs. Static uniform quantization leads to improperly represented capacity with overemphasis on outliers and underrepresenting crucial ranges.
Which step gets optimized with only weights quantization? And same for activations quantization? ––
How does reducing activation quantization improve prefill speedup?
How is quantization performed? –– Two ways to perform:
Compute in FP16: Input remain in FP16, dequantize weights back to FP16, and compute in native quantized state. All activation stay in FP16.
Compute in INT8, accumulate in INT32: Quantize inputs to FP4 x≈sxqx by computing integer multiplication in INT4 Wx=(sWsx)(qWqx), i.e. accumulate in INT32 qWqx, and add rescaling factors back (sWsx).
Uniform or Logarithmic codes.
Uniform – quantize the points to space evenly, i.e. gap between 100-100.1 and 0.1-0.2 is same. Find the scale factor s for the final quantization accuracy, (16 for INT4), compute q=⌊x/s⌋, and recompute by x^=q∗s. There’s a asymmetric variant that adds a offset q=⌊x/s⌋+z,x^=s∗(q−z), useful to cover asymmetric ranges like ReLU. Weights quantization is generally symmetric while activations are asymmetric.
Logarithmic – Compute information-theoretically optimal 4-bit grid such that expected quantization error is minimized. It identifies the statistic of the distribution of the quantity, and spaces the grid such that ranges more likely to occur are more tightly packed and conversely, heavy-tail is loosely packed. Example: Given normally distributed weights, resulting grid represents range near 0.95 confidence interval tightly (near zero), and rest of the range loosely.
Scale factor granularity – Scale factor can be applied in a spectrum of values. At one end is no scale factor s=1; all values persist in native values, with gradual refinement of scale factor. Another end is single scale factor for the whole layer which can lead to severe loss in accuracy for some values.
Similar to normalization logic, next step is per-feature dimension scale factor (row or column). For example: for an input of N×D, we can either have N or D scale factors corresponding to each input or each feature-dimension respectively.
Further refinement adds block-wise scale which further divides a channel in blocks of varying sizes (from 128 to 16) depending on tolerable accuracy loss. For example: A block of 32 four-bit along with a FP16 scale factor occupies 3232∗4+16=4.5 bits per value. So, 1 bit or 2-bit models represent extreme quantized models capable of compressing model parameters even further.
Sensitive layers – MLP layers gets quantized to 4 bits while important layers (embedding, first or final layer, layer norms) remain at FP16.
What are the training strategies for quantizing a pre-trained model –– Post training quantization enables quantization through post training a model using a small calibration dataset. Available in three levels:
Rounding to nearest integer
Minimize layer wise error: Run a calibration dataset through the network, compute the optimal scale factor as per the granularity (layer or block-level), and minimize the difference between quantized and non-quantized layer outputs.
Minimize model output error: Use Hessian of the loss with respect to each weight to make the rounding decision that minimizes end-to-end output error. GPTQ processes weights column by column using inverse hessian for each rounding decision. AWQ and GGUF are other tools used for quantization at model level.
Quantization strategies for scratch model –– PTQ works at 8 bits for weights-only quantization but hits a wall for 2,4 bit quantized models or for quantizing sensitive layers like activation layers. QAT (Quantized aware training) simulates quantization during pretraining phase of the model to learn model weights that are robust to rounding errors.
How does PTQ work? –– Run the forward pass as if the model behaves like a quantized model, but backpropagation enters into complicated territory because quantization function is a staircase function which is non-differentiable, thus producing no gradient signal. To make it work, STE (Straight-through estimator) proposes to simulate fake quantization where Quantization is only enabled for forward pass, and backward pass assumes quantization as an identity function. To visualize these, suppose x is the INT8 input with W in FP32, we add a wt_quant layer qW=wt_quant(W) that quantizes the input and gradient of the layer is identity.
What’s the problem with STE? –– Weight sitting above a point always round back to the same point, nullifying update by the gradient. Solution is to add a stochastic rounding that uses a Bernoulli distribution to output the discrete rounded value. For example: 2.3 has 0.7 probability of rounding off to 2 and 0.3 to 3.
KV cache quantization –– For transformer based models, KV cache is a special case of activation quantization which is computed for every input token, persists across entire context length and grows quadratically with sequence length (exceeding model weights for long sequences). KVQuant showed 3-bit quantization with 0.1 perplexity loss, using per-channel quantization and careful outlier handling. NVIDIA’s Blackwell architecture supports NVFP4 for KV cache natively.