Spartan: Differentiable Sparsity via Regularized Transportation
Authors: Kai Sheng Tai
Taipeng Tian
Ser-Nam Lim
Venue: NeurIPS 2022 Accept (2022)
Keywords: sparsity
optimal transport

Abstract: We present Spartan, a method for training sparse neural network models with a predetermined level of sparsity. Spartan is based on a combination of two techniques: (1) soft top-k masking of low-magnitude parameters via a regularized optimal transportation problem and (2) dual averaging-based parameter updates with hard sparsification in the forward pass. This scheme realizes an exploration-exploitation tradeoff: early in training, the learner is able to explore various sparsity patterns, and as the soft top-k approximation is gradually sharpened over the course of training, the balance shifts towards parameter optimization with respect to a fixed sparsity mask. Spartan is sufficiently flexible to accommodate a variety of sparsity allocation policies, including both unstructured and block-structured sparsity, global and per-layer sparsity budgets, as well as general cost-sensitive sparsity allocation mediated by linear models of per-parameter costs. On ImageNet-1K classification, we demonstrate that training with Spartan yields 95% sparse ResNet-50 models and 90% block sparse ViT-B/16 models while incurring absolute top-1 accuracy losses of less than 1% compared to fully dense training.

Spartan: Differentiable Sparsity via Regularized Transportation
Kai Sheng Tai Meta Al
Taipeng Tian
Ser-Nam Lim Meta Al
Meta Al
Abstract
We present Spartan, a method for training sparse neural network models with predetermined level of sparsity. Spartan is based combination of two techniques: (I) soft top-k masking of low-magnitude parameters via a regularized optimal transportation problem and (2) dual averaging-baat based parameter updates with hard sparsification in the forward pass. This scheme realizes an exploration- exploitation tradeoff: early in training, the learner is able to explore various sparsity patterns, and as the soft top-k approximation is gradually sharpened over the course of training. the balance shifts towards parameter optimization with respect to a fixed sparsity mask. Spartan is sufficiently flexible to accommodate a variety sparsity allocation policies, including both unstructured and block structured sparsity, as well general cost-sensitive sparsity allocation mediated by linear models of per-parameter costs. On ImageNet-I K classification, Spartan yields 95% sparse ResNet-50 models and 90% block sparse VIT-B/16 models while incurring absolute top-1 accuracy losses of less than 1% compared to fully dense training.
Introduction
Sparse learning algorithms search for model parameters that minimize training loss while retaining only small fraction of non-zero entries. Parameter sparsity yields benefits along several axes: reduced model storage costs, greater computational and energy efficiency during training and infer- ence, and potentially improved model generalization [53]. However, sparse training is challenging optimization problem: for the general problem of learning parameters E I subject to the constraint that 0 is k-sparse, the feasible set is the union of (+) axis-aligned hyperplanes intersecting at the origin, each of dimension k. This complex, nonconvex geometry of the feasible set compounds the existing difficulty of optimizing deep neural network models.
This work focuses improving the generalization error of sparse neural network models. To this end, we introduce Spartan, or Sparsity via Regularized Transportation- a sparse learning algorithm that leverages optimal transportation-basee based top-k masking operation induce parameter sparsity during training. Spartan belongs to the family of "parameter dense" training algorithms that maintains dense parameter vector 0 Rd throughout training 4 24], in contrast "parameter sparse" training algorithms that adhere
Ō(k) memory budget for representing the parameters of k- sparse model B 31 B 16. While computational cost and memory usage at training time are important considerations, Spartan primarily optimizes for performance inference time.
Intuitively, Spartan aims to achieve controlled transition between the exploration and exploitation of various sparsity patterns during training. In the exploration regime, the goal is for the learner to poor
easily transition between differing sparsity patterns order to escape those that correspond minima of the loss function. On the other hand, the goal of exploitation is for the learner to optimize model parameters given a fixed sparsity pattern, or a small set of similar patterns. This latter setting
Correspondence Kai Sheng Tai (kst@meta.com
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Soft top-k masking (k 30)
Soft top-k masks
1.0
|0]
Os41
10
[1(00+1)
[ox(0)]
0p-1
Hg(00-1)
(9:)
(1,(8)
0.0
100
(a)
(b)
(c)
Figure 1: Soft top-k masking. (a. b) The soft top-k masking operator oj computes approximately k-sparse outputs, with the sharpness of the mask controlled by the parameter B. (c) Small updates iterates far from the k-sparse feasible set (0: 0++1) can correspond to a large perturbation in parameter space after projection by I Updates to iterates 9g in the approximately sparse region (shaded in grey) correspond smaller post-projection perturbations.
avoids frequent oscillations between disparate sparsity patterns, which may be detrimental to the optimization process.
We operationalize this intuition through the use of differentiable soft top-k masking operation
of R + Rdd This function maps parameters 0 to an approximately sparse output that suppresses soft top-k mask with sharpness parameter 3: at = 0, of simply scales the input by constant, and as - 20, the mapping reduces to hard top-k magnitude-basee selection (Figure [ (a, b)). Soft top-k masking therefore constrains the iterates 0 to be close to the set of exactly k-sparse vectors, with the strength of this constraint mediated by the parameter B. Figures (c) and give some geometeric intuition for the effect of this mapping. We implement soft top-k masking using regularized optimal transportation formulation [I] 42 and
low-magnitude entries
We parameterize
demonstrate that this technique scales networks the order of 10 parameters in size.
We evaluate Spartan on ResNet-50 ResNet-50, find that sparse models trained with Spartan achieve higher generalization accuracies than those trained with existing methods at sparsity levels of 90% and above. In particular, we train ResNet-50 models to 76.5% top-1 validation accuracy 95% sparsity and to 74.2% accuracy at 97.5% sparsity, improving on the previous state-of-the-an by 0.6% and 1.6% respectively. Our sparse ViT-B/16 models reduce model storage size by 10x and inference FLOPs by 7.4x the cost of 0.6% accuracy reduction relative to DeiT-B [S We further demonstrate that Spartan is effective for block structured pruning, form of structured sparsity that is more amenable acceleration current GPU hardware than unstructured sparsity (18). On a ViT-B/16 model with 16 x 16 block structured pruning, Spartan achieves 79. % top-1 accuracy 90% sparsity, compared the baseline of 74.1% at the same sparsity level.
and ViT 1S] models trained on the ImageNet- IK dataset.
To summarize, we make the following contributions in this paper:
We present Spartan, sparsity-inducing training algorithm based on a soft top-k masking op- eration. We show that Spartan interpolates between two existing sparse learning algorithms: iterative magnitude pruning AS and Top-K Always Sparse Training (24).
We empirically evaluate Spartan using ResNet-50 and VIT models trained on the ImageNet-IK dataset, demonstrating consistent improvements over the prior oate-of-the-art
We study the effect of Spartan's 's hyperparameters on an exploration-exploiiee during training and on the final accuracy of the trained models.
tradeoff
We provide an open source implementation axcebokrresearch/ppaatta
Spartan at https://githun com/
Notation. We use R2., and Re++ to denote the nonnegative real numbers and the strictly positive real numbers respectively. 1, is the all-ones vector of dimension d. Given a vector E RRd, I|pp denotes

---

.Mo,0)
B
B =
f = 10
Figure 2: A 2D example soft masking, We plot probability densities corresponding to the action of the soft top-1 operator of on the 2D standard Gaussian distribution (darker colors indicate higher density). Specifically, we visualize the densities of of (3)) ar N 12) for range of sharpness parameters 3. Higher values of constrain iterates to closer to the k-sparse feasible set.
the p-norm of x, || the number of nonzero entries in æ, and the vector of elementwise absolute values. When z and y are vectors, O denotes elementwise multiplication and x/ denotes elementwise division. The operator II, (z) == arg m || ar--a -yll denotes Euclidean projection onto the set of k-sparse vectors.
Related Work
Neural network pruning. Our use of a magnitude-basse criterion for sparse learning draws on early work in the area of neural network pruning [ Mlggituueebbsee pruning is computationally cheap relative to alternative criteria that rely on first- or second-order information [+]] 27 20 and
perhaps surprisingly performant option despite its simplicity [7 More generally, Spartan builds on substantial body of previous work that aims to jointly optimize sparsity patterns and model
parameters for deep neural networks 43
29
4g
inter alia; see, e.g., 22 for survey].
Spartan as a generalization of existing methods. We highlight two particularly relevant sparse training methods in the literature: the iterative magnitude pruning (IMP) method (Algorithm and Top- K Always Sparse Training, or KAST (Algorithm
48
(4
Algorithm 1 Iterative magnitude pruning update
Algorithm Dual averaging Top-KAST update
Input: parameters 0r, loss function L(0), spar- sity budget k, step size n
Input: parameters Or, loss function L(0), spar- sity budget k, step size ne
Output: parameters B2+1
Output: parameters f,41
1: 0, 1T (0.)
1: 0, [T; (0.)
2: 0.+1 = 02 The VO. vO.)))
2: 0.+1 = Or The L(.).)
3: return f.+1
3: return f:-1
Note that both Algorithms and 2 compute the forward pass through the network using the sparsified parameters 0, obtained by setting all but the top-k entries of 0 by magnitude to zero. Algorithms | and 2 differ in the presence of the Vo, term in line 2. This is a d x d diagonal matrix where (Vo,) iff Or,i was not set to zero by the projection IT, and (Vo,) ) 0 otherwise.
At the extreme points of the sharpness parameter Spartan's parameter update procedure reduces to those of IMP and Top-KAST. We can thus view Spartan as a method that generalizes and smoothly interpolates between this pair of algorithms. As 3 + oo, Spartan is equivalent to IMP: this approach sparsifies parameters using top-k magnitude-base forward pass receive zero gradient in the backward pass. At 8 this method again sparsifies parameters by magnitude with binary mask, but unlike IMP, all entries are updated in the backward pass using the gradient of the loss with respect to masked parameters. The Top-KAST update is thus an application of the straight-through gradient method
binary mask, and entries that were masked in the
Spartan reduces Top-KAS KAST:
[5 (5, otherwise See Sec. for
known as lazy projection or dual averaging (DA) optimization further discussion on this point.
41

---

Smooth approximations. Our use of the soft top-k operation is related to prior methods that use the logistic sigmoid function as a differentiable approximation the step function for sparse training 330 36 These approaches similarly regulate the sharpness of the approximation with temperature parameter that scales the input logits to the sigmoid function. A distinctive characteristic of Spartan is that the parameter & directly controls the degree of sparsity of the mask; this differs from previous approaches involving the logistic sigmoid approximation that only indirectly control the degree of sparsity using an L1 penalty term.
Traneforrer--peecffi methods. Several approaches are specialized to transformer based architec- tures. These include structured pruning of pre-trained transformers for NLP tasks [39 2 40 and sparse training methods specifically applied to vision transformers 77 In contrast, Spartan is a general
general-purpose purpose sparse training algorithm that is designed to be agnostic to the model architecture.
3 Sparsity via Regularized Transportation
We begin by motivating our proposed approach. A key advantage of the dual averaging method (Algorithm{) over IMP (Algorithm is that it mitigates the issue of gradient sparsity. In the IMP update, only those parameters that were not masked in the forward pass receive a nonzero gradient, resulting in slow progress at high levels of sparsity. In contrast. dual averaging applies dense update to the parameter vector 0 in each prrration-intutiven parameters that are below the top-k magnitude threshold and are consequently masked to zero. Dual averaging iterates are thus able to more quickly explore the space of possible sparsity patterns.
these dense updates can help "revive"
On the other hand, large variations in the sparsity patterns realized post-projection can lead to instability in optimization, thus hampering the final performance of the model. For instance, Top- KAST empirically benefits from the application of additional sparsity
regularization on the parameters [4 (Sec. 2.3), 577. This issue motivates our use of soft top-k approximation as mechanism for controlling the stability of our training iterates.
Each iteration of Spartan consists of the following two steps (Algorithm 3]: (1) an approximately
k-sparse masking of the model parameters using the soft top-k operator, and (2) dual averaging-based
parameter updates with the set of exactly k-sparse vectors as the feasible set. This procedure aims to combine the advantages of Algorithms [ and within single parameter update scheme. Note that the Spartan parameter update is essentially identical to Algorithm D save for the inclusion of the intermediate soft top-k masking step.
Algorithm 3 Spartan parameter update
Input: parameters 9,, loss function L(0), sparsity budget k, sharpness parameter B, step size T)s Output: parameters f/+1
I: of (0)) = 92 @ softtopk(111 k,)
apply soft masking i> project onto k-sparse set I compute dual averaging update
2: 6, II; (o% (0,)
3: 8i+1 = 0₂ T, Vo (L.)
4: return 0,+1
In the following, we begin by describing the soft top-k masking operation and address the issues of scalability and of applying soft top-k masking to induce structured sparsity. We then discuss how the update procedure outlined in Algorithm 3 can be incorporated within a complete training loop.
3.1 Soft Top-k Masking
Our soft top-k masking scheme is based on the soft top-k operator softtopk(r, k, B8) described by Xie et al.
42 softtopk takes as input a vector v E RR budget parameter k > 0 and a sharpness
parameter 3 2 O, and outputs m [0, 1] smoothed version of the top-k indicator vector. By using the magnitudes |0] of the parameter vector measure of the "value" of each entry, we obtain soft top-k magnitude pruning operator of (-) by masking the entries of the parameter vector with the output of the soft top-k operator:
o (0) = 0 @ (ofttopk(10),k..

---

Algorithm 4 Soft top-k forward pass
Input: values v E IR costs с E R budget E (0, 12 sharpness parameter / 2 0, max. itera- tions T, tolerance initial dual variable Mo
Output: mask m E [0, 1] dual variables E RR v E IRd
1: Bv/c, mmg a
normalize values by costs
2: for = 1,...., , do
3: 1/- loge c log( L. + exp(2 + (t-1))
>> normalize mask entries > normalize sum to be equal to k compute mask
4:
Ht = log k log 2; exp(2; + Pr,i)
5:
m. exp(= + μr 14 log c)
6:
if [e [|eT(nn (m, m4-1)l < c/-'m,111 m then
check stopping criterion
7:
return m4, Pe, Vt
8: return m.r, pr, MT
Algorithm 5 Soft top-k backward pass
Input: gradient w.r.t outputs E Rd, mask m E [0, 11 costs c, budget k, sharpness parameter / Output: gradient w.r.t. imputs
1: a1 = E; g.m, (1 m²), a2 = E, c.m?
2: return 3m (I. m) @ (g/c 01 01 a2))
We introduce a mild generalization the soft top-k operator from Xic et al. 442 by incorporating a strictly positive cost vector c. In particular, require that the output mask [0, 1] satisfies the budget constraint c" m = k. This abstraction is useful for modeling several notions of heterogeneous parameter costs within model; for instance, parameters that repeatedly invoked within network have a higher associated computational cost. As a concrete example, Evci et al. [d observed that the FLOP count of ResNet-50 model at fixed global sparsity can vary by factor of over 2 due to differences in the individual sparsities of convolutional layers with differing output dimensions.
To derive this cost-sensitive soft top-k operator, begin by expressing the mask as the solution to the following linear program (LP):
maximize m MER
v" subject 0 m 3 1d:
cf F'm
Deferring the derivation to the Appendix, we can rewrite this LP as the following optimal transporta- tion (OT) problem:
minimize YEREX
CigYes
subject to Y1, =c 1,Y = [k: 1.c k],
(2)
with the cost matrix C = (-v/c, 0j. We can recover the mask from Y as m, = Yn/e. Now following [2] I. we introduce entropic regularization obtain smoothed version of the top-k operator, with the degree of smoothing controlled by 3:
minimize Y
- -
(3)
subject to Y1, -, 1.Y = [k: cTc-
where = Yij (log Yı, 1) is the entropy of Y. In this standard form, we can efficiently solve this regularized OT problem using the Sinkhorn-Knopp algorithm
Algorithms 4 and 5 describe the forward and backward passes of the resulting cost-sensitive soft top-k operator (see Appendices B and for their derivations). The expression for the gradient in Algorithm 5 follows from Theorem in Xie et al. 1 with some algebraic simplification. We remark that this closed-form backward pass is approximate in the sense that assumes that we obtain the optimal dual variables μ., in the forward pass; in practice, do not encounter issues when using estimates of the dual variables obtained with tolerance e 0.01 and a maximum iteration count of T 100 in Algorithm 2
Scalability. Since we apply soft top-k masking to high-dimensional parameters each iteration of training, it is important to minimize the computational overhead of Sinkhorn iteration in the
5

---

correlation with final mask
spoch-to-epoch correlation
1.0
Pearson correlation
0.6
0.6
Spartan ma 20)
0.4
Spartan (8mo 10)
Spartan (Amax
0.2
Top-KAST
Iter. Pruning
0.0
0
20
40
60 80 100
0
20
40
60
80
100
epoch
epoch
Figure 3: Spartan parameterizes an exploration-exploittioon tradeoff. For each run of ResNet-50 training on ImageNet-1K IK, we plot Pearson correlation coefficients between the sparsity mask at the end of each epoch and the mask obtained at the end training (left), and between sparsity masks at the end of subsequent epochs (right). Kinks in the correlation curves at epochs 20 and 80 are respectively due to the end of sparsity annealing and the start of fine-tuning with fixed masks (see for details on the training schedule).
Section
forward pass. We note that in order to compute the hard top-k projection step in dual averaging (as in Algorithm 3) it is necessary to find the index ig of the kth largest entry of [0]/c. Given the index ik, we the convergence of Sinkhorn iteration. This heuristic is based on the observation that (6 l/c, is the threshold value of the normalized value vector 0l|c in hard top-k masking. Concretely, demonstrate in Sec. 4.3 that our implementation of Spartan incurs a per-iteration runtime overhead of approximately 5% over standard dense ResNet-50 training.
the heuristic initialization eo -3 0, //a, for the dual variable μ to accelerate
Structured sparsity. To implement block structured sparsity, we instantiate one mask variable per block and mask all parameters within that block with the same mask variable. To compute the mask, we use the sum of the magnitudes standard pruning setting with uniform costs across all parameters, the corresponding cost c is simply the total number of entries in the block.
the entries each block as the corresponding value o, In
3.2 Training with Spartan
Training schedule. When training with Spartan, we divide the training process into a warmup phase, an intermediate phase, and a fine-tuning phase. In our experiments, the warmup phase consists of the first 20% of epochs, the intermediate phase the next 60%, and the fine-tuning phase the final 20%. In the warmup phase, we linearly anneal the global sparsity of the model from value of at the start of training to the target level of sparsity. Throughout the intermediate phase, we maintain the target sparsity level, and during the fine-tuning phase, we fix the sparsity mask that used the start of fine-tuning. This training schedule is similar to those used in prior work 37
Exploration vs. exploitation as function of B. From the start of training up the start of the fine-tuning phase, we linearly ramp the sharpness parameter 33 from an initial value of to the final value 3max- We interpret the spectrum of updates parameterized by of an exploration-exploiaatio tradeoff. In Figure g we illustrate this phenomenon by plotting Pearson correlation coefficients between sparsity masks at different stages of training. We observe that iterative magnitude pruning converges on the final mask relatively carly in training, which is indicative of insufficient exploration of different sparsity patterns-consequee lower validation accuracy than the remaining models. The three models trained with Spartan each ramp 3 from 1 the start of training to Bmxx at epoch 80. The value of 9max correlates well both with the Pearson correlation with the final mask, and with the Pearson correlation between masks high-exploration regime
being characteristic
this model achieves
at the end of subsequent epochs. Spartan thus interpolates between of standard dual averaging and the low-exploration regime of iterative magnitude pruning. In this example, the intermediate setting 8max = achieves the highest top-1 validation accuracy of 74.6%, with
at 73.5% and IMP at 68.2%.

---

Table 1: Top-1 accuracies on ImageNet-IK validation set with fully dense training.
Method
Epochs
Accuracy (%)
Method
Epochs
Accuracy (%)
100
77.08
20.11
100
77.78 -.22
ResNet-50
200
77.46
+0000
ViT-B/16
400
77.17 10.04
300
80.06 + 0.11
Table 2: ResNet-50 top-1 accuracies on ImageNet- IK validation set at varying levels of sparsity and epochs of training. When available, we report means and standard deviations over trials.
Sparsity
Method
Epochs
80 %
90%
95%
97.5%
99%
100
75.1 0.05
73.0 + n.00
69.7 =0.77
6
RigL (ERK)
500
77.1 à 0.06
76.4 4 0.05
74.5 = 0.09
25
STR
100
76.19
74.31
70.40
62.84
51.82
17
ProbMask
100
74.68
71.50
66.837
61.07
40
OptG
100
74.28
72.38
62.10
IMP
100
75.3
73.7 0.144
70.6 0.06
Top-KAST
100
76.08
= 0.02
75.13
1.0.3
73.19
2.00
(with PP)
100
76.24
0.07
75.23
1.002
73.25
0.00
(with ERK)
100
76.42
0.03
75.51
0000
(with PP
100
76.76
= 0.08
75.74
0.08
& FRK)
200
77.51
0.03
76.94
+.0.0
300
77.64
0.05
77.16 1.119
100
76.66
75.48
0.155
73.51
20.16
70.23
1.005
Top-KAST
200
77.55
0.08
76.84
+0111
75.20
=011
71.72
#0.0
400
77.75
0.09
77.37
+ 0.07
75.90
0001
72.58
10.13
100
76.89
0.09
76.17
0.10
74.68
0.21
71.95
20.11
63.87 0.12
Spartan
200
77.56
0.19
77.06
0.13
75.92
0.01
73.41
+0.55
65.77
0.03
400
77.57
77.40
76.48
0.20
74.15
10100
66.80 0.03
Repooted accuracics for models closest to the listed sparsity level. *Model tranndd 98% sparsity.
Empirical Evaluation
In this section, we report the results of our sparse training experiments on the ImageNet- dataset with two standard architectures: ResNet-50 [21] and ViT-B/16 [5), consisting of 25.6M and 86.4M parameters respectively. On ResNet-50, we evaluate both unstructured and block structured pruning. We subsequently present empirical studies of the sensitivity of Spartan to the value of 3, the effect of running Spartan without the dual averaging step, and the computational overhead of soft top-k masking.
evaluate only unstructured pruning, whereas on ViT-B/16.
4.1 ImageNet-IK Classification
In our experiments, we run Spartan with the training schedule described in Section 3.2 and evaluate our models on the ImageNet- I dataset with the standard validation split and report means and standard deviations over 3 independent trials. We provide additional details on our experimental setup in the Appendix.
We train
ResNet-50 experimental setup. For consistency with our baselines, use standard data augmenta- tion with horizontal flips and random crops at 224 x 224 resolution. For all Spartan runs, we use 3max 10, which we selected based on models trained at 95% accuracy. We sparsify the parameters of all linear and convolutional layers with global sparsity budget, excluding bias parameters and the parameters of batch normalization layers. Our baselines are iterative magnitude pruning 48), Rigl. with the Erdos-Renyis Kernell (ERK) sparsity distribution [16] Soft Threshold Weight Repa- rameterization (STR) 25). probabilistic masking (ProbMask) 47), OptG 46, and Top-KAST with Powerpropagation and ERK [Z] [7) We additionally re-run the most performant bascline method,
7

---

Table 3: ViT-B/16 top-1 accuracies on ImageNet- IK validation set at 90% sparsity with unstructured sparsity and block structured pruning of attention layers.
Sparsity Structure
Method
Epochs
Unstructured
16 x 16 blocks
32 x 32 blocks
100
78.05 0.07
67.69 10223
64.79 1014
Top-KAST
300
80.86 20.03
74.11 10.20
70.67 1.96
100
79.88
0.16
75.62 0.00
74.50 00223
Spartan
300
81.18
1.0.4
79.12 1.1.6
78.43
0.10
Top-KAST, using reimplementation in our own codebase. For Top-KAST, we exclude the first convolutional layer from pruning (following [37 [6) and we use fully dense backward passes (i.e. with backward sparsity of 0), since this setting achieved the highest accuracy in prior work [24] We use mixed precision training with a batch size of 4096 on NVIDIA A100 GPUs.
VIT experimental setup. We use the ViT-B architecture with 16 X 16 patches at 224 x 224 input resolution. We augment the training data using RandAugment IT., MixUp 45 and CutMix 44
Our ViT models are trained from random initialization, without any pretraining. We set maxx 20 for Spartan with unstructured sparsity, and Smax = 320 and fmax = 640 for Spartan with 16 16 and 32 x 32 blocks respectively. These latter settings are the values of Smax for the unstructured case scaled up by factors of 16 and 32 since each block averages the magnitudes of B entries, we expect the variance of the block values to correspondingly decrease by factor of B², and we thus compensate by scaling up 3 by B. In the block structured case, we exclude the input convolutional layer and the output classification head from pruning since their parameter dimensions are not divisible by the block size. We use mixed precision training with a batch size of 4096 on 16 NVIDIA A100 GPUs across 2 nodes.
Results. Table lists the top-1 validation accuracies achieved by fully dense ResNet-50 and ViT-B/16 models. Table preports validation accuracies for ResNet-50 models at 80%, 90%, 95%, 97.5% and 99% sparsity, and Table reports validation accuracies for ViT at 90% sparsity. In the Appendix, we report additional measurements of inference FLOP costs for our sparse models and the results of experiments with FL.OP-sensiti pruning.
For ResNet-50 models, we find that Spartan outperforms all our baselines across all training durations sparsity levels of 90% and above. In particular, Spartan achieves mean top-1 accuracy within 1% of fully dense training at 95% sparsity. We observe that additional epochs of sparse training consistently improves the final generalization accuracy; in contrast, validation accuracy peaks at 200 epochs for dense training. This trend persists at 800 training epochs, where Spartan achieves 67.18 + 0.13 top-1 accuracy at 99% sparsity. For the Top-KAST baseline, we omit results 99% sparsity due to training instability. We note that the accuracy improvements in the ResNet-50 models relative to Top-KAST are not a result of increased FLOP counts at a given level of sparsity, as can be seen from the FLOP measurements reported in Appendix E
For VIT-B/16 models, Spartan outperforms Top-KAST for both unstructured and block structured pruning. We observe a particularly large improvement over Top-KAST in the block structured case, where Spartan improves absolute validation accuracy by 7.8% for 32 x 32 block structured pruning. For unstructured pruning, Spartan achieves comparable accuracy to SVITE ? (81.2% for Spartan vs. 81.3% for SViTE), but with 30% higher sparsity (90% for Spartan vs. 60% for SViTE). In exchange for 0.6% reduction in accuracy relative to DeiT-B 3S Spartan reduces model storage cost by 10x and the FLOP cost of inference by 7.4x.
4.2 Sensitivity and Ablation Analysis
Figure (left) shows the effect of ablating the dual averaging step Spartan- iie. omitting hard top-k projection in the forward pass over range of settings of fmax for 95% sparse ResNet-50 models trained for 100 epochs. The dashed line shows top-1 accuracy with TAS- For Spartan training without dual averaging, we compute a hard top-k mask at the end of epoch 80, and as with standard Spartan training, we fix this mask until the end of training. In the absence of top-k projection, we find that accuracy increases with increasing 3max up Smax 80; at lower settings of Bmax.
8

---

74
17.5
Sinkhorn only
+ dual caching
top-1 accuracy
overhead
15.0
+ sorting
72
12.5
10.0
70
Spartan (no DA)
Spartan
68
5.0
Bmax
Figure 4: Effect of and dual averaging (left). Top-1 I ImageNet- IK validation accuracies for Spartan with and without dual averaging (DA) and with varying 9max 95% sparsity. Computational overhead (right), Percentage increase in per-iteration wall clock runtime over dense training for Spartan with standard Sinkhorn iteration, Sinkhorn with dual caching, and Sinkhorn with sorting.
the higher approximation error soft masking is detrimental to the final accuracy of the model. In contrast, the use of top-k projection in the forward pass mitigates this mismatch between training and inference and improves the final accuracy of the sparse model.
4.3 Computational Overhead
We evaluate the computational overhead of Spartan over standard dense training by measuring wall clock training times for a ResNet-50 model (Figure Dl right). Our benchmark measures runtime single NVIDIA A 100 GPU over 50 iterations of training with batch size of 256. We use random inputs and labels to avoid incurring data movement costs. We compare three approaches: standard Sinkhorn iteration, Sinkhorn iteration with the dual variable μ initialized using its final value from the previous iteration ("dual caching"), and Sinkhorn iteration with μ initialized using the value - -B l/c.. where we compute the index is the kth largest entry the normalized value vector |0|c by sorting its entries in each iteration ("sorting").
Standard Sinkhorn iteration incurs higher overheads as 3 increases-this is due to the additional iterations required to reach convergence the regularized OT problem more closely approximates the original I... We find that dual caching and sorting both prevent this growth runtime over the range of B values that we tested. In our remaining experiments, we use the simple aatinn-based approach since corresponds to the lowest relative overhead over standard dense training (approximately 5%). We note that this relative overhead decreases as the batch size increases since the cost of computing the mask independent of the batch size. Spartan is therefore best suited to training with large batch sizes since amortize the cost of mask updates over the size of the batch.
Discussion
Extensions. As an alternative to magnitude- based pruning, we may also apply Spartan conjunction with learned value parameters, as in methods such Movement Pruning [35.. In this approach, we would compute masks using a set auxiliary parameters & E g instead of the magnitudes 10]: of (0; Φ) = 0 @ softtopk(), k, 9), and similarly for hard projection. We remark that while this variant requires additional memory during training to store the value parameters, there is additional cost during inference since the sparsity mask is fixed at the end of training.
Limitations. Since Spartan retains a dense parameter vector and computes dense backward passes during training, it incurs higher memory and computational costs each iteration than methods like RigL [T6] that use both sparse parameters and sparse backward passes. Nevertheless, terms total computational or energy cost over full training run, Spartan may remain a competitive option it requires fewer iterations of training to reach given accuracy threshold relative to these methods. However, we do not compare total training FLOP costs in our empirical evaluation.
note that in
further limitation that cost-sensitive pruning with Spartan is only compatible with relatively crude linear cost models of the form cf m. where is the cost vector and is the mask. This restriction

---

arises due to the requirements of the regularized OT formulation used to compute the soft top-k mask. In particular, this limitation precludes the use of cost models involving interaction terms such those arising from spatially coherent sparsity patterns. For example, the cost models that we consider here cannot encode the notion that coherently pruning an entire row or column of a parameter matrix will yield more cost savings than pruning random entries of the matrix.
Optimization with Dual Averaging. As discussed in Section Spartan and Top-KAST are both applications of the dual averaging method 34 41 The empirically observed effectiveness of dual averaging sparse training is reminiscent of its success in a related domain where continuous optimization is complicated by discrete constraints: neural network quantization. In this area, the algorithm M which trains binary neural networks with parameters in {+1]*, d is considered to be a foundational method. As observed by Bai et al. 22, BinaryConnect is itself a special case of dual averaging, and the theoretical underpinnings of this connection to dual averaging were further analyzed by Dockhorn et al. [14]. These parallels with neural network quantization suggest that similar ideas may well apply towards deepening our understanding of sparse learning.
Societal Impacts. Inference with deep neural network models is a computationally intensive process. At present, the total energy footprint associated with serving these models in production is expected to continue growing in tandem with the rising prevalence of large transformer-based architectures in vision and NLP applications. Research towards improving the energy efficiency of deep neural networks is therefore an important counterbalancing force against increasing resource usage by these models. The development of sparse learning algorithms is particularly relevant these efforts, and we expect that the impact of these approaches will further increase sparsity-aware hardware acceleration becomes more widely available.
Conclusions & Future Work
In this work, we describe a sparse learning algorithm that interpolates between two parameter update schemes: standard stochastic gradient updates with hard masking and the dual averaging method. We show that there exists intermediate regime between these two methods that yields improved generalization accuracy for sparse convolutional and transformer models, particularly at higher levels of sparsity. While we have demonstrated promising empirical results with our proposed method, the learning dynamics of stochastic optimization for deep networks under sparsity constraints remains relatively poorly understood from theoretical standpoint. There thus exists ample opportunity for further work towards better understanding sparse learning algorithms, which may in turn inspire future algorithmic advances in this area.
Acknowledgements and Disclosure of Funding
We are grateful to Trevor Gale for his feedback on this work. We also thank our anonymous reviewers for their valuable suggestions on improving the manuscript. This work was funded by Meta.
References
[1] Kambiz Azarian, Yash Bhalgat, Jinwon Lee, and Tijmen Blankevoort. Learned threshold pruning. arXiv preprint arXiv:2003.0007. 2020.
[2] Yu Bai, Yu-Xiang Wang, and Edo Liberty. ProxQuant: Quantized Neural Networks via Proximal Operators. In International Conference on Learning Representations, 2018.
[3] Guillaume Bellec, David Kappel, Wolfgang Maass, and Robert Legenstein. Deep Rewiring: Training very sparse deep networks. In International Conference on Learning Representations, 2018.
[4] Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré. Iterative Bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing. 37(2), 2015.
[5] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.344 2013.
10

---

[6] Arnay Chavan, Zhiqiang Shen. Zhuang Liu, Zechun Liu, Kwang- Ting Cheng, and Eric Xing. Vision transformer slimming: Multi-dimension searching continuous optimization space. arXiv preprini arXiv:2201.008// 2022.
[7] Tianlong Chen, Yu Cheng, Zhe Gan, Lu Yuan, Lei Zhang, and Zhangyang Wang. Chasing sparsity in vision transformers: An end-to-end exploration. In Advances in Neural Information Processing Systems, 2021.
[8] Maxwell D Collins and Pushmeet Kohli. Memory bounded deep convolutional networks. arXiv preprint 4424v.1412.1442, 2014.
[9] Matthieu Courbariaux, Yoshua Bengio, and Jean-Pierre David. BinaryConnect: Training deep neural networks with binary weights during propagations. In Advances in Neural Information Processing Systems, 2015.
[10] Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le. Randaugment: Practical automated data augmentation with reduced search space. In IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, 2020.
[11] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, 2013.
[12] Marco Cuturi, Olivier Teboul, and Jean- Philipph Vert. Differentiable ranking and sorting using optimal transport. In Advances in Neural Information Processing Systems, 2019.
[13] Tim Dettmers and Luke Zettlemoyer. Sparse networks from scratch: Faster training without losing performance. arXiv preprint arXiv:1907.04844 2019..
[14] Tim Dockhorn, Yaoliang Yu, Eyyüb Sari, Mahdi Zolnouri, and Vahid Partovi Nia. Demystifying and generalizing binaryconnect. 2021.
[15] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[16] Utku Evci, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen. Rigging the lottery: Making all tickets winners. In International Conference on Machine Learning. 2020.
[17] Trevor Gale, Erich Elsen, and Sara Hooker, The state of sparsity in deep neural networks. arXiv e-prints, weX.iv1102.095555 2019.
[18] Scott Gray, Alec Radford, and Diederik P Kingma. GPU kernels for block-sparse weights. arXiv preprint arXiv: 1711.09224 2017.
[19] Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. Advances in Neural Information Processing Systems, 2015.
[20] Babak Hassibi and David Stork. Second order derivatives for network pruning: Optimal brain surgeon. Advances in Neural Information Processing Systems, 1992.
[21] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recognition. IEEE Conference on Computer Vision and Pattern Recognition, 2016.
[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. Journal of Machine Learning Research, 2021.
[23] Steven A Janowsky. Pruning versus clipping in neural networks. Physical Review A, 1989.
[24] Siddhant Jayakumar, Razvan Pascanu, Jack Rae, Simon Osindero, and Erich Elsen. Top-KAST: Top-K Always Sparse Training. Advances in Neural Information Processing Systems, 2020.
11

---

[25] Aditya Kusupati, Vivek Ramanujan, Raghav Somani, Mitchell Wortsman, Prateek Jain, Sham Kakade, and Ali Farhadi. Soft threshold weight reparameterization for learnable sparsity. In International Conference on Machine Learning, 2020.
[26] François Lagunas, Ella Charlaix, Victor Sanh, and Alexander M Rush. Block pruning for faster transformers. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, 2021.
[27] Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. Advances in Neural Information Processing Systems, 2, 1989.
[28] Namhoon Lee, Thalaiyasingam Ajanthan, and Philip Torr SNIP: Single-shot Network Pruning based on Connection Sensitivity. In International Conference on Learning Representations. 2018.
(29] C Louizos, M Welling, and DP Kingma. Learning sparse neural networks through Lo regular- ization. In International Conference Learning Representations, 2018.
[30] Jian-Hao Luo and Jianxin Wu. Autopruner: An end-to-end trainable filter pruning method for efficient deep model inference. Pattern Recognition, 2020.
[31] Decebal Constantin Mocanu, Elena Mocanu, Peter Stone, Phuong H Nguyen, Madeleine Gibescu, and Antonio Liotta. Scalable training of artificial neural networks with adaptive sparse connectivity inspired by network science. Nature Communications, 2018.
[32] Hesham Mostafa and Xin Wang. Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization. In International Conference on Machine Learning. 2019.
[33] Michael C Mozer and Paul Smolensky. Skeletonization: A technique for trimming the fat from a network via relevance assessment. In Advances in Neural Information Processing Systems, 1988.
[34] Yurii Nesterov. Primal-dual subgradient methods for convex problems. Mathematical Program- ming, 2009.
[35] Victor Sanh, Thomas Wolf, and Alexander Rush. Movement pruning: Adaptive sparsity by fine-tuning. In Advances in Neural Information Processing Systems, 2020.
[36] Pedro Savarese, Hugo Silva, and Michael Maire. Winning the lottery with continuous sparsifi- cation. Advances in Neural Information Processing Systems, 2020.
[37] Jonathan Schwarz, Siddhant Jayakumar, Razvan Pascanu, Peter Latham, and Yee Teh. Power- propagation: A sparsity inducing weight reparameterisation. In Advances in Neural Information Processing Systems, 2021.
[38] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning, 2021.
[39] Ziheng Wang, Jeremy Wohlwend, and Tao Lei. Structured pruning of large language models. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2020.
[40] Mengzhou Xia, Zexuan Zhong, and Danqi Chen. Structured pruning learns compact and accurate models. arXiv preprint arXiv:2204.0000 00408, 2022.
[41] Lin Xiao. Dual averaging method for regularized stochastic learning and online optimization. Advances in Neural Information Processing Systems, 2009.
[42] Yujia Xie, Hanjun Dai, Minshuo Chen, Bo Dai, Tuo Zhao, Hongyuan Zha Wei Wei, and Tomas Pfister. Differentiable top-k with optimal transport. In Advances Neural Information Processing Systems, 2020.
12

---

[43] Dong Yu, Frank Seide, Gang Li, and Li Deng. Exploiting sparseness in deep neural networks for large vocabulary speech recognition. IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2012.
[44] Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. In IEEE/CVF International Conference on Computer Vision, 2019.
[45] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018.
[46] Yuxin Zhang, Mingbao Lin, Mengzhao Chen, Zihan Xu, Fei Chao, Yunhan Shen, Ke Li, Yongjian Wu, and Rongrong Ji. Optimizing gradient-driven criteria in network sparsity: Gradient is all you need. arXiv preprint arXiv:2201.12822 2022.
[47] Xiao Zhou, Weizhong Zhang, Hang Xu, and Tong Zhang. Effective sparsification of neural networks with global sparsity constraint. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, 2021.
[48] Michael Zhu and Suyog Gupta. To Prune, or Not to Prune: Exploring the Efficacy of Pruning for Model Compression. International Conference on Learning Representations, 2018.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See
Section 5
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Section and the Appendix.
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [Yes]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any assets either in the supplemental material or as URL? [Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you're rou're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "yhLedqzE9YA", "reviewer": "Reviewer_4MMR", "summary": "The paper presents Spartan - a \"dense-to-sparse\" algorithm to train sparse neural networks in which the sparsity of parameters can be enforced directly by a predetermined budget.\n\nThe paper has $3$ main contributions:\n- By incorporating the soft top-$k$ operator in the parameter update and by introducing the sharpness parameter, Spartan allows to interpolate between Iterative magnitude pruning and Top-$K$ Always sparse training.\n\n- Spartan shows very competitive and consistent performance over the existing methods and also the fully dense training.\n\n- The authors also study the effect of the sharpness parameter on the accuracy, as a trade-off between exploration and exploitation.", "questions": "- The problem $(2)$ is indeed an optimal transport problem, but it is not formulated in the standard form, i.e. what is the cost matrix, the admissible coupling (transport plan), the two marginal constraints? This makes it difficult to calculate and justify the update of Sinkhorn iterations.\n\n- Please specify which kind of entropic regularization you are using. More precisely, $[4]$ use KL divergence but $[11]$ use the negative entropy. The two types are closely related, but mathematically, they result in a bit different Sinkhorn iterations. Based on the form of the mask $m$, I guess you are using the negative entropy, but not the KL divergence.\n\n- In the Appendix, please consider to provide the details of the:\n   + Soft top-$k$ operator in $[41]$. It is the core component of Spartan, so I believe it deserves some brief description.\n   + Calculation of the backward pass. Even though the authors cite theorem $3$ in $[41]$, it is complicated for outsider, thus difficult to justify if the calculation in the Algorithm 5 is correct or not. \n  + Calculation of the forward pass. Please also specify the entropic-regularization problem to which the authors apply.\n\n- What is the dimension $d$ of the input vectors in the applications?\n\n- How do the authors come up with the specific form of the initialization $-\\beta |\\theta_{i_k}| / c_{i_k}$?\nDoes it come from the empirical observation, or some heuristic?\n\n- I see the definition of the projection $\\Pi_k$ but wonder how is it implemented in practice? Is it simply keeping the largest $k$ elements and setting all others to zero?", "limitations": "The authors do discuss the limitations and potential negative societal impact of their work.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Amongst core components of Spartan, the algorithm 4 leaves me much doubt.\n\n- It appears to me that the algorithm 4 is nothing but the Sinkhorn algorithm. In that case, it is not clear to me the motivation of the stopping criteria in the line 7. In practice (and also in the optimal transport community), it is enough to control the marginal violation (see Remark 4.6 in [CompOT]).\n  \n - The calculation of the Sinkhorn iterations does not look right to me. Here, I present my own derivation of the Sinkhorn update. First, let us define $Y := [y, y'] \\in \\mathbb R^{d \\times 2}$ by stacking two vectors in $\\mathbb R^d$ and $C := [\\frac{V}{c}, 0_d] \\in \\mathbb R^{d \\times 2}$, where $0_d$ is the zero vector in $\\mathbb R^d$. Then clearly, $\\langle C, Y \\rangle = (\\frac{V}{c})^T y$, where $\\langle C, Y \\rangle = \\sum_{ij} C_{ij} Y_{ij}$ denotes the Frobenius product. We can rewrite the problem $(2)$ as \n\n$$\n\\min_{Y \\in \\mathbb R^{d \\times 2}_+} \\langle -C, Y \\rangle \\text{ subject to } Y 1_2 = c \\text{ and } Y^T 1_d = (k, 1_d^T c - k)^T.\n$$\n\nThe two marginals are not yet normalized. So, denote $s = 1_d^T c$, then define $c^* := \\frac{c}{s}, C^* := sC = [\\frac{V}{c^*}, 0_d]$ and $k^* := \\frac{1}{s} (k, 1_d^T c - k)^T = ( \\frac{k}{s}, 1 - \\frac{k}{s})^T$. The above problem is equivalent to\n\n$$\n\\min_{Y \\in \\mathbb R^{d \\times 2}_+} \\langle -C^*, Y \\rangle \\text{ subject to } Y 1_2 = c^* \\text{ and } Y^T 1_d = k^*.\n$$\n\nNow, this is a properly defined optimal transport problem. If we define the entropic regularization problem as: for $\\beta > 0$,\n\n$$\n\\min_{Y \\in \\mathbb R^{d \\times 2}_+} \\langle -C^*, Y \\rangle + \\frac{1}{\\beta} H(Y) \\text{ subject to } Y 1_2 = c^* \\text{ and } Y^T 1_d = k^*.\n$$\n\nwhere $H(Y) = \\sum_{ij} Y_{ij} (\\log Y_{ij} - 1)$. The Sinkhorn update now reads: for dual vectors $u \\in \\mathbb R^2$ and $v \\in \\mathbb R^d$, we have: $v = \\log c^* - \\log \\sum_j \\exp(\\beta C^*_{\\cdot, j} + u_j)$ and \n$u = \\log k^* - \\log \\sum_i \\exp(\\beta C^*_{i, \\cdot} + v_i)$, and the optimal plan $Y$ is given by: $Y_{ij} = \\exp(v_i + u_j + \\beta C^*_{ij})$. \n\nComparing $v, u, Y$ to $\\nu, \\mu, m$ from the lines 4,5,6 in the Algorithm 4, respectively, I do not see why they are equivalent.\n\nReference:\n\n[CompOT] Gabriel Peyré and Marco Cuturi (2019), \"Computational Optimal Transport: With Applications to Data Science\", Foundations and Trends in Machine Learning: Vol. 11: No. 5-6, pp 355-607.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "sfGhMusYpZT", "reviewer": "Reviewer_1csR", "summary": "This article introduces a new method, Spartan, which allows training neural networks with sparsity constraints. Spartan belongs to the \"dense-to-sparse\" family of methods: it maintains a dense parameter during training and makes it sparse little by little. Spartan controls the level of sparsity thanks to a parameter $\\beta$ which performs an interpolation between two popular methods to train sparse neural networks (Top-KAST ref [23] of the paper at $\\beta = 0$ and IMP [47] at $\\beta = +\\infty$). The central idea is to use, during training, a \"mask\" on the network parameters based on a soft-topk operator (implemented via regularized optimal transport).\n", "questions": "1° Relations with Top-KAST: \n\nIt is said in the introduction that the Top-KAST method [47] is a \"dense-to-sparse\" method, i.e. it maintains in memory a more or less dense parameter during the iterations. This point seems to me quite surprising because in [47] the authors specify that the ultimate goal of Top-KAST is to maintain sparsity during training (\"In this work we propose Top-KAST, a method that preserves constant sparsity throughout training (in both the forward and backward-passes)\"). Can the authors clarify this point?\n\n\n2° About notations:\n\nThere is an important notation which comes up regularly but which is not explained and difficult to catch. In Algorithm 1 and Algorithm 3 at the backward pass the gradient with respect to a vector is used and multiplied with the gradient of the loss. \n\nFor example in Algorithm 1 line 2 $\\theta_{t+1} = \\theta_t - \\eta_t \\nabla \\tilde{\\theta}_t \\nabla L(\\tilde{\\theta}_t)$. I don't think this notation is really standard. If I read between the lines, I understand this notation as $\\nabla \\tilde{\\theta}_t=$ \"support of $\\tilde{\\theta}_t$\" (i.e. $0$ or $1$) and the multiplication with $\\nabla L(\\tilde{\\theta}_t)$ as being a point-to-point multiplication by the support mask. More precisely $\\nabla \\tilde{\\theta}_t \\nabla L(\\tilde{\\theta}_t) = \\nabla L(\\tilde{\\theta}_t)$ if $\\tilde{\\theta}_t \\neq 0$ and $0$ otherwise. \n\nIs this correct? I think this point deserves more details because it is difficult to understand  what is done in practice.\n\n3° About Algorithm 3:\n\nI have trouble understanding the interest of line 2 in algorithm 3: after finding $\\sigma^{\\beta}_k(\\theta_t)$ by a soft-topk operator on $\\theta_t$ (and as shown in figure 2) an additional step of projection onto $k$-sparse set is added $\\tilde{\\theta_t} = \\Pi_k(\\sigma^{\\beta}_k(\\theta_t))$ and the gradient of the loss is computed with respect to this $\\tilde{\\theta_t}$. It seems to me that the two are quite redundant. \n\nIt is argued in the article that this is to \"mitigate the issue of gradient sparsity\" and an ablation study is conducted on this point in section 4.2 to justify this step. I find these arguments quite empirical and it is difficult to feel the underlying intuition. Do the authors have some theoretical points to support them? For example, can we show for a simple function that the optimization \"goes wrong\" if we do not add this top-k step?\n\nAlso, it seems to me that if this step was not there, then when $\\beta = 0$ Spartan would no longer be equivalent to Top-KAST: am I correct ?\n\n4° Experimental section:\n\nThe experiments are in my opinion convincing. The idea of alternating an exploration and exploitation phase is sounded and very well illustrated by Figure 3. The experiments show well the interest of Spartan from the point of view of performance vs sparsity. Although Spartan does not outperform Top-KAST much on ResNet-50 (of the order of the variance in most cases) the ViT experiment shows the opposite and Spartan has very good performance especially in the structured case. \n\nThe computational overhead compared to dense training seems reasonable to me. On this point I think it is also important to put the Top-KAST one in perspective on the same figure: how do they compare? \n\n\n5° Small remarks: \n- The \"sparsity mask\" is never clearly defined. We understand that it is $m \\in [0,1]^{d}$ but I think it is important to say it clearly. \n- Figure 1 (c) doesn't look good in pdf because I think it's a .png or .jpg : it could be interesting to recompile the figure in pdf so that it's more readable. \n- I think that the link between equation (2) and the optimal transport could be more detailed (e.g. in appendix). Having the formulation \"min on a coupling\" could help understanding for readers who are not familiar with optimal transport. \n- What are \"crude linear cost models\" (Section 5: discussion) ?\n\nTo conclude I would like to note that it is quite rare to have a \"Limitations\" paragraph and I find this a good approach that I welcome (especially since the remarks made by the authors are good remarks that also help in understanding the method).\n\n\n------ AFTER REBUTTAL ------\n\nI thank the authors for their response which addressed most of my concerns, therefore I change my score to 8.", "limitations": "NA", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall I find this article very well written, especially the introduction, the background and the presentation of the contributions. The authors have made a notable effort of pedagogy which greatly facilitates the reading and the understanding of the article. In particular, figures 1 and 2 are very clear, didactic and I thank the authors for this effort. On the other hand, the proposed method seems to me quite relevant, and well supported by many experiments. The fact that Spartan is linked to and generalizes IMP and Top-KAST strengthens the contribution. The method seems to me quite new: it skillfully combines the soft-topk operator of Xie et al [41] with dual averaging ideas. \n\nHowever, there are a few points that are, in my opinion, unclear and would need to be clarified. These prevent me from putting a very clear accept on this article, but I would gladly change my mind depending on the authors' answers.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "mU3xynnpOKx", "reviewer": "Reviewer_X3Jw", "summary": "This paper proposes a new method for pruning DNN parameters. The proposed method uses a differential top-K operator so that the proposed method interpolates between IMP and Top-KAST and balances exploitation and exploration well.", "questions": "How was $-\\beta |\\theta_{i_k}| / c_{i_k}$ (i.e., the initial value) derived?", "limitations": "The authors mentioned the limitations and social impacts well in Section 5.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n\n- The proposed method is simple and easy to use.\n- The proposed method outperforms existing methods.\n- The experiments include the modern ViT architecture.\n\n# Weaknesses\n\n- While I grasp the intuitive idea of the proposed method, e.g., the proposed method interpolates IMP and Top-KAST, I don't find the rationale behind the specific design of the proposed update rule. More detailed discussions and hopefully principled interpretation, e.g., the update of the proposed method can be interpreted as a descent step of some objective function, should be given.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "cxZrhQIVPMB", "reviewer": "Reviewer_HMyH", "summary": "This paper proposes a gradual neural network pruning algorithm which utilizes soft top-k mechanism for the weight update of masked/unmasked parameters. In particular, authors adopt the differentiable top-k mechanism of Xie et al. (2020), which smooths the top-k problem in a principled manner via viewing top-k as an optimal transport problem and adds the entropic regularizer to make the allocation smooth. The model applies both soft and hard top-k for the evaluation of the model, and only soft for the gradient computation. When combined with an appropriate scheduling of sparsity and entropic regularization intensity, the method achieves SOTA-like performance on ResNet-50 and ViT-B trained on ImageNet-1K.", "questions": "- Up to my knowledge, Top-KAST has two different types of sparsity; forward sparsity and backward sparsity. I wonder what backward sparsity the authors used for reproducing the top-KAST results. Are you using fully dense backwards?\n\n- I am curious where the inference FLOP benefits of Spartan comes from. Could authors provide a layerwise sparsity plot? Also, how would the inference FLOPs of top-KAST be if one uses ERK layerwise sparsities?\n\n- I believe that the inference FLOP gain is a very strong benefit of the proposed method. I suggest the authors to move them to the main text if possible.", "limitations": "Authors adequately addressed these, in my humble opinion.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "__Strength.__ First of all, I must say that the paper is very clearly written; I like how algorithms 1,2 are placed, which makes algorithm 3 very sensible and straightforward thing to do. Also, the \"dual averaging\"-based perspective toward the top-KAST update was quite new and fresh to me (although I would have appreciated slightly more detailed discussion on what dual averaging is), and I think this could be inspirational for many future works. Another big strength of the proposed method is its inference efficiency, having much lower (theoretical) inference FLOP than top-KAST.\n\n__Weakness.__ My main concern is on the _practicality_ of the proposed method. Although the central idea of the algorithm is interesting and it seems like it gives slight boost in terms of the final sparsity-to-performance tradeoff, I am slightly worried about the practicality of the method. Spartan has additional hyperparameter $\\beta$ which needs to be carefully selected and scheduled for the performance gain over the existing methods (which may require many training runs to be tuned). Also, unlike magnitude pruning (whose active parameters decrease as the training proceeds) and top-KAST (which has decreased number of peak #parameters via backward sparsity), Spartan has all parameters active throughout the training, having very little benefit in terms of resource used for training (training FLOPs, memory).\n\nHere are some nitpicks (or much minor-er concerns):\n- In line 60, the paper mentions that \"We show that Spartan interpolates between ...\" but I do not think this point has been rigorously shown. In fact, a key idea of top-KAST is using backward sparsity, which is slightly less than forward sparsity but not dense, and Spartan uses the fully dense(-but-rescaled) backward.\n- In line 137, authors seem to argue that the entropic regularization is necessary in order to efficiently solve Eq.(2) (or equivalently Eq. (1)). I do not think this is necessarily true. In fact, the solution of Eq.(1) is actually quite easy to get; for uniform c, one would simply need to perform the magnitude-based pruning (or cost-rescaled magnitude whenever c is non-uniform). Thus, I view the entropic regularization as an artificially introduced (but nevertheless neat) artifact to make the allocation softer, instead of a requirement.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept