operty Importance: The filters are pruned based onrinsic properties of CNNs. These pruning approaches
modify the network training loss. After the pruning,odel performance is enhanced through fine-tuning.
g these methods, Hu et al. [14] utilized the sparoutputs in a large network to remove filters withpercentage of zero activations. The !1 -norm basedg [18] assumes that parameters or features with small
are less informative, and thus should be pruned first.
anov et al. [27] considered the first-order gradientuate the importance of filters and removed the leastant ones. In [34], the importance scores of the final
ses are propagated to every filter in the network andNNs are pruned by removing the one with the least imce. He et al. [12] calculated the geometric median iners, and the filters closest to this are pruned. Most deor filter evaluation functions are ad-hoc, which brings
vantage of low time complexity, but also limits the acion and compression ratios.

aptive Importance: Different from the property imce based approaches, another direction embeds theg requirement into the network training loss, and ema joint-retraining optimization to generate an adaptive
g decision. Liu et al. [24] and Zhao et al. [36] ima sparsity constraint on the scaling factor of the batch
lization layer, allowing channels with lower scaling
to be identified as unimportant. Huang et al. [16]in et al. [23] introduced a new scaling factor par (also known as a mask) to learn a sparse strucruning where filters corresponding to a scaling faczero are removed. Compared with property imporbased filter pruning, adaptive importance based methually yield better compression and acceleration reue to their joint optimization. However, because thechanged, the required retraining step is heavy in bothne time and human labor, usually demanding anotherof hyper-parameter tuning. For some methods, e.g.,
ask-based scheme, the modified loss even requires

Differentiable Transportation Pruning

ETH Zurich…

Feature Map
Generation

Input

…

…

Filter 5!

Filter 4

Target

0

rank

1⁄5

SmoothTransportation
Plan

Gradients…

Fine-Tuning/Inference Phase…

…

Filter 1

1

Soft mask

Filter 3

Filter 2

Source3⁄5Filter Selection

rank

Deep learning algorithms are increasingly employed at
the edge. However, edge devices are resource constrained
and thus require efficient deployment of deep neural networks. Pruning methods are a key tool for edge deployment
as they can improve storage, compute, memory bandwidth,and energy usage. In this paper we propose a novel accurate pruning technique that allows precise control over
the output network size. Our method uses an efficient optimal transportation scheme which we make end-to-end differentiable and which automatically tunes the explorationexploitation behavior of the algorithm to find accurate
sparse sub-networks. We show that our method achieves
state-of-the-art performance compared to previous pruning
methods on 3 different datasets, using 5 different models,across a wide range of pruning ratios, and with two typesof sparsity budgets and pruning granularities.

3

Soft Mask Learning

TU Delft

Figure 1. Framework of HRank. In the left column, we first use
images to run through the convolutional layers to get the featuremaps. In the middle column, we then estimate the rank of eachfeature map, which is used as the criteria for pruning. The rightcolumn shows the pruning (the red filters), and fine-tuning where
the green filters are updated and the blue filters are frozen.

…

Figure 1. During training we multiply the filters in a layer with
soft masks mϵ . The soft masks are obtained by solving an optimal transport problem: minimizing a transportation cost between auniform source distribution over trainable importance scores s anda Bernoulli target distribution defined by sparsity ratio i.e. 3/5.
During training the soft masks gradually converge to hard masks.
…

the pruning complexity. Moreover, compared with existing
property importance based methods, HRank also leads tosignificant improvements in acceleration and compression[18, 34] , achieving even better results than state-of-the-art
adaptive importance based methods [13, 24, 26, 16, 23, 36].
The motivation behind enforcing the high rank of featuremaps comes from an empirical and quantitative observation. As shown in Fig. 2, we have found that the averagerank of feature maps generated by a single filter is always
the same, regardless of how much data the CNN has seen.
It suggests that the ranks of feature maps in deep CNNs canbe accurately estimated using only a small portion of the input images, and thus can be highly efficient. Based on thisidea, we mathematically prove that feature maps with lower
ranks contribute less to accuracy. Thus the filters generatingthese feature maps can be removed first.
We conduct extensive experiments on two benchmarks,
CIFAR-10 [17] and ImageNet [31], using many representative large CNN models, including VGGNet [32],GoogLeNet [33], ResNet [11] and DenseNet [15]. The results demonstrate the superior performance of HRank overexisting filter pruning methods, including property importance based approaches[18, 34] and state-of-the-art adaptive
importance based models [13, 24, 26, 16, 23, 36]. We also

Abstract

2…

Axelera AI…

1

Jan C. van Gemert2
Torsten Hoefler3
Evangelos Eleftheriou1 Bram-Ernst Verhoef1

Conv Layer

Yunqiang Li1
Bert Moons1

In this paper, as illustrated in Fig. 1, we propose afully differentiable transportation pruning method that allows precise control of the output network size. We draw
inspiration from Xie et al. (2020) [65] who formulate a soft
top-k operator based on entropic optimal transport. Recent
work [59] shows good pruning results with the soft top-k
operator. This soft top-k operator, however, is computationally expensive [10] and relies on implicit, and therefore lessaccurate, gradients [11] for back-propagation. Instead, we
speed up the soft top-k operator with an approximate bilevel optimization which is optimized using accurate automatic differentiation [48]. Moreover, in contrast to [59] we
can automatically tune the temperature hyper-parameter.
We make the following contributions: i) A fully differentiable pruning method that allows precise control over
output network size; ii) An efficient entropic optimal transportation algorithm that significantly reduces the computational complexity of bi-level optimization; iii) We increase
the ease-of-use of the algorithm by means of an automatic
temperature annealing mechanism instead of a manuallychosen decay schedule; iv) We show state-of-the-art results
on 3 datasets with 5 models for structured and unstructured
pruning using layer-wise and global sparsity budgets.

1. Introduction
As the world is getting smaller, its edge is getting larger:
more compute-limited edge devices are used. To unlock
deep learning on the edge, deep networks need to be efficient and compact. There is a demand for accurate networks
that fit exactly in pre-defined memory constraints.
Pruning is an effective technique [20, 41, 67] to find asparse, compact model from a dense, over-parameterized
deep network by eliminating redundant elements such as filters or weights while retaining accuracy. Pruning is a hard
problem because finding a good-performing sub-network
demands selecting which part of a network to keep —and
which part to prune—, requiring an expensive search over alarge discrete set of candidate architectures.
Recent work [23, 52, 40, 25] finds that optimizing acontinuously differentiable mask to sparsify network connection while simultaneously learning network parameters
facilitates search space exploration and enables gradientbased optimization. Prior methods [23, 40], however, control the sparsity, i.e. pruned network size, by an additional
penalty term in the loss function that is difficult to optimize
and tune, and hampering a precise control over the sparsity.
12. Related Work
Pruning in deep learning. Network pruning makes models
smaller by removing elements that do not contribute significantly to accuracy. Pruning methods can be subdivided into
methods that promote structured (i.e. blocked) [42, 27] or
unstructured (i.e. fine-grained) sparsity [6, 13, 43, 61, 68],
see [22] for a review. Filter pruning is a form of structured
sparsity by removing entire filters from the network’s layers [34, 30]. Filter pruning often achieves practical network
compression and significant acceleration as entire featuremaps are no longer computed. For these reasons, we here
focus on filter pruning.
Selection criteria for pruning. Network pruning can bephrased as a form of neural architecture search [38]. Suchan architecture search may involve training and evaluating
several random subnetworks based on leave-some-out approaches [5, 58, 32], but this can be quite expensive for large
models. More efficient approaches assign an importancescore as selection criteria. This score can be based on L1/L2
norm of weights [67], geometric median [20] of filters,
Kullback–Leibler divergence [41], Wasserstein barycenterof channel probability vectors [55], prediction error [45],or empirical sensitivity of a feature map [33]. These methods first prune a pretrained model based on the importancescores and subsequently fine-tune the pruned model so themodel can adapt to the reduced set of parameters. Oncepruned, however, the model is not allowed to explore othersubnetworks. Here, we also use importance scores, which
we learn as latent variables from which soft masks are computed through optimal transport. The soft mask is automatically annealed to a hard mask during training, allowing themodel to explore different subnetworks in the process.
Differentiable continuous relaxation. Recent work aims
to select the best sub-architecture by training a binary
mask on the network weights [15]. This binary mask,
however, is non-differentiable and applying a hard binary
mask during the forward pass may impair network performance [23]. Other methods approximate binary masks
using a sigmoid function, which gradually converges to abinary step function by means of a decaying temperature
hyper-parameter [23, 25, 40, 52]. To control the sparsity ratio, these methods usually introduce a sparsity penalty termin the loss function, which requires an additional difficultto-tune hyper-parameter if an exact pruning ration is required. Our method mitigates these issues as it allows foran exact pruned network size using optimal transport anddecays the temperature variable automatically.
Learning soft mask via entropic optimal transport. Our
method draws inspiration from the work of Xie et al.
(2020) [65], who formulate a soft top-k operator based on
entropic optimal transport [51], optimized via Sinkhorn’salgorithm [8]. The authors apply their method to predictivemodeling applications such as information retrieval. As re-marked by [10], this method can be computationally expensive if applied to pruning, as hundreds of Sinkhorn iterations are needed in the forward pass of each update step. In
this paper we phrase network pruning as an efficient optimaltransport problem. Compared to [65] our method significantly reduces the number of Sinkhorn iterations requiredin the forward pass. Furthermore, in contrast to [65], who
employ implicit gradients [11], our method relies on accurate automatic differentiation [48]. Finally, we replace themanually-chosen decay schedule of the entropic hyperparameter in [65] with an automatic schedule, thereby simplifying the optimization process.
Bregman divergence based proximal method. Previous
work [26, 53] has applied to optimal transport the proximal
point algorithm [47] based on the Bregman divergence [4]to alleviate the potential numerical instability of Sinkhorn
iterations. Xie et al. [66] use the proximal algorithm to optimize optimal transport in deep generative models [2, 16]
and achieve good performance, but still performs a large
amount of proximal point iterations per mini-batch step,
making it computational inefficient for our pruning purposes. In addition, its back-propagation pass [66] relieson the envelope theorem [1] that does not go into proximal
point iterations to accelerate training, which causes inaccurate gradients. Instead, we only use a single proximal point
iteration, i.e. a single Sinkhorn iteration, per mini-batch stepand compute gradients via automatic differentiation [48].
Budget aware pruning. Budgeted pruning focuses oncompressing the network subject to the prescribed explicit
resource constraints, such as targeted amount of layer-wisefilters [28], global filters [60], FLOPs [29, 17], or execution
latency [54, 24]. Our method can directly control the pruned
network size given a sparsity budget. Moreover, we show
how it can be easily extended to general budgeted pruning
by learning budget-constrained sparsity across layers [46].
Our technique improves the accuracy while using less hyperparameters, e.g. ChipNet [60] adds 6 extra hyperparameters while we only introduce a single temperature hyperparameter, which eases hyperparameter tuning effort.

3. Differentiable Transportation Pruning
Given a neural network f (x; w) with input x and a set ofn filters forming the columns of weight matrix w ∈ Rm×n ,filter pruning removes filters by applying a binary maskm ∈ {0, 1}n to the n filters (1 = kept, 0 = pruned).
If we need to keep exactly k out of n filters, then thesparse learning problem can be formulated as:arg min Ltrain f (x; w, m) ,
m,w
(1)
Xn
s.t.
mi = k, m ∈ {0, 1}n ,i=1where Ltrain is the training loss, e.g. cross entropy loss.

m3

m3

m3

m3

m3

(0, 0, 1)
(1,
0,
0)

0.5

s3
100

101
102
iteration

103(f) importance score

0

m1
100

101
102
iteration(g) soft mask

By hard m

0.5

m2
103

0

(e) ℓ = 103

m3

m3

s2
s1

100

1

(d) ℓ = 102

2

m2 m1
100

101
102
iteration(h) hard mask

m1

1

102

10 2

(c) ℓ = 1012mm1(b) ℓ = 1002mm1(a) no entropic relaxation

2
mm1

2
mm1

m2

(0, 1, 0)

103

1

By soft m

100

101
102
iteration

103(i) training loss

Figure 2. Illustrative example: The figure is best understood after Eq. (8) and it illustrates how
P our method can prune 2 weights outof 3 by learning a binary mask m ∈ {0, 1}3 . We can visualize in 2D because the simplex 3 can be shown as a triangle in twoTdimensions. In this example we chose to favor the 2nd weight, and thus define some loss
P on the weights as Ltrain = w m where we
define w = [2, 1, 3]. To keep just a single weight, we set k to 1, yielding the constraint 3i=1 mi = 1 from Eq. (1), where the optimal
solution is m = (0, 1, 0). The top row (a)-(e) illustrates the solution space, where the corners of the triangle are the hard masks: (1, 0, 0),(0, 1, 0) and (0, 0, 1). Specifically, (a), (b) denote the sparse and non-sparse solution space without or with entropy regularization at ℓ = 1.
The color bar indicates low (red) and high (violet) loss values. A green dot in the center is the initial position; and the black dot is thecurrent position. the bottom row (f)-(i): training curve of importance score, soft mask, discrete mask and training loss for one thousand
training steps. The hard mask in (h) is obtained by Eq. (6). (a): the discrete optimal transportation loss at ℓ = 1 just after initializing
the importance scores s. Note the ordering of s1 , s2 , s3 in (f) at ℓ = 1, which explains the low loss value for m3 . (b)-(e): Snapshots at
iterations ℓ = 100 , 101 , 102 , 103 after making the discrete optimal transport problem differentiable by entropy regularization in Eq. (7)where ε = 10. Note how the black dot moves to the correct solution. Note the training loss Ltrain in (i), as it pushes s2 towards top-1 in (f),and minimizing transportation loss pushes the soft mask m2 towards 1 in (g). The minimum value of the training loss is 1, due to how wechose the dot product loss Ltrain .

