Technical Report: NEMO DNN Quantization for Deployment Model

💡 Research Summary

This technical report defines a formal framework for layer‑wise quantization of deep neural networks (DNNs) with a focus on final deployment, documenting the NEMO (NEural Minimization for PyTorch) framework. Four representations are described: FullPrecision, FakeQuantized, QuantizedDeployable, and IntegerDeployable. FullPrecision corresponds to the standard PyTorch model using 32‑bit floating‑point arithmetic. Layers are modeled as a sequence of Linear operators (convolutions or fully‑connected matrices), optional Batch‑Normalization, and a non‑linear activation (e.g., ReLU). Linear operators are expressed as ϕ = b + h·w·x, where bias b can be absorbed into subsequent operators if desired.

FakeQuantized introduces a restricted set of quantized values for weights and activations, but the forward pass still operates on real‑valued tensors. Quantization is defined as t_i = α_t + ε_t·q_i with q_i∈ℤ, where ε_t is the quantum (step size) and α_t is the offset. The quantization function Q_t maps real numbers to integer images and is monotonic, piecewise‑constant. NEMO adopts a PACT‑like scheme for ReLU: the activation is clipped at a learned upper bound β, and the step size ε_y is set to β/(2^Q − 1). During training, the Straight‑Through Estimator (STE) is used so that gradients are computed on the underlying floating‑point tensors, effectively ignoring the quantization function in the backward pass.

QuantizedDeployable extends FakeQuantized by ensuring that every operator consumes quantized inputs and produces quantized outputs. Batch‑Normalization parameters are also quantized, and Linear weights are replaced by their quantized counterparts b_w = ⌊w/ε_w⌉·ε_w. A global scaling factor ε is propagated through the network via net.set_deployment(eps_in=1./255), guaranteeing consistent scaling across layers. Consequently, all intermediate tensors become integer images, allowing inference without any real‑valued arithmetic.

IntegerDeployable goes a step further: the model can be executed using only integer arithmetic. PACT_IntegerAct replaces the standard PACT activation, and a requantization step converts tensors from one quantized space to another. Requantization is defined mathematically: given two quantized spaces Z_a and Z_b with steps ε_a and ε_b, the integer image is approximated as Q_b ≈ ⌊(ε_a·D/ε_b)·Q_a⌉ ≫ d, where D = 2^d is a large integer chosen to bound the relative error by 1/D. This operation reduces to a right‑shift, enabling efficient hardware implementation.

The report provides formal definitions (Definitions 2.1, 2.2, 3.1) for tensor quantization, quantized tensors, and requantization, and derives linear quantization as an affine clipping operation. It also details how the quantized space of a linear combination (ϕ) is formed from the product of weight and activation quantized spaces, leading to a combined quantum ε_ϕ = ε_w·ε_x.

Implementation in NEMO follows a clear pipeline:

1. net = nemo.transform.quantize_pact(net, dummy_input) – creates the FakeQuantized model.
2. net.harden_weights() – freezes weights to their quantized values.
3. net.set_deployment(eps_in=1./255) – propagates the scaling factor.
4. net = nemo.transform.integerize_pact(net, eps_in=1./255) – produces the IntegerDeployable model.

Currently NEMO still stores data in float32, so integer‑only inference on GPUs incurs a small penalty due to extra conversion operators. The authors note that dedicated integer hardware would eliminate this overhead.

Key insights include: (1) a rigorous, layer‑wise mathematical treatment of quantization that separates quantization, quantization‑aware training, and deployment stages; (2) the use of STE to keep training pipelines unchanged while still learning quantized parameters; (3) a practical requantization scheme that leverages bit‑shifts to map between quantized spaces with bounded error. These contributions enable memory‑ and compute‑constrained environments (e.g., edge devices) to run DNN inference using only integer arithmetic while preserving accuracy, and they lay a foundation for future integration with specialized integer accelerators.