[Paper Review] Only Train Once: A One-Shot Neural Network Training And Pruning Framework

Tianyi Chen, Bo Ji+8 more|arXiv (Cornell University)|Jul 15, 2021

Advanced Neural Network ApplicationsComputer Science81 references43 citations

TL;DR

OTO trains and compresses full neural networks from scratch in one pass, using zero-invariant groups and a novel Half-Space Stochastic Projected Gradient to prune without fine-tuning.

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.

Motivation & Objective

• Motivate the need for efficient model compression without multi-stage fine-tuning in pruning workflows.
• Introduce a framework that can train and prune full models simultaneously from scratch.
• Define zero-invariant groups (ZIGs) to enable safe, output-preserving pruning.
• Develop a novel optimization algorithm, Half-Space Stochastic Projected Gradient (HSPG), to promote structured sparsity.
• Demonstrate state-of-the-art or competitive compression and accuracy on standard benchmarks across CNNs and NLP models.

Proposed method

• Partition model parameters into zero-invariant groups (ZIGs) so zero groups can be pruned without affecting output.
• Formulate a structured-sparsity regularization with a mixed l1/l2 norm over ZIGs to induce group sparsity.
• Solve the non-smooth regularized problem using Half-Space Stochastic Projected Gradient (HSPG), which has an Initialization Stage (SGD) and a Group-Sparsity Stage (Half-Space Projection).
• Introduce a Half-Space Projection operator that zeros entire groups when their descent condition is satisfied, enabling aggressive sparsity.
• Prune the full model by removing zero ZIGs to obtain a slimmer architecture without additional fine-tuning.
• Provide automatic or prescribed ZIG partition guidance for common architectures (Conv-BN, Residual blocks, Fully Connected and Multi-Head Attention).

Experimental results

Research questions

• RQ1 Can a single training run from scratch yield a compressed, high-performance model without fine-tuning?
• RQ2 How can parameters be partitioned into groups (ZIGs) such that zero-valued groups can be pruned without changing outputs?
• RQ3 Does a structured-sparsity approach with HSPG promote zero groups more effectively than standard proximal methods?
• RQ4 What are the empirical gains in FLOPs, parameter counts, and accuracy across standard benchmarks when applying OTO?
• RQ5 Is the approach architecture-agnostic and scalable to CNNs and NLP models like BERT?

Key findings

• OTO achieves significant FLOPs and parameter reductions while maintaining competitive accuracy across benchmarks.
• VGG16 and VGG16-BN on CIFAR-10 show substantial FLOPs and parameter reductions with maintained Top-1 accuracy.
• ResNet-50 on CIFAR-10 and ImageNet achieve notable FLOPs reductions and parameter reductions with minimal accuracy loss.
• BERT on SQuAD demonstrates competitive results under the one-shot pruning framework.
• HSPG outperforms standard proximal methods in promoting group sparsity and maintains convergence properties.
• The pruning process yields a compressed model that computes identical outputs to the full model for the same inputs without fine-tuning.