The binary mask m is discrete and therefore difficult to
optimize with gradient descent. To overcome this, we makethe binary mask depend on a continuous latent importancescore s which can be optimized by gradient descend. By
sorting these importance scores we allow exact control overthe pruned output network size by keeping only the topk importance scores. Specifically, the top-k operator firstsorts the importance scores s, and then assigns the top kimportance scores to 1 and the remaining importance scores
to 0, that is:(1, if si is in the top-k of s,mi =(2)0, otherwise.
The top-k operator can be parameterized in terms of the solution of an optimal transport problem [65]. This allows
us to parameterize the pruning mask in terms of optimaltransport, which is the problem of efficiently moving probability mass from a source distribution to a target distribution. In our case, we know the target Bernoulli distribution
over the binary values, which is given by the fraction of thek filters out of the n total filters to keep: P (1) = nk andP (0) = n−k
n . Then, the optimization problem reduces tominimizing the transportation costs of moving probability
mass from the trainable importance scores s to this target
distribution. See Fig. 1 for an illustration.
The two probability distributions used in optimal transport are over the importance scores s and over the two binary values. WeP
define them by considering
P2 two discrete
n
a
δ
and
b
=
measures a =i=1 i si
j=1 bj δqj supported on {si }ni=1 and {qj : {0, 1}} respectively whereδx is the Dirac at location x, intuitively a unit of mass
which is infinitely concentrated at location x, the ai andbj are the probability mass. Following [65], we define aas an empirical discrete uniform distribution (source) withai = 1/n. Our binary mask variable is distributed asb = [1 − (k/n), k/n]T (target).
These two distributions now allow optimal transportation optimization where probability mass is moved according to the cost it takes to move one unit of probability mass.
We define the cost matrix C ∈ Rn×2 which aligns n filters (source) to the 2 binary options (target). The elementsof the cost matrix are squared Euclidean distances and are
denoted as:
Ci1 = s2i , Ci2 = (si − 1)2 , i = 1, 2, ..., n(3)

where the targets are binary (0 or 1) and thus the target value
for Ci1 is 0, and the target value for Ci2 is 1.
With the source distributions, target distribution and cost
matrix defined, the optimal transportation plan can be formulated as:P∗ = arg min ⟨C, P⟩ ,(4)This objective is an ε-strongly convex function which has aunique solution. As ε → 0 the unique solution Pε of (7)converges to the original optimal sparse solution of (4), i.e.
with hard masks.
Our soft masks mϵ are now obtained by inserting P∗ε into
(6), as:mϵ = nP∗ε · [0, 1]T .
(8)Interlude: illustrative example. We have now described
the main ingredients of our method. This allows us to givean illustrative example of the main setting before we describe how to solve the optimization problem. In Fig. 2
we show a 2D visualization for pruning n = 3 importancescores to just a single score.
Dual formulation for optimization. To facilitate the optimization problem in (7), we make use of Lagrangean duality. That is, by introducing two dual variables f ∈ Rn andg ∈ R2 for each marginal constraint of the transport plansin U(a, b), the Lagrangian of (7) is optimized, resulting inan optimal Pε to the regularized problem given by:P∈U (a,b)

where U(a, b) = {P ∈ Rn×2 : P12 = a, PT 1n = b} isa set of coupling measures satisfying marginal constraints
P12 = a and PT 1n = b, 1n and 12 denote vectors with nones and 2 ones, and P ∈ Rn×2 denotes the general probabilistic coupling matrix, i.e. the transportation plan, and⟨·, ·⟩ is the Frobenius dot-product.
The optimal transportation plan P∗ can be computed as:((
1/n, if i ≤ k0,
if i ≤ k∗∗
Pσi ,1 =, Pσi ,2 =0,otherwise
1/n, otherwise(5)with σ being the sorting permutation, i.e. sσ1 < sσ2 < ··· <sσn . Given P∗ , the binary mask can be parameterized as afunction of the optimal transport plan:m = nP∗ · [0, 1]T .

(6)

Differentiable transportation relaxation. So far, the derived hard mask in Eq. (6) is not differentiable everywherewith respect to importance scores, which can be seen fromthe sorting permutation applied in Eq. (5) and also in thelack of a gradient when using masking with a hard zero.
To make the method differentiable, we derive a soft maskby smoothing the optimal transportation plan with entropic
regularization. By adding an entropy regularization term we
extend the solution space by allowing non-sparse solutionsof the transportation plan. That is, we now allow soft masks.
Define thePdiscrete entropy of a coupling matrix as
H(P) = − ij Pij (log(Pij ) − 1), an entropic regularization is used to obtain smooth solutions to the originaltransport problem in Eq. (4):P∗ε = arg min ⟨C, P⟩ − εH(P).
P∈U (a,b)(7)

Pε = ef /ε ⊙ e−C(s)/ε ⊙ eg/ε .

(9)where ⊙ is element wise product (see details in A.1). Notethat the Pε in Eq. (9) is differentiable with respect to s.
Bi-level optimization. Now that we have relaxed the problem, we aim to jointly learn the dual variables f , g and themodel variables s, w. The dual variables are optimized byminimizing the Sinkhorn divergence as in Eq. (7). Thisis equivalent to maximizing its dual problem, denoted as
Lεdual (f , g, s) = ⟨f , a⟩ + ⟨g, b⟩ − ε ef /ε , e−C(s)/ε · eg/ε(See details in A.2). Finally the model variables are obtained by minimizing the training loss Ltrain .
This implies a bi-level optimization problem [7], i.e. an
optimization problem which contains another optimization
problem as a constraint. The model variables s, w appear
as the upper-level variables, and the dual variables f , g as
the lower-level variables. We rewrite the original problem
in Eq. (1) as the following bi-level optimization problem:min Ltrain f x; w, m(P∗ε (s)) ,
(10)
s,w

s.t. f ∗ , g∗ = arg maxf ,g Lεdual (f , g, s),(11)where P∗ε (s) is obtained by inserting the outcome from (11)into (9), and the mask mε in (10), dubbed soft mask, iscomputed by inserting P∗ε (s) into Eq. (8).
Nonetheless, this bi-level optimization problem is expensive to train with Sinkhorn’s algorithm: Firstly, per minibatch the inner optimization in Eq. (11) needs to perform
Sinkhorn’s iterative algorithm for hundreds of iterations toconverge, which is computationally inefficient. Secondly,
during the back-propagation pass, the gradient of P∗ε (s)with respect to the importance scores s is computed by differentiating [48] through all Sinkhorn iterations, which isAlgorithm 1: Differentiable Transportation Pruning
Probabilities a = 1n /n, b = [p, 1 − p]T on supports {si }ni=1 , {mj }2j=1 , weights: w, training loss: Ltrain , training
iterations: L, learning rate: α, initialization: ε, g(1) ← 12 , P(1) ← n1 1n 1T
2for ℓ = 1, 2, ..., L doC(ℓ)
Cost matrix C with element Cij = (si − mj )2 ; Gibbs kernel K = e− ε ⊙ P(ℓ)
(ℓ)
(ℓ+1)
/ε
Update dual variables: f (ℓ+1) = ε log a − ε log (Keg /ε ); g(ℓ+1) = ε log b − ε log (KT ef
)
(ℓ+1)
(ℓ+1)
/ε
/ε
Update transportation plan P(ℓ+1) = ef
⊙ K ⊙ eg
Update model variables with SGD: s(ℓ+1) = s(ℓ) − α∇s Ltrain ; w(ℓ+1) = w(ℓ) − α∇w Ltrain
Derive pruned architecture based on soft mask m = nP(L) · [0, 1]T for finetuning.
expensive (see details in A.3). Finally, we have also introduced an additional hyperparameter ε to decay, whichmakes the problem even harder to optimize.
Approximate bi-level optimization. We note that inour bi-level optimization the target distribution is a fixedBernoulli distribution with parameter defined by the sparsity. The source distribution changes only when the trainable importance score updates via SGD. This differs fromprevious approaches [65, 66] involving a stochastic random sampling of the target or source distribution and where
the inner optimization in Eq. (11) is expected to converge
per mini-batch step. We therefore propose to approximateP∗ε (s) in Eq. (11) using only a single Sinkhorn step, i.e.,
without optimizing the inner optimization problem untilconvergence per mini-batch step.
Given the ℓ-th mini-batch update step, we approximate
the training loss in Eq. (10) with:
Ltrain f x; w, m(P∗ε (s))

(12)

≈Ltrain f x; w, m(P(ℓ+1)(s)) ,ε(ℓ+1)(s) is obtained by applying a single Sinkhorn
where Pε
update step. Similar techniques for solving bi-level optimization have been used in architecture search [36], hyperparameter tuning [39] and meta-learning [12].
To enhance the convergence guarantee using just a single Sinkhorn inner iteration, we use generalized proximal
point iteration based on Bregman divergence [53, 66] (seedetails in A.4), which can be viewed as applying iteratively
Cthe Sinkhorn algorithm [51] with a e− ε/ℓ kernel (see details in A.5). Thus, a proximal point iteration keeps track of
the previous Sinkhorn update like Momentum [57] that incorporating a moving average of previous gradients, which
accelerates the convergence and improves optimization sta(ℓ+1)
bility. We therefore compute Pε(s) in Eq. (12) as,P(ℓ+1)(s) = eεf (ℓ+1)εC(s)g(ℓ+1)⊙ e− ε ⊙ P(ℓ) ⊙ e ε , (13)using the transportation plan P(ℓ) from the previous step.
Let K = e−C(s)/ε ⊙ P(ℓ) denote the Gibbs kernel in the

Sinkhorn algorithm, and the dual variables can be computed
using a single Sinkhorn iteration, i.e. a proximal point iteration with block coordinate ascent as (see details in A.6):
f (ℓ+1) = ε log a − ε log (Keg
g

(ℓ+1)

(ℓ)

= ε log b − ε log (K e

T f

/ε

(14)

);

(ℓ+1)

/ε

).

(15)

A side advantage of the updates in Eq. (13) is that it can iteratively decay the regularization parameter through ε/ℓ (seedetails in A.5). This means we can perform an automatic
decaying on the regularization temperature, and as ℓ → ∞the solution gradually approaches the optimal sparse transportation plan of the original objective (1).
The gradient of Eq. (13) with respect to the importancescores s can be obtaind using automatic differentiation. We
can therefore use normal SGD to jointly train the model
weights w and importance scores s in Eq. (12).
Our method results in fast training and converges well
as shown by our experiments in Section 4. After training,
we derive the sparse architecture based on the soft masks
m = nP(L) · [0, 1]T and finetune the model. We outline theiterative training procedure in Algorithm 1.

4. Experiments
We validate our proposed method mainly on filter pruning using different models and datasets, and then extend to
unstructured pruning. Filter pruning is a form of structured
pruning that removes entire filters from the layers of a network, thereby compressing the network, and potentially accelerating the network by reducing the number of FLOPs
needed to execute it.
Datasets and networks. We first conduct our experiments
on CIFAR-10 with ResNet-56 and on CIFAR-100 with
VGG-19. Following [63], for the CIFAR datasets we train
our baseline models with accuracy comparable to those in
the original papers. We then evaluate our method on thelarge-scale ImageNet dataset with ResNet-34 and ResNet50. Following [63], we use the PyTorch [49] pretrainedmodels as our baseline models. We also apply our method
to prune lightweight architecture MobileNetV2.

1.0

filter1
filter2
filter3

0.50

values@layer 24

soft mask

1.00

score

0.00
filter1
filter2
filter3

0.30
0.15

0.6
0.4
0.2
0.0

L1 norm

filter1
filter2
filter3

0.03
0.00
101

Accuracy(%)

0.00
0.06

102

104

103

iteration

105

soft mask
L1 norm

0.8

0

50

100

75
70
65
60
55
50
45
0

25

50

150
filter index

75

200

100
Epoch

125

250

150

175

(a) On CIFAR-10(b) On ImageNet
Figure 3. Convergence analysis experiments. (a). Evolution of the soft masks associated with three filters from ResNet-56 trainedon CIFAR-10 with ε = 1. Early in training our method explores different filter compositions. In later training phases the soft masks
have converged to hard masks. For each mask, the soft mask, importance score and L1 norm of the respective filter are correlated. (b).
Experiments with ResNet-50 on ImageNet with ε = 0.25. After training the soft mask has converged to a discrete hard mask solution.
The L1 norms of the filters are highly correlated with the soft masks. Bottom: Validation accuracy during training with differential
transportation pruning and during finetuning. The star point shows the test accuracy for the derived pruned architecture by our soft mask
before finetuning.
Table 1. Various pruning ratios with ResNet-56/VGG-19 on CIFAR-10/CIFAR-100: Accuracy comparison with GReg-1 [63] and L1 +one-shot [30]. Baseline accuracy for ResNet-56
is 93.36% and for VGG-19 is 74.02%. Each experiment is repeated 3 times with random
initialization, mean and std accuracies are reported.
Pruning ratio (%)
ResNet-56(CIFAR-10)

Sparsity (%) / Speedup

VGG-19

70

90

92.5

95

49.82/1.99×70.57/3.59×90.39/11.41×93.43/14.76×95.19/19.31×

L1 +one-shot (Acc.%) [30] 92.97±0.15 91.88±0.09 87.34±0.21 87.31±0.28 82.79±0.22
GReg-1 (Acc. %) [63]93.06±0.09 92.23±0.21 89.49±0.23 88.39±0.15 85.97±0.16

Ours (Acc. %)93.46±0.18 92.46±0.10 89.84±0.14 88.03±0.33 86.18±0.06

Pruning ratio (%)(CIFAR-100)

50

Sparsity (%) / Speedup

50

60

70

80

90

74.87/3.60×84.00/5.41× 90.98/8.84× 95.95/17.30×98.96/44.22×

L1 +one-shot (Acc.%) [30] 71.49±0.14 70.27±0.12 66.05±0.04 61.59±0.03 51.36±0.11

Table 2. Unstructured pruning on ImageNet for ResNet-50.
MethodBaseline Pruned Droped Sparsity
Acc.(%) Acc.(%) Acc.(%) (%)

GSM [9]Sparse VD [44]DPF [35]
WoodFisher [56]GReg-1 [63]GReg-2 [63]
Ours

75.72
76.69
75.95
75.98
76.13
76.13
76.13

1.42
1.41
1.40
0.78
0.68
0.86
0.63

80.00
80.00
82.60
82.70
82.70
82.70
82.70

Table 3. MobileNetV2 on CIFAR-10/100
with global sparsity ratio over network.
MethodSparsity (%)

Unpruned

GReg-1 (Acc.%) [63]71.50±0.12 70.33±0.12 67.35±0.15 63.55±0.29 57.09±0.03

ChipNet [60]Ours

Ours (Acc. %)71.56±0.19 70.41±0.06 67.74±0.13 63.98±0.25 57.21±0.09

ChipNet [60]Ours

Training settings. The importance scores corresponding toeach layer’s filters are initialized to the L2-norm of the respective pretrained filter weights. For fair comparisons, we
adopt the same pruning rate schedules as [63]. Also following [63], for each pruning schedule we denote the model
speedup as the pruning-induced reduction in the numberof floating point operations (FLOPs) relative to the originalmodel. The whole training procedure consists of two stages:

74.30
75.28
74.55
75.20
75.45
75.27
75.50

Acc.
CIFAR-10 CIFAR-100

0

93.70

72.80

60

91.83
91.93

66.61
70.45

80

90.41
91.49

52.96
68.90

optimal transportation pruning and post-pruning finetuning.
Our main focus is on how to train the soft masks infirst stage. An important hyperparameter in this respectis the regularization constant ε. For ResNet-56 and MobileNetV2, we use ε = 1, for all other networks we set
ε = 0.25 (see A.7). The model variables and importancescores are trained with standard Stochastic Gradient Descent (SGD) with a momentum of 0.9. For CIFAR datasets,

4.1. Effect of training soft masks
How do soft masks converge? We show how our soft
masks converge to discrete masks during training in thissection. In Fig. 3(a), we trained ResNet-56 on CIFAR-10
with a regularization hyperparameter ε = 1 and a pruning
rate of 0.5, and visualize the evolution during training of thesoft mask, importance score and L1 norm values for threefilters. Early in training, all filters have similar L1 norms
and soft masks. Subsequently, however, the algorithm starts
to explore different filter compositions: The soft mask of filter1 first converges to one, then fluctuates between zero andone, after which it moves back towards one, its final value.
The soft masks corresponding to filter2 and filter3 display
similar exploratory behavior. In later training phases, thesoft masks converge to hard masks (≈0 or ≈1), and the L1
norm of filters with mask values near zero have gradually
decreased to zero, a result of the weight decay term in theloss. Since the mask encodes the neurons’ connection information, one could interpret the exploration-exploitation
process as performing architecture search.
Training overhead with vs without pruning masks. We
measure the wall-clock time for a ResNet-50 on a single
A100 GPU for 300 iterations using a batch size of 64 withor without training a pruning mask. The training overhead
increases by ∼0.5%, which is generally negligible.

4.2. Accuracy vs. number of Sinkhorn steps
To improve the computational efficiency of our method,
we propose an optimal transportation algorithm that requires only a single Sinkhorn iteration to converge. To verify that a single Sinkhorn iteration indeed suffices, we performed experiments on CIFAR-10 with ResNet-56 using
a single GPU. In Fig. 4 we plot the accuracy and training
time as a function of the number of innner Sinkhorn steps.
A larger number of Sinkhorn steps significantly increases
training time without markedly affecting accuracy. Thus asingle Sinkhorn iteration suffices.

4.3. Results for ResNet-56/VGG-19 on CIFAR
We first explore the effect of applying different pruning schedules on the accuracy of ResNet-56 on CIFAR-10
and VGG-19 on CIFAR-100. We compare our method to
two other pruning methods: L1-norm based one-shot pruning [30] and growing regularization ”GReg-1” [63] where

9930

Accuracy
Training cost96

GPU-hours

Accuracy(%)

a weight decay of 5×10−4 and a batch size of 256 are used.
On ImageNet we use a weight decay of 10−4 and train on
4 Tesla V100 GPUs with a batch size of 64 per GPU. In all
experiments except for MobileNetV2 that follows the training settings of [60], the initial learning rate is set as 0.1 anda cosine learning rate decay is applied. In finetuning stage,
we employed the same training settings as [63].

15
93

90

1

10
50
Number of Sinkhorn steps

100

0

Figure 4. Training cost as a function of the number of inner
Sinkhorn steps on CIFAR-10 with ResNet-56 using a single GPU.
Each experiment is repeated 3 times with random initialization.
The training cost for a Sinkhorn step is measured in GPU-hours.
Increasing the number of Sinkhorn steps, significantly increases
the training cost, but hardly affects the accuracy. A single
Sinkhorn step is sufficient for our method.

the results are taken from the literature [63]. We use a simple uniform pruning ratio scheme, where the pruning ratio is
identical across all pruned layers. Following common practice [14], we do not prune the first layer, and for ResNet56 we only prune the first layer in a residual block (as aresult of the restriction imposed by the residual addition)[63, 30]. We explore a broad range of pruning ratios, covering Speedups between 2× and 44×. For fair comparisons,the finetuning scheme (e.g., learning rate schedule, numberof epochs, etc.) is identical across all methods.
The results are presented in Table 1. Note that all pruning methods result in exactly the same pruned network size.
Therefore different observed accuracies between methods
result from how the filters are pruned. Our proposed method
first performs exploration followed by exploitation. In doing so, our method searches for best architecture while gradually removing the effect of pruned filters on the network
output. Our method outperforms L1+one-shot in all examined conditions and GReg-1 for all conditions except one.

4.4. Results for ResNet-34/50 on ImageNet
We next evaluate our method with ResNet-34 andResNet-50 on the large-scale ImageNet dataset. We compare our results to previous pruning methods, including
L1 Norm [30], TaylorFO [45], IncReg [62], SFP [19],HRank [34], Factorized [31], DCP [69], CCP-AC [50],
LFPC [18], GReg-1 and GReg-2 [63], SRR [64], MetaPruning [37], AutoPruner [40], PGMPF [3] and Random Pruning [32]. For fair comparisons, following [63], we use theofficial PyTorch ImageNet training example to assure that
implementation details such as data augmentation, weightdecay, momentum, etc. match between methods.

Table 4. Acceleration comparison for structured filter pruning on
ImageNet. Speedup reflects the reduction on FLOPs.
MethodBackboneBaseline Pruned
Acc
Speed
Acc (%) Acc (%) Drop (%) Up

L1 (pruned-B) [30]Taylor-FO [45]ResNet-34
GReg-1 [63]GReg-2 [63]
Ours

73.23
73.31
73.31
73.31
73.31

72.17
72.83
73.54
73.61
74.28

1.06
0.48
-0.23
-0.30
-0.97

1.32×
1.29×
1.32×
1.32×
1.32×

SFP [19]
HRank [34]Factorized [31]
DCP [69]
CCP-AC [50]
SRR [64]
ResNet-50
AutoPruner [40]MetaPruning [37]
PGMPF [3]Random Prune [32]GReg-1 [63]GReg-2 [63]
Ours

76.15
76.15
76.15
76.01
76.15
76.13
76.15
76.60
76.01
76.13
76.13
76.13
76.13

74.61
74.98
74.55
74.95
75.32
75.11
74.76
75.40
75.11
75.13
75.16
75.36
75.55

1.54
1.17
1.60
1.06
0.83
1.02
1.39
1.20
0.90
1.00
0.97
0.77
0.58

1.72×
1.78×
2.33×
2.25×
2.18×
2.27×
2.05×
2.05×
2.20×
2.04×
2.31×
2.31×
2.31×

LFPC [18]GReg-1 [63]GReg-2 [63]Ours

ResNet-50

76.15
76.13
76.13
76.13

74.46
74.85
74.93
75.24

1.69
1.28
1.20
0.89

2.55×
2.56×
2.56×
2.56×

ResNet-50

75.60
76.18
76.13
76.13
76.13

71.07
71.69
73.75
73.90
74.26

4.53
4.49
2.38
2.23
1.87

3.00×
3.05×
3.06×
3.06×
3.06×

IncReg [62]Taylor-FO [45]GReg-1 [63]GReg-2 [63]Ours

We first explore the behavior of our method by training ResNet-50 on ImageNet with the pruning setting corresponding to 3.06× speedup [63] (see detail in A.7). In
the top row of Fig. 3(b), we visualize the soft mask andL1 norm of 256 filters from layer 24 of the trained model.
Note how all soft masks have converged to hard masks andthat the L1 norms of the filters are highly correlated with
their corresponding mask values. In the bottom row ofFig. 3(b), we plot the validation accuracy of ResNet-50 during training. We first use our method to prune the model
for 90 epochs and, following previous work [63], finetune
the pruned model for another 90 epochs. During the firststage, our method searches for the best architecture and reduces the importance of the pruned filters. After 90 epochs
we can safely prune the filters without a large impact on theaccuracy, in the star point we test the accuracy for the derived sparse model before finetuning, after which finetuning
helps to further improve the accuracy. At the beginning offinetuning we observe a small accuracy drop since we use alarger learning rate to start with.
Table 4 gives an overview of all comparisons. We groupthe methods with similar speedup together for easy comparison. We observe that our method consistently achieves

the best result among all approaches, for various speedup
comparisons on different architectures. On ResNet-34, ourmethod improves the baseline (unpruned) model by 0.97%
accuracy. Previous work [22] explained how pruning can
improve generalization performance to improve accuracy
for unseen data drawn from the same distribution. Such improved accuracy after pruning has been observed in earlier
work [21], but has been more apparent for smaller datasets
like CIFAR. In agreement with [63] we also observe improvements over the baseline accuracy on the much morechallenging ImageNet benchmark. With larger speedups,
the advantage of our method becomes more obvious. For
example, our method outperforms Taylor-FO [45] by 1.05%
top-1 accuracy at the 2.31× setting, while at 3.06×, ours is
better by 2.57% top-1 accuracy. Previous work has madeconsiderable progress defining importance scores based on
the networks weights [30] [18]. In contrast to those studies
our method learns the importance scores used to prune the
network in an end-to-end fashion. Our results indicate thiscan improve the accuracy of pruned networks.
Unstructured pruning on ImageNet. Thus far we focusedon structured filter pruning, but previous studies have also
explored unstructured pruning. In Table 2, we show that
our method can also be applied to unstructured pruning use
cases. Our method compares favorably to other advanced
unstructured pruning methods for ResNet-50 on ImageNet.
Global sparsity for MobileNetV2. Thus far we havefocused on normal convolutions, but networks based on
depthwise convolutions may be harder to prune. In Table 3,
we compare to ChipNet [60] using the global sparsity ratio
for the lightweight MobileNetV2 architecture on CIFAR10/100. Learning with a global sparsity ratio may distributedifferent sparsity ratios across layers, while respecting thepreset sparsity across the network. For both datasets, wererun the experiments of ChipNet using its released hyperparameters and training settings. Under different sparsity
ratios, our method consistently outperforms ChipNet.
We also further extend our method to general budgetedpruning with e.g. a certain FLOPs budget (see detail in A.9).

5. Conclusion
In this paper we propose a differentiable sparse learning algorithm based on entropy regularized optimal transportation that achieves exact control over the sparsity ratio.
We formulate a differentiable bi-level learning objective to
jointly optimize the soft masks and the network parameters using standard SGD methods combined with Sinkhorn
optimization. Optimizing this bi-level problem is computational intensive, but we show that by using a Bregman divergence we can improve the efficiency, i.e. requiring only asingle Sinkhorn iteration, while still guaranteeing algorithmic convergence. We demonstrated the effectiveness of our
method on different datasets for different pruning rates.

Acknowledgements. This work is funded by EU’s Horizon
Europe research and innovation program under grant agreement No. 101070374.

References[1] SN Afriat. Theory of maxima and the method of lagrange. SIAM Journal on Applied Mathematics, 20(3):343–
357, 1971.
[2] Martin Arjovsky, Soumith Chintala, and Léon Bottou.
Wasserstein generative adversarial networks. In International conference on machine learning, pages 214–223.
PMLR, 2017.
[3] Linhang Cai, Zhulin An, Chuanguang Yang, Yangchun Yan,and Yongjun Xu. Prior gradient mask guided pruning-awarefine-tuning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 1, 2022.
[4] Yair Censor and Stavros Andrea Zenios. Proximal minimization algorithm with d-functions. Journal of Optimization
Theory and Applications, 73(3):451–464, 1992.
[5] Soravit Changpinyo, Mark Sandler, and Andrey Zhmoginov. The power of sparsity in convolutional neural networks.
arXiv preprint arXiv:1702.06257, 2017.
[6] Daiki Chijiwa, Shin’ya Yamaguchi, Yasutoshi Ida, Kenji
Umakoshi, and Tomohiro Inoue. Pruning randomly initialized neural networks with iterative randomization. Advancesin Neural Information Processing Systems, 34:4503–4513,
2021.
[7] Benoı̂t Colson, Patrice Marcotte, and Gilles Savard. An
overview of bilevel optimization. Annals of operations research, 153(1):235–256, 2007.
[8] Marco Cuturi. Sinkhorn distances: Lightspeed computationof optimal transport. Advances in neural information processing systems, 26, 2013.
[9] Xiaohan Ding, Xiangxin Zhou, Yuchen Guo, Jungong Han,
Ji Liu, et al. Global sparse momentum sgd for pruning very
deep neural networks. Advances in Neural Information Processing Systems, 32, 2019.
[10] Konstantin Ditschuneit and Johannes S Otterbach. Autocompressing subset pruning for semantic image segmentation. In Pattern Recognition: 44th DAGM German Conference, pages 20–35. Springer, 2022.
[11] Marvin Eisenberger, Aysim Toker, Laura Leal-Taixé, Florian Bernard, and Daniel Cremers. A unified framework
for implicit sinkhorn differentiation. In Proceedings ofthe IEEE/CVF Conference on Computer Vision and Pattern
Recognition, pages 509–518, 2022.
[12] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Modelagnostic meta-learning for fast adaptation of deep networks.
In International conference on machine learning, pages
1126–1135. PMLR, 2017.
[13] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. ICLR,
2019.
[14] Trevor Gale, Erich Elsen, and Sara Hooker. The stateof sparsity in deep neural networks.
arXiv preprint
arXiv:1902.09574, 2019.

[15] Shangqian Gao, Feihu Huang, Jian Pei, and Heng Huang.
Discrete model compression with resource constraint for
deep neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages
1899–1908, 2020.
[16] Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics,
pages 1608–1617. PMLR, 2018.
[17] Shaopeng Guo, Yujie Wang, Quanquan Li, and Junjie Yan.
Dmcp: Differentiable markov channel pruning for neuralnetworks. In Proceedings of the IEEE/CVF conference on
computer vision and pattern recognition, pages 1539–1547,
2020.
[18] Yang He, Yuhang Ding, Ping Liu, Linchao Zhu, Hanwang
Zhang, and Yi Yang. Learning filter pruning criteria for deep
convolutional neural networks acceleration. In Proceedingsof the IEEE/CVF conference on computer vision and patternrecognition, pages 2009–2018, 2020.
[19] Yang He, Guoliang Kang, Xuanyi Dong, Yanwei Fu, and Yi
Yang. Soft filter pruning for accelerating deep convolutional
neural networks. IJCAI, 2018.
[20] Yang He, Ping Liu, Ziwei Wang, Zhilan Hu, and Yi
Yang. Filter pruning via geometric median for deep convolutional neural networks acceleration. In Proceedings ofthe IEEE/CVF conference on computer vision and patternrecognition, pages 4340–4349, 2019.
[21] Yihui He, Xiangyu Zhang, and Jian Sun. Channel pruning
for accelerating very deep neural networks. In Proceedingsof the IEEE international conference on computer vision,
pages 1389–1397, 2017.
[22] Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden,
and Alexandra Peste. Sparsity in deep learning: Pruning
and growth for efficient inference and training in neural networks. J. Mach. Learn. Res., 22(241):1–124, 2021.
[23] Zehao Huang and Naiyan Wang. Data-driven sparse structure selection for deep neural networks. In Proceedings ofthe European conference on computer vision (ECCV), pages
304–320, 2018.
[24] Ryan Humble, Maying Shen, Jorge Albericio-Latorre, Eric
Darve, and Jose M Alvarez.
Soft masking for costconstrained channel pruning. ECCV, 2022.
[25] Minsoo Kang and Bohyung Han. Operation-aware soft channel pruning using differentiable masks. In International Conference on Machine Learning, pages 5122–5131. PMLR,
2020.
[26] Jeffrey J Kosowsky and Alan L Yuille. The invisible handalgorithm: Solving the assignment problem with statisticalphysics. Neural networks, 7(3):477–490, 1994.
[27] Ronny Krashinsky, Olivier Giroux, Stephen Jones, Nick
Stam, and Sridhar Ramaswamy. Nvidia ampere architecture
in-depth. NVIDIA blog: https://devblogs. nvidia. com/nvidiaampere-architecture-in-depth, 2020.
[28] Carl Lemaire, Andrew Achkar, and Pierre-Marc Jodoin.
Structured pruning of neural networks with budget-aware
regularization. In Proceedings of the IEEE/CVF Conference
on Computer Vision and Pattern Recognition, pages 9108–
9116, 2019.

[29] Bailin Li, Bowen Wu, Jiang Su, and Guangrun Wang. Eagleeye: Fast sub-net evaluation for efficient neural network
pruning. In ECCV 2020, pages 639–654. Springer, 2020.
[30] Hao Li, Asim Kadav, Igor Durdanovic, Hanan Samet, andHans Peter Graf. Pruning filters for efficient convnets. ICLR,
2017.
[31] Tuanhui Li, Baoyuan Wu, Yujiu Yang, Yanbo Fan, Yong
Zhang, and Wei Liu. Compressing convolutional neural networks via factorized convolutional filters. In Proceedings ofthe IEEE/CVF Conference on Computer Vision and Pattern
Recognition, pages 3977–3986, 2019.
[32] Yawei Li, Kamil Adamczewski, Wen Li, Shuhang Gu, Radu
Timofte, and Luc Van Gool. Revisiting random channel
pruning for neural network compression. In Proceedings ofthe IEEE/CVF Conference on Computer Vision and Pattern
Recognition, pages 191–201, 2022.
[33] Lucas Liebenwein, Cenk Baykal, Harry Lang, Dan Feldman,
and Daniela Rus. Provable filter pruning for efficient neuralnetworks. arXiv preprint arXiv:1911.07412, 2019.
[34] Mingbao Lin, Rongrong Ji, Yan Wang, Yichen Zhang,
Baochang Zhang, Yonghong Tian, and Ling Shao. Hrank:
Filter pruning using high-rank feature map. In Proceedingsof the IEEE/CVF conference on computer vision and patternrecognition, pages 1529–1538, 2020.
[35] Tao Lin, Sebastian U Stich, Luis Barba, Daniil Dmitriev, andMartin Jaggi. Dynamic model pruning with feedback. arXiv
preprint arXiv:2006.07253, 2020.
[36] Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts:
Differentiable architecture search. ICLR, 2019.
[37] Zechun Liu, Haoyuan Mu, Xiangyu Zhang, Zichao Guo, Xin
Yang, Kwang-Ting Cheng, and Jian Sun. Metapruning: Metalearning for automatic neural network channel pruning. In
Proceedings of the IEEE/CVF international conference on
computer vision, pages 3296–3305, 2019.
[38] Zhuang Liu, Mingjie Sun, Tinghui Zhou, Gao Huang, andTrevor Darrell. Rethinking the value of network pruning.
ICLR, 2019.
[39] Jelena Luketina, Mathias Berglund, Klaus Greff, and Tapani
Raiko. Scalable gradient-based tuning of continuous regularization hyperparameters. In International conference on
machine learning, pages 2952–2960. PMLR, 2016.
[40] Jian-Hao Luo and Jianxin Wu. Autopruner: An end-to-end
trainable filter pruning method for efficient deep model inference. Pattern Recognition, 107:107461, 2020.
[41] Jian-Hao Luo and Jianxin Wu. Neural network pruning withresidual-connections and limited-data. In Proceedings ofthe IEEE/CVF Conference on Computer Vision and Pattern
Recognition, pages 1458–1467, 2020.
[42] Fanxu Meng, Hao Cheng, Ke Li, Huixiang Luo, Xiaowei
Guo, Guangming Lu, and Xing Sun. Pruning filter in filter. Advances in Neural Information Processing Systems,
33:17629–17640, 2020.
[43] Lu Miao, Xiaolong Luo, Tianlong Chen, Wuyang Chen,
Dong Liu, and Zhangyang Wang. Learning pruning-friendly
networks via frank-wolfe: One-shot, any-sparsity, and no retraining. In International Conference on Learning Representations, 2022.

[44] Dmitry Molchanov, Arsenii Ashukha, and Dmitry Vetrov.
Variational dropout sparsifies deep neural networks. In International Conference on Machine Learning, pages 2498–
2507. PMLR, 2017.
[45] Pavlo Molchanov, Arun Mallya, Stephen Tyree, Iuri Frosio, and Jan Kautz. Importance estimation for neural network pruning. In Proceedings of the IEEE/CVF Conference
on Computer Vision and Pattern Recognition, pages 11264–
11272, 2019.
[46] Xuefei Ning, Tianchen Zhao, Wenshuo Li, Peng Lei, Yu
Wang, and Huazhong Yang. Dsa: More efficient budgetedpruning via differentiable sparsity allocation. In ECCV 2020,
pages 592–607. Springer, 2020.
[47] Neal Parikh, Stephen Boyd, et al. Proximal algorithms.
Foundations and trends® in Optimization, 1(3):127–239,
2014.
[48] Adam Paszke, Sam Gross, Soumith Chintala, Gregory
Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automaticdifferentiation in pytorch. 2017.
[49] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer,
James Bradbury, Gregory Chanan, Trevor Killeen, Zeming
Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[50] Hanyu Peng, Jiaxiang Wu, Shifeng Chen, and Junzhou
Huang. Collaborative channel pruning for deep networks.
In International Conference on Machine Learning, pages
5113–5122. PMLR, 2019.
[51] Gabriel Peyré, Marco Cuturi, et al. Computational optimaltransport: With applications to data science. Foundations
and Trends® in Machine Learning, 11(5-6):355–607, 2019.
[52] Pedro Savarese, Hugo Silva, and Michael Maire. Winningthe lottery with continuous sparsification. Advances in Neural Information Processing Systems, 33:11380–11390, 2020.
[53] Bernhard Schmitzer. Stabilized sparse scaling algorithms
for entropy regularized transport problems. arXiv preprint
arXiv:1610.06519, 2016.
[54] Maying Shen, Hongxu Yin, Pavlo Molchanov, Lei Mao,
Jianna Liu, and Jose Alvarez. Structural pruning via latencysaliency knapsack. In Advances in Neural Information Processing Systems, 2022.
[55] Yucong Shen, Li Shen, Hao-Zhi Huang, Xuan Wang, andWei Liu. Cpot: Channel pruning via optimal transport. arXiv
preprint arXiv:2005.10451, 2020.
[56] Sidak Pal Singh and Dan Alistarh. Woodfisher: Efficientsecond-order approximation for neural network compression. Advances in Neural Information Processing Systems,
33:18098–18109, 2020.
[57] Ilya Sutskever, James Martens, George Dahl, and Geoffrey
Hinton. On the importance of initialization and momentum
in deep learning. In International conference on machinelearning, pages 1139–1147. PMLR, 2013.
[58] Kenji Suzuki, Isao Horiba, and Noboru Sugie. A simple neural network pruning algorithm with application to filter synthesis. Neural processing letters, 13(1):43–53, 2001.

[59] Kai Sheng Tai, Taipeng Tian, and Ser Nam Lim. Spartan:
Differentiable sparsity via regularized transportation. Advances in Neural Information Processing Systems, 35:4189–
4202, 2022.
[60] Rishabh Tiwari, Udbhav Bamba, Arnav Chavan, and Deepak
Gupta. Chipnet: Budget-aware pruning with heaviside continuous approximations. In International Conference on
Learning Representations, 2021.
[61] Marc Aurel Vischer, Robert Tjarko Lange, and Henning
Sprekeler. On lottery tickets and minimal task representations in deep reinforcement learning. ICLR, 2022.
[62] Huan Wang, Xinyi Hu, Qiming Zhang, Yuehai Wang, Lu Yu,
and Haoji Hu. Structured pruning for efficient convolutional
neural networks via incremental regularization. IEEE Journal of Selected Topics in Signal Processing, 14(4):775–788,
2019.
[63] Huan Wang, Can Qin, Yulun Zhang, and Yun Fu. Neuralpruning via growing regularization. ICLR, 2021.
[64] Zi Wang, Chengcheng Li, and Xiangyang Wang. Convolutional neural network pruning with structural redundancy
reduction. In Proceedings of the IEEE/CVF Conference
on Computer Vision and Pattern Recognition, pages 14913–
14922, 2021.
[65] Yujia Xie, Hanjun Dai, Minshuo Chen, Bo Dai, Tuo Zhao,
Hongyuan Zha, Wei Wei, and Tomas Pfister. Differentiable
top-k with optimal transport. Advances in Neural Information Processing Systems, 33:20520–20531, 2020.
[66] Yujia Xie, Xiangfeng Wang, Ruijia Wang, and Hongyuan
Zha. A fast proximal point method for computing exact
wasserstein distance. In Uncertainty in artificial intelligence,
pages 433–453. PMLR, 2020.
[67] Jianbo Ye, Xin Lu, Zhe Lin, and James Z Wang. Rethinking the smaller-norm-less-informative assumption inchannel pruning of convolution layers. arXiv preprint
arXiv:1802.00124, 2018.
[68] Shuai Zhang, Meng Wang, Sijia Liu, Pin-Yu Chen, and Jinjun Xiong. Why lottery ticket wins? a theoretical perspectiveof sample complexity on sparse neural networks. Advancesin Neural Information Processing Systems, 34:2707–2720,
2021.
[69] Zhuangwei Zhuang, Mingkui Tan, Bohan Zhuang, Jing Liu,Yong Guo, Qingyao Wu, Junzhou Huang, and Jinhui Zhu.
Discrimination-aware channel pruning for deep neural networks. Advances in neural information processing systems,
31, 2018.

A. Supplementary Material

A.3. Expensive to train with Sinkhorn’s algorithm

A.1. Optimal coupling to regularized problem

The bi-level optimization problem in Eq. (10) and Eq. (11)is expensive to train with Sinkhorn’s algorithm: During forward pass, per mini-batch the inner optimization in Eq. (11)needs to perform Sinkhorn’s iterative algorithm for hundreds of iterations to converge, which is computationallyinefficient. During the back-propagation pass, the gradientof P∗ε (s) with respect to the importance scores s is computed by differentiating [48] through all Sinkhorn iterations,which is expensive. We display the forward and backwardof Sinkhorn algorithm in Fig. 5.

Here we prove the optimal coupling to the regularized problem in Eq. (7) is:Pε = ef /ε ⊙ e−C/ε ⊙ eg/ε .

(16)

Proof. Introducing two dual variables f ∈ Rn and g ∈R2 for marginal constraints P12 = a and PT 1n = b,given
P the discrete entropy of a coupling matrix H(P) =− ij Pij (log(Pij ) − 1), the Lagrangian of Eq. (7) is optimized as follows:

A.4. Bregman divergence based optimization

ξ(P, f , g) = ⟨C, P⟩−εH(P)−⟨f , P12 − a⟩− g, P 1n − b . In this paper we use similar Bregman Divergence Dh as
P
(17)
[66]
 based on entropy function H(x) = − i xi log(xi )−First order conditions then yield:
1 as:
X
xi X
∂ξ(P, f , g)
xi .
(21)
DH (x, y) = −xi log += C + ε log(Pij ) − fi − gj = 0,
(18)
yi
∂Pij
i
i
T

which result for an optimal coupling to the regularized problem, in the matrix expression shown in Eq. (16).

Based on the defined Bregman Divergence, a single proximal point iteration for problem (4) can be written as:

A.2. Dual problem computation

P(ℓ+1) =Here we prove that minimizing the regularized optimaltransport distance in Eq. (7) is equivalent to maximizing its
dual problem:DEf /ε −C(s)/ε
g/ε
max
⟨f
,
a⟩
+
⟨g,
b⟩
−
ε
e
,
e
·
e
n
2
f ∈R ,g∈R(19)

Proof. We start from the result in Eq. (16), and substitute itin the Lagrangian ξ(P, f , g) of Eq. (17), where the optimal
P is a function of f and g, we obtain that the Lagrange dual
function equals:DEf , g 7→ ef /ε , e−C(s)/ε ⊙ C(s) eg/ε(20)
− εH ef /ε ⊙ e−C(s)/ε ⊙ eg/εThe negative entropy of P scaled by ε, namelyε ⟨P, logP − 1n×2 ⟩, can be stated explicitly as a functionof f , g, C:ε ⟨P, logP − 1n×2 ⟩DET= ef /ε ⊙ e−C(s)/ε ⊙ eg/ε , f 1T
2 + 1n g − C(s) − ε1n×2DE= − ef /ε , e−C(s)/ε ⊙ C(s) eg/ε + ⟨f , a⟩ + ⟨g, b⟩ −DEε ef /ε , e−C(s)/ε · eg/εtherefore, the first term in Eq. (20) cancels out with the firstterm in the entropy above. The remaining terms are those
appearing in Eq. (19).

=minP∈U (a,b)min

P∈U (a,b)⟨C, P⟩ − εDH (P, P(ℓ) )DE(22)
C − ε log P(ℓ) , P − εH(P).

Denote C′ = C − ε log P(ℓ) . Note that for optimization
problem 22, P(ℓ) is a fixed value that is not relevant to optimization variable P. Comparing to Eq. (7), the problem
in Eq. (22) can be solved by Sinkhorn iteration by replacing
CC′
Gibbs kernel K by K′ = e− ε = e− ε ⊙P(ℓ) . With this reorganization, we can solve it with Sinkhorn algorithm. [66]have shown both theoretically and empirically that a single Sinkhorn inner iteration is sufficient to converge undera large range of fixed ε, we therefore compute P(ℓ+1) as:

(ℓ+1)
(ℓ+1)
C
/ε
/ε
P(ℓ+1) = ef
⊙ e− ε ⊙ P(ℓ) ⊙ eg
, (23)
which is the expression of Eq. (13).

A.5. Proximal point iteration as iterative Sinkhorn
We refer to the proof in Chapter 4.2 of the book [51].
The general setting for proximal point iteration is to defineC
Gibbs kernel as K = e− ε ⊙ P(ℓ) , the proximal point iterations thus have the form:

(ℓ+1)
(ℓ+1)
C
/ε
/ε
P(ℓ+1)(s) =ef⊙ e− ε ⊙ P(ℓ) ⊙ egε(ℓ+1)C
(ℓ+1)
(1)
/ε
(24)
= ef
⊙ · · ⊙ef /ε ⊙ e− ε


(ℓ+1)
(1)
/ε
⊙ · · ⊙eg /ε .
⊙ P(ℓ) ⊙ eg

The proximal point iteration iteratively applies Sinkhorn’sCalgorithm with a e− ε/ℓ kernel, i.e. with a decaying regularization parameter ε/ℓ, therefore it can perform an automatic

ForwardC𝑠!

……

Eq. (3)Eq. (24)

Eq. (23)

Eq. (24)

Transportation
Plan

……

𝑠#

0

Score

Target

1

𝑔(")

𝑓 ($)

𝑔($)

𝑔(%&")

𝑓 (%)

……

𝑠"

Eq. (23)……

…

K

Backward…

𝑔(%)

Sinkhorn 𝑡 = 1, … , 𝑇

Figure 5. The display of training with Sinkhorn’s algorithm. We show the forward and backward pass over one single layer: given
importance score s, the cost matrix is computed by Eq. (3), we thus get K; in forward pass we input K to iteratively compute f (t+1) with
Eq. (27) and g(t+1) with Eq. (28) using Sinkhorn algorithm for T iterations; during back-propagation pass, the gradient of soft mask withrespect to the importance scores is computed by differentiating through all Sinkhorn iterations. A large T makes it expensive to train.
Table 5. Training setting: For the SGD optimizer, in the parentheses are the momentum and weight decay. For ImageNet, batch size is 64
per GPU. We learn the soft mask using cosine learning rate schedule.
Settings

CIFAR

ImageNet

Optimizer

SGD (0.9, 5e-4)

SGD (0.9, 1e-4)

LR schedule (soft mask)
LR schedule (finetune)

Cosine LR schedule (0.1)
Multi-step (0:1e-2, 60:1e-3, 90:1e-4)

Multi-step (0:1e-2, 30:1e-3, 60:1e-4, 75:1e-5)

120 + 120

90 + 90256

256

Training Epoch
Batch size

decaying schedule on the regularization as ℓ → ∞ to gradually approach the optimal discrete plan.

A.6. Sinkhorn iterations on dual problem
A simple approach to solving the unconstrained maximization problem in Eq. (19) is to use an exact block coordinate
ascent strategy, namely to update alternatively f and g tocancel the respective gradients in these variables of the objective of (19). Indeed, one can notice after a few elementary computations that, writing Ldual (f , g, s) for the objective of (19), we have the gradients, w.r.t., f and g as follows:∇f Ldual (f , g, s) = a − ef /ε ⊙ (Keg/ε ),∇g Ldual (f , g, s) = b − eg/ε ⊙ (Kef /ε ),

(25)
(26)

where K = e−C(s)/ε is defined as the Gibbs kernel in
Sinkhorn’s algorithm. Block coordinate ascent can therefore be implemented in a closed form by applying successively the following updates, starting from any arbitraryg(1) , for t ≥ 1:
f (t+1) = ε log a − ε log (Keg
g

(t+1)

(t)

/ε

= ε log b − ε log (K e

(27)

);

T f (t+1) /ε

).

(28)

A.7. Experimental setting details
How is the ε selected. The ε controls a trade-off between
exploration and exploitation. A large ε leads to a greater
exploration by a softer mask. We verify that soft masks
converge to hard masks by training for few epochs and measuring the average of the squared differences between themfor different ε candidates. As is common to hyperparametertuning, this requires some effort. Therefore for similar architectures we used the same ε, e.g., ε = 1 for lightweightResNet-56 and MobileNetV2.
Training setting. In Table 5 we summarize the detailed
training settings of this work. For easier comparison, we
use same training settings as [63] in finetuning. In our soft
mask learning phase, we use a cosine learning rate schedule for updating the importance scores and network weights
with SGD.
Pruning ratio. We use same pruning ratio as in [63] for faircomparison. In Table 6 we give the specific pruning ratio
used for our experiments in the paper. We here briefly explain how we set the pruning ratio. We use two different architectures: single-branch (i.e. VGG-19) and multi-branch(i.e. ResNet). (i) For VGG19, we use the following pruning ratio setting. As an example, “[0:0, 1-9:0.3, 10-15:0.5]”
means “for the first layer (index starting from 0), the pruning ratio is 0; for layer 1 to 9, the pruning ratio is 0.3; forTable 6. Training setting: For the SGD optimizer, in the parentheses are the momentum and weight decay. For ImageNet, batch size is 64
per GPU. We learn the soft mask using cosine learning rate schedule.
Datasets

Backbone

SpeedupPruning ratio

CIFAR-10

ResNet-56−−[0, p, p, p], p ∈ {0.5, 0.7, 0.9, 0.925, 0.95}[0, 0.50, 0.60, 0.40, 0]

CIFAR-100

VGG-19ImageNetResNet-34

1.32×

ImageNetResNet-50

2.31×[0, 0.60, 0.60, 0.60, 0.21]

ImageNetResNet-50

2.56×[0, 0.74, 0.74, 0.60, 0.21]

ImageNetResNet-50

3.06×[0, 0.68, 0.68, 0.68, 0.50]

Table 7. Training recipes for ResNet50 under structured pruning on ImageNet-1K. Our method is flexible for different training
recipes.
RecipeBaseline
Stonger

Speed Up

Acc.

2.31×

75.54
77.16

[0:0, 1-15:p], p ∈ {0.5, 0.6, 0.7, 0.8, 0.9}−−

Table 8. FLOPs budgets for WideResNet-26 on CIFAR-100
MethodBudget (%)

Acc.

Unpruned

100

80.21

ChipNet [60]
Ours

40

76.88
78.63

ChipNet [60]
Ours

20

77.15
78.17

Sparsity ratio

1layer 10 to 15, the pruning ratio is 0.5”. (ii) For a ResNet,
if it has N stages, we will use a list of N floats to represent
its pruning ratios for the N stages. For example, ResNet-56
has 4 stages in conv layers, then “[0, 0.7, 0.7, 0.7]” means“for the first stage (the first conv layer), the pruning ratio
is set as 0; the other three stages have pruning ratio of 0.7”.
Besides, since we do not prune the last conv layer in a residual block, which means for a two-layer residual block (forResNet-56), we only prune the first layer; for a three-layer
bottleneck block (for ResNet-34 and ResNet-50), we only
prune the first and second layers.
Unstructured pruning. For unstructured pruning, each individual weight needs to couple an importance score which
may cause higher memory cost, therefore we follow previous works that use the magnitude of the weight parameteras an importance score.

Initialization
3rd epoch
20th epoch

0.5
0

0

5

10

15
20
25
Layer index
Figure 6. Learned sparsity ratios over layers under FLOPs budget.

A.9. Extension to FLOPs and latency budgets
The FLOPs or latency budget pruning can be viewed as alearnable differentiable sparsity allocation problem. We formulate it in the following.
Formulation. Given budget BF , the network weights as w,
the optimization problem of kept ratios A = {α(l) }l=1,...,L(i.e., 1−sparsity ratio) of L layers can be written as:arg min Ltrain f (x; w, m, A) ,m,w,A

A.8. Training recipe variants
Our primary goal was to compare pruning methods as fairly
as possible. The reported results for ResNet-50 used thebaseline training recipe from CReg for fair comparisons.
Here, we performed additional analyses, pruning ResNet50 on ImageNet with Label Smoothing, TrivialAugment,
Random Erasing, Mixup, and Cutmix to form a stronger
training recipe. We did not apply other effective techniquessuch as Long Training (e.g. 600 epochs), LR optimizations,
EMA, Weight Decay tuning, and Inference Resize tuning,which may further improve by ∼3% accuracy as shown in
PyTorch. We trained on a large batch size of 4,096 on 8
NVIDIA A100 GPUs, scaling up the learning rate by 5×.
As expected, Table 7 shows increased pruning accuracy,
showing that our method can be augmented with different
training recipes.

s.t.

1 Xn
(l)
mi = α(l) , m(l) ∈ {0, 1}n ,i=1nF(A) ≤ BF , 0 ≤ A ≤ 1.

(29)

where Ltrain is training loss. F(A) is the consumed resource
corresponding to the kept ratios. If A is predefined layerwise or global kept ratio, the formulation degrades to our
original formulation in Eq. (1). For FLOPs or latency budgets, we learn the kept ratios. Note F(A) can be denoted
as a function of kept ratio, see [60] and [54]. By setting
A = Sigmoid(Θ) where Θ is a group of learnable parameters where 0 ≤ A ≤ 1 is satisfied naturally, we optimize
Θ to learn kept ratio A by minimizng Ltrain and a budgets
penalty loss ||F(A) − BF ||2 . Our optimal transport method
dynamically aligns to learned ratios A.
We add the experimental comparisons in Table 8 with
FLOPs budget for WideResNet-26 for CIFAR-100.

FLOPs budget. We use the same way as [60] to calculate
the FLOPs budget that assumes a sliding window is used toachieve convolution and the nonlinear computational overhead is ignored. We define FLOPs budget as:
PN
F(A) =(j−1)+ 1) · nj · α(j) · Aj
j=1 (Kj · nj−1 · α,
PN
j=1 (Kj · nj−1 + 1) · nj · Aj

(30)
where N denotes the number of convolutional layers in the
network, Kj denotes area of the kernel, and Aj and nj denote area of the feature maps and the channel count, respectively, in the j th layer.
We note that the FLOPs budget is formulated as a function of kept ratios α(j) . Also the the latency cost of eachlayer can be approximated by the kept ratios of previous
layer and current layer as shown in [24].