Train faster, perform better : modular adaptive training in over-parameterized models

Shi, Yubin and Chen, Yixuan and Dong, Mingzhi and Yang, Xiaochen and Li, Dongsheng and Wang, Yujiang and Dick, Robert and Lv, Qin and Zhao, Yingying and Yang, Fan and Lu, Tun and Gu, Ning and Shang, Li (2023) Train faster, perform better : modular adaptive training in over-parameterized models. In: Advances in Neural Information Processing Systems, 2023-12-10 - 2023-12-16.

[thumbnail of Shi-etal-NeurIPS-2023-Train-faster-perform-better] Text. Filename: Shi-etal-NeurIPS-2023-Train-faster-perform-better.pdf
Accepted Author Manuscript
Restricted to Repository staff only until 10 December 2024.
License: Strathprints license 1.0
Download (2MB) | Request a copy

Abstract

Despite their prevalence in deep-learning communities, over-parameterized models convey high demands of computational costs for proper training. This work studies the fine-grained, modular-level learning dynamics of over-parameterized models to attain a more efficient and fruitful training strategy. Empirical evidence reveals that when scaling down into network modules, such as heads in self-attention models, we can observe varying learning patterns implicitly associated with each module’s trainability. To describe such modular-level learning capabilities, we introduce a novel concept dubbed modular neural tangent kernel (mNTK), and we demonstrate that the quality of a module’s learning is tightly associated with its mNTK’s principal eigenvalue λmax. A large λmax indicates that the module learns features with better convergence, while those miniature ones may impact generalization negatively. Inspired by the discovery, we propose a novel training strategy termed Modular Adaptive Training (MAT) to update those modules with their λmax exceeding a dynamic threshold selectively, concentrating the model on learning common features and ignoring those superfluous ones. Unlike most existing training schemes with a complete BP cycle across all network modules, MAT can significantly save computations by its partially-updating strategy and can further improve performance. Experiments show that MAT nearly halves the computational cost of model training and outperforms the accuracy of baselines.

CORE (COnnecting REpositories)