Neural Network Pruning for Invariance Learning

摘要

Model scaling laws have driven the development of neural networks with more-and-more parameters. As a result, there is a growing need for neural network pruning to enable the efficient deployment during inference. We take a deeper look at neural network pruning from the lens of invariance preservation. We argue that successfully pruned neural networks should be invariant to transformations which do not alter the data's underlying semantics (ex. translations). To this goal, we first show existing post-pruning algorithms do not perserve desired invariances. We then propose a framework to discover novel architectures that do capture desired invariances from data via pruning. Specifically, we show contrastive learning with small initialization can effectively transfer invariance preservation encoded in the model weights to the pruning mask. Our approach consistently outperform traditional pruning algorithms on fully-connected, convolutional, and transformer networks across 3 vision, 40 tabular, and 1 natural language datasets.