Paper: SWAP: Sparse Entropic Wasserstein Regression for Robust Network Pruning (ICLR-2024)

Reviewer comment (Evaluation):
Is it possible to show an analysis in term of computational expense incurred, as the scale of the model increases?

Author response:
We thank the reviewer for this good suggestion. We have added a new section (A.8 Algorithm Scalability) in the appendix of the updated manuscript. The content is as follows.

*In this section, we analyze the scalability of Algorithm 1 with respect to the number of model parameters involved in pruning. 
The result is shown in Figure 8. It can be observed that the execution time scales linearly with the number of pruning parameters. The extra cost of solving the OT is marginal. Theoretically, one could derive this linear scalability by inspecting Line 8, which is the most time-consuming step. The required operations can be decomposed as a sequential operations: matrix-vector multiplications $\vec{Gw}$ and $\vec{G\bar{w}}$ in $O(np)$, the vector subtraction $\vec{Gw}-\vec{G\bar{w}}$ in $O(n)$, a matrix-matrix multiplication $\vec{\Pi}(\vec{Gw}-\vec{G\bar{w}})$ in $O(n^2)$, a matrix transposition and multiplication with $\vec{G}$ in $O(np)$, and the vector subtraction and scalar multiplication $\lambda(\vec{w}-\vec{\bar{w}})$ in $O(p)$. Thus, the overall complexity is $O(np)$, with $p$ significantly larger than $n$ practically. Given fixed fisher sample size $n$, the loop of Algorithm 1 scales linearly with the number of pruning parameters $p$.*

Reviewer comment (Presentation):
Suggestion to change some of the cited references to be included in parentheses for better presentation. Eg in page 3 Computing the distance between .... (Nadjahi et al., 2020).

Author response:
Thanks for the suggestion! We have changed the cited references to be included in parentheses as you suggested.

Reviewer comment (Experiments):
In the introduction, GPT-4 is used an example of model with substantial size and complexity, Is it possible to show case such a pruning in a language model?

Author response:
In principle, the proposed method can be applied to prune a language model. Namely, the proposed method has no assumptions on the model’s architecture. The formulations of (3) (i.e. the LR formulation) and (5) (i.e. the sparse EWR formulation) originate from approximating the Taylor Expansion of the loss function shown in (1). Both the LR and the EWR methods can be used as long as the loss function is differentiable, which is the common case for the majority class of deep learning models including language models that are based on neural networks and transformers. 

In the revised version, we conducted an extra experiment on a larger model ResNet50, shown in Table 1. Our conclusions in the original paper still hold for the newly obtained results. We hence anticipate that the proposed method would be applicable for larger models including transformer architectures. This is definitely a direction worth exploring with significant efforts.

Reviewer comment (Theory):
Is it possible to conduct a convergence analysis for algorithm 1? If it is difficult, does the main challenge lie in the simultaneous optimization of $\Pi_t$?

Author response:
Yes, we agree with the reviewer that the convergence analysis for problem (5) is difficult, and the main challenge lie in the simultaneous optimization of $\Pi_t$ and $\vec{w}$. Given fixed sparsity $k$, the convergence analysis of using block coordinated descent (i.e. alternating optimization) in solving (5) essentially lies on proving the block coordinate descent together with the iterative hard thresholding in Line 10 leads to convergence. To the best of our knowledge, this question remains open.

The most relevant work we’ve identified so far are:
- [1] Peng, Liangzu, and René Vidal. "Block Coordinate Descent on Smooth Manifolds." arXiv preprint arXiv:2305.14744 (2023).
- [2] Eisenmann, Henrik, et al. "Riemannian thresholding methods for row-sparse and low-rank matrix recovery." Numerical Algorithms 93.2 (2023): 669-693.

Literature [1] shows that convergence can be guaranteed with Riemannian gradient descent. However it does not cover the case of using IHT. Literature [2] proposed a Riemannian version of IHT with local convergence. Whether integrating the Riemannian IHT into the block coordinate descent on manifold leads to local convergence remains open. This is definitely a direction worth exploring for future work.

However, we would like to point out that for our particular problem in network pruning, the convergence of an algorithm solving problem (5) with given $k$ is not as relevant. The main reason is that in Algorithm 1 (or any gradual pruning strategies developed from (5)), the number of iteration steps is fixed. In each iteration, the target sparsity is different (i.e. $k$ is different) and we are solving different problems with one SGD iteration. 

