Only Train Once: A One-Shot Neural Network Training And Pruning Framework

2021/07/15 by Tianyi Chen, Bo Ji, Chen, Tianyi +18 · 1 voice · 43 citations
Computer Science · Mathematics · #Advanced Neural Network Applications #Algorithm #Artificial intelligence #Artificial neural network #Computer science #Deep neural networks #Domain Adaptation and Few-Shot Learning #FLOPS #Heuristic #Inference #Invariant (physics) #Machine learning #Mathematics #Parallel computing #Partition (number theory) #Pruning #Speedup #Stochastic Gradient Optimization Techniques #cs.LG

paper · pdf · doi:10.48550/arxiv.2107.07467

published in arXiv (Cornell University) (Cornell University) · Accepted by NeurIPS 2021

openalex publication_date 2021/07/15 · openalex created_date 2021/07/19 · arxiv created 2021/11/12 · arxiv updated 2021/11/15 · openalex updated_date 2026/08/05

Abstract

Structured pruning is a commonly used technique in deploying deep neural networks (DNNs) onto resource-constrained devices. However, the existing pruning methods are usually heuristic, task-specified, and require an extra fine-tuning procedure. To overcome these limitations, we propose a framework that compresses DNNs into slimmer architectures with competitive performances and significant FLOPs reductions by Only-Train-Once (OTO). OTO contains two keys: (i) we partition the parameters of DNNs into zero-invariant groups, enabling us to prune zero groups without affecting the output; and (ii) to promote zero groups, we then formulate a structured-sparsity optimization problem and propose a novel optimization algorithm, Half-Space Stochastic Projected Gradient (HSPG), to solve it, which outperforms the standard proximal methods on group sparsity exploration and maintains comparable convergence. To demonstrate the effectiveness of OTO, we train and compress full models simultaneously from scratch without fine-tuning for inference speedup and parameter reduction, and achieve state-of-the-art results on VGG16 for CIFAR10, ResNet50 for CIFAR10 and Bert for SQuAD and competitive result on ResNet50 for ImageNet. The source code is available at https://github.com/tianyic/only_train_once.

Citations

• Nearly unbiased variable selection under minimax concave penalty

Cited by

• NIRVANA: Structured Pruning Reimagined for Large Language Model Compression
• ModHiFi: Identifying High Fidelity predictive components for Model Modification
• C-SWAP: Explainability-Aware Structured Pruning for Efficient Neural Networks Compression
• Pruning Overparameterized Multi-Task Networks for Degraded Web Image Restoration
• Optimally Deep Networks -- Adapting Model Depth to Datasets for Superior Efficiency
• Structured Output Regularization: a framework for few-shot transfer learning
• SOLAR: Switchable Output Layer for Accuracy and Robustness in Once-for-All Training
• Low-bit Model Quantization for Deep Neural Networks: A Survey
• Local Representative Token Guided Merging for Text-to-Image Generation
• Pruning by Block Benefit: Exploring the Properties of Vision Transformer Blocks during Domain Adaptation
• Loss-Aware Automatic Selection of Structured Pruning Criteria for Deep Neural Network Acceleration
• WINA: Weight Informed Neuron Activation for Accelerating Large Language Model Inference

Discussions

• Only Train Once: A One-Shot Neural Network Training and Pruning Framework [hn, 65 points, 11 comments]

Related

• Auto-Train-Once: Controller Network Guided Automatic Network Pruning from Scratch
• HESSO: Towards Automatic Efficient and User Friendly Any Neural Network Training and Pruning
• A Framework For Pruning Deep Neural Networks Using Energy-Based Models
• Accelerating DNN Training with Structured Data Gradient Pruning
• Full deep neural network training on a pruned weight budget