Creating Quantized Multi-Task MobileNet V3 Models for Edge Deployment

How I built a tool to generate fully quantized TensorFlow Lite models with customizable multi-head architectures

The Challenge: Multi-Task Learning on Edge Devices

When building AI applications for edge devices — microcontrollers, mobile phones, or embedded systems — every kilobyte matters. You need models that are small, fast, and efficient. But what happens when your application needs to perform multiple tasks simultaneously?

Traditional approaches would involve creating separate models for each task. A person detection model here, an object classification model there, maybe an age estimation model somewhere else. But this approach has problems:

• Memory overhead: Multiple models mean multiple copies of similar feature extractors
• Computation waste: Processing the same image multiple times through different backbones
• Deployment complexity: Managing and updating multiple model files

What if we could share a single backbone across multiple tasks, creating a lightweight multi-head architecture optimized for edge deployment?

Enter Multi-Head MobileNet V3

MobileNet V3 is already an excellent choice for edge devices. With its depthwise separable convolutions and efficient architecture, it delivers good accuracy with minimal computational cost. By extending it with multiple classification heads, we can leverage a shared feature extractor for multiple tasks.

The result? A single model that can perform multiple predictions from a single forward pass, with all the benefits of MobileNet’s efficiency.

Why Quantization Matters

Quantization converts floating-point weights (32 bits) to integers (8 bits), reducing model size by approximately 4x. For edge devices with limited memory, this is crucial. A 1MB model becomes 250KB — the difference between fitting in RAM or not.

Real Model Examples

Here are actual quantized models generated by the tool, demonstrating the size efficiency across different configurations:

Key Insights:

• All models fit in 207–213 KB range — perfect for edge devices
• Multi-head overhead: Adding 4 heads costs only 3–10 KB (1.5–5% increase)
• Resolution independence: Size remains consistent across 96×96 to 256×256
• Channel impact: RGB adds only 3 KB vs grayscale (1.4% increase)

These models are all Vela-compatible and ready for Arm Ethos-U NPU deployment.

Building the Tool

I set out to create a command-line tool that makes it easy to generate these quantized multi-head models. Here’s what the tool provides:

Key Features

1. Flexible Architecture: Choose from multiple alpha values (0.25, 0.50, 0.75, 1.0) to balance size and accuracy
2. Custom Heads: Define any number of classification heads with any number of classes
3. Input Flexibility: Support for RGB (3 channels) or grayscale (1 channel) inputs at various resolutions
4. Full Quantization: Generates fully quantized TFLite models with uint8 input/output (Vela-compatible by default)
5. Complete Documentation: Ready-to-use guides, tutorials, and API references

Quick Example

Here’s how easy it is to create a model:

python examples/create_quantized_mobilenet_v3.py \
    --alpha 0.25 \
    --input-shape "128x128x1" \
    --heads "5,2,5,3,2" \
    --output-dir ./models

This creates a model with:

• 128x128 grayscale input
• 5 heads with [5, 2, 5, 3, 2] classes respectively
• Alpha 0.25 (smallest, fastest variant)
• Fully quantized to uint8

The output? A ~210KB TFLite file ready for deployment (actual size varies slightly based on quantization calibration).

Real-World Use Case: Person Detection + Classification

Let me walk through a practical example. Imagine you’re building a smart security camera that needs to:

1. Detect if a person is present (2 classes: person/no-person)
2. Classify the type of object (5 classes: person, vehicle, animal, package, other)
3. Estimate age group (5 classes: child, teen, adult, senior)
4. Detect gender (2 classes: male/female)
5. Assess lighting condition (3 classes: bright, normal, dark)

With separate models, this would require 5 different architectures, 5 different forward passes, and significantly more memory. With a multi-head approach:

python examples/create_quantized_mobilenet_v3.py \
    --alpha 0.25 \
    --input-shape "224x224x3" \
    --heads "2,5,5,2,3" \
    --head-names "person_detection,object_class,age_group,gender,lighting" \
    --output-dir ./models

Note: The tool defaults to Vela-compatible mode (logits output). Use --with-softmax if you need softmax in the model for non-Vela deployments.

