Stochastic Subnetwork Annealing: A Regularization Technique for Fine Tuning Pruned Subnetworks
This addresses the challenge of fine-tuning pruned subnetworks for researchers and practitioners in model compression, though it is an incremental improvement over existing iterative pruning methods.
The paper tackles the problem of overfitting and convergence issues in iterative pruning of deep neural networks by introducing Stochastic Subnetwork Annealing, a regularization technique that uses stochastic masks to gradually evolve subnetworks, resulting in smoother optimization and improved robustness at high sparsity levels.
Pruning methods have recently grown in popularity as an effective way to reduce the size and computational complexity of deep neural networks. Large numbers of parameters can be removed from trained models with little discernible loss in accuracy after a small number of continued training epochs. However, pruning too many parameters at once often causes an initial steep drop in accuracy which can undermine convergence quality. Iterative pruning approaches mitigate this by gradually removing a small number of parameters over multiple epochs. However, this can still lead to subnetworks that overfit local regions of the loss landscape. We introduce a novel and effective approach to tuning subnetworks through a regularization technique we call Stochastic Subnetwork Annealing. Instead of removing parameters in a discrete manner, we instead represent subnetworks with stochastic masks where each parameter has a probabilistic chance of being included or excluded on any given forward pass. We anneal these probabilities over time such that subnetwork structure slowly evolves as mask values become more deterministic, allowing for a smoother and more robust optimization of subnetworks at high levels of sparsity.