Finding Lottery Tickets in Vision Models Via Data-driven Spectral Foresight Pruning

Recent advances in neural network pruning have shown how it is possible toreduce the computational costs and memory demands of deep learning modelsbefore training. We focus on this framework and propose a new pruning atinitialization algorithm that leverages the Neural Tangent Kernel (NTK) theoryto align the training dynamics of the sparse network with that of the denseone. Specifically, we show how the usually neglected data-dependent componentin the NTK's spectrum can be taken into account by providing an analyticalupper bound to the NTK's trace obtained by decomposing neural networks intoindividual paths. This leads to our Path eXclusion (PX), a foresight pruningmethod designed to preserve the parameters that mostly influence the NTK'strace. PX is able to find lottery tickets (i.e. good paths) even at highsparsity levels and largely reduces the need for additional training. Whenapplied to pre-trained models it extracts subnetworks directly usable forseveral downstream tasks, resulting in performance comparable to those of thedense counterpart but with substantial cost and computational savings. Codeavailable at: https://github.com/iurada/px-ntk-pruning