Pruning-as-Reconstruct: Masked Autoencoders are Efficient Importance Indicators
Primary Area: representation learning for computer vision, audio, language, and other modalities
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Keywords: Network Pruning, Deep Neural Networks, Masked Autoencoder, Reconstruction
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.
Abstract: Network pruning has emerged as an effective technique for reducing the size and computational complexity of neural networks, thereby addressing the challenges of deployment on resource-limited devices. However, existing pruning criteria are predominantly based on handcrafted heuristics or calculated statistics, hindering their generality and effectiveness. In this paper, we reveal that masked autoencoder (MAE) can exploit the hidden semantic information within structured parameters, thereby functioning as a learnable pruning criterion. Specifically, to address the dimension inconsistency problem between layers, we propose a parallel training pipeline, facilitating stable and efficient MAE training on weight matrices. Based on the 'harder-reconstructed-more-important' assumption, we explore diverse pruning strategies and formulate structured pruning as a sample-without-replacement problem that strikes a balance between algorithm complexity and performance. Extensive experiments on benchmark datasets, including CIFAR-10 and ImageNet, demonstrate that our method can efficiently compress both convolutional neural networks and transformers. Furthermore, the trained MAE exhibits transferability across various structures and datasets, avoiding repetitive training from scratch and highlighting its potential as a universal pruning criterion. To the best of our knowledge, this is the first work that establishes a connection between structured pruning and self-supervised learning.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 941
Loading