On-Device Training Under 256KB Memory

TL;DR

This work introduces a novel algorithm‑system co‑design that enables on‑device training within extremely tight memory limits on microcontrollers. By combining real quantized graph optimization with Quantization‑Aware Scaling (QAS), memory‑efficient Sparse Update, and a lightweight compile‑time Tiny Training Engine (TTE) that prunes backward graphs and supports in‑place updates, the approach achieves training within $256\ \mathrm{KB}$ SRAM and $1\ \mathrm{MB}$ Flash without auxiliary memory. QAS stabilizes 8‑bit training, while automated contribution analysis selects sparse update schemes that maximize downstream accuracy under memory constraints; TTE further reduces memory and accelerates training through code generation and graph optimization. The framework matches or surpasses cloud training accuracy on tinyML benchmarks (e.g., VWW) with substantial memory reductions (over 1000×) and speedups (up to 25×), enabling lifelong on‑device learning for privacy‑preserving personalization. The work lays a foundation for practical on‑device lifelong learning on ultra‑low‑memory hardware, with potential extensions to other modalities and models.

Abstract

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.

Figures (15)

• Figure 1: Algorithm and system co-design reduces the training memory from 303MB (PyTorch) to 141KB with the same transfer learning accuracy, leading to 2300$\times$ reduction. The numbers are measured with MobilenetV2-w0.35 sandler2018mobilenetv2, batch size 1 and resolution 128$\times$128. It can be deployed to a microcontroller with 256KB SRAM.
• Figure 2: Real quantized graphs (our optimized graph, designed for efficiency) vs.fake quantized graphs (for QAT, designed for simulation). The fake quantize graphs cannot provide memory saving due to floating-point operations. We need to use real quantized graph to fit the tight memory constraint.
• Figure 3: The quantized model has a very different weight/gradient norm ratio (i.e., $\lVert\mathbf{W}\rVert/\lVert\mathbf{G}\rVert$) compared to the floating-point model at training time. QAS stabilizes the $\lVert\mathbf{W}\rVert/\lVert\mathbf{G}\rVert$ ratio and helps optimization. For example, in the highlighted area, the ratios of the quantized model fluctuate dramatically in a zigzag pattern (weight, bias, weight, bias, ...); after applying QAS, the pattern stabilizes and matches the fp32 counterpart.
• Figure 4: Different update paradigms of two linear layers in a deep neural network.
• Figure 5: Contribution analysis of updating biases and weights. (a) For bias update, the accuracy generally goes higher as more layers are updated, but plateaus soon. (b) For updating the weight of a specific layer, the later layers appear to be more important; the first point-wise conv (pw1) in an inverted bottleneck block sandler2018mobilenetv2 appears to be more important; and the gains are bigger with more channels updated. (c) The automated selection based on contribution analysis is effective: the actual downstream accuracy shows a positive correlation with $\sum\Delta\text{acc}$.