Train Flat, Then Compress: Sharpness-Aware Minimization Learns More Compressible Models

Abstract

Model compression by way of parameter pruning, quantization, or distillation has recently gained popularity as an approach for reducing the computational requirements of modern deep neural network models for NLP. Inspired by prior works suggesting a connection between simpler, more generalizable models and those that lie within wider loss basins, we hypothesize that optimizing for flat minima should lead to simpler parameterizations and thus more compressible models. We propose to combine sharpness-aware minimization (SAM) with various task-specific model compression methods, including iterative magnitude pruning (IMP), structured pruning with a distillation objective, and post-training dynamic quantization. Empirically, we show that optimizing for flatter minima consistently leads to greater compressibility of parameters compared to vanilla Adam when fine-tuning BERT models, with little to no loss in accuracy on the GLUE text classification and SQuAD question answering benchmarks. Moreover, SAM finds superior winning tickets during IMP that 1) are amenable to vanilla Adam optimization, and 2) transfer more effectively across tasks.

arXiv commentEMNLP 2022 Findings, 28 pages

References in corpus (19)

• Adam: A Method for Stochastic Optimization
• Distilling the Knowledge in a Neural Network
• HuggingFace's Transformers: State-of-the-art Natural Language Processing
• The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks
• What is being transferred in transfer learning?
• Comparing Rewinding and Fine-tuning in Neural Network Pruning
• Fantastic Generalization Measures and Where to Find Them
• Towards Understanding Generalization of Deep Learning: Perspective of Loss Landscapes
• The De-democratization of AI: Deep Learning and the Compute Divide in Artificial Intelligence Research
• Linear Mode Connectivity and the Lottery Ticket Hypothesis
• Non-Vacuous Generalization Bounds at the ImageNet Scale: A PAC-Bayesian Compression Approach
• Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers
• An Empirical Investigation of the Role of Pre-training in Lifelong Learning
• Lost in Pruning: The Effects of Pruning Neural Networks beyond Test Accuracy
• Efficient Sharpness-aware Minimization for Improved Training of Neural Networks
• Sharpness-Aware Training for Free
• Robustness to Pruning Predicts Generalization in Deep Neural Networks
• Low-Pass Filtering SGD for Recovering Flat Optima in the Deep Learning Optimization Landscape
• Unmasking the Lottery Ticket Hypothesis: What's Encoded in a Winning Ticket's Mask?