Paper: Learning Invariances via Neural Network Pruning (ICLR-2024)

Reviewer comment (Theory):
As far as I understand, how connectivity models achieve invariance / equivariance does not only rely on the architecture but also on certain constraints on weights (especially weight sharing), for example, the translation invariance / equivariance with convolution. This work only emphasizes on the architecture, but there are no discussions on weights. I am doubtful of the claim (hypothesis) that the subnetwork after retraining "preserves desired invariances and hence improves the inductive bias" and I wonder if this is empirically the case.

Author response:
**[On Architectural Emphasis]**

Our goal is to improve the **network architecture’s inductive bias**.

We agree weights play an important role in preserving invariance; however, the **network architecture’s inductive bias** is **independent of its weights** and refers to the inductive bias captured by **the structure of the network**. Improving the network architecture’s structure will improve its prior and optimization trajectory[1] **leading to better trained models**. In fact convolutional nets can be used as feature extractors **even under random initialization**[1,2,3]. Learning invariance by solely updating the weights will not achieve a better architectural inductive bias.

We empirically verified IUNet’s pruned network has **a better structure and inductive bias in our ablation studies, Table 4.** Our results show the pruned subnetwork achieves **substantially better** downstream accuracy than both the supernetwork trained normally and the supernetwork trained with our novelties, PIs and ILO.

It is *“empirically the case” IUNet’s captures better such invariances. As detailed in the [On Consistency Experiment] rebuttal section, network pruning is *“necessary for obtaining an invariance-preserving network”* under maximum likelihood training.

We will update our main text to emphasize our motivation for improving the **network architecture’s inductive bias** rather than just the weights.

Reviewer comment (Experiments):
Rotation, as one of the most explored transformation in the literature of invariant/equivariant deep learning, is however not included in the study of this work.

Author response:
**[On Rotation Invariance]**

We did not introduce rotations into our pipeline since the **SVHN dataset is a digit classification** task. Specifically, digits are not rotationally invariant since 6 and 9 fall under the same space group yet are distinct digits.

Reviewer comment (Reproducibility):
After pruning, is ILO or data augmentation used in training the subnetwork?

Author response:
**[On Transformations and Data Augmentations]**

In Tables 3 and 4, the set of transformations used are mentioned in Section 4.3. For computer vision, these are (1) resize crops, (2) horizontal flips, (3) color jitter, and (4) random grayscale. For tabular data, this is (5) SCARF[8], randomly corrupting features by drawing the corrupted versions from its empirical marginal distribution.

In Tables 3 and 4, all methods adopt some **standard dataset augmentation**. On vision datasets, we follow ShrinkBench[9], which adopts random horizontal flips and crops. On tabular datasets, we follow MLP+C[4], which adopts no augmentation by default.

**For fairness, the maximum likelihood branch is trained on the same dataset, with or without ILO.** The contrastive learning branch applies transformations (1-5) on training samples to create positive and negative pairs, but the underlying data used to train all models are the same. We will add these details to Section E.

Reviewer comment (Experiments):
In table 3, are the weights of pruned IUNets also re-initialised or re-trained before evaluating the consistency? And if re-trained is ILO or augmentation used in training? As I mentioned, the inductive bias of invariance / equivariance does not only exist in the architecture but also in the shared weights. So if the IUNet weights are re-initialised then the results are too good to be true, otherwise an unfair comparison.

Author response:
**[On Consistency Experiment 1/2]**

The weights are re-initialized to right after the first epoch fairly for all models in Table 3. 

We greatly appreciate your suggestion to deeply re-analyze these results. We found an edge case that the consistency metric did not properly capture. Although IUNet’s consistency metric is >98%, **the model predicts the same class for nearly all inputs**! This is because IUNet has majority small weights and not been fully trained, so the encoder consistently predicts a near-zero vector representation, regardless of input image.

