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

By Shishir G. Patil, Paras Jain, Prabal Dutta, Ion Stoica, Joseph E. Gonzalez. 2022

54 citations recorded by the source corpus.

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

Subfields: Artificial Intelligence, Computer Networks and Communications, Computer Vision and Pattern Recognition, Model Architectures, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Song Han
2. Chuang Gan
3. Ji Lin
4. Ligeng Zhu
5. Joseph E. Gonzalez
6. Ion Stoica
7. Wei-Ming Chen
8. Quoc V. Le

In Computer Networks and Communications

1. Joseph E. Gonzalez
2. Ion Stoica
3. Song Han
4. Chuang Gan
5. Ji Lin
6. Scott Shenker
7. Mosharaf Chowdhury
8. Hector Garcia-Molina

In Computer Vision and Pattern Recognition

1. Song Han
2. Chuang Gan
3. Ji Lin
4. Luc Van Gool
5. Trevor Darrell
6. Dacheng Tao
7. Quoc V. Le
8. Andrew Zisserman

In Model Architectures

1. Song Han
2. Chuang Gan
3. Ji Lin
4. James Hegarty
5. Haotong Zhang
6. Ligeng Zhu
7. Mostafa Elhoushi
8. Lianmin Zheng

In Efficiency & Compression

1. Chuang Gan
2. Song Han
3. Ji Lin
4. Ligeng Zhu
5. Joseph E. Gonzalez
6. Wei-Ming Chen
7. Wei-Chen Wang
8. Ion Stoica

Citations in the corpus

• Efficient Memory Management for Large Language Model Serving with PagedAttention (2023, methodology): “Swapping [23, 42, 55], recomputation [7, 24] and their combination [40] have been utilized to reduce the peak memory of training.”
• OLLA: Optimizing the Lifetime and Location of Arrays to Reduce the Memory Usage of Neural Networks (2022, background): “• More recently, combining several of these techniques has been proposed to increase their effectiveness and mitigate their drawbacks [4, 62] without fully eliminating them.”
• On-Device Training Under 256KB Memory (2022, methodology): “Recently, POET [57] utilizes rematerialization and paging to train on microcontrollers, but it relies on a large external memory.”