PaperlayerView exact revision on arXiv

Source-linked AI summary

On-Device Training Under 256KB Memory

Ji Lin, Ligeng Zhu, Wei-Ming Chen, Wei-Chen Wang, Chuang Gan, Song Han

arXiv:2206.15472v4cs.CV

TL;DR

Tiny IoT devices cannot normally train models within their memory constraints, despite the privacy and personalization benefits of local adaptation. The paper combines quantization-aware optimization, sparse updates, and system co-design to enable training under 256KB SRAM; on VWW, it matches or surpasses cloud-trained accuracy while reducing memory by more than 1000×. The demonstrated scope is currently limited to vision recognition with CNNs.

  • Problem

    Tiny IoT devices have limited SRAM, while quantized graphs and full back-propagation make on-device training difficult despite its privacy and customization benefits.

  • Method

    The framework combines Quantization-Aware Scaling for quantized gradients with Sparse Update and algorithm-system co-design to reduce training memory.

  • Results

    The framework enables CNN training under 256KB SRAM and 1MB Flash, reducing memory by more than 1000× versus PyTorch and TensorFlow while matching or surpassing cloud-trained accuracy on VWW.

  • Takeaways & Limitations

    Tiny IoT devices can perform inference and continuously adapt to new data for lifelong on-device learning.

  • Takeaways & Limitations

    The study is limited to vision recognition with CNNs; extending it to audio and models such as RNNs and Transformers remains future work.

Abstract

from arXiv · show

On-device training enables the model to adapt to new data collected from the sensors by fine-tuning a pre-trained model. Users can benefit from customized AI models without having to transfer the data to the cloud, protecting the privacy. However, the training memory consumption is prohibitive for IoT devices that have tiny memory resources. We propose an algorithm-system co-design framework to make on-device training possible with only 256KB of memory. On-device training faces two unique challenges: (1) the quantized graphs of neural networks are hard to optimize due to low bit-precision and the lack of normalization; (2) the limited hardware resource does not allow full back-propagation. To cope with the optimization difficulty, we propose Quantization-Aware Scaling to calibrate the gradient scales and stabilize 8-bit quantized training. To reduce the memory footprint, we propose Sparse Update to skip the gradient computation of less important layers and sub-tensors. The algorithm innovation is implemented by a lightweight training system, Tiny Training Engine, which prunes the backward computation graph to support sparse updates and offload the runtime auto-differentiation to compile time. Our framework is the first solution to enable tiny on-device training of convolutional neural networks under 256KB SRAM and 1MB Flash without auxiliary memory, using less than 1/1000 of the memory of PyTorch and TensorFlow while matching the accuracy on tinyML application VWW. Our study enables IoT devices not only to perform inference but also to continuously adapt to new data for on-device lifelong learning. A video demo can be found here: https://youtu.be/0pUFZYdoMY8.

1 Introduction

On-device training supports private, local model adaptation but is difficult on tiny IoT devices because training exceeds their memory and computation budgets. The paper addresses this gap through algorithm-system co-design for quantized training and constrained back-propagation.

  • Motivation: On-device training adapts pre-trained models to newly collected sensory data locally, supporting user customization, lifelong learning, and privacy.Examples include continual learning from typing data and recognizing new camera objects.
  • Challenges: 256KB SRAM is barely sufficient for inference on tiny IoT devices, while training additionally requires backward computation and intermediate activations.These requirements make training substantially more demanding than inference on microcontrollers.
  • Challenges: >1000× memory gaps make current cloud-oriented training frameworks impractical for tiny IoT devices.PyTorch, TensorFlow, JAX, and MXNet do not target tight edge-resource constraints.
  • Challenges: Quantized graphs are difficult to optimize because they combine low-precision tensors with fused graphs lacking Batch Normalization, while full back-propagation exceeds tiny-device resources.Updating only the final layer also causes poor accuracy.
  • Approach: Quantization-Aware Scaling stabilizes quantized training, and Sparse Update reduces memory by skipping gradients for less important layers and sub-tensors.The framework combines these algorithmic techniques with system co-design to enable tiny on-device training.
  • Results: The framework enables convolutional-network training under 256KB SRAM and 1MB Flash, matching cloud-training accuracy on VWW while reducing memory by more than 1000× versus PyTorch and TensorFlow.It also reduces per-iteration time by more than 20× and was deployed on a Cortex M7 microcontroller.

2 Approach

