Quantization for On-Device Machine Learning: Optimizing Feature Encoding for Speed, Size, and Accuracy

Quantization for on-device machine learning balances model size, inference speed, and accuracy by encoding features and weights in lower-precision formats. This article walks through practical options, trade-offs, and concrete examples to help you pick a quantization strategy that fits resource-constrained environments without guessing blindly.

Why quantization matters for on-device ML

On-device models often face limits on memory, CPU cycles, battery, and download size. Quantization reduces numeric precision (for example from 32-bit floating point to 8-bit integer), which usually cuts memory footprint and speeds up inference. The trade-off is potential accuracy loss, but with the right pipeline that loss can be small or recoverable.

Lower model size — smaller weights reduce APK or binary size and memory use.
Faster inference — integer math can be faster on many mobile NPUs and CPUs.
Reduced power — lower-precision ops can use less energy.

Core quantization strategies

Choose among these common approaches based on constraints and tolerance for accuracy shifts.

Post-training dynamic quantization — maps weights to int8, quantizes activations on the fly. Fast to apply, works well for transformer and LSTM dense layers.
Post-training static quantization (a.k.a. calibration) — extracts activation ranges from a representative dataset and stores quantization parameters. Often better accuracy for CNNs and conv-heavy models.
Quantization-aware training (QAT) — simulates quantization during training so the model learns to resist low-precision errors. Useful when accuracy sensitivity is high.

Feature encoding and per-channel vs per-tensor

Feature encoding matters before numeric quantization. Typical steps include scaling numeric inputs, bucketing or embedding categorical features, and normalizing continuous features. For weights, per-channel quantization (each output channel gets its own scale) often preserves accuracy for convolutions. Per-tensor quantization has simpler storage and may be adequate for many dense layers.

Practical pipeline: prepare, calibrate, convert

A minimal post-training static quantization pipeline looks like this:

Export or load a trained model in CPU mode.
Define representative calibration data that mirrors on-device inputs.
Run the model on calibration data to collect activation ranges.
Convert weights and activations to integer with saved scales and zero-points.
Validate size, latency, and accuracy on a holdout set.

The following example uses PyTorch eager-mode quantization to show both dynamic and static cases. Replace placeholders with your model and dataloader.

import torch
import torch.quantization as tq

# assume model is a pre-trained torch.nn.Module on CPU
model.eval()
# Dynamic quantization example: quick, often used for transformer/dense
qmodel_dyn = tq.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
torch.save(qmodel_dyn.state_dict(), \'qmodel_dynamic.pth\')

# Static post-training quantization (eager)
# 1) set qconfig for backend that your device/runtime supports
tq.backends.quantized.engine = \'fbgemm\'
model.qconfig = tq.get_default_qconfig(\'fbgemm\')
tq.prepare(model, inplace=True)

# 2) Calibration: run representative data
for batch in calibration_loader:   # calibration_loader must reflect on-device inputs
    inputs, _ = batch
    model(inputs)

# 3) Convert to int8
tq.convert(model, inplace=True)
torch.save(model.state_dict(), \'qmodel_int8.pth\')

# Profiling: compare sizes and inference time
print("FP32 size:", os.path.getsize(\'model_fp32.pth\'))
print("INT8 size:", os.path.getsize(\'qmodel_int8.pth\'))

Tips to preserve accuracy

Some practical adjustments tend to reduce accuracy regression after quantization:

Representative dataset quality — use diverse samples that reflect real usage for calibration; small biased sets can cause large errors.
Per-channel for convs — enable per-channel weight quantization where supported.
Mixed precision — keep sensitive layers (like first/last) in FP16 or FP32 while quantizing the rest.
Use QAT when modest fine-tuning with simulated quantization can recover performance.

Measuring impact: size, speed, and accuracy

Track a small set of metrics before and after quantization:

Model binary size (bytes).
Inference latency (cold and warm starts) on target device.
Memory footprint at peak during inference.
Task accuracy and class-wise metrics to spot regressions.

Profilers and on-device benchmarks vary by platform. For Android, try Android Profiler and the NNAPI delegate. For iOS, Instruments and Core ML delegates help identify bottlenecks.

Deployment options: runtime considerations

Several runtimes support quantized models with different performance characteristics:

TFLite — strong mobile support for int8 with both post-training and QAT flows.
PyTorch Mobile — supports eager-mode quantization and mobile optimization.
ONNX Runtime — conversion to ONNX with quantization tooling can leverage runtime accelerations.
NNAPI / Core ML — platform delegates that can exploit device-specific accelerators.

When converting across formats, confirm that quantization parameters (scale, zero_point) survive the conversion and that the runtime uses hardware-accelerated int8 kernels where possible.

Case study: small vision model

Imagine a small MobileNet-style classifier used for on-device image tagging. After static int8 quantization with per-channel weights and a 500-sample calibration set drawn from deployment images, size can drop by around 4x and latency may decrease depending on CPU vectorization. If accuracy dropped too much, a short QAT run of a few epochs with a low learning rate often recovers most of the lost performance.

Checklist before shipping

Verify representative calibration data used in the pipeline.
Measure latency and memory on the target device under realistic load.
Run edge-case tests for out-of-distribution inputs to detect quantization instability.
Consider mixed precision for critical layers.
Document the quantization method and exact runtime used so future builds are reproducible.

Quantization is not a one-size-fits-all cure, but with a clear pipeline and careful validation it can unlock meaningful size and performance gains for on-device ML. Start with quick dynamic quantization for early iterations, move to calibrated static quantization for vision models, and apply QAT when higher fidelity is required. Experimentation and profiling on the target hardware remain central to a successful deployment.