The computational cost in Algorithm 1 is mainly incurred by Line 8 (the stochastic gradient descent step). Notably, the cost of computing the OT plan $\Pi_t$ is marginal. We added a section (A.8 Algorithm Scalability) in the Appendix. The content is as follows.

*The result is shown in Figure 8. It can be observed that the execution time scales linearly with the number of pruning parameters. The extra cost of solving the OT is marginal. Theoretically, one could derive this ... The required operations can be decomposed as a sequential operations: matrix-vector multiplications $\vec{Gw}$ and $\vec{G\bar{w}}$ in $O(np)$, the vector subtraction $\vec{Gw}-\vec{G\bar{w}}$ in $O(n)$, a matrix-matrix multiplication $\vec{\Pi}(\vec{Gw}-\vec{G\bar{w}})$ in $O(n^2)$, a matrix transposition and multiplication with $\vec{G}$ in $O(np)$, and the vector subtraction and scalar multiplication $\lambda(\vec{w}-\vec{\bar{w}})$ in $O(p)$. Thus, the overall complexity is $O(np)$, with $p$ significantly larger than $n$ practically. Given fixed fisher sample size $n$, the loop of Algorithm 1 scales linearly with the number of pruning parameters $p$.*

Reviewer comment (Experiments):
It would be more interesting to compare with more baselines, in addition to the LR method.

Author response:
We thank the reviewer for pointing out the weakness, however the choice of comparing to mainly compare to the LR method is twofold. The main motivation is that we would like to investigate the effect of the optimal transport that we introduced to the problem formulation. Since LR can be seen as a naive case of EWR, it is natural to use LR as the baseline for EWR to see how effective the Wasserstein regression would be for combating the noise. The second reason is that LR is the current state of the art that outperforms other methods within the area of post-training network pruning , and our proposed method will not get worse even for noise-free scenarios. It will converge to LR with a large number of clean training data and therefore also outperform other methods.  

Moreover, we have also listed the performance of other methods (MP, WF, and CBS) in Table 1. Since the results we reproduce for LR are consistent with the original paper, we believe the results for the other methods should be consistent as well.

Reviewer comment (Experiments):
The experiments are weak, without test on state-of-the-art architectures like transformers, or larger models like ResNet50, making it suspicious that the proposed approach does not work well on larger model sizes.

Author response:
To make the experiment more comprehensive, we have added the pruning performance benchmarking for ResNet50 in Table 1. One could see that EWR outperforms LR for both ResNet20 and ResNet50 on the same dataset CIFAR10. Correspondingly, we revised our results and analysis for Table 1:
The advantages of EWR over the others are reflected by the three more challenging tasks: ResNet20 and ResNet50 on CIFAR10 and MobileNetV1 on ImageNet, especially in the presence of noisy gradients.

Reviewer comment (Experiments):
Why using 0.84, 0.74, 0.75, 0.63 ... values in Table 2 and 3? This is very uncommon and even in Table 1, the results are following traditional sparsity levels.

Author response:
regarding the sparsity levels in Tables 2 and 3, in contrast to Table 1, is as follows.

In Table 1, all sparsity levels are "target sparsity" and all the pruning algorithms are fine-tuned on the training set after every pruning. This is to make our implemented LR and EWR align with the performance benchmarking results of MP, WF, and CBS, such that all algorithms are comparable to each other. The “Sparsity” column is the target sparsity, such that every line in Table 1 is the final pruning result. 
Tables 2 and 3 showcase the performance during one pruning process, where the loss and the accuracy values are taken immediately after pruning (without imposing any fine-tuning). The sparsity levels are intermediate levels. This was clarified in the appendix A.5 of the original version:
“In Table 2, Table 3, and additional results provided in the Appendix, sparsity is set using a linear gradual pruning strategy, progressing from 0 to 0.75 or 0.95 across ten distinct stages for MLPNet and ResNet, and from 0 to 0.75 across eight distinct stages for MobileNetV1.”