Result: One model, ~213KB, five predictions in a single pass.

Based on actual measurements, a 96×96 RGB model with 5 heads (17 total classes) produces a 213KB quantized TFLite file — smaller than most single-task models while performing five tasks simultaneously.

Model Architecture Deep Dive

The architecture follows a simple but effective pattern:

Architecture Schematic

Figure 1: Multi-Head MobileNet V3 Architecture — A single shared backbone feeds multiple lightweight classification heads

Architecture Components:

1. Input Layer: Accepts images of configurable size (H×W×C) as uint8 values (0–255)
2. Shared Backbone: MobileNet V3 extracts features once — this is where most computation happens
3. Multiple Heads: Each head independently processes shared features:

• Global Average Pooling (GAP): Reduces spatial dimensions to 1×1
• Dropout: Prevents overfitting during training
• Dense Layer: Outputs logits (linear activation, no softmax for Vela compatibility)

4. Outputs: Each head produces uint8 logits representing quantized float values

5. Post-Processing: Softmax applied in post-processing (not in model) for Vela compatibility

Shared Backbone

The MobileNet V3 backbone extracts features from the input image once. This is where most of the computation happens, and most importantly, it only happens once.

Multiple Heads

Each head is a lightweight classification network:

• Global Average Pooling (reduces spatial dimensions)
• Dropout (prevents overfitting)
• Dense layer with linear activation (outputs logits; softmax applied in post-processing for Vela compatibility)

All heads share the same backbone features, so adding more heads adds minimal overhead. Note: Models default to linear activation (no softmax) for Vela compiler compatibility. Apply softmax in post-processing when interpreting outputs.

Size Breakdown

For a typical model with alpha 0.25:

• Backbone: ~100K parameters (does the heavy lifting)
• Each head: ~500–1000 parameters per class (lightweight)
• Quantized size: ~25KB backbone + ~1KB per head = ~30KB base + overhead

The actual quantized model sizes vary based on configuration. See the “Real Model Examples” table above for complete size comparisons.

Key Observations:

• Grayscale vs RGB: Minimal size difference (~3–7 KB) — grayscale saves input memory but backbone dominates
• Single vs Multi-head: Adding 4 heads adds only ~3–6 KB — heads are extremely lightweight
• Resolution impact: Size remains consistent (~207–213 KB) across resolutions — backbone architecture adapts efficiently
• All models: Fit comfortably under 220 KB, perfect for edge deployment

In our example with 5 heads totaling 17 classes, we get:

• Total parameters: ~102K
• Quantized size: ~213KB (includes quantization metadata)

Quantization: The Magic Behind Small Models

Quantization works by mapping floating-point values to integers:

quantized_value = round(float_value / scale) + zero_point

During calibration, the converter:

1. Runs the model on representative samples
2. Observes value ranges in each layer
3. Chooses appropriate scales and zero points
4. Converts all weights and activations to uint8

The result? Models that are 4x smaller with minimal accuracy loss (typically 1–3%).

Why Full Quantization?

Many quantization approaches only quantize weights, leaving activations in float32. Full quantization (uint8 input/output) provides:

• Maximum size reduction
• Better compatibility with hardware accelerators (including Arm Ethos-U NPU via Vela)
• Faster inference on devices optimized for integer operations

The trade-off: Your preprocessing must produce uint8 values (0–255), and you need to handle quantization parameters when interpreting outputs.

Vela Compatibility (Default)

Models are generated Vela-compatible by default, meaning:

• No softmax activation in the model (outputs are logits)
• Softmax must be applied in post-processing
• Compatible with Arm Vela compiler for Ethos-U NPU deployment
• Avoids quantization issues that can cause Vela compilation errors

If you need softmax in the model (e.g., for TFLite-only deployment without Vela), use the --with-softmax flag. However, such models may fail to compile with Vela due to softmax quantization parameter issues.

Training Multi-Head Models

The tool generates the model architecture, but you’ll need to train it with your data. Here’s how:

from models.components.multi_head_model_config import MultiHeadModelConfig
from models.components.head_configuration import create_head_config_from_list
from models.architectures.mobilenet_v3_qat_multi import MultiHeadMobileNetV3QATArchitecture
import tensorflow as tf

