Low-rank Gradient Approximation For Memory-Efficient On-device Training of Deep Neural Network
This work addresses memory constraints for on-device training, enabling more advanced optimization techniques on mobile devices, though it is incremental as it builds on existing gradient methods.
The paper tackles the memory limitation problem in on-device training of deep neural networks by proposing a low-rank gradient approximation method, which reduces training memory by about 33.0% for Adam optimization and achieves a 4.5% relative lower word error rate on an ASR personalization task.
Training machine learning models on mobile devices has the potential of improving both privacy and accuracy of the models. However, one of the major obstacles to achieving this goal is the memory limitation of mobile devices. Reducing training memory enables models with high-dimensional weight matrices, like automatic speech recognition (ASR) models, to be trained on-device. In this paper, we propose approximating the gradient matrices of deep neural networks using a low-rank parameterization as an avenue to save training memory. The low-rank gradient approximation enables more advanced, memory-intensive optimization techniques to be run on device. Our experimental results show that we can reduce the training memory by about 33.0% for Adam optimization. It uses comparable memory to momentum optimization and achieves a 4.5% relative lower word error rate on an ASR personalization task.