Using python code as an example:
```{python} 
target_sparsity = 0.95
pruning_stage = np.array([0,1,2,3,4,5,6,7,8,9]) 
```
And remark that in stage 0 we record the target weights $\bar{w}$ and no actual pruning is conducted:
```{python} 
total_stages = 10 - 1
```
Hence:
```{python} 
(pruning_stage / total_stages * target_sparsity).round(2)
```
yields the sparsity levels during this pruning process as
```{python}
array([0.  , 0.11, 0.21, 0.32, 0.42, 0.53, 0.63, 0.74, 0.84, 0.95])
```
To clarify it further, we added the following content after the above cited paragraph:

*The values are computed with linear incremental steps, from zero to the target sparsity.*

Additionally, we further clarified this in the captions of Tables 2, 3, 6, and 7.

Reviewer comment (Writing):
The writing and clarity of the paper could be significantly improved, as many phrases and statements are unclear, beginning with the abstract sentence: "_This study unveils a cutting-edge technique for neural network pruning that judiciously addresses noisy gradients during the computation of the empirical Fisher Information Matrix (FIM)._"

Author response:
We thank the review for the comprehensive review and acknowledgement of our paper’s work. We feel sorry for the unclear phrases and statements. We have done polishing and rephrasing in the recent submitted version. Careful proofreading will be conducted further.

Reviewer comment (Theory):
Some choices in the method could be better motivated, for example, the pruning step (step 10 of Algorithm 1) as a projection onto the l-0 norm ball. While the authors reference relevant literature, the analysis of the optimization problem lies outside the scope of the paper.

Author response:
We agree with the reviewer that the pruning step in Algorithm 1, especially the iterative hard-thresholding (IHT), should be better motivated and explained. Towards this goal, we added the following text after two paragraphs below equation (14):  
Line 10 in Algorithm 1 employs the IHT method that is commonly used in sparse learning, which together with line 9, forms a projected gradient descent algorithm. It finds a sparse representation of the updated gradient in line 9. Intuitively, IHT keeps the dominant model weights and essentially preserves the most impactful aspects of the to-be-trimmed model.

Reviewer comment (Writing):
The manuscript contains many typos and grammatical errors. Although readable in its current form, the reviewer suggests a careful review of the manuscript to improve writing. Examples provided include:  
- Now let’s comparing the covariance between…  
- Intuitively, a large dataset of high-quality training samples diminishes concerns over gradient noise, making the empirical fisher a close approximation to the true fisher.  
- Importantly, this seamless trade-off eludes the combination of the Euclidean distance with gradient averaging.  
- Remark that LR is a special case of EWR…

Author response:
We thank the review for the comprehensive review and acknowledgement of our paper’s work. We feel sorry for the unclear phrases and statements. We have done polishing and rephrasing in the recent submitted version. Careful proofreading will be conducted further.

Reviewer comment (Experiments):
It is not obvious what kind of assumption is made regarding the noise, and it is unclear what kind of noise is introduced in the experiments.

Author response:
**Answer 1** regarding the added noise is as follows:

The noise is added to the data as a Gaussian distribution, with zero mean and a specified standard deviation σ. The value of σ is calibrated in the following way:
- For a well-trained neural network (before pruning), we compute its derivative as a vector, and obtain the standard deviation of the elements in the vector. Denote it by σ.
- We then add Gaussian noise to the data. Denote the standard deviation of this Gaussian noise by σ_{data}.  Denote the standard deviation of the computed derivative after adding this Gaussian noise as σ’. We have σ’ > σ.
- We tune and calibrate σ_{data}, such that σ’ = 2 σ = σ + σ. This is the case of the noise level being “σ”. If we calibrate it to be σ’ = 2 σ = σ + 2σ, then it is the case of the noise level to be “2 σ”.

To clarify this in the paper, we added in Appendix A.5 the content below:
In calibrating noise for data in neural networks, we start with a well-trained network. First, we calculate the standard deviation $\sigma$ of the network's derivative. Then, we add Gaussian noise with zero mean to the data. After adding the noise, the standard deviation of the network's derivative changes to a new value, $\sigma'$, which is always greater than sigma. The goal is to adjust the standard deviation of the Gaussian noise so that $\sigma'$ becomes $\sigma'=\sigma+\sigma$ (referred to as noise level being $\sigma$) or $\sigma'=\sigma+2\sigma$  (referred to as noise level being $2\sigma$).

Additionally, we modified the term “noisy gradient” to be “noisy data” in tables and the text descriptions  to avoid ambiguity.