**To correct for this edge case, we rerun Table 3 and additional consistency experiments on CIFAR10, CIFAR100, and SVHN.** Because IUNet lacks weight sharing, we expect randomly-initialized CNN to outperform both it and MLPs. With maximum likelihood training, we expect IUNet to effectively bridge the invariance consistency gap between MLPs and CNNs, by learning CNN-like weights from its improved inductive bias. Overall, we expect IUNet to outperform MLPs in accuracy by learning a network structure with better inductive bias.

Reviewer comment (Reproducibility):
Also in table 3 and 4, what transformations are used?

Author response:
**[On Transformations and Data Augmentations]**

In Tables 3 and 4, the set of transformations used are mentioned in Section 4.3. For computer vision, these are (1) resize crops, (2) horizontal flips, (3) color jitter, and (4) random grayscale. For tabular data, this is (5) SCARF[8], randomly corrupting features by drawing the corrupted versions from its empirical marginal distribution.

In Tables 3 and 4, all methods adopt some **standard dataset augmentation**. On vision datasets, we follow ShrinkBench[9], which adopts random horizontal flips and crops. On tabular datasets, we follow MLP+C[4], which adopts no augmentation by default.

**For fairness, the maximum likelihood branch is trained on the same dataset, with or without ILO.** The contrastive learning branch applies transformations (1-5) on training samples to create positive and negative pairs, but the underlying data used to train all models are the same. We will add these details to Section E.

Reviewer comment (Experiments):
In table 4, what if data augmentation is applied to all models, no matter with and without ILO? I wonder how ILO can improve over simple augmentation (especially with a pruned model trained without ILO)

Author response:
**[On Transformations and Data Augmentations]**

In Tables 3 and 4, the set of transformations used are mentioned in Section 4.3. For computer vision, these are (1) resize crops, (2) horizontal flips, (3) color jitter, and (4) random grayscale. For tabular data, this is (5) SCARF[8], randomly corrupting features by drawing the corrupted versions from its empirical marginal distribution.

In Tables 3 and 4, all methods adopt some **standard dataset augmentation**. On vision datasets, we follow ShrinkBench[9], which adopts random horizontal flips and crops. On tabular datasets, we follow MLP+C[4], which adopts no augmentation by default.

**For fairness, the maximum likelihood branch is trained on the same dataset, with or without ILO.** The contrastive learning branch applies transformations (1-5) on training samples to create positive and negative pairs, but the underlying data used to train all models are the same. We will add these details to Section E.

Reviewer comment (Novelty):
Most of the ideas are derived from different papers and combined together.

Author response:
**[On Novelty and Existing Work]**

We highlight our contributions over the existing works and will update our writing to make this more clear.

On task novelty:

**IUNet is the first general-purpose invariance learning via pruning approach**. Our approach tackles vision, text, and tabular data across MLPs, CNNs, and transformers. Previous work[8] only explored computer vision where the authors state they hope to *“see if [network pruning for invariance learning] can be used in other contexts”*.


On method novelty:

1) ILO: **We are the first to combine maximum likelihood (cross entropy) with contrastive learning (NCE) for network pruning. This combination improves network pruning, which contradicts previous work[7] that found contrastive learning hurts network pruning.** The authors of [7] claim they hope their negative results *“spark interest in developing more representation-friendly pruning methods for supervised contrastive learning”.*  We successfully demonstrate previous efforts failed because they overfit the contrastive objective[14]. Combining both maximum likelihood and contrastive learning is necessary to boost pruning performance.

2) PIs: **We are the first to propose small weight initialization for network pruning.** Small initial weights force the neural network to only grow weights that are truly important. Otherwise, unimportant weights will stay large, causing important weights to be pruned (i.e. lazy training phenomenon[1, 15]). This issue occurs more on deeper networks[15]. Hence, unlike earlier invariance learning for pruning work[8], **IUNet improves the scalability of said approach to deep networks**.

