A Systematic DNN Weight Pruning Framework using Alternating Direction Method of Multipliers

Tianyun Zhang; Shaokai Ye; Kaiqi Zhang; Jian Tang; Wujie Wen; Makan Fardad; Yanzhi Wang

Source-linked AI summary

A Systematic DNN Weight Pruning Framework using Alternating Direction Method of Multipliers

Tianyun Zhang, Shaokai Ye, Kaiqi Zhang, Jian Tang, Wujie Wen, Makan Fardad, Yanzhi Wang

arXiv:1804.03294v3Apr 10, 2018cs.NEcs.CVcs.LG

TL;DR

Prior DNN pruning work is mainly heuristic and iterative, motivating a systematic approach with specified sparsity. The paper applies ADMM to decompose constrained nonconvex pruning into tractable subproblems and reports substantial compression without accuracy loss.

Problem
Prior DNN pruning methods are mainly heuristic and iterative, lacking guarantees on weight-reduction ratio and convergence time.
Method
The framework formulates pruning as constrained nonconvex optimization with combinatorial sparsity constraints and solves it using ADMM.
Results
The method consistently outperformed prior work, achieving 71.2× LeNet-5 weight reduction and 21× AlexNet weight reduction without accuracy loss.
Takeaways & Limitations
The framework provides systematic sparsity control and supports substantial weight and computation reduction in the evaluated DNNs.
Takeaways & Limitations
The framework requires tuning the parameters l_i, and ADMM has no general guarantee of convergence to an optimal point for nonconvex problems.

Abstract

from arXiv · show

Weight pruning methods for deep neural networks (DNNs) have been investigated recently, but prior work in this area is mainly heuristic, iterative pruning, thereby lacking guarantees on the weight reduction ratio and convergence time. To mitigate these limitations, we present a systematic weight pruning framework of DNNs using the alternating direction method of multipliers (ADMM). We first formulate the weight pruning problem of DNNs as a nonconvex optimization problem with combinatorial constraints specifying the sparsity requirements, and then adopt the ADMM framework for systematic weight pruning. By using ADMM, the original nonconvex optimization problem is decomposed into two subproblems that are solved iteratively. One of these subproblems can be solved using stochastic gradient descent, the other can be solved analytically. Besides, our method achieves a fast convergence rate. The weight pruning results are very promising and consistently outperform the prior work. On the LeNet-5 model for the MNIST data set, we achieve 71.2 times weight reduction without accuracy loss. On the AlexNet model for the ImageNet data set, we achieve 21 times weight reduction without accuracy loss. When we focus on the convolutional layer pruning for computation reductions, we can reduce the total computation by five times compared with the prior work (achieving a total of 13.4 times weight reduction in convolutional layers). Our models and codes are released at https://github.com/KaiqiZhang/admm-pruning

1 Introduction

DNNs impose substantial storage and computation burdens, motivating compression methods. The paper proposes ADMM-based systematic pruning and reports higher compression without accuracy loss.

Large-scale DNNs create significant model-size and computational burdens, especially for embedded and IoT systems.
Prior compression approaches include weight pruning, sparsity regularization, weight clustering, and low-rank approximation.
Existing pruning methods iteratively remove relatively unimportant weights and retrain the network to maintain accuracy.
The proposed framework formulates pruning as constrained nonconvex optimization and applies ADMM to solve it systematically.ADMM decomposes the problem into two iteratively solved subproblems, using stochastic gradient descent for one and an analytical solution for the other.
71.2× weight reduction was achieved on LeNet-5 and 21× on AlexNet without accuracy loss.Convolutional-layer pruning reduced total computation by five times compared with prior work, with 13.4× convolutional-layer weight reduction.

2 Related Work on Weight Reduction/Model Compression

Related work reduces DNN storage and computation by exploiting redundancy, but prior pruning approaches are largely heuristic and retraining-intensive. The paper positions its framework as systematic, faster-converging, and applicable to broader compression settings.

DNN weight-reduction methods exploit redundancy across filters and channels to reduce model storage.
Low-rank approximation methods such as SVD can be difficult to use without accuracy degradation in very large DNNs.
Prior pruning heuristically removes small-magnitude weights and retrains the DNN, achieving 9× AlexNet weight reduction with virtually zero accuracy degradation.
The proposed method is described as systematic, with higher compression, faster convergence, and generality for structured pruning and weight clustering.

3 Background of ADMM

ADMM uses variable splitting to turn difficult optimization problems into separate subproblems. Its decomposition supports efficient treatment of differentiable and structured terms.

ADMM is presented as a tool for solving regularized convex optimization and some nonconvex problems, including problems with combinatorial constraints.For nonconvex problems, it may converge to a solution that is not globally optimal but is sufficiently good for many applications.
Variable splitting rewrites an optimization problem with separate variables linked by the constraint x = z.
The augmented Lagrangian decomposes the split problem into two subproblems that can be solved separately and efficiently.
When the differentiable term and quadratic penalty are differentiable, the first subproblem can be solved by gradient descent.
A quadratic subproblem involving a structured term may be solved analytically by exploiting that structure.

4 Problem Formulation and Proposed Framework

