Quantization

Quantization techniques perform computations and store tensors in lower precision data types like 8-bit integer instead of floating point precision. There are multiple approaches to quantize a deep learning model categorized as:

• Post-training quantization (PTQ)
• Quantization aware training (QAT)

In post-training quantization, the model is trained in floating point precision and later converted to the lower precision data type. There are two types of post-training quantization:

1. Static quantization: quantizes the weights and activations of the model. Quantizing the activations statically requires data to be calibrated (i.e., recording the activation values to compute the optimal quantization parameters with representative data).
2. Dynamic quantization: quantized the weights ahead of time (like static quantization) but the activations are dynamically at runtime.

Sometimes post-training quantization is not able to achieve acceptable task accuracy. In general, this is where quantization-aware training (QAT) can be used: during training, fake-quantization modules are inserted in the forward and backward passes to simulate quantization effects, allowing the model to learn representations that are more robust to reduced precision.

Burn does not currently support QAT. Only post-training quantization (PTQ) is implemented at this time.

Module Quantization

Quantizing the weights of your model after training is quite simple. We have access to the weight tensors and can collect their statistics, such as the min and max value when using MinMaxCalibration, to compute the quantization parameters.

use burn::module::Quantizer;
use burn::tensor::quantization::{Calibration, QuantLevel, QuantParam, QuantScheme, QuantValue};

// Quantization config
let scheme = QuantScheme::default()
    .with_level(QuantLevel::Block(32))
    .with_value(QuantValue::Q4F)
    .with_param(QuantParam::F16);
let mut quantizer = Quantizer {
    calibration: Calibration::MinMax,
    scheme,
};

// Quantize the weights
let model = model.quantize_weights(&mut quantizer);

Calibration

Calibration is the step during quantization where the range of all floating-point tensors is computed. This is pretty straightforward for weights since the actual range is known at quantization-time (weights are static), but activations require more attention.

To compute the quantization parameters, Burn supports the following Calibration methods.

Quantization Scheme

A quantization scheme defines how an input is quantized, including the representation of quantized values, storage format, granularity, and how the values are scaled.

#![allow(unused)]
fn main() {
let scheme = QuantScheme::default()
    .with_mode(QuantMode::Symmetric)         // Quantization mode
    .with_level(QuantLevel::block([2, 16]))  // Granularity (per-tensor or per-block)
    .with_value(QuantValue::Q8S)             // Data type of quantized values, independent of how they're stored
    .with_store(QuantStore::Native)          // Storage format for quantized values
    .with_param(QuantParam::F16);            // Precision for quantization parameters
}

Quantization Mode

Quantization Level

Quantization Value

Quantization Store

Native storage is not supported for sub-byte quantization values.

Quantization Parameters Precision