On-device learning gets faster with smart batch size tuning
TASTE: Throughput-Aware Batch Size Tuning for On-Device Edge Learning
Machine LearningArtificial IntelligenceHardware ArchitecturePerformance
Summary
Training AI models directly on devices like smartphones or Raspberry Pi can protect user privacy but is slow because of limited hardware. The authors found a way to adjust the number of data samples processed at once (batch size) to speed up training without hurting accuracy. They used a smart search method to find the best batch size and combined it with other tricks to double training speed on small devices. Their approach also works well when the AI learns new things over time, keeping old knowledge stable while learning fast.
on-device learningbatch sizethroughputBayesian optimizationsupervised learningcontinual learninggradient accumulationlearning rate scalingedge devicescatastrophic forgetting
Authors
Avik Bhatnagar, Federico Nicolas Peccia, Oliver Bringmann
Abstract
The rise of privacy-preserving artificial intelligence (AI) has shifted the focus of model adaptation and personalization towards on-device learning, where deep learning models are finetuned directly on edge hardware using local user data. However, this shift requires optimization of deep learning training on resource-constrained hardware to maximize throughput while maintaining predictive accuracy. This paper introduces a novel technique for on-device model training that incorporates an efficient Bayesian optimization-based batch size tuning approach to maximize hardware throughput. To evaluate the impact of this hyperparameter on the learning dynamics, we investigated two distinct paradigms: standard supervised learning (SL) and online continual learning (CL). Experimental results across various edge devices demonstrate a throughput ceiling, beyond which increasing the batch size yields no additional throughput gains. The proposed tuning approach identifies the optimal batch size, which, when combined with gradient accumulation and linear learning rate scaling, achieves up to a 2X increase in training throughput on platforms such as Raspberry Pi 4 compared to maximum batch sizes, without compromising model accuracy. Furthermore, in the CL paradigm, we demonstrate that optimal batch sizes maintain the stability-plasticity balance required for incremental learning, effectively mitigating catastrophic forgetting while maximizing computational efficiency on edge-hardware.