The approach combines Quantization-Aware Scaling for stable real-quantized training with sparse updates and compile-time graph transformations to fit training within tight device memory.

  • 2.1 Optimizing Real Quantized Graphs: Real quantized graphs keep tensors in int8 for deployment and updates, unlike fake-quantized QAT graphs that retain floating-point operations.The distinction is central to realizing memory and computation savings on microcontrollers.
  • 2.1 Optimizing Real Quantized Graphs: Quantized training is unstable because mixed bit-precisions and fused-away Batch Normalization distort weight-to-gradient norm ratios.Quantization can make ratios larger and produce different patterns across weights and biases.
  • 2.1 Optimizing Real Quantized Graphs: Quantization-Aware Scaling compensates quantized gradients by scaling factors, stabilizing weight-to-gradient ratios and matching floating-point training accuracy.QAS enables fully quantized int8 forward and backward computation.
  • 2.2 Memory-Efficient Sparse Update: Sparse Update skips gradients for less important layers and sub-tensors, reducing memory and computation when full-model back-propagation exceeds the hardware budget.The search selects bias updates and partial weight updates under a memory constraint.
  • 2.2 Memory-Efficient Sparse Update: Contribution analysis prioritizes updates by estimated accuracy gain, while evolutionary search selects layer subsets and update ratios under the memory constraint.The approximation that tensor contributions can be summed is reported as effective.
  • 2.3 Tiny Training Engine (TTE): Tiny Training Engine prunes backward nodes for frozen weights, reorders gradient updates, and compiles forward and backward graphs into deployable code.These transformations address runtime overhead and reduce intermediate-tensor storage.

3 Experiments

Experiments evaluate quantized training, sparse update, and Tiny Training Engine on tinyML models and downstream vision datasets. The combined approach improves optimization while reducing memory and latency enough for microcontroller training, with VWW accuracy matching or exceeding cloud-trained baselines.

  • Experimental setup: Experiments use MobileNetV2, ProxylessNAS, and MCUNet on STM32F746, with quantized ImageNet-pretrained models fine-tuned on downstream datasets.Accuracy statistics are simulated on GPUs and verified against microcontroller training at the same level.
  • Quantization-aware scaling: QAS closes the accuracy gap between int8 and floating-point fine-tuning without extra memory cost.It stabilizes optimization of real quantized graphs by calibrating gradient scales.
  • Sparse update: Sparse update achieves higher downstream accuracy than last-k-layer updates while using a smaller memory footprint.Updating only selected layers, biases, and sub-tensors preserves more transfer-learning capacity than classifier-only or bias-only updates.
  • TinyML results: 206KB measured SRAM supports 89.1% top-1 accuracy on VWW, exceeding MCUNet’s 88.7% cloud-trained result and the 80% MLPerf Tiny requirement.The on-device result also surpasses TF-Lite Micro+MobileNetV2’s 86.2% inference-only result under 256KB.
  • Memory and latency: 20-21× lower measured peak memory results when sparse update is combined with TTE graph optimization.Sparse update alone provides 7-9× savings, while operator reordering produces the larger total reduction.
  • Memory and latency: 23-25× faster training is achieved by sparse update with TTE operators than full update with TF-Lite Micro operators.The comparison measures per-image training latency on the STM32F746 microcontroller.

4 Related Work

Related work reduces transfer-learning or deployment costs through parameter-efficient adaptation, model compression, architecture search, and lightweight inference systems. These approaches generally lack the combination of training support and resource efficiency targeted by this paper.

  • Efficient transfer learning: Parameter-efficient transfer learning updates classifiers, biases, normalization parameters, or small parallel branches instead of the full model.Classifier-only updates can lose accuracy when domain shift is large.
  • Systems for deep learning: Cloud training frameworks carry substantial host-language and runtime overhead, while lightweight inference libraries do not support training.This leaves a gap between resource-constrained deployment and on-device adaptation.
  • Tiny deep learning on microcontrollers: TinyML research reduces resource requirements through pruning, quantization, neural architecture search, and microcontroller-oriented systems.Examples include TF-Micro, CMSIS-NN, TinyEngine, MicroTVM, and CMix-NN.

5 Conclusion

The paper presents a co-designed solution for on-device transfer learning on microcontrollers within tight memory and Flash budgets. It reports large memory and latency reductions while supporting continual adaptation, but its demonstrated scope remains limited to vision CNNs.

  • 5 Conclusion: The framework enables on-device training on microcontrollers with 256KB SRAM and 1MB Flash without auxiliary memory.It reduces training memory by more than 1000× versus PyTorch and TensorFlow and latency by more than 20× versus TensorFlow-Lite Micro.
  • 5 Conclusion: Tiny IoT devices can perform inference and continuously adapt to new data for lifelong learning.This conclusion is stated within the demonstrated scope of the framework.
  • 5 Conclusion: The current study is limited to vision recognition with CNNs, with audio and RNN or Transformer extensions left for future work.The paper also notes electricity consumption from experiments on many downstream datasets.

