Sparse Training of Neural Networks based on Multilevel Mirror Descent

TL;DR

The paper tackles the challenge of training sparse neural networks efficiently by marrying linearized Bregman iterations with a multilevel (coarse-to-fine) optimization framework. It introduces adaptive freezing of nonzero parameters to exploit sparsity and enable computational savings, while ensuring convergence under a Polyak–Łojasiewicz-type condition. The proposed multilevel LinBreg method achieves highly sparse yet accurate models on standard vision benchmarks and reports substantial theoretical FLOPs reductions relative to SGD. Experimental results on CIFAR-10 and TinyImageNet show competitive accuracy with higher sparsity compared to state-of-the-art sparse training methods, and the approach yields real CPU speedups with sparse kernels. The work provides a principled convergence analysis and highlights practical implications for resource-efficient training of large neural networks.

Abstract

We introduce a dynamic sparse training algorithm based on linearized Bregman iterations / mirror descent that exploits the naturally incurred sparsity by alternating between periods of static and dynamic sparsity pattern updates. The key idea is to combine sparsity-inducing Bregman iterations with adaptive freezing of the network structure to enable efficient exploration of the sparse parameter space while maintaining sparsity. We provide convergence guaranties by embedding our method in a multilevel optimization framework. Furthermore, we empirically show that our algorithm can produce highly sparse and accurate models on standard benchmarks. We also show that the theoretical number of FLOPs compared to SGD training can be reduced from 38% for standard Bregman iterations to 6% for our method while maintaining test accuracy.

Figures (7)

• Figure 1: Achieved accuracy and sparsity of different optimizers on CIFAR-10 for different architectures.
• Figure 2: Timing breakdown for training ResNet18 for 10 epochs on an AMD Ryzen 5 4500U CPU when employing SparseProp modules to exploit unstructured sparsity. The sparsity induced by ML LinBreg can lead to real-time computational savings.
• Figure 3: Achieved accuracy and sparsity of different optimizers on TinyImageNet for different architectures.
• Figure 4: Ablation study comparing accuracy and sparsity across different choices of the hyperparameters $\lambda$ and $m$ for ResNet18 on CIFAR-10.
• Figure 5: Mean and standard deviation of validation accuracy, sparsity, and train loss of the model parameters over training epochs, shown for different regularization parameters, $\lambda = 0.005$ and $\lambda = 0.01$ and for LinBreg with $\lambda=0.2$.

Theorems & Definitions (20)

• Proposition 3.1
• Proposition 3.2
• Theorem 4.1
• Corollary 4.2
• proof : Proof of Proposition \ref{['subgradients1']}
• proof : Proof of Proposition \ref{['subgradients2']}
• Lemma 1.1
• proof
• Lemma 1.2
• proof