On-Device Learning

Deploying ML on microcontrollers (MCUs) represents the extreme end of resource-constrained inference. MCUs typically have tens to hundreds of kilobytes of SRAM, megabytes of flash storage, and operate at clock speeds of tens to hundreds of megahertz. Yet despite these severe constraints, on-device ML enables always-on sensing applications that are impossible with cloud-based approaches.

Inference Frameworks for MCUs

The primary framework for MCU deployment is TensorFlow Lite Micro (TFLite Micro), which provides a stripped-down inference engine designed for bare-metal environments. TFLite Micro uses a static memory allocation scheme, operates without an operating system, and supports a subset of TensorFlow operations optimized for MCU architectures.

Memory Management Challenges

Memory management is the primary challenge for MCU deployment. The model weights must fit in flash memory, while the intermediate activations during inference must fit in SRAM. The arena memory allocation strategy pre-allocates a fixed buffer for all intermediate tensors, avoiding dynamic allocation and heap fragmentation.

Power Management

Power consumption is a critical constraint for battery-powered MCU applications. Always-on keyword detection models must consume microwatts to enable months or years of battery life.

• Duty cycling — Alternate between deep sleep and wake states to reduce average power.
• Event-driven inference — Only run the ML model when a simple threshold detector indicates interesting input.
• Hardware power gating — Shut down unused peripherals and memory banks during inference.
• Cascaded models — Use a tiny always-on model as a trigger for a larger, more accurate model.

Edge deployment places ML inference on devices closer to the data source, ranging from powerful edge servers to smartphones to IoT sensors. The primary motivations are latency reduction, privacy preservation, bandwidth savings, and reliability.

• Latency reduction — Eliminate the network round-trip to the cloud, achieving sub-millisecond inference.
• Privacy preservation — Sensitive data (face images, health signals) never leaves the device.
• Bandwidth savings — Process raw sensor data locally, transmitting only results or anomalies.
• Reliability — Operate continuously even without network connectivity.

Model Conversion and Optimization

Models trained in frameworks like PyTorch or TensorFlow must be converted to edge-optimized formats. Each target platform has its own optimized inference runtime.

Split Inference

Split inference distributes model computation between edge devices and the cloud, placing early layers on-device and later layers in the cloud. The split point is chosen to minimize the data transmitted between device and cloud.

Operational Challenges at Fleet Scale

Edge deployment introduces unique operational challenges. Model updates must be distributed to potentially millions of devices across diverse hardware and software versions. Monitoring and debugging are harder without direct access to devices.

On-device inference optimization encompasses all techniques for making ML models run efficiently on resource-constrained devices. The goal is to meet latency, memory, and power requirements while maintaining acceptable model quality. This often requires a combination of model-level and system-level optimizations.

Quantization: The Most Impactful Optimization

Quantization is the single most impactful optimization for on-device inference. Converting models from FP32 to INT8 typically yields 4x reduction in model size, 2-4x speedup, and significant energy savings.

Efficient Model Architectures

Model architecture choice is equally important. Architectures designed for efficiency achieve better accuracy per FLOP than larger architectures compressed to the same size.

• MobileNetV3 — Depthwise separable convolutions with squeeze-and-excitation for mobile GPUs.
• EfficientNet-Lite — Compound-scaled architecture optimized for on-device inference.
• SqueezeNet — Fire modules that minimize parameters while maintaining accuracy.
• MobileViT — Lightweight vision transformer designed for mobile deployment.
• TinyBERT — Distilled BERT variant for on-device NLP tasks.

Runtime Optimizations

Runtime optimizations include operator fusion, memory planning, and hardware-specific kernel selection. Inference runtimes implement these optimizations automatically during model conversion.

On-device learning goes beyond inference to enable models to adapt and improve directly on the device using local data. This is particularly valuable for personalization and for privacy-sensitive applications where training data should never leave the device.

Federated Learning

Federated learning enables distributed model improvement without centralizing data. Each device trains a local model update on its private data and sends only the model update to a central server. The server aggregates updates from many devices to improve the global model.

class="tok-comment"># Simplified Federated Averaging (FedAvg) pseudocode
def federated_averaging(global_model, devices, rounds):
    for r in range(rounds):
        local_updates = []
        class="tok-comment"># Select subset of devices for this round
        selected = random.sample(devices, k=fraction * len(devices))
        for device in selected:
            local_model = copy(global_model)
            class="tok-comment"># Train locally for E epochs on device data
            for epoch in range(E):
                local_model.train(device.data)
            local_updates.append(local_model.weights - global_model.weights)
        class="tok-comment"># Aggregate: weighted average of updates
        avg_update = weighted_average(local_updates, weights=device_data_sizes)
        global_model.weights += avg_update
    return global_model

On-Device Fine-Tuning

On-device fine-tuning is a simpler alternative where a pre-trained model is adapted to local conditions. The key challenge is performing gradient computation and weight updates within the device's memory and compute constraints.

Catastrophic Forgetting

Continuous adaptation raises concerns about model stability and catastrophic forgetting. Without careful regularization, a model adapted to recent local data may lose its ability to handle less frequent inputs.

Extreme optimization pushes the boundaries of what is possible on the most resource-constrained devices. When working with only 256 KB of memory and a few MHz of compute, every byte and every cycle counts. This requires rethinking standard ML practices from first principles.

Binary and Ternary Neural Networks

Binary and ternary neural networks represent weights and activations with just 1 or 2 bits, replacing multiplication with simple bitwise operations. While these extreme quantization approaches sacrifice significant accuracy, they enable ML on devices that could not otherwise support any neural network computation.

Model-Hardware Co-Design

Model-hardware co-design for MCUs involves selecting architectures that match the specific capabilities of the target hardware. ARM Cortex-M processors have SIMD instructions (CMSIS-NN) optimized for certain tensor shapes.

Knowledge Distillation for MCUs

Knowledge distillation takes on heightened importance for MCU deployment. A large teacher model trained in the cloud is distilled into an extremely compact student model designed specifically for the target MCU.