Advancing Dynamic Sparse Training by Exploring Optimization Opportunities

Dynamic Sparse Training (DST) has been effectively addressing the substantial training resource requirements of increasingly large Deep Neural Networks (DNNs). Characterized by its dynamic "train-prune-grow" schedule during training, DST implicitly develops a bi-level structure for training the weights while discovering a subnetwork topology. However, such a structure is consistently overlooked by the current DST algorithms for further optimization opportunities, and these algorithms, on the other hand, solely optimize the weights while determining masks heuristically. In this paper, we extensively study DST algorithms and argue that the training scheme of DST naturally forms a bi-level problem in which the updating of weight and mask is interdependent. Based on this observation, we introduce a novel efficient training framework called BiDST, which for the first time, introduces bi-level optimization methodology into dynamic sparse training domain. Unlike traditional partialheuristic DST schemes, which suffer from suboptimal search efficiency for masks and miss the opportunity to fully explore the topological space of neural networks, BiDST excels at discovering excellent sparse patterns by optimizing mask and weight simultaneously, resulting in maximum 2.62% higher accuracy, 2.1× faster execution speed, and 25× reduced overhead. Code available at https://github.com/jjsrf/ BiDST-ICML2024 .