The framework formulates weight pruning as a cardinality-constrained nonconvex optimization problem and applies ADMM to split it into iteratively solved subproblems. Stochastic gradient descent handles the differentiable subproblem, while analytical projection enforces sparsity, followed by retraining and rapid convergence.

4.1 Problem Formulation of Weight Pruning: Each DNN layer’s weights and biases are defined explicitly, with convolutional weights represented as four-dimensional tensors and fully connected weights as two-dimensional matrices.
4.1 Problem Formulation of Weight Pruning: The training objective minimizes DNN loss while constraining each layer’s weight cardinality to a desired limit l_i.card(·) counts nonzero elements, and l_i specifies the desired number of weights in layer i.
4.1 Problem Formulation of Weight Pruning: Hard cardinality constraints directly target sparsity and can produce a higher degree of sparsity than regularization in the objective.
4.2 Systematic Weight Pruning Framework using ADMM: The framework can constrain the total number of weights across the whole DNN, not only the number in each layer.
4.2 Systematic Weight Pruning Framework using ADMM: ADMM reformulates the constrained problem using variable splitting, indicator functions, equality constraints, an augmented Lagrangian, and dual variables.The constraint sets are nonconvex, with S_i = {W_i | card(W_i) ≤ l_i}.
4.2 Systematic Weight Pruning Framework using ADMM: The first ADMM subproblem combines DNN loss with a differentiable quadratic regularizer and is solved using stochastic gradient descent.
4.2 Systematic Weight Pruning Framework using ADMM: The second subproblem projects onto the cardinality-constrained set by retaining the l_i largest-magnitude weights and setting the rest to zero.Although the set is nonconvex, its special structure makes the Euclidean projection analytically computable.
4.2 Systematic Weight Pruning Framework using ADMM: With warm starts, AlexNet convergence takes approximately one-tenth of the iterations required for original DNN training, while ADMM subproblems add negligible computational time.

5 Experimental Results

Experiments on MNIST and ImageNet evaluate ADMM-based pruning for weight and computation reduction without accuracy loss. The framework consistently exceeds the prior iterative pruning heuristic in compression and convolutional-layer computation reduction.

Experimental setup: Experiments evaluate LeNet-300-100, LeNet-5, and AlexNet on MNIST and ImageNet ILSVRC-2012 for comparison with prior heuristic pruning.LeNet models use MNIST, while AlexNet uses ImageNet ILSVRC-2012.
AlexNet computation reduction: The framework reduces total computation by five times compared with the prior work when pruning focuses on convolutional layers.Its main gains come from pruning conv2 through conv5, while the first convolutional layer is difficult to prune.
LeNet results: 22.9× weight reduction on LeNet-300-100 exceeds the prior method’s 12× parameter reduction without accuracy loss.The reported reduction is 90% higher than the prior work.
LeNet results: 71.2× weight reduction on LeNet-5 is achieved without accuracy loss, compared with 12× parameter reduction from the prior heuristic.LeNet-5 achieves 99.2% test accuracy on MNIST.
AlexNet weight reduction: 21× weight reduction on AlexNet is achieved without accuracy loss, with 80.2% top-5 validation accuracy, more than twice the prior heuristic’s reduction.The AlexNet experiment focuses purely on model-size reduction.
AlexNet computation reduction: 13.4× weight reduction in AlexNet convolutional layers is achieved without accuracy loss, compared with 2.7× in the prior work.The convolutional-layer experiment targets computation reduction because AlexNet computation is dominated by its five convolutional layers.

6 Discussion

The framework addresses nonconvex weight pruning through ADMM, but convergence depends on initialization, penalty parameters, and stopping choices. Experiments show that suitable settings produce effective pruning and retraining restores the loss level.

ADMM has no general guarantee of convergence to an optimal point for nonconvex problems and may reach different points under different initializations and penalty parameters.
Parameters and initialization of ADMM: The method initializes ADMM from a pretrained model, retains the largest-magnitude weights in each layer for Z0_i, and sets U0_i to zero.
Parameters and initialization of ADMM: Penalty parameters require balancing accuracy and sparsity: values that are too small may slow or prevent convergence, while values that are too large may degrade DNN accuracy.
Parameters and initialization of ADMM: The pruning levels l_i are tuned using existing results, proportional reduction, and binary search to find the smallest values without accuracy loss.
Discussion of the proposed framework: After ADMM convergence, near-zero weights are removed and retraining decreases the increased loss back to its pre-pruning level.
Discussion of the proposed framework: ADMM bypasses the cardinality function's nondifferentiability by decomposing the problem, while dynamically updating the regularization target during iterations.

7 Conclusions and Future Work

The paper presents an ADMM-based systematic framework for DNN weight pruning and reports substantial compression without accuracy loss. It also identifies extensions toward structured pruning and a unified approach combining weight pruning, activation reduction, and weight clustering.

The framework formulates DNN weight pruning as a constrained nonconvex optimization problem and solves it with ADMM-based iterative subproblems.
22.9× weight reduction on LeNet-300-100 and 71.2× on LeNet-5 were achieved without accuracy loss.
21× weight reduction was achieved on AlexNet without accuracy loss, while convolutional-layer pruning reached 13.4× weight reduction and five-times computation reduction compared with prior work.
Future Work: Future work will incorporate structure and regularity into pruning and unify weight pruning, activation reduction, and weight clustering.