Paper: Distributional Structured Pruning by Lower bounding the Total Variation Distance using Witness functions (ICLR-2024)

Reviewer comment (Experiments):
The paper seems to lack extensive testing across a diverse range of neural network architectures and data distributions. This limitation raises concerns about the applicability of the proposed methods in various real-world scenarios and different neural network models.

Author response:
We will provide a new battery of experiments highlighting the efficacy of our algorithm to other datasets and architectures within the next couple of days. Specifically, we will present results evaluating our algorithm on ResNet50 models trained on ImageNet. We'll post those results here, and please check out the general response as well!

Reviewer comment (Theory):
The paper's methodologies rely on certain assumptions regarding witness functions and distribution moments, necessary for theoretical formulations but potentially unrealistic in practice.

Author response:
The assumptions made on (a) the distributions, and (b) the witness functions themselves, are very minimal. In Theorems 1 and 2, we only assume we have access to $k$ moments of  the distribution. Similarly, the witness functions themselves only need to have bounded first and second order moments; that is $\mathbb{E}[u^\top \varphi(X)] <\infty$ and $\mathbb{E}[(u^\top \varphi(X))^2] < \infty$.

Reviewer comment (Experiments):
VGG-16 is a very old NN architecture. It would be better if the authors could demonstrate that their approach works for newer NN architecture.

Author response:
Further experiments are being conducted, and are coming soon! In particular, we compare the performance of WitnessPrune with CHIP [2] and TVSPrune [1] for pruning ResNet50 on the ImageNet dataset. We'll post those results both here and in the general response.

Reviewer comment (Evaluation):
How do these methods scale with increasingly complex and deep neural networks? Is there an in-depth analysis of the computational overhead and efficiency trade-offs when applying the proposed pruning techniques to larger-scale networks, especially compared to traditional pruning methods?

Author response:
We thank the reviewer for this insightful question. We answer the question in 2 parts.  
1. Traditional structured pruning strategies that require gradients often require expensive backpropagation in order to compute saliency scores (which are used to identify which filters are important, and which are not). **In this work, we focus on the challenging use case where the training data and the loss function are unavailable to us**. As a benefit to this, we no longer need to compute the gradients.  
2. The key computational cost involved in this algorithm is the storage of moments required to compute the $r^l_j$ scoress defined in Equation 7. Unlike previous works such as TVSPrune [1] and CHIP [2], we do not need to store the moments of entire feature maps - we need only store the first two moments of the witness function $f(X) = u^\top \varphi(X)$. Thus, prior to computing the $r^l_j$ scores, we only  need to store $2n_l$ real numbers per class for each layer. This is in contrast to [1] and [2], where moments of the entire (flattened) feature maps are stored. Furthermore, unlike prior work such as TVSPrune, for classification tasks where the dataset has $K$ classes, we need only store $2K$ moments, whereas TVSPrune requires the storage of ${K \choose 2}$ moments.

Reviewer comment (Experiments):
Could the authors provide further evidence on how the proposed pruning strategies perform across a wider range of scenarios? For instance, how do these methods fare with different types of neural network architectures, and varying data modalities?

Author response:
We'll post additional experimental results both here and in the general response, in the next 48 hours. Our experiments will focus on ResNet50 model trained on the ImageNet dataset, and compare with some standard baselines [1],[2],[3].

Reviewer comment (Presentation):
Formatting of references: The paper chooses to use references without separations from text. While this is suitable for some scenarios, in other situations it obstructs the flow of the reading, and it may be beneficial to update the references in these situations. This can likely be changed quite easily.

Author response:
**Clarity issues: In certain parts of the paper, clarity issues cause confusion when reading. Some examples include:**

**-   Formatting of references. The paper chooses to use references without separations from text. While this is suitable for some scenarios, in other situations it obstructs the flow of the reading, and it may be beneficial to update the references in these situations. This can likely be changed quite easily.**
**-   Unclear definitions. There are multiple sections of the paper that can be revised to make definitions clearer. One example is in Section 2.2: the very first sentence abruptly ends with a "..., and let." Formatting issues and typos throughout the paper should be revised and updated.**

We have updated the text with corrections to the main document. Please do check it out!

Reviewer comment (Writing):
Unclear definitions: There are multiple sections of the paper that can be revised to make definitions clearer. One example is in Section 2.2: the very first sentence abruptly ends with a "..., and let."

Author response:
**Clarity issues: In certain parts of the paper, clarity issues cause confusion when reading. Some examples include:**

**-   Formatting of references. The paper chooses to use references without separations from text. While this is suitable for some scenarios, in other situations it obstructs the flow of the reading, and it may be beneficial to update the references in these situations. This can likely be changed quite easily.**
**-   Unclear definitions. There are multiple sections of the paper that can be revised to make definitions clearer. One example is in Section 2.2: the very first sentence abruptly ends with a "..., and let." Formatting issues and typos throughout the paper should be revised and updated.**

We have updated the text with corrections to the main document. Please do check it out!

Reviewer comment (Writing):
Formatting issues and typos throughout the paper should be revised and updated.

Author response:
**Clarity issues: In certain parts of the paper, clarity issues cause confusion when reading. Some examples include:**

**-   Formatting of references. The paper chooses to use references without separations from text. While this is suitable for some scenarios, in other situations it obstructs the flow of the reading, and it may be beneficial to update the references in these situations. This can likely be changed quite easily.**
**-   Unclear definitions. There are multiple sections of the paper that can be revised to make definitions clearer. One example is in Section 2.2: the very first sentence abruptly ends with a "..., and let." Formatting issues and typos throughout the paper should be revised and updated.**

