Model optimization

Edge devices often have limited memory or computational power. Various optimizations can be applied to models so that they can be run within these constraints. In addition, some optimizations allow the use of specialized hardware for accelerated inference.

For model optimization, LiteRT provides post-training quantization (PTQ) with AI Edge Quantizer (AEQ). AEQ is a framework-agnostic post-training quantization toolkit designed for LiteRT and LiteRT-LM that converts unquantized models exported from PyTorch, JAX, TensorFlow, or Keras into hardware-optimized integer formats across CPUs, GPUs, and NPUs.

For comprehensive documentation, advanced algorithms (such as GPTQ, OCTAV, and Hadamard transforms), and full recipe references, visit the AI Edge Quantizer GitHub repository.

Why models should be optimized

There are several main ways model optimization can help with application development.

Size reduction

Quantization can be used to reduce the size of a model. Smaller models have the following benefits:

  • Smaller storage size: Smaller models occupy less storage space on your users' devices.
  • Smaller download size: Smaller models require less time and bandwidth to download to users' devices.
  • Less memory usage: Smaller models use less RAM when they are run, which frees up memory for other parts of your application to use, and can translate to better performance and stability.

Latency reduction

Latency is the amount of time it takes to run a single inference with a given model. Quantization can reduce the amount of computation required to run inference using a model, resulting in lower latency and reduced power consumption.

Accelerator compatibility

Some hardware accelerators, such as NPUs, GPUs, and the Edge TPU, can run inference extremely fast with models that have been correctly optimized.

Generally, these types of devices require models to be quantized in a specific way. See each hardware accelerator's documentation to learn more about their requirements.

Trade-offs

Optimizations can potentially result in changes in model accuracy, which must be considered during the application development process.

The accuracy changes depend on the individual model being optimized, and are difficult to predict ahead of time. Generally, models that are optimized for size or latency will lose a small amount of accuracy. Depending on your application, this may or may not impact your users' experience. In rare cases, certain models may gain some accuracy as a result of the optimization process.

To address accuracy loss, AI Edge Quantizer supports advanced techniques:

  • Selective quantization: Quantize only specific layers or subgraphs while keeping accuracy-sensitive operations in floating-point precision.
  • Mixed-precision quantization: Combine different bit widths (such as INT4/INT8 mixed weights).
  • Blockwise quantization: Quantize weights in sub-channel blocks (block sizes of 32, 64, 128, or 256) to better preserve dynamic range and model accuracy in low-bit schemes like INT4.
  • Advanced quantization algorithms: Apply algorithms such as GPTQ, OCTAV, and Hadamard transforms to minimize quantization error.

For details and implementation examples, see the AI Edge Quantizer GitHub repository.

Types of quantization

AI Edge Quantizer supports three main post-training quantization techniques:

Technique Example recipe Requires calibration data Weights Activations Inference Model size Accuracy
Weight-only "weight_only_wi8_afp32" No int float32 float32 Small Highest
Dynamic "dynamic_wi8_afp32" No int float32 int Small Small accuracy loss
Static "static_wi8_ai8"/"static_wi8_ai16" Yes int int int Smallest Small to moderate loss

Generally, we recommend dynamic quantization for CPU/GPU deployment and static quantization for NPU deployment.

For migration from legacy quantization with the TFLite Converter (TFLQ), see the TFLQ to AEQ Migration Guide.

Installation

Install AI Edge Quantizer using pip:

pip install ai-edge-quantizer

Quick start

AI Edge Quantizer takes an unquantized source model (.tflite or .litertlm) and a quantization recipe to produce a hardware-optimized model.

Command Line (aeq)

Quantize directly from your terminal:

aeq --model_file="path/to/input.tflite" \
    --recipe=dynamic_wi8_afp32 \
    --output_dir="/path/to/output"

Python API

from ai_edge_quantizer import quantizer

# 1. Initialize quantizer with an unquantized .tflite model
qt = quantizer.Quantizer("path/to/model.tflite")

# 2. Load a ready-to-use recipe
qt.load_quantization_recipe("dynamic_wi8_afp32")

# 3. Quantize and export the optimized model
qt.quantize().export_model("path/to/quantized_model.tflite")

References