Reviewer comment (Evaluation):
The annexure table 7 for tabular data is a bit unclear as it highlights only a few good instances, while the other pruning methods still work better than the proposed method. The reason is not mentioned. While the idea is to capture and preserve the invariance during the pruning, this means somewhere something is still missing in the proposed method, it solves the problem in some of the datasets but not all.

Author response:
**[On Tabular Dataset Results]**
The competing baseline “MLP+C” performs time-consuming (few days for all 40 datasets) hyperparameter optimization over regularization techniques, whereas our approach does not (few hours for all 40 datasets). Thus, there is a trade-off between MLP+C's more time-consuming good results compared to IUNet’s more economical, yet still quite good results.

It is well known that tabular data has tricky results as the dataset quality can drastically vary: from sparse to dense inputs and from few to many samples[9, 10, 11, 12]. Some datasets are naturally biased towards more sample efficient algorithms like XGBoost[10, 13]. Other datasets may be bottlenecked by regularization rather than architectural design[9, 11, 12]. Despite this, **IUNet boosts MLP performance on tabular data in aggregate and achieves the best performance without expensive hyperparameter tuning.**

Reviewer comment (Writing):
The paper mentions the encoder-decoder architecture of the proposed method, but throughout the paper, the convention is not followed properly.

Author response:
**[Clarifying Encoder and Decoder Blocks]**

Thank you for pointing out this detail! We will update Section 3.2 to be consistent. In terms of empirical results, as mentioned in Appendix E.4, we follow the encoder-decoder architecture, where only the encoder is pruned and the decoder (final linear layer following conventions) remains dense.

Reviewer comment (Experiments):
The authors mention that their method improves compression performance on existing models, how about a completely new model? Does it still work the same? Or do the pre-trained models need to be trained extensively?

Author response:
**[On Compressing Pretrained vs Randomly Initialized Networks]**

Our experiments are run on **randomly initialized networks** (i.e. *“completely new models”*) using **“existing” architectures** (CNNs and MLPs). IUNet (1) **randomly initializes** the supernetwork with weights $\theta_{M^\dagger}^{(0)}$, (2) scales the random weights with PIs, $\theta_M^{(0)} = \kappa \theta_{M^\dagger}^{(0)}$, (3) trains the supernetwork **from scratch** using ILO to get $\theta_M^{(T)}$, (4) prunes the trained supernetwork into a subnetwork, (5) **randomly re-initializes** the pruned subnetwork as $\theta_{P}^{(0)}$ with Lottery Ticket Reinitialization (i.e. resetting weights to that of $\theta_{M^\dagger}^{(0)}$), and (6) finetunes the **randomly re-initialized** subnetwork **from scratch** using maximum likelihood to get $\theta_P^{(T)}$. 

Yes, **IUNet works on pretrained networks with minimal adaptation!** We provide **new experiments** on **pretrained transformers**, following the SMC-Benchmark[1], which prunes the pretrained RoBERTa[2] on the Common Sense QA dataset[3]. Using the official Github repo, *https://github.com/VITA-Group/SMC-Bench/*, we reran SMC’s provided pruning algorithms OMP (after), OMP (before), and SNIP[4]. Additionally, we implemented and ran IUNet and GRaSP[5] on this benchmark. We report the accuracy under different sparsities.

For a fair comparison, we leave out ILO from IUNet and adopt the same finetuning protocol and hyperparameters as SMC. We tried 2 different settings for kappa: [0.125, 0.0625]. Because SMC uses pretrained models: (1) IUNet scales the pretrained weights rather than random initialization by kappa, (2) IUNet reinitializes the pruned models weights to the fine-tuned model initialized from pretrained weights, not random initializations.

Note, IUNet is a **post-tuning pruning algorithm** (similar to OMP (after) in the benchmark paper), where the model is first pretrained, then fine-tuned (either with maximum likelihood or ILO), then pruned, then sparsely fine-tuned to convergence. OMP(before), SNIP, and GRaSP are **pruning at initialization (PaI)**[1, 6] algorithms, where the model first pretrained, then pruned (iterating through the train set once) then sparsely fine-tuned to convergence.

