Sparse Training of Neural Networks based on Multilevel Mirror Descent

TL;DR

This paper presents a novel dynamic sparse training algorithm for neural networks that combines Bregman iterations with multilevel optimization, achieving high sparsity, reduced computational cost, and maintained accuracy.

Contribution

It introduces a new sparse training method based on mirror descent and multilevel optimization, with proven convergence and significant efficiency improvements.

Findings

Achieves high sparsity with maintained accuracy on benchmarks.

Reduces FLOPs from 38% to 6% compared to SGD.

Cuts training time by about 50% with a sparsity-aware CPU implementation.

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.We additionally show a training time reduction by about 50%, when using a sparsity-aware CPU implementation of our method.