Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
review_content: 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.
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.

__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).

Here are some nitpicks (or much minor-er concerns):
- 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.
- 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.
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?

- 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?

- 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.
review_title: null

final_score: rating: 8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.
review_content: 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).

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. 

However, 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.
questions: 1° Relations with Top-KAST: 

It 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?


2° About notations:

There 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. 

For 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. 

Is this correct? I think this point deserves more details because it is difficult to understand  what is done in practice.

3° About Algorithm 3:

I 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. 

It 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?

Also, 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 ?

4° Experimental section:

The 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. 

The 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? 


5° Small remarks: 
- 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. 
- 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. 
- 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. 
- What are "crude linear cost models" (Section 5: discussion) ?

To 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).


------ AFTER REBUTTAL ------

I thank the authors for their response which addressed most of my concerns, therefore I change my score to 8.
limitations: NA
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: 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.
strengths_and_weaknesses: # Strengths

- The proposed method is simple and easy to use.
- The proposed method outperforms existing methods.
- The experiments include the modern ViT architecture.

# Weaknesses

- 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.
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.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: 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.

The paper has $3$ main contributions:
- 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.

- Spartan shows very competitive and consistent performance over the existing methods and also the fully dense training.

- The authors also study the effect of the sharpness parameter on the accuracy, as a trade-off between exploration and exploitation.
strengths_and_weaknesses: Amongst core components of Spartan, the algorithm 4 leaves me much doubt.

- 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]).
  
 - 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 

$$
\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.
$$

The 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

$$
\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^*.
$$

Now, this is a properly defined optimal transport problem. If we define the entropic regularization problem as: for $\beta > 0$,

$$
\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^*.
$$

where $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 
$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})$. 

Comparing $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.

Reference:

[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.
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.

- 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.

- In the Appendix, please consider to provide the details of the:
   + Soft top-$k$ operator in $[41]$. It is the core component of Spartan, so I believe it deserves some brief description.
   + 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. 
  + Calculation of the forward pass. Please also specify the entropic-regularization problem to which the authors apply.

- What is the dimension $d$ of the input vectors in the applications?

- How do the authors come up with the specific form of the initialization $-\beta |\theta_{i_k}| / c_{i_k}$?
Does it come from the empirical observation, or some heuristic?

- 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.
review_title: null

Meta-review: metareview: All reviewers agree that the paper is clearly written and proposes an algorithm which is both novel and efficient.

The rebuttal has clarified a number of points, and thereby adressed most of the concerns of the reviewers. The authors are thus strongly encouraged to take into account the comments of the reviewers and to add some of the clarifications that they provided in this discussion in the paper and supplementary materials.

confidence: Certain
Decision: Accept