```
Method\Sparsity   | 0.2   | 0.36  | 0.488 | 0.590 | 0.672
—-----------------+-------+-------+-------+-------+-------
OMP (Before)      | 26.54 | 58.89 | 25.88 | 19.57 | 19.57
SNIP              | 18.67 | 20.48 | 18.84 | 19.00 | 19.16
GRaSP             | 21.79 | 19.25 | 18.43 | 17.77 | 20.63
—-----------------+-------+-------+-------+-------+-------
OMP (After)       | 76.09 | 73.87 | 30.14 | 19.57 | 19.57
—-----------------+-------+-------+-------+-------+-------
IUNet* (k=0.125)  | 75.10 | 74.53 | 34.89 | 19.57 | 19.57
IUNet* (k=0.0625) | 75.02 | 74.37 | 53.80 | 19.57 | 19.57
*no-ILO
```

The densely fine-tuned model reaches an accuracy of 75.76%.

At low sparsities (sparsity= [0.2, 0.36]), **all post-tuning pruning algorithms (OMP(After) and IUNet) outperform PaI algorithms (OMP(Before), SNIP, GRaSP)**. At these sparsities, *OMP(After) and IUNet are within experimental uncertainty of each other*. This agrees with prior work[1, 6] which finds PaI is less effective than post-tuning algorithms.

At middle sparsities (sparsity=0.488), **IUNet outperforms all baselines (OMP (after), OMP (before), SNIP, GRaSP),** because PIs alleviates the lazy learning phenomenon which is a key bottleneck in existing pruning algorithms[4]. This is a notable finding, showing **IUNet scales to larger sparsities than baselines.**

At high sparsities (sparsity= [0.5, 0.672]), **all methods fail**, which agrees with SMC[4], suggesting at least 50% of weights are crucial to the transformer models on these tasks.

**Hence, as shown in results above, IUNet generalizes to pre-trained models and outperforms strong baselines.**

Reviewer comment (Theory):
It seems that many presentations of the paper should be revised. For example, in Section 3.2, the explanation for relations of parameters of supernetwork and subnetwork are quite confusing. Since $\theta_P^{(0)}$ and $\theta_M^{(T)}$ are parameter "vectors", hence the set inclusion operation is not suitable. Also, the expression $|\theta_P^{(0)}| < |\theta_M^{(T)}|$ does not imply that $\theta_P^{(0)}$ is more sparse than $\theta_M^{(T)}$, which does not give any information about the network capacity. Rather, the authors should express the sparsity or network capacity via $\ell_0$-norm (number of non-zeros, nnz). In the same sense, in Section 3.2.1, it seems difficult to say that the higher weight value of criterion (3) is a criterion of network capacity. Even if the parameter value is small, if it is non-zero, it will eventually be counted in the FLOPs calculation, so (for example) the number of non-zeros should be used rather than weight magnitude.

Author response:
**[On Sensitivity to Initialization]**

We ran all experiments 3 times from scratch under different random seeds for all methods. Hence, the trends we report between IUNet and OMP already consider OMP’s sensitivity to initialization.

**[On Paper Presentation]**

Thank you for pointing out this confusing notation. In Section 3.2, we will update our notation from $\theta$ to $\omega$, where $\omega$ denotes the set of parameters in each network: $|\omega_P|<<|\omega_M|$ and $\omega_P \subset \omega_M$.

Your understanding of the l0-norm is correct, because the l0-norm of parameters where pruned parameters are zero is equivalent to the size of the set of remaining parameters. More formally:

$$
nnz(\theta_M) = \sum_{\theta_{Mi} \in \theta_{M}} 1[\theta_{Mi} \ne 0] = \sum_{\omega_{Pi} \in \omega_P} 1 = |\omega_P|
$$