Reviewer comment (Evaluation):
The reviewer asks what kinds of noise the proposed method can be a good choice for.

Author response:
**Answer 2** regarding what kinds of noise can this method be a good choice for, is as follows.

In our experiments, we discovered that if the noise (random fluctuations) is either very small or very large, the performance of two methods, LR and EWR (our method), tend to be similar. It is straightforward that small noise does not lead to visible difference between LR and EWR. When the noise is very large, gradients would be heavily contaminated such that few true covariance information could be recovered in Hessian approximation, and hence LR and EWR stay on par. Regarding the question on whether the noise level is considered to be “small” or “very large”, it depends on the size of both the model and the dataset. We don’t have a general answer to this, and it is mainly based on experimental observations.  However, in all cases, EWR never performed worse than LR, no matter the level of noise. We also think that if the mean of the noise is centered around zero (meaning it's equally likely to be positive or negative), our method could be a particularly good choice. This is because the way EWR interpolates, or blends data, might help in balancing out the noise.

Reviewer comment (Reproducibility):
“The noise level σ is set to be the standard deviation of the original gradients”. Why is this the noise level for both gradients and data? I would like to see a more detailed explanation how how the noise is added to data and gradients.

Author response:
We thank the reviewer for acknowledging the contributions of our work. We appreciate the reviewer for pointing out the ambiguity here. 

**Answer 1** regarding the noise level is as follows.

The noise is added to the data as a Gaussian distribution, with zero mean and a specified standard deviation σ. The value of σ is calibrated in the following way:

- For a well-trained neural network (before pruning), we compute its derivative as a vector, and obtain the standard deviation of the elements in the vector. Denote it by σ.
- We then add Gaussian noise to the data. Denote the standard deviation of this Gaussian noise by σ_{data}.  Denote the standard deviation of the computed derivative after adding this Gaussian noise as σ’. We have σ’ > σ.
- We tune and calibrate σ_{data}, such that σ’ = 2 σ = σ + σ. This corresponds to the case of the noise level being “σ”. If we calibrate it to be σ’ = 3 σ = σ +2 σ, then it corresponds to the case of the noise level to be “2 σ”.
\end{itemize}

To clarify this in the paper, we added in Appendix A.5 the content below:

*In calibrating noise for data in neural networks, we start with a well-trained network. First, we calculate the standard deviation $\sigma$ of the network's derivative. Then, we add Gaussian noise with zero mean to the data. After adding the noise, the standard deviation of the network's derivative changes to a new value, $\sigma'$, which is always greater than sigma. The goal is to adjust the standard deviation of the Gaussian noise so that $\sigma'$ becomes $\sigma'=\sigma+\sigma$ (referred to as noise level being $\sigma$) or $\sigma'=\sigma+2\sigma$  (referred to as noise level being $2\sigma$).*

Additionally, we modified the term “noisy gradient” to be “noisy data” in tables and the text descriptions  to avoid ambiguity.

Reviewer comment (Evaluation):
Can the authors also provide an accuracy table for Table 2 and Table 3?

Author response:
**Answer 2** regarding the accuracy of Table 2 and Table 3 is as follows.

Thank you for this good advice! We have added Table 6 and Table 7 showing the testing accuracy results, as compliment to the results of Table 2 and Table 3. 

Remark that these four tables differ with the benchmarking board of Table 1: In Table 1, all the pruning algorithms are fine-tuned on the training set after every pruning. This is to make our implemented LR and EWR align with the performance benchmarking results of MP, WF, and CBS, such that all algorithms are comparable to each other. The “Sparsity” column is the target sparsity, such that every line in Table 1 is the final pruning result. 

Tables 2, 3, 6 and 7 showcase the performance during one pruning process, where the loss and the accuracy values are taken immediately after pruning (without imposing any fine-tuning). This was clarified in the appendix A.5 of the original version:
“ In Table 2, Table 3, and additional results provided in the Appendix, sparsity is set using a linear gradual pruning strategy, progressing from 0 to 0.75 or 0.95 across ten distinct stages for MLPNet and ResNet, and from 0 to 0.75 across eight distinct stages for MobileNetV1. Notably, all recorded loss values are captured immediately post-pruning, devoid of any subsequent fine-tuning.”