We have updated the text with corrections to the main document. Please do check it out!

Reviewer comment (Theory):
The TV lower bound is stated only in terms of populational quantities, such as exact moments of the distribution. In practice, these quantities have to be *estimated* from finitely many samples. The regime the authors seem to have in mind is the fixed-dimension infinite-samples setting. However, for neural network pruning, their main application, a *high-dimensional* scaling in which the data dimension and the number of samples grow proportionally would be a more pertinent setting to analyze. This greatly complicates the analysis for their lower bound since one would need to analyze the estimation error in their moment functional. Even the simple task of estimating moment tensors of high-dimensional distributions can be tricky. Hence, a finite sample analysis of their lower bound is warranted.

Author response:
**The most glaring issue is that the TV lower bound is stated only in terms of populational quantities, such as exact moments of the distribution. In practice, these quantities have to be estimated from finitely many samples. The regime the authors seem to have in mind is the fixed-dimension infinite-samples setting.** 
We would like to clarify that in the setting of this paper - that is, using the TV-separation between the class-condtional distributions of the **feature maps generated by individual filters** as defined in Equation (1) on page 2-  the dimensions of feature maps, particularly in later layers, are not high-dimensional. That is, the number of samples of each class is actually greater than the dimension of the individual feature maps.  For models trained on CIFAR10, we typically choose 100 samples per class to compute the moments of the witness function $u^\top \varphi(X)$.

**However, for neural network pruning, their main application, a high-dimensional scaling in which the data dimension and the number of samples grow proportionally would be a more pertinent setting to analyze.**  
In the final layers, of most neural network architectures, the size of the feature maps is often *low-dimensional*, and moreover, these layers have been shown to be most amenable to pruning [1],[7]. We illustrate this with a VGG-16 model trained on CIFAR10, where $n=100$ is the number of samples. We see that the dimension of the feature maps is relatively small.


|Layers | 4,5,6 | 7,8,9 | 10,11,12 |
|---|---|---|---|
|  Dimension of feature map $d$ |  64 |  16 | 4  | 
| Ratio of dimension to sample size $d/n$ |  .64 |  .16 |  .04 | 

We see from this table that the majority of feature maps that can be effectively pruned are *not* high dimensional. Note that this trend can be observed over a variety of architectures as well.

**This greatly complicates the analysis for their lower bound since one would need to analyze the estimation error in their moment functional. Even the simple task of estimating moment tensors of high-dimensional distributions can be tricky. Hence, a finite sample analysis of their lower bound is warranted.**

For feature maps that are high dimensional, we note that in Corollaries 1 and 2, we derive lower bounds on the TV distance with respect to the Fisher Linear Discriminant [5] and the Minimax Probability Machine (MPM) [6].  We also recall that works such as [1],[7] observed that initial layers possessing high-dimensional feature maps are typically the most difficult to prune effectively. In this case, we can use existing methodologies, such as those proposed in  [8],[6] for robust estimation of Fisher and MPM based lower bounds on the TV distance. Specifically, if we choose $u^\top \varphi(X) = u^\top X$, we can directly apply Equation 15 proposed in *Robust Fisher Discriminant Analysis*, by Kim et al [8].

 However, we feel that a broader investigation of these aspects is beyond the scope of our current work, which focuses on using the lower bounds for pruning, which it does effectively, as evidenced by the empirical results presented in Section 7.3.

Reviewer comment (Writing):
Can the authors explain why a “witness function” is named as such? The reviewer’s guess is that it comes from the general form of the lower bound from A.1 and the authors view it as the function that achieves the supremum (a “witness” of the supremum), but it would be helpful to explain the term in the introduction.

Author response:
**Can the authors explain why a “witness function” is named as such? My guess is that it comes from the general form of the lower bound from A.1 and the authors view it as the function that achieves the supremum (a “witness” of the supremum), but it would be helpful to explain the term in the introduction.**

The term "witness function" is a term commonly used in the literature to describe functions that "witness" the effect of distributions. Some recent and highly cited works that utilize this terminology are [2],[3],[4].

Reviewer comment (Writing):
The use of descriptor “robust” when the authors say “robust lower bound” is somewhat misleading. “Robust” is typically used in the context of estimation, but there is no estimation analysis in this paper.

Author response:
**-   The use of descriptor “robust” when they say “robust lower bound” is somewhat misleading. “Robust” is typically used in the context of estimation, but there is no estimation analysis in this paper.**

In this work, we use the term "robust" in the sense of *distributional robustness*. As stated in a previous response, the proposed bound is a worst-case lower bound. Thus, given any distributions $\mathcal{P}$ and $\mathcal{Q}$  with collections of moments **P** and **Q** respectively, Theorem 1 provides a worst-case lower bound on the TV distance between them.

Reviewer comment (Theory):
Can \varphi be non-polynomial? The defining condition for \varphi: R^d \to R^n is that its expectation is a function of only the first k moments. Are there any non-polynomial functions that satisfy this condition?

Author response:
**-   Can \varphi be non-polynomial? The defining condition for \varphi: R^d \to R^n is that its expectation is a function of only the first k moments. Are there any non-polynomial functions that satisfy this condition?**

$\varphi$ can indeed be non-polynomial. For instance, if we choose $u^\top \varphi(X)$ to be the moment generating function, there are distributions for which closed-form expressions of the MGF exist. An example of this would be the univariate Gaussian distribution $\mathcal{N}(\mu,\sigma)$, where we can simply choose $\varphi(X) =  \mathrm{exp}(\mu t + \frac{1}{2}\sigma^2 t^2)$, for a given real $t$.