Quantization Folding

Quantization folding merges linear operations—most commonly Batch Normalization (BN) layers—into the preceding Convolutional or Fully Connected layers. This fusion creates a single, equivalent layer with adjusted weights and bias. The primary benefit is a simpler graph with fewer operations, which is easier to quantize and execute efficiently. For example, a Conv2D → BatchNorm sequence becomes a single fused Conv2D layer.

Folding operations before quantization reduces cumulative quantization error. When BN layers are quantized separately, the quantization noise from both the preceding layer's weights and the BN layer's parameters (scale, offset) compounds. By folding BN, its linear transformation is absorbed into integer-friendly weights, eliminating a source of error. This often results in higher accuracy for Post-Training Quantization (PTQ) without requiring retraining.

By eliminating entire layers (like BatchNorm) from the execution graph, quantization folding directly reduces computational and memory overhead during inference. This is critical for on-device deployment where every operation counts. Benefits include:

• Fewer kernel launches and memory accesses.
• Lower latency and power consumption.
• Simplified operator fusion in subsequent compiler passes for target hardware like NPUs or mobile CPUs.

A folded, linearized graph is ideal for mapping to highly optimized integer compute kernels. Hardware accelerators like Neural Processing Units (NPUs) and DSPs have dedicated units for operations like INT8 convolution. Folding ensures these kernels process a single, consolidated operation rather than a chain, maximizing hardware utilization and throughput. This is a prerequisite for achieving peak performance on edge AI silicon.

Quantization folding is not an optional step but a foundational optimization in modern ML compilers and converters. It is automatically applied by frameworks during export to efficient runtimes. Key tools that perform folding include:

• TensorFlow Lite Converter (TFLite)
• ONNX Runtime graph optimizations
• PyTorch's torch.ao.quantization.fuse_modules
• TVM and Apache MXNet's NNVM compiler. This standardization ensures optimized models across diverse deployment targets.

While most beneficial for inference, the concept also interacts with Quantization-Aware Training (QAT). During QAT, the forward pass simulates quantization on the folded graph to train robust parameters. However, the training graph itself often keeps BN layers separate for proper gradient calculation. The compiler then folds the trained BN layers away during the final export to an integer model, ensuring the training-time simulation matches the deployed, folded architecture.

Batch Normalization Folding is the foundational operation that quantization folding builds upon. It is a graph transformation that merges the parameters of a batch normalization layer into the weights and bias of a preceding linear layer (e.g., a convolution or fully connected layer).

• Mechanism: The batch normalization operation (scale, shift, mean, variance) is mathematically absorbed into the preceding layer's weights (W_folded) and bias (b_folded).
• Purpose: This eliminates the batch normalization layer from the inference graph, reducing computational overhead and memory accesses.
• Prerequisite for Quantization: Folding batch normalization before quantization is critical. It ensures the quantizer sees a single, simpler linear operation, leading to more accurate calibration of quantization ranges and a more efficient integer kernel.

Compute Graph Optimization encompasses a suite of transformations applied to a neural network's computational graph to improve inference latency, memory usage, and hardware efficiency. Quantization folding is a prime example.

• Common Optimizations: Includes operator fusion (e.g., Conv+ReLU), constant folding, dead code elimination, and layout optimizations.
• Hardware-Aware: These optimizations are often specific to a target backend, such as a mobile NPU or GPU, to map the graph optimally to the hardware's execution units.
• Compiler Role: Frameworks like TensorFlow Lite, ONNX Runtime, and TVM perform these optimizations during the model conversion or compilation phase, preparing the model for deployment.

Quantization-Aware Training (QAT) is a training paradigm that simulates quantization during the forward pass, allowing the model to adapt its parameters for lower precision before deployment. It often incorporates folding.

• Process: Fake quantization nodes are inserted into the training graph. During training, weights and activations are quantized and dequantized, but gradients are calculated with respect to the full-precision weights using a Straight-Through Estimator (STE).
• Interaction with Folding: In QAT, batch normalization folding is typically performed before inserting the fake quantizers. This creates a training graph that closely mirrors the final, folded, and quantized inference graph, leading to superior accuracy compared to Post-Training Quantization.
• Outcome: Produces a model whose weights are already robust to the quantization error that will be applied during inference.

Static Quantization is a Post-Training Quantization (PTQ) method where the quantization parameters (scale and zero-point) for all tensors, including activations, are determined once during a calibration phase and remain fixed during inference.

• Calibration: Requires a small, representative dataset to profile activation ranges. This is when quantization folding is applied.
• Requires Folding: For static quantization to be effective and accurate, operations like batch normalization must be folded beforehand. An unfused batch norm layer would require dynamic per-batch statistics, breaking the static assumption.
• Efficiency: Enables the most aggressive runtime optimizations, as the entire computation graph uses pre-defined integer math, minimizing control flow overhead.

Operator Fusion is a graph optimization that combines multiple sequential operators into a single, compound kernel. It is a broader category that includes quantization folding.

• Goal: Reduce kernel launch overhead, intermediate tensor memory writes/reads, and improve data locality.
• Examples: Fusing a Convolution with a following ReLU activation, or a MatMul with a BiasAdd. BatchNorm folding into Conv is a specific, mathematically exact form of fusion.
• Hardware Benefit: Fused operators can be executed by a single, highly optimized kernel on the target accelerator (e.g., a Conv2D + BiasAdd + ReLU kernel on a GPU or NPU), dramatically speeding up inference.