# On-Device Training Under 256KB Memory

arXiv: https://arxiv.org/abs/2206.15472 · HTML: https://mlsanity.com/paper/2206.15472/

## Key facts

- Cited by 160 papers in the map (data through 14 Oct 2026).
- Builds on 49 of its 76 references that are in the map.
- First posted on arXiv 30 Jun 2022.

## Builds on (most cited first, 30 of 49)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [MobileNetV2: Inverted Residuals and Linear Bottlenecks](https://mlsanity.com/paper/1801.04381/) (2018)
- [Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding](https://mlsanity.com/paper/1510.00149/) (2015)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference](https://mlsanity.com/paper/1712.05877/) (2017)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [BitFit: Simple Parameter-efficient Fine-tuning for Transformer-based Masked Language-models](https://mlsanity.com/paper/2106.10199/) (2021)
- [Learning Transferable Architectures for Scalable Image Recognition](https://mlsanity.com/paper/1707.07012/) (2017)
- [MnasNet: Platform-Aware Neural Architecture Search for Mobile](https://mlsanity.com/paper/1807.11626/) (2018)
- [XNOR-Net: ImageNet Classification Using Binary Convolutional Neural Networks](https://mlsanity.com/paper/1603.05279/) (2016)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Big Transfer (BiT): General Visual Representation Learning](https://mlsanity.com/paper/1912.11370/) (2019)
- [Channel Pruning for Accelerating Very Deep Neural Networks](https://mlsanity.com/paper/1707.06168/) (2017)
- [Learning Efficient Convolutional Networks through Network Slimming](https://mlsanity.com/paper/1708.06519/) (2017)
- [ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware](https://mlsanity.com/paper/1812.00332/) (2018)
- [Do Better ImageNet Models Transfer Better?](https://mlsanity.com/paper/1805.08974/) (2018)
- [Large Batch Training of Convolutional Networks](https://mlsanity.com/paper/1708.03888/) (2017)
- [Ray: A Distributed Framework for Emerging AI Applications](https://mlsanity.com/paper/1712.05889/) (2017)
- [PACT: Parameterized Clipping Activation for Quantized Neural Networks](https://mlsanity.com/paper/1805.06085/) (2018)
- [HAQ: Hardware-Aware Automated Quantization with Mixed Precision](https://mlsanity.com/paper/1811.08886/) (2018)
- [DeCAF: A Deep Convolutional Activation Feature for Generic Visual Recognition](https://mlsanity.com/paper/1310.1531/) (2013)
- [AMC: AutoML for Model Compression and Acceleration on Mobile Devices](https://mlsanity.com/paper/1802.03494/) (2018)
- [CNN Features off-the-shelf: an Astounding Baseline for Recognition](https://mlsanity.com/paper/1403.6382/) (2014)
- [MCUNet: Tiny Deep Learning on IoT Devices](https://mlsanity.com/paper/2007.10319/) (2020)
- [MetaPruning: Meta Learning for Automatic Neural Network Channel Pruning](https://mlsanity.com/paper/1903.10258/) (2019)

## Built on by (most cited first, 30 of 160)

- [On-Device Language Models: A Comprehensive Review](https://mlsanity.com/paper/2409.00088/) (2024)
- [Offsite-Tuning: Transfer Learning without Full Model](https://mlsanity.com/paper/2302.04870/) (2023)
- [Sensitivity-Aware Visual Parameter-Efficient Fine-Tuning](https://mlsanity.com/paper/2303.08566/) (2023)
- [Tiny Machine Learning: Progress and Futures](https://mlsanity.com/paper/2403.19076/) (2024)
- [Q-GaLore: Quantized GaLore with INT4 Projection and Layer-Adaptive Low-Rank Gradients](https://mlsanity.com/paper/2407.08296/) (2024)
- [COAT: Compressing Optimizer states and Activation for Memory-Efficient FP8 Training](https://mlsanity.com/paper/2410.19313/) (2024)
- [How Far Are We From AGI: Are LLMs All We Need?](https://mlsanity.com/paper/2405.10313/) (2024)
- [Galaxy: A Resource-Efficient Collaborative Edge AI System for In-situ Transformer Inference](https://mlsanity.com/paper/2405.17245/) (2024)
- [Streaming Deep Reinforcement Learning Finally Works](https://mlsanity.com/paper/2410.14606/) (2024)
- [PockEngine: Sparse and Efficient Fine-tuning in a Pocket](https://mlsanity.com/paper/2310.17752/) (2023)
- [Adapting to Distribution Shift by Visual Domain Prompt Generation](https://mlsanity.com/paper/2405.02797/) (2024)
- [On-device Training: A First Overview on Existing Systems](https://mlsanity.com/paper/2212.00824/) (2022)
- [AI-native Interconnect Framework for Integration of Large Language Model Technologies in 6G Systems](https://mlsanity.com/paper/2311.05842/) (2023)
- [Multisize Dataset Condensation](https://mlsanity.com/paper/2403.06075/) (2024)
- [TinyTrain: Resource-Aware Task-Adaptive Sparse Training of DNNs at the Data-Scarce Edge](https://mlsanity.com/paper/2307.09988/) (2023)
- [Empirical Guidelines for Deploying LLMs onto Resource-constrained Edge Devices](https://mlsanity.com/paper/2406.03777/) (2024)
- [Unleashing the Potential of Spiking Neural Networks by Dynamic Confidence](https://mlsanity.com/paper/2303.10276/) (2023)
- [Computation-efficient Deep Learning for Computer Vision: A Survey](https://mlsanity.com/paper/2308.13998/) (2023)
- [Efficient On-device Training via Gradient Filtering](https://mlsanity.com/paper/2301.00330/) (2023)
- [On-Device Domain Learning for Keyword Spotting on Low-Power Extreme Edge Embedded Systems](https://mlsanity.com/paper/2403.10549/) (2024)
- [Asteroid: Resource-Efficient Hybrid Pipeline Parallelism for Collaborative DNN Training on Heterogeneous Edge Devices](https://mlsanity.com/paper/2408.08015/) (2024)
- [Tensor-Compressed Back-Propagation-Free Training for (Physics-Informed) Neural Networks](https://mlsanity.com/paper/2308.09858/) (2023)
- [Towards Green AI in Fine-tuning Large Language Models via Adaptive Backpropagation](https://mlsanity.com/paper/2309.13192/) (2023)
- [ElasticTrainer: Speeding Up On-Device Training with Runtime Elastic Tensor Selection](https://mlsanity.com/paper/2312.14227/) (2023)
- [LifeLearner: Hardware-Aware Meta Continual Learning System for Embedded Computing Platforms](https://mlsanity.com/paper/2311.11420/) (2023)
- [Training Machine Learning models at the Edge: A Survey](https://mlsanity.com/paper/2403.02619/) (2024)
- [On-Device Unsupervised Image Segmentation](https://mlsanity.com/paper/2303.12753/) (2023)
- [Reduced Precision Floating-Point Optimization for Deep Neural Network On-Device Learning on MicroControllers](https://mlsanity.com/paper/2305.19167/) (2023)
- [To Spike or Not To Spike: A Digital Hardware Perspective on Deep Learning Acceleration](https://mlsanity.com/paper/2306.15749/) (2023)
- [Efficient Low-rank Backpropagation for Vision Transformer Adaptation](https://mlsanity.com/paper/2309.15275/) (2023)

## Abstract (arXiv)

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.