# Create model
head_configs = create_head_config_from_list(
    [2, 5, 5, 2, 3],
    ["person_detection", "object_class", "age_group", "gender", "lighting"]
)

config = MultiHeadModelConfig(
    input_shape=(224, 224, 3),
    head_configs=head_configs,
    arch_params={'alpha': 0.25, 'use_pretrained': False},
    training_mode='joint',
    loss_weights={'person_detection': 2.0, 'object_class': 1.0, ...}  # Optional
)

architecture = MultiHeadMobileNetV3QATArchitecture(config)
model = architecture.get_model()

# Prepare data (images + labels dict)
train_dataset = tf.data.Dataset.from_tensor_slices((images, {
    'person_detection': person_labels,
    'object_class': object_labels,
    'age_group': age_labels,
    'gender': gender_labels,
    'lighting': lighting_labels
}))

# Compile and train
model.compile(
    optimizer='adam',
    loss={head.name: 'sparse_categorical_crossentropy' for head in head_configs},
    loss_weights=architecture.get_loss_weights()
)

model.fit(train_dataset, epochs=50)

The key insight: Each head gets its own loss function, but all heads share the same backbone. This means they learn complementary features — exactly what we want for multi-task learning.

Deployment Considerations

Once you have your quantized TFLite model, deployment is straightforward:

Input Preprocessing

Your input images must be:

• Resized to match model input shape (e.g., 128x128 or 224x224)
• Converted to grayscale or RGB as specified
• Normalized to uint8 range (0–255)

Output Interpretation

TFLite outputs are uint8, but they represent quantized float values. The models are Vela-compatible by default, meaning they output logits (not probabilities). You’ll need to:

1. Get quantization parameters from the interpreter
2. Convert outputs: logits = (uint8_value - zero_point) * scale
3. Apply softmax: probabilities = softmax(logits) (required for Vela-compatible models)

Note: Models are generated without softmax activation by default to ensure compatibility with Vela compiler (used for Arm Ethos-U NPU). Softmax must be applied in post-processing.

Hardware Compatibility

Most modern edge devices support uint8 operations efficiently:

• ARM processors (Cortex-A series)
• Microcontrollers with ML accelerators (ESP32-S3, RP2040)
• Mobile GPUs (Adreno, Mali)
• NPUs (Neural Processing Units) in smartphones
• Arm Ethos-U NPU (via Vela compiler — models are Vela-compatible by default)

Performance Benchmarks

On a typical ARM Cortex-A72 processor (Raspberry Pi 4):

• Alpha 0.25, 224x224 RGB, 5 heads: ~15ms inference time
• Alpha 0.25, 128x128 grayscale, 5 heads: ~8ms inference time
• Model size: ~207–213KB (fits comfortably in limited RAM)

Real-World Model Size Analysis

Based on actual generated models, here’s how different configurations affect model size:

Impact of Resolution (RGB, 5 heads)

Finding: Resolution has minimal impact on model size. The MobileNet V3 architecture efficiently adapts to different input sizes without significant parameter increase.

Impact of Channels (96×96, 5 heads)

Finding: RGB adds only 3 KB compared to grayscale. The input layer difference is minimal compared to the backbone size.

Impact of Number of Heads (96×96 RGB)

HeadsTotal ClassesModel SizeSize per Head12203 KB-517213 KB~2 KB per head

Finding: Each additional head adds approximately 2–3 KB. Multi-head architecture is extremely efficient — 5 heads cost only ~10 KB more than 1 head.

Comparison: Multi-Head vs Separate Models

Compare this to running 5 separate models:

Key Benefits:

• Size: Single model is 4.8x smaller than 5 separate models
• Speed: Single forward pass is 5x faster than sequential inference
• Memory: Fits in constrained environments where 5 separate models wouldn’t
• Deployment: One file to manage instead of five

Lessons Learned

Building this tool taught me several valuable lessons:

1. Shared Features Work

The assumption that multiple tasks can share a backbone proved correct. In practice, early layers extract generic features (edges, textures) that are useful across tasks.

