LocalProp cuts memory use in deep learning by updating weights locally
LocalProp: Neuro-Localized Memory-Efficient Backpropagation
Artificial IntelligenceComputer Vision and Pattern RecognitionMachine Learning
Summary
Training deep learning models usually involves adjusting all the parts of the model at once, which uses a lot of computer memory and is not how brains learn. The authors propose LocalProp, a method where only small sections of the model are updated at a time, making training more memory efficient. They use a step called pruning to pass learning signals from higher to lower parts of the model and finish with a final fine-tuning. Experiments show LocalProp performs well on big image datasets while using much less GPU memory than traditional methods.
What this means in practice
- •For deep learning engineers: Enable training large neural networks with reduced GPU memory usage by updating weights locally and pruning between updates.
- •For cloud ai service providers: Reduce compute resource costs for large-scale model training by employing memory-efficient local update techniques.$Commercial implications: Allows selling AI training services with lower hardware requirements and costs, enhancing scalability.
Authors
Diana-Nicoleta Grigore, Iuliana Georgescu, Radu Tudor Ionescu
Abstract
The current deep learning training paradigm employs end-to-end backpropagation, regardless of the training stage, i.e. pre-training or fine-tuning. However, backpropagating through the entire model is neither biologically plausible nor memory efficient, since learning inside the brain is highly localized. Therefore, we propose LocalProp, a training procedure that locally updates the weights of a model. Our neuro-localized weight updates follow the "pre-training then fine-tuning" paradigm, where the pre-training is based on I-JEPA. After locally updating the weights, a pruning operation is performed, followed by a short final fine-tuning phase. Pruning helps by sending the learning signal from higher blocks to lower blocks. We perform experiments on several datasets, including large-scale benchmarks such as ImageNet, and empirically show that LocalProp reaches good performance at a fraction of GPU peak memory. By varying the number of jointly optimized blocks, we identify gradient-propagation span as a practical control over the accuracy-memory trade-off.