SWAP: Sparse Entropic Wasserstein Regression for Robust Network Pruning

Question 1: 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).

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

Question 2: Is it possible to show an analysis in terms of computational expense incurred, as the scale of the model increases?

Answer 2: 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$.

Question 3: 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?

Answer 3: 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.

Question 1: 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$?

Answer 1: 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$.

Question 2: It would be more interesting to compare with more baselines, in addition to the LR method.

Answer 2: 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.

Question 3: Just to be curious, as the main contribution lies in better estimating the gradient under the existence of noise, is it possible to extend the proposed method to other applications beyond network pruning?

Answer 3: Yes, this is possible, but we haven’t tested all the possible applications and we are listing one of them here. A closely related scenario is when there is noise or corruption in computing the gradient of the trained network, as opposed to the noisy data that we are considering in this work. This naturally arises in federated learning scenarios where the gradient updates from the clients are not reliable.

We appreciate the reviewer for the acknowledgement of our contributions.

Question 1: 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.

Answer 1: 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.

Question 2: 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.

Answer 2 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:

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:

total_stages = 10 - 1

Hence:

(pruning_stage / total_stages * target_sparsity).round(2)

yields the sparsity levels during this pruning process as

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.

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

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.

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.

Answer 3 regarding what would be an appropriate choice of the regularization weight for different kinds / magnitudes of noise, is as follows.

The resilience of noise originates from the neighborhood interpolation mechanism. Intuitively, if the gradients are really noisy, then we would trust each single data point less, because it may likely be contaminated with noise. Instead, we may want to enlarge the neighborhood size of performing the interpolation to combat the noise. A large value of $\varepsilon$ does this favor for us. On the contrary, a small $\varepsilon$ leads to few points participating effectively in interpolation. An extreme case is $\varepsilon=0$ such that the entropic regularization term is entirely removed, indicating that no interpolation is performed when computing the Euclidean distance between two points. Yet, this case is not equivalent to the LR case, because the optimal transportation does not necessarily happen between data points in their originally assigned order. They are equivalent only if we impose the transportation plan to be $\Pi=*diag*(1/n)$ consistently. An illustration of the neighborhood size can be seen in Figure 10 (a) - (d). The lighter the color, the more points are involved in neighborhood interpolation in computing the Euclidean distance.

Thank you for the comprehensive response and clarification re LR & EWR. I have read over the responses by the authors and have taken a look at the revision, which looks much better. I re-iterate that the topic of this work falls outside my area of expertise, but I would increase my score to an 8.

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.

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