Reviewer comment (Theory):
In the case of proactive initialization (PI), it is as if only the parameter size was simply adjusted, but looking at the results in Figure 3, the parameters were initialized small with $\kappa = 0.125$. It is natural that the result of $\kappa = 0.125$ has many weights with smaller values. There needs to be some mathematical evidence or experimental results for proactive initialization that can further demonstrate motivation and reasoning about it.

Author response:
**[On PIs]**

PIs will scale the original weights by multiplying the Kaiming-initialized weights by a kappa scalar as stated in Section 3.2.2:

$\theta_M^{(0)} = \kappa \theta_{M \dagger}^{(0)} $

**We agree it is natural thus for the learned weights to also have smaller magnitudes!** This accomplishes our goal of filtering out unimportant weights, because most weights will have small magnitudes, while only important weights will have large magnitudes. In Figure 3, **we empirically verify** that the magnitude of important weights are the same (large) with and without PIs, yet the magnitude of most weights are small **only with PIs.**

**Encouraging most weights to be small matters**, because the change in output logits is proportional to the magnitude of pruned weights. This relation between the magnitude of pruned weights and output logits can be seen in a simple 2 layer neural network with k-Lipschitz activation:

Let the 2-layer supernetwork be: $f(x) = \sigma(\sigma(xW_{1})W_2), x\in R^d, W_1 \in R^{d\times d}, W_2 \in R^{d\times 1}$. 

Let the pruned subnetwork be: $g(x) = \sigma(\sum_{i\in P_2'}\sigma(\sum_{j,i\in P_1'}xW_{1,j,i})W_{2,i}))$ and pruned weight indices be $P_2 = argtopk(-|W_2|), P_1 = argtopk(-|W_1|)$ and unpruned weight indices be $P_2' = ${$ 1,...,d $}$ \setminus P_2, P_1' = ${$ 1,...,d $}$ ^2 \setminus P_1$. 

Then the difference in the pruned network’s and original network’s output logits is:

$|f(x) - g(x)| \le k|\sigma(xW_{1})W_2 - \sum_{i\in P_2'}\sigma(\sum_{j,i\in P_1'}xW_{1,j,i})W_{2,i}|$

$$|
f(x) - g(x)| \le k C_1 + k^2 C_2
$$

where $C_1 = |\sum_{i \in P_2} \sigma(xW_{1})_i|$

and $C_2 = \sum_{i\in P_2'} (W_{2,i}\sum_{j,i\in P_1}|xW_{1,j,i}|)$

$|f(x) - g(x)| = \mathcal{O}(\sum_{i\in P_2}|W_{2,i}|+\sum_{j,k\in P_1}|W_{1,j,k}|)$

**Figure 2abc presents empirical evidence of PIs efficacy**. As we decrease kappa, there is a **tradeoff** between performance of the unpruned network and performance improvement in the pruned network. Interestingly, **kappa has a sweet spot, where the pruned network’s performance grows higher than that of the Kaiming-initialized (kappa=1.0) supernetwork’s!** We highlight this is **a substantial finding among the network pruning literature**[10] and supports our hypothesis that through pruning, the subnetwork can improve the inductive bias of the supernetwork. 

**Table 4, ablation studies, presents empirical evidence of PIs efficacy**. If we train a model without PIs, performance is worse on computer vision datasets, as the pruning stage cannot differentiate important from unimportant weights. Tabular datasets show less consistent results since they are much more varied in size and feature type [14, 15, 16, 17]. In some tabular datasets, regularization becomes the bottleneck rather than model architecture [15, 16, 17]. Nonetheless, our ablation studies demonstrate **PIs will improve the downstream accuracy of IUNet**.

**“Our [On “SNIP” and “GRaSP”] rebuttal section also presents empirical evidence of PIs efficacy.** By simply adjusting the initialization, **PIs scales transformers to larger sparsities than that of baselines**.

We will update section 5.4 to make this empirical evidence more salient.