Learning Invariances via Neural Network Pruning

Thank you for your insightful comments! We provided additional experiments and clarifications to address your concerns. Please let us know if you have any remaining doubts!

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

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

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

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

[Clarifying Figure 2]

Thank you for pointing out this errata! We will updated the paper with the correct x- and y-labels. Results with kappa=0.0625 are too poor (below the visible y-axis). We will update our figures to clarify this as well.

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

[Learning different Pruning Masks for different transformation types]

We thank the reviewer for this excellent suggestion, but leave this as future work.

Thank you for your response! We provided additional experiments and clarifications to address your concerns. Please let us know if you have any remaining doubts!

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

I appreciate the authors for their kind rebuttals, but they only partially address my concerns.

I understand why Proactive initialization helps the network pruning in theory (though the authors simply justify only for 2-layer networks) and practice. Also, the descriptions on Figure 4 helps me understand the points than before reading the rebuttals.

However, I have still concern about "how and why pruning is required for learning invariance" (I think it is the most important part since it is the motivation of this paper). The authors say that inductive bias heavily depends on the architecture itself rather than on the model parameter, but this paper fully utilizes the pretrained models, in which one trains model at "good initialization".

According to author's claim, I think that this paper should conduct the experiments without pretrained models, i.e., training from scratch with Invariance Learning Objective (ILO). Furthermore, in my view, ILO still seems far away from network pruning, and applying ILO to dense models (without pruning) is expected to better performance than pruned networks.

Considering these points, I will maintain my score in the current shape.

[On Figure 4]

We provide further discussion on the visualizations below and will update our text accordingly.

In Figure 4, we visualize randomly-chosen filters from the first layer of each network. We randomly chose the filters for fairness, and picked the first layer because it is the easiest to interpret. The main conclusion from Figure 4 is: IUNet bridges the gap between MLPs and CNNs and decision trees architecturally. We make the following specific observations:

1) Locally Connected Regions: For Figures 4efg, IUNet finds locally connected regions with high RGB values. Local connections help preserve translation and scale invariance as shown in Table 3 and verified by previous works[1, 2, 3, 11]. The presence of locally-connected regions visually confirms that IUNet can rediscover CNN architectures from MLPs.

2) Color Invariance: In Figures 4eg (CIFAR10/SVHN), IUNet assigns high values to locally connected regions uniformly across all RGB channels, which improves color invariance. In Figure 4f (CIFAR100), this is not as consistent due to a more challenging learning objective, though we find it still improves color invariance over the MLP supernetwork.

3) Discovering Substructures: Nearly all IUNet output neurons favor a clear pattern of input neurons (the exception being the top left filter of Figure 4f), whereas nearly no OMP output neurons exhibit any pattern. This suggests IUNet does a good job recovering invariance preserving substructures across the supernetwork’s first layer.

4) Axis-Aligned Tree Structures: For Figure 4h, IUNet discovers subnetworks that focus on primarily 1 input feature for each output neuron. This behavior mirrors tree-based models where neurons are activated where single features reach some threshold [12, 13, 14]. Such network architecture follows existing work which found the axis-aligned inductive bias in tree-based models is critical to XGBoost’s performance over standard MLP setups [14].

[References]

[1] Ulyanov, Dmitry, Andrea Vedaldi, and Victor Lempitsky. "Deep image prior." Proceedings of the IEEE conference on computer vision and pattern recognition. 2018.

[2] Cao, Yun-Hao, and Jianxin Wu. "A random cnn sees objects: One inductive bias of cnn and its applications." Proceedings Of The AAAI Conference On Artificial Intelligence. Vol. 36. No. 1. 2022.

[3] Rosenfeld, Amir, and John K. Tsotsos. "Intriguing properties of randomly weighted networks: Generalizing while learning next to nothing." 2019 16th conference on computer and robot vision (CRV). IEEE, 2019.

[4] Liu, Shiwei, et al. "Sparsity May Cry: Let Us Fail (Current) Sparse Neural Networks Together!." arXiv preprint arXiv:2303.02141 (2023).

