TinyTL: Reduce Activations, Not Trainable Parameters for Efficient On-Device Learning

Han Cai 1 , Chuang Gan 2 , Ligeng Zhu 1 , Song Han 1
Massachusetts Institute of Technology, MIT-IBM Watson AI Lab

Abstract

On-device learning enables edge devices to continually adapt the AI models to new data, which requires a small memory footprint to fit the tight memory constraint of edge devices. Existing work solves this problem by reducing the number of trainable parameters. However, this doesn't directly translate to memory saving since the major bottleneck is the activations, not parameters. In this work, we present Tiny-Transfer-Learning (TinyTL) for memory-efficient on-device learning. TinyTL freezes the weights while only learns the bias modules, thus no need to store the intermediate activations. To maintain the adaptation capacity, we introduce a new memory-efficient bias module, the lite residual module, to refine the feature extractor by learning small residual feature maps adding only 3.8% memory overhead. Extensive experiments show that TinyTL significantly saves the memory (up to 6.5x) with little accuracy loss compared to fine-tuning the full network. Compared to fine-tuning the last layer, TinyTL provides significant accuracy improvements (up to 34.1%) with little memory overhead. Furthermore, combined with feature extractor adaptation, TinyTL provides 7.3-12.9x memory saving without sacrificing accuracy compared to fine-tuning the full Inception-V3.

On-Device Learning, not Just Inference

Efficient On-device Learning Requires Small Training Memory

Activation is the Main Bottleneck, not Parameters

Tiny Transfer Learning

Results on Transfer Learning Datasets

Combining with Batch Size 1 Training

Citation

 @inproceedings{
    cai2020tinytl,
    title={TinyTL: Reduce Memory, Not Parameters for Efficient On-Device Learning},
    author={Cai, Han and Gan, Chuang and Zhu, Ligeng and Han, Song},
    booktitle={Advances in Neural Information Processing Systems},
    volume={33},
    year={2020}
} 

Acknowledgments: We thank MIT-IBM Watson AI Lab, NSF CAREER Award #1943349 and NSF Award #2028888 for supporting this research. We thank MIT Satori cluster for providing the computation resource.