STARFISH: faST Accuracy Recovery in pruned networks From Internal State Healing

Pruning is a process designed to reduce the number of weights in a large neural network. This can substantially speed up inference but might cause a considerable reduction in the model's accuracy, and thus it is usually followed by a healing process that regains some of the lost accuracy. In this paper, we propose a new healing method, *STARFISH*, that can recover (most of) the accuracy of any pruned network efficiently. The main idea of STARFISH is to optimize the pruned network to align with the original network's internal state representations using a tiny calibration set of unlabeled examples. For the common case of removing $$\begin{gathered} 50 \\ \% \end{gathered}$$ of the weights, STARFISH healing improves the recovered accuracy by up to $$\begin{gathered} 22 \\ \% \end{gathered}$$ over the state-of-the-art methods on ViT-based networks. Its advantage is even more pronounced under aggressive pruning. For example, after eliminating $$\begin{gathered} 75 \\ \% \end{gathered}$$ of the weights in a DeiT-B network for ImageNet, STARFISH uses only $$\begin{gathered} 0.4 \\ \% \end{gathered}$$ of the number of training images as a calibration set and recovers $$\begin{gathered} 82 \\ \% \end{gathered}$$ of the original dense accuracy, whereas competing recovery techniques reach only $$\begin{gathered} 40 \\ \% \end{gathered}$$ of the dense model accuracy. The paper includes numerous ablations demonstrating the advantages of STARFISH over other common healing and distillation methods. Open-source code will be made available upon acceptance.