[5] Liu, Yinhan, et al. "Roberta: A robustly optimized bert pretraining approach." arXiv preprint arXiv:1907.11692 (2019).

[6] Talmor, Alon, et al. "Commonsenseqa: A question answering challenge targeting commonsense knowledge." arXiv preprint arXiv:1811.00937 (2018).

[7] Lee, Namhoon, Thalaiyasingam Ajanthan, and Philip HS Torr. "Snip: Single-shot network pruning based on connection sensitivity." arXiv preprint arXiv:1810.02340 (2018).

[8] Wang, Chaoqi, Guodong Zhang, and Roger Grosse. "Picking winning tickets before training by preserving gradient flow." arXiv preprint arXiv:2002.07376 (2020).

[9] Frankle, Jonathan, et al. "Pruning neural networks at initialization: Why are we missing the mark?." arXiv preprint arXiv:2009.08576 (2020).

[10] Ma, Xiaolong, et al. "Sanity checks for lottery tickets: Does your winning ticket really win the jackpot?." Advances in Neural Information Processing Systems 34 (2021): 12749-12760.

[11] Neyshabur, Behnam. "Towards learning convolutions from scratch." Advances in Neural Information Processing Systems 33 (2020): 8078-8088.

[12] Marton, Sascha, et al. "GRANDE: Gradient-Based Decision Tree Ensembles." arXiv preprint arXiv:2309.17130 (2023).

[13] Marton, Sascha, et al. "GradTree: Learning Axis-Aligned Decision Trees with Gradient Descent." NeurIPS 2023 Second Table Representation Learning Workshop. 2023.

[14] Grinsztajn, Léo, Edouard Oyallon, and Gaël Varoquaux. "Why do tree-based models still outperform deep learning on typical tabular data?." Advances in Neural Information Processing Systems 35 (2022): 507-520.

[15] Gorishniy, Yury, et al. "Revisiting deep learning models for tabular data." Advances in Neural Information Processing Systems 34 (2021): 18932-18943.

[16] Kadra, Arlind, et al. "Well-tuned simple nets excel on tabular datasets." Advances in neural information processing systems 34 (2021): 23928-23941.

[17] McElfresh, Duncan, et al. "When Do Neural Nets Outperform Boosted Trees on Tabular Data?." arXiv preprint arXiv:2305.02997 (2023).

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

[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}|$

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.

Thank you for your review! We provided additional experiments and clarifications to address your concerns. Please let us know if you have any remaining doubts!

[Motivation for Pruning]

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 (i.e. applying ILO “without network pruning”) will not achieve such properties, hence pruning is necessary.

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. Hence, IUNet indeed finds a pruned subnetwork with better inductive bias (or prior) than that of the supernetwork.

We will update our main text to emphasize the necessity of pruning clearer.

[On “SNIP” and “GRaSP”]

We agree that comparing IUNet with SNIP and GraSP can strengthen our claim that IUNet improves general pruning performance.

To this end, we compare IUNet against SNIP and GRaSP on the SMC-Benchmark[4], which prunes the pretrained RoBERTa[5] on the Common Sense QA dataset[6]. 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[7]. Additionally, we implemented and ran IUNet and GRaSP[8] 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)[4, 9] 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[4, 9] 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 both SNIP and GRaSP.

That's my confusion, thanks for the correction to the authors. Also, I appreciate the authors for their response.

Thank you for your timely response! We want to make an important clarification.

As stated in Section 3.2, our paper experiments are run on randomly initialized networks, not pretrained models.

Specifically, in the paper, 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)}$.

Tables 1 and 2 compare the randomly initialized supernetwork and the randomly initialized pruned network (IUNet) both trained under the same maximum likelihood objective, where we find "the inductive bias heavily depends on the architecture itself rather than on the model parameters". Table 4 shows the pruned networks (IUNet) outperforms applying ILO to dense models without pruning (IUNet (No-Prune)).

We will update our notation and wording in Section 3.2 to make this clearer. Please let us know if you have any remaining questions. Thanks again!