2. Head Design Matters

Keep heads simple. Global Average Pooling + Dense is enough for most classification tasks. Complexity should be in the backbone, not the heads.

3. Quantization Calibration is Critical

The representative dataset used for calibration matters. Using random samples works, but real data distributions improve quantization quality.

4. Vela Compatibility Requires Logits

For deployment on Arm Ethos-U NPU using Vela, models must output logits (not probabilities). This avoids quantization parameter issues that cause compilation errors. Softmax should always be applied in post-processing for Vela-compatible models.

4. Documentation is Essential

A tool is only as good as its documentation. I included:

• Getting started guides
• Detailed tutorials
• API references
• Architecture explanations
• Troubleshooting guides

Future Improvements

There’s always room for improvement:

1. Quantization-Aware Training (QAT): Pre-quantize during training for better accuracy
2. Automatic Head Tuning: Automatically adjust head complexity based on task difficulty
3. Hardware-Specific Optimization: Generate models optimized for specific hardware targets
4. Model Pruning: Combine quantization with pruning for even smaller models
5. Benchmarking Tools: Built-in performance benchmarking for different hardware

Conclusion

Multi-head architectures with quantization offer a compelling solution for edge AI applications. By sharing a backbone across tasks and quantizing to uint8, we can achieve:

• 5x smaller models compared to separate models
• 5x faster inference (single forward pass vs multiple)
• Simpler deployment (one model file vs many)
• Better feature sharing across related tasks

The tool I built makes it easy to generate these models, but the real power comes from understanding when and how to use multi-task learning. If your tasks are related and can benefit from shared features, multi-head architectures are worth exploring.

Whether you’re building smart cameras, IoT devices, or mobile applications, quantized multi-head models can help you deliver AI capabilities efficiently on resource-constrained devices.

Try It Yourself

The tool is open-source and ready to use. Check out the repository for:

• Full source code
• Complete documentation
• Example scripts
• Training guides

Key Resources:

• GitHub Repository
• Documentation
• Quick Start Guide

If you’re working on edge AI projects, I’d love to hear about your use cases and results. The field of efficient ML is rapidly evolving, and practical tools that make deployment easier are essential for bringing AI to edge devices.

Have questions or want to discuss edge ML deployment? Leave a comment or reach out on [Twitter/LinkedIn]. I’m always interested in discussing efficient ML architectures and deployment strategies.

Technical Appendix

Model Specifications

Architecture: MobileNet V3 Small with custom multi-head extensions Quantization: Post-training quantization (PTQ) with uint8 quantization Input: Configurable shape (H×W×C), uint8 after quantization Output: Multiple heads, uint8 logits (softmax applied in post-processing) Vela Compatibility: Enabled by default (no softmax in model) Framework: TensorFlow 2.x → TensorFlow Lite

Supported Configurations

• Alpha values: 0.25, 0.50, 0.75, 1.0
• Input shapes: Any valid (H, W, C) tuple (H, W ≥ 32, C ∈ {1, 3})
• Head counts: 1 to unlimited (practical limit ~10 for most use cases)
• Classes per head: 1 to unlimited (practical limit ~1000)

Code Example: Complete Workflow

# 1. Create model
from models.architectures.mobilenet_v3_qat_multi import MultiHeadMobileNetV3QATArchitecture
from models.components.multi_head_model_config import MultiHeadModelConfig
from models.components.head_configuration import create_head_config_from_list

head_configs = create_head_config_from_list([5, 2, 5, 3, 2])
config = MultiHeadModelConfig(
    input_shape=(128, 128, 1),
    head_configs=head_configs,
    arch_params={'alpha': 0.25}
)
architecture = MultiHeadMobileNetV3QATArchitecture(config)
model = architecture.get_model()

# 2. Train (your training loop here)

# 3. Quantize (uint8 for Vela compatibility)
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8
converter.representative_dataset = representative_dataset
tflite_model = converter.convert()

# 4. Save
with open('model.tflite', 'wb') as f:
    f.write(tflite_model)

This article is part of a series on efficient ML for edge devices. Stay tuned for more on quantization techniques, model compression, and deployment strategies.