weight quantization aware training
**Quantization-Aware Training (QAT)** is the **training technique that simulates the effects of low-bit quantization during the forward pass while maintaining full-precision gradients** — by inserting fake quantization operations that round weights and activations to discrete values during training, the model learns to compensate for quantization error, producing quantized models with significantly higher accuracy than post-training quantization (PTQ), especially critical for aggressive quantization like INT4 and INT2 where PTQ causes unacceptable quality degradation.
**QAT vs. PTQ (Post-Training Quantization)**
| Aspect | PTQ | QAT |
|--------|-----|-----|
| Training required | No | Yes (fine-tune or full train) |
| Accuracy loss (INT8) | 0.1-0.5% | <0.1% |
| Accuracy loss (INT4) | 1-5% | 0.1-0.5% |
| Accuracy loss (INT2) | 20-40% (unusable) | 2-10% (usable) |
| Cost | Minutes | Hours-days |
| Use case | INT8 deployment | INT4/INT2, edge devices |
**Fake Quantization**
```python
def fake_quantize(x, scale, zero_point, num_bits=8):
"""Simulates quantization during training"""
qmin, qmax = 0, 2**num_bits - 1
# Quantize
x_q = torch.clamp(torch.round(x / scale + zero_point), qmin, qmax)
# Dequantize (back to float for computation)
x_dq = (x_q - zero_point) * scale
return x_dq
# Forward: discrete values (simulates INT arithmetic)
# Backward: straight-through estimator (gradient flows as if identity)
```
**Straight-Through Estimator (STE)**
```
Forward: x → round(x) → x_q (non-differentiable!)
Backward: ∂L/∂x ≈ ∂L/∂x_q (pretend round() is identity)
STE enables gradient-based optimization despite discrete rounding:
- Forward pass: Exact quantization behavior
- Backward pass: Gradients pass through as if no quantization
- Result: Weights learn to cluster near quantization grid points
```
**QAT Training Process**
```
1. Start with pretrained FP32 model
2. Insert fake-quantize nodes:
- After each weight tensor (weight quantization)
- After each activation tensor (activation quantization)
3. Calibrate quantization ranges (min/max or percentile)
4. Fine-tune for 5-20% of original training steps
5. Export truly quantized model (replace fake-quant with real INT ops)
```
**Advanced QAT Techniques**
| Technique | Description | Benefit |
|-----------|------------|--------|
| Learned step size (LSQ) | Backprop through scale factor | Better scale calibration |
| Mixed precision QAT | Different bits per layer | Accuracy-efficient tradeoff |
| PACT | Learnable clipping range for activations | Reduces outlier impact |
| DoReFa | Quantize gradients too | Enables low-bit training |
| Binary/Ternary QAT | 1-2 bit weights | Extreme compression |
**QAT for LLMs**
| Model | QAT Method | Bits | Quality Retention |
|-------|-----------|------|------------------|
| Llama-2-7B + QAT | GPTQ-aware fine-tune | INT4 | 99% of FP16 |
| BitNet b1.58 | 1.58-bit QAT (ternary) | ~2bit | 90-95% of FP16 |
| QuIP# | Incoherence QAT | INT2 | 85-90% of FP16 |
| SqueezeLLM | Sensitivity-aware QAT | Mixed 3-4 bit | 98% of FP16 |
**Deployment**
- INT8 QAT: Supported everywhere (TensorRT, ONNX Runtime, CoreML).
- INT4 QAT: Requires specific kernels (CUTLASS, custom CUDA).
- Binary/Ternary: Specialized hardware (XNOR-net accelerators).
- QAT → ONNX export: Most frameworks support fake-quant → real quantized graph conversion.
Quantization-aware training is **the gold standard for deploying neural networks at reduced precision** — while post-training quantization works well for moderate compression (INT8), QAT's ability to learn compensation for quantization error makes it essential for aggressive compression (INT4 and below) that enables deployment on edge devices, mobile phones, and cost-efficient inference servers where every bit of precision reduction translates directly to memory savings and throughput improvements.