Spending Your Winning Lottery Better After Drawing It

AJAY KUMAR JAISWAL, Haoyu Ma, Tianlong Chen, Ying Ding, Zhangyang Wang

Keywords: Lottery Ticket Hypothesis, Sparse Network Optimization, Sparse Training

Abstract:

Lottery Ticket Hypothesis (LTH) (Frankle & Carbin, 2019) suggest suggests that a dense neural network contains a sparse sub-network that can match the performance of the original dense network when trained in isolation from scratch. Most works retrain the sparse sub-network with the same training protocols as its dense network, such as initialization, architecture blocks, and training recipes. However, till now it is unclear that whether these training protocols are optimal for sparse networks.
In this paper, we demonstrate that it is unnecessary for spare retraining to strictly inherit those properties from the dense network. Instead, by plugging in purposeful "tweaks" of the sparse subnetwork architecture or its training recipe, its retraining can be significantly improved than the default, especially at high sparsity levels. Combining all our proposed "tweaks" can yield the new state-of-the-art performance of LTH, and these modifications can be easily adapted to other sparse training algorithms in general. Specifically, we have achieved a significant and consistent performance gain of 1.05% - 4.93% for ResNet18 on CIFAR-100 over vanilla-LTH. Moreover, our methods are shown to generalize across datasets (CIFAR10, CIFAR100, TinyImageNet) and architectures (Vgg16, ResNet-18/ResNet-34, MobileNet). All codes will be publicly available.

One-sentence Summary: In this work, we demonstrate that purposely re-tweaking the sparse subnetwork re-training step of LTH by incorporating smoothness can achieve better performance.