# POET: Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging

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

## Key facts

- Cited by 26 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 40 references that are in the map.
- First posted on arXiv 15 Jul 2022.

## Builds on (most cited first, 20 of 20)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [EfficientNetV2: Smaller Models and Faster Training](https://mlsanity.com/paper/2104.00298/) (2021)
- [SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and <0.5MB model size](https://mlsanity.com/paper/1602.07360/) (2016)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware](https://mlsanity.com/paper/1812.00332/) (2018)
- [What is the State of Neural Network Pruning?](https://mlsanity.com/paper/2003.03033/) (2020)
- [Federated Learning for Mobile Keyboard Prediction](https://mlsanity.com/paper/1811.03604/) (2018)
- [HAWQ: Hessian AWare Quantization of Neural Networks with Mixed-Precision](https://mlsanity.com/paper/1905.03696/) (2019)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [The Carbon Footprint of Machine Learning Training Will Plateau, Then Shrink](https://mlsanity.com/paper/2204.05149/) (2022)
- [Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization](https://mlsanity.com/paper/1910.02653/) (2019)
- [Federated Evaluation and Tuning for On-Device Personalization: System Design & Applications](https://mlsanity.com/paper/2102.08503/) (2021)
- [ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training](https://mlsanity.com/paper/2104.14129/) (2021)
- [Dynamic Tensor Rematerialization](https://mlsanity.com/paper/2006.09616/) (2020)
- [E2-Train: Training State-of-the-art CNNs with Over 80% Energy Savings](https://mlsanity.com/paper/1910.13349/) (2019)
- [On-Device Neural Net Inference with Mobile GPUs](https://mlsanity.com/paper/1907.01989/) (2019)
- [Memory Optimization for Deep Networks](https://mlsanity.com/paper/2010.14501/) (2020)

## Built on by (most cited first, 26 of 26)

- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Tiny Machine Learning: Progress and Futures](https://mlsanity.com/paper/2403.19076/) (2024)
- [On-device Training: A First Overview on Existing Systems](https://mlsanity.com/paper/2212.00824/) (2022)
- [Asteroid: Resource-Efficient Hybrid Pipeline Parallelism for Collaborative DNN Training on Heterogeneous Edge Devices](https://mlsanity.com/paper/2408.08015/) (2024)
- [Training Machine Learning models at the Edge: A Survey](https://mlsanity.com/paper/2403.02619/) (2024)
- [OLLA: Optimizing the Lifetime and Location of Arrays to Reduce the Memory Usage of Neural Networks](https://mlsanity.com/paper/2210.12924/) (2022)
- [SlimPipe: Memory-Thrifty and Efficient Pipeline Parallelism for Long-Context LLM Training](https://mlsanity.com/paper/2504.14519/) (2025)
- [Colossal-Auto: Unified Automation of Parallelization and Activation Checkpoint for Large-scale Models](https://mlsanity.com/paper/2302.02599/) (2023)
- [On-Device Training of Fully Quantized Deep Neural Networks on Cortex-M Microcontrollers](https://mlsanity.com/paper/2407.10734/) (2024)
- [Rockmate: an Efficient, Fast, Automatic and Generic Tool for Re-materialization in PyTorch](https://mlsanity.com/paper/2307.01236/) (2023)
- [Peer-to-Peer Deep Learning for Beyond-5G IoT](https://mlsanity.com/paper/2310.18861/) (2023)
- [FlexInfer: Breaking Memory Constraint via Flexible and Efficient Offloading for On-Device LLM Inference](https://mlsanity.com/paper/2503.03777/) (2025)
- [FedEL: Federated Elastic Learning for Heterogeneous Devices](https://mlsanity.com/paper/2509.16902/) (2025)
- [XEngine: Optimal Tensor Rematerialization for Neural Networks in Heterogeneous Environments](https://mlsanity.com/paper/2212.09290/) (2022)
- [Towards Federated Learning Under Resource Constraints via Layer-wise Training and Depth Dropout](https://mlsanity.com/paper/2309.05213/) (2023)
- [Breaking On-device Training Memory Wall: A Systematic Survey](https://mlsanity.com/paper/2306.10388/) (2023)
- [Resource-Efficient Personal Large Language Models Fine-Tuning with Collaborative Edge Computing](https://mlsanity.com/paper/2408.10746/) (2024)
- [TrainDeeploy: Hardware-Accelerated Parameter-Efficient Fine-Tuning of Small Transformer Models at the Extreme Edge](https://mlsanity.com/paper/2603.09511/) (2026)
- [Moccasin: Efficient Tensor Rematerialization for Neural Networks](https://mlsanity.com/paper/2304.14463/) (2023)
- [Privacy-Enhanced Training-as-a-Service for On-Device Intelligence: Concept, Architectural Scheme, and Open Problems](https://mlsanity.com/paper/2404.10255/) (2024)
- [Online Resource Allocation for Edge Intelligence with Colocated Model Retraining and Inference](https://mlsanity.com/paper/2405.16029/) (2024)
- [MUDAS: Mote-scale Unsupervised Domain Adaptation in Multi-label Sound Classification](https://mlsanity.com/paper/2506.11331/) (2025)
- [FTTE: Enabling Federated and Resource-Constrained Deep Edge Intelligence](https://mlsanity.com/paper/2510.03165/) (2025)
- [What changes after deployment? A survey on On-device Learning in TinyML](https://mlsanity.com/paper/2605.31226/) (2026)
- [QScheduler: Adaptive Gradient Sampling for Zeroth-Order On-Device Training on INT8 NPUs](https://mlsanity.com/paper/2607.18802/) (2026)
- [GaLe: memory-efficient Global Approximate and Local Exact features](https://mlsanity.com/paper/2609.02689/) (2026)

## Abstract (arXiv)

Fine-tuning models on edge devices like mobile phones would enable privacy-preserving personalization over sensitive data. However, edge training has historically been limited to relatively small models with simple architectures because training is both memory and energy intensive. We present POET, an algorithm to enable training large neural networks on memory-scarce battery-operated edge devices. POET jointly optimizes the integrated search search spaces of rematerialization and paging, two algorithms to reduce the memory consumption of backpropagation. Given a memory budget and a run-time constraint, we formulate a mixed-integer linear program (MILP) for energy-optimal training. Our approach enables training significantly larger models on embedded devices while reducing energy consumption while not modifying mathematical correctness of backpropagation. We demonstrate that it is possible to fine-tune both ResNet-18 and BERT within the memory constraints of a Cortex-M class embedded device while outperforming current edge training methods in energy efficiency. POET is an open-source project available at https://github.com/ShishirPatil/poet