A Video Demo

A video demo deploys the framework on an STM32F746 microcontroller to adapt an ImageNet-pretrained MCUNet for person detection. Training uses camera-fed VWW samples and completes within the device’s tight memory budget.

  • A Video Demo: The demo deploys on an STM32F746 with 320KB SRAM and 1MB Flash for on-device learning.The deployed model is MCUNet, pre-trained on ImageNet, adapted to classify whether a person is in front of the camera.
  • A Video Demo: The demo reports decent accuracy within the tight memory budget for a tiny on-device learning application.Figure 12 provides a screenshot of the video demo.
  • A Video Demo: 100 VWW camera images, split evenly between positive and negative samples, are used for training.The complete pipeline, including camera I/O overhead, takes around minutes.

B Variance of Different Runs

Across three independent runs, the sparse-update training results show very small variance, particularly for average accuracy.

  • ±0.07% standard deviation was observed across three independent runs of MobileNetV2-w0.35 with sparse updates.The scheme was searched under a 100KB analytic memory constraint.
  • Using the same data-loader random seed produces exactly the same results across multiple runs.Weight quantization after each iteration almost eliminates GPU-training nondeterminism.

C Training Setups & Discussions

The training setup balances transfer-learning accuracy, optimization stability, and microcontroller memory constraints through floating-point classifiers, batch-size approximations, and gradient accumulation.

  • The experiments use SGD with QAS, zero weight decay, cosine learning-rate decay, and warm-up of 1 epoch on VWW or 5 epochs on other datasets.Initial learning rates were extensively tuned for each run.
  • Data type of the classifier: 2.3% average downstream-accuracy improvement comes from keeping the classifier floating-point rather than quantized.The classifier is only 0.3% of model size for 10 classes, so the overhead is marginal.
  • Single-batch training & momentum: Batch size 128 with momentum reaches 74.4% average accuracy versus 72.4% without momentum, whereas batch size 1 reaches 71.5% with momentum versus 72.3% without it.Momentum doubles parameter-update memory requirements when its buffer is safely quantized.
  • Single-batch training & momentum: Without momentum, normal-batch and single-batch training achieve similar average accuracy: 72.4% versus 72.3%.This supports using batched training results as an approximation for evaluation.
  • Gradient accumulation: 2.5% lower accuracy occurs on VWW when simulated training approximates actual microcontroller training, because binary labels provide limited information and unstable gradients.Gradient accumulation closes this gap at an 11% extra memory cost, achieving 89.1% top-1 accuracy within 256KB.

D Evolutionary Search vs. Random Search

Evolutionary search is more sample-efficient than random search for optimizing sparse-update schemes under a memory constraint and finds a better final solution.

  • Evolutionary search finds a better final sparse-update solution than random search under a 100KB analytic memory constraint.The comparison uses MCUNet-5FPS and evaluates solution quality by higher sum of ∆acc.
  • Evolutionary search explores the sparse-update search space with better sample efficiency than random search.The search is performed for MCUNet-5FPS under the same 100KB analytic memory constraint.

E Amount of Compute

The study evaluates eight downstream training schemes with modest GPU cost, while ImageNet pre-training is a larger one-time expense.

  • 3.2 GPU hours are required to train one run across all 8 downstream datasets, excluding hyper-parameter tuning.Experiments were simulated on NVIDIA GeForce RTX 3090 GPUs.
  • 31.5 GPU hours are required for ImageNet pre-training over 300 epochs, performed once for each model.

F More Contribution Analysis Results

Contribution analysis on Cars shows that later layers matter more for accuracy improvement, while the first point-wise convolution in each block contributes most within that block.

  • Later layers contribute more to accuracy improvement in MobileNetV2-w0.35 and ProxylessNAS-w0.3 on Cars.
  • Within each block, the first point-wise convolutional layer contributes most to accuracy improvement.
  • Figures 14 and 15 present contribution analyses for MobileNetV2-w0.35 and ProxylessNAS-w0.3, respectively.

G Other Partial Update Methods That Did Not Work

Several partial-update alternatives were rejected because they either consumed too much memory, performed poorly with quantized weights, or significantly reduced transfer-learning accuracy.

  • LoRA was unsuitable because merging its update with quantized weights produces a full-size weight that consumes substantial SRAM.
  • Randomly initialized low-rank updates were difficult to apply to quantized weights and led to worse performance empirically.
  • Replacing convolutions with grouped convolutions reduced update cost but significantly hindered transfer-learning accuracy.
  • Changing model architectures for lighter convolutional alternatives also introduced unwanted complexity.
Loading 2206.15472v4…