Unveiling Invariances via Neural Network Pruning
Invariance describes transformations that do not alter data's underlying
semantics. Neural networks that preserve natural invariance capture good
inductive biases and achieve superior performance. Hence, modern networks are
handcrafted to handle well-known invariances (ex. translations). We propose a
framework to learn novel network architectures that capture data-dependent
invariances via pruning. Our learned architectures consistently outperform
dense neural networks on both vision and tabular datasets in both efficiency
and effectiveness. We demonstrate our framework on multiple deep learning
models across 3 vision and 40 tabular datasets.
Derek Xu, Yizhou Sun, Wei Wang
# Unveiling Invariances via Neural Network Pruning

###### Abstract

Invariance describes transformations that do not alter data's underlying semantics. Neural networks that preserve natural invariance capture good inductive biases and achieve superior performance. Hence, modern networks are handcrafted to handle well-known invariances (ex. translations). We propose a framework to learn novel network architectures that capture data-dependent invariances via pruning. Our learned architectures consistently outperform dense neural networks on both vision and tabular datasets in both efficiency and effectiveness. We demonstrate our framework on multiple deep learning models across 3 vision and 40 tabular datasets.

## Introduction

Preserving invariance is a key property in successful neural network architectures. Invariance occurs when the semantics of data remains unchanged under a set of transformations [1]. For example, an image of a cat can be translated, rotated, and scaled, without altering its underlying contents. Neural network architectures that represent data passed through invariant transformations with the same representation inherit a good inductive bias [20, 17, 18] and achieve superior performance [21, 22].

Convolutional Neural Networks (CNNs) are one such example. CNNs achieve translation invariance by operating on local patches of data and weight sharing. Hence, early CNNs outperform large multilayer perceptrons (MLP) in computer vision [14, 15]. Recent computer vision works explore more general spatial invariances, such as rotation and scaling [23, 24, 25, 26, 27, 28, 29, 30, 31, 32]. Other geometric deep learning works extend CNNs to non-Euclidean data by considering additional data-type specific invariances, such as permutation invariance [33, 25, 26].

Designing invariant neural networks requires substantial human effort: both to determine the set of invariant transformations and to handcraft architectures that preserve said transformations. In addition to being labor-intensive, this approach has not yet succeeded for all data-types [24, 25, 26, 27]. For example, designing neural architectures for tabular data is especially hard because the set of invariant tabular transformations is not clearly-defined. Thus, the state-of-the-art deep learning architecture on tabular data remains MLP [1, 26].

Existing invariance learning methods operate at the data augmentation level [10, 25, 26, 27], where a model is trained on sets of transformed samples rather than individual samples. This makes the network resilient to invariant transformations at test time. Contrastive learning (CL) is shown to be an effective means of incorporating invariance [25], and has seen success across various tasks [24, 25, 26, 27, 28, 29], including tabular learning [1]. While these approaches train existing neural networks to capture new data-dependent invariances, the model architecture itself still suffers from a weak inductive bias.

In contrast, existing network pruning works found shallow MLPs can automatically be compressed into sparse subnetworks with good inductive bias by pruning the MLP itself [26]. Combining pruning and invariance learning has largely been unsuccessful [27]. Furthermore, pruning for invariance does not scale to deep MLPs, possibly due to issues in the lazy training regime [24, 25] where performance improves while weights magnitudes stay near static over training. Combining invariance learning with network pruning remains an open question.

We propose Invariance **U**nveiling **N**eural **N**etworks, IUNet, a pruning framework that discovers invariance-preserving subnetworks from deep and dense supernetworks. We hypothesize pruning for invariances fails on deep networks due to the lazy training issue [12], where performance improvement decouples from weight magnitudes. We address this with a proactive initialization scheme (PIs), which prevents important weights from being accidentally pruned by assigning low magnitudes to major
ity of weights. To capture useful invariances, we propose a novel invariance learning objective (ILO), that successfully combines CL with network pruning by regularizing CL with the supervised objective.

To the best of our knowledge, we are the first to automatically design deep architectures that incorporate invariance using pruning. We summarize our contributions below:

* Designing architectures from scratch is difficult when desired invariances are either unknown or hard to incorporate. We automatically discover an invariance-preserving subnetwork that outperforms an invariance-agnostic supernetwork on both computer vision and tabular data.
* Network pruning is used to compress models for mobile devices. Our approach consistently improves compression performance for existing vision and tabular models.
* Contrastive learning traditionally fails with network pruning. We are the first to successfully combine contrastive learning with network pruning by regularizing it in our invariance learning objective.
* In the lazy training regime, performance improves drastically while weight magnitudes stay relatively constant, hence weights important for downstream performance may not have large magnitudes and hence be falsely pruned. We provide a simple yet effective approach that encourages only important weights to have large magnitudes before the lazy training regime begins.

## Related Work

### Learning Invariances

Most invariant networks are handcrafted for spatial invariances [1, 1, 1, 2, 3, 4, 16, 17, 18, 19]. Interestingly, subnetworks rarely outperform the original supernetwork, which has been dubbed the "Jackpot" problem [13]. In contrast to existing works, we successfully combine OMP with contrastive learning, alleviate the lazy learning issue, and outperform the original supernetwork.

## Proposed Method: IUNet

### Problem Setting

We study the classification task with inputs, \(x\in\mathcal{X}\), class labels, \(y\in\mathcal{Y}\), and hidden representations, \(h\in\mathcal{H}\). Our neural network architecture, \(f(x,\theta):\mathcal{X}\rightarrow\mathcal{Y}\) is composed of an encoder, \(f_{\mathcal{E}}(\cdot,\theta):\mathcal{X}\rightarrow\mathcal{H}\) and decoder, \(f_{\mathcal{D}}(\cdot,\theta):\mathcal{H}\rightarrow\mathcal{Y}\)

Figure 1: Overview for the IUNet Framework. The supernetwork, \(f^{P}(\cdot,\theta_{M})\), is initialized using PIS and trained on the ILO objective to obtain \(\theta_{M}^{(T)}\). Magnitude-based pruning is used to get a new architecture \(f^{P}=\mathcal{P}(\theta_{M}^{(T)})\). The new architecture, \(f^{P}(\cdot,\theta_{P})\), is initialized via lottery ticket reinitialization and finetuned with supervised maximum likelihood loss.
where \(\theta\in\Theta\) are the weights and \(f=f_{\mathcal{E}}\circ f_{\mathcal{D}}\). In the context of training, we denote the weights after \(0<t<T\) iterations of stochastic gradient descent as \(\theta^{(t)}\).

First, we define our notion of invariance. Given a set of invariant transformations, \(\mathcal{S}\), we wish to discover a neural network architecture \(f^{*}(x,\theta)\), such that all invariant input transformations map to the same representation, shown in Equation 1. We highlight our task focuses on the discovery of novel architectures, \(f^{*}(\cdot,\theta)\), not weights, \(\theta\), because good architectures improves the inductive bias [20].

\[f^{*}_{\mathcal{E}}(x,\theta)=f^{*}_{\mathcal{E}}(g(x),\theta),\forall g\in \mathcal{S},\forall\theta\in\Theta. \tag{1}\]

### Framework

We accomplish this by first training a dense supernetwork, \(f^{M}(\cdot,\theta_{M})\), with enough representational capacity to capture the desired invariance properties, as shown in Equation 2. A natural choice for \(f^{M}(\cdot,\theta_{M})\) is a deep MLP, which is a universal approximator [17].

\[\exists\theta^{*}_{M}\in\Theta_{M}:f^{M}_{\mathcal{E}}(x,\theta^{*}_{M})=f^{M }_{\mathcal{E}}(g(x),\theta^{*}_{M}),\forall g\in\mathcal{S}. \tag{2}\]

Next, we initialize the supernetwork's weights, \(\theta^{(0)}_{M}\), using our Proactive Initialization Scheme, PIS, and train the supernetwork with our Invariance Learning Objective, ILO, to obtain \(\theta^{(T)}_{M}\). We discuss both PIS's and ILO's details in following sections.

We construct our new untrained subnetwork, \(f^{P}(\cdot,\theta^{(0)}_{P})\), from the trained supernetwork, \(f^{M}(\cdot,\theta^{(T)}_{M})\), where the subnetwork contains a fraction of the supernetwork's weights, \(\theta^{(0)}_{P}\subset\theta^{(T)}_{M}\) and \(|\theta^{(0)}_{P}|\ll|\theta^{(T)}_{M}|\), and is architecturally different from the supernetwork, \(f^{P}(\cdot,\cdot)\neq f^{M}(\cdot,\cdot)\). For this step, we adopt standard One-shot Magnitude-based Pruning (OMP), where the smallest magnitude weights and their connections in the supernetwork architecture are dropped. We adopt OMP because of its success in neural network pruning [16, 1]. We represent this step as an operator mapping supernetwork weights into subnetwork architectures \(\mathcal{P}:\Theta_{M}\rightarrow\mathcal{F}_{P}\), where \(\mathcal{F}_{P}\) denotes the space of subnetwork architectures.

We hypothesize the trained subnetwork, \(f^{P}(\cdot,\theta^{(T)}_{P})\), can outperform the trained original supernetwork, \(f^{M}(\cdot,\theta^{(T)}_{M})\), if it learns to capture the right invariances and hence achieving a better inductive bias. The ideal subnetwork, \(f^{P*}(\cdot,\theta_{P*})\), could capture invariances without training, as shown in Equation 3.

\[f^{P*}_{\mathcal{E}}(x,\theta_{P*})=f^{P*}_{\mathcal{E}}(g(x),\theta_{P*}), \forall g\in\mathcal{S},\forall\theta_{P*}\in\Theta_{P*} \tag{3}\]

Finally, we re-initialize the subnetwork's weights, \(\theta^{(0)}_{P}\), using the Lottery Ticket Re-initialization scheme [16] then finetune the subnetwork with maximum likelihood loss to obtain \(\theta^{(T)}_{P}\). We call this framework, including the ILO objective and PIS initialization scheme, IUNet1, as shown in Figure 1

Footnote 1: IUNet prunes an ineffective supernetwork into an efficient effective subnetwork. OMP prunes an inefficient effective supernetwork into an efficient but slightly less effective subnetwork.

\begin{table}
\begin{tabular}{|c|c|c c c|} \hline Dataset & \(\text{MLP}_{\text{vis}}\) & \(\text{OMP}^{(\text{MLP}_{\text{vis}})}\) & \(\beta\text{-LASSO}^{(\text{MLP}_{\text{vis}})}\) & IUNet (\(\text{MLP}_{\text{vis}}\)) \\ \hline CIFAR10 & 59.266 \(\pm\) 0.050 & 59.668 \(\pm\) 0.171 & 59.349 \(\pm\) 0.174 & **64.847 \(\pm\) 0.121** \\ CIFAR100 & 31.052 \(\pm\) 0.371 & 31.962 \(\pm\) 0.113 & 31.234 \(\pm\) 0.354 & **32.760 \(\pm\) 0.288** \\ SVHN & 84.463 \(\pm\) 0.393 & 85.626 \(\pm\) 0.026 & 84.597 \(\pm\) 0.399 & **89.357 \(\pm\) 0.156** \\ \hline Dataset & ResNet & OMP (ResNet) & \(\beta\text{-LASSO}^{(\text{ResNet})}\) & IUNet (ResNet) \\ \hline CIFAR10 & 73.939 \(\pm\) 0.152 & 75.419 \(\pm\) 0.290 & 74.166 \(\pm\) 0.033 & **83.729 \(\pm\) 0.153** \\ CIFAR100 & 42.794 \(\pm\) 0.133 & 44.014 \(\pm\) 0.163 & 42.830 \(\pm\) 0.412 & **53.099 \(\pm\) 0.243** \\ SVHN & 90.235 \(\pm\) 0.127 & 90.474 \(\pm\) 0.192 & 90.025 \(\pm\) 0.201 & **94.020 \(\pm\) 0.291** \\ \hline \end{tabular}
\end{table}
Table 1: Comparing different pruning approaches to improve the inductive bias of \(\text{MLP}_{\text{vis}}\) and ResNet on computer vision datasets. Notice, IUNet performs substantially better than existing pruning-based methods by discovering novel architectures that better capture the inductive bias. IUNet flexibly boosts performance of off-the-shelf models.

\begin{table}
\begin{tabular}{|c|c|c c|} \hline Dataset & \(g(\cdot)\) & \(\text{MLP}_{\text{vis}}\) & IUNet (\(\text{MLP}_{\text{vis}}\)) \\ \hline \multirow{4}{*}{CIFAR10} & resize. & 44.096 \(\pm\) 0.434 & **97.349 \(\pm\) 4.590** \\  & horiz. & 80.485 \(\pm\) 0.504 & **99.413 \(\pm\) 1.016** \\  & color. & 56.075 \(\pm\) 0.433 & **98.233 \(\pm\) 3.060** \\  & graysc. & 81.932 \(\pm\) 0.233 & **99.077 \(\pm\) 1.598** \\ \hline Dataset & \(g(\cdot)\) & \(\text{MLP}_{\text{TAB}}\) & IUNet (\(\text{MLP}_{\text{TAB}}\)) \\ \hline mfeat. & feat. & 46.093 \(\pm\) 1.353 & **51.649 \(\pm\) 4.282** \\ \hline \end{tabular}
\end{table}
Table 2: Comparing the pruned IUNet (\(\text{MLP}_{\text{vis}}\)) model to an equivalent CNN. Although IUNet (\(\text{MLP}_{\text{vis}}\)) cannot outperform CNN, it bridges the gap between MLP and CNN architectures without any human design intervention.

\begin{table}
\begin{tabular}{|c|c|c c|} \hline Dataset & \(\text{MLP}_{\text{vis}}\) & OMP (\(\text{MLP}_{\text{vis}}\)) & \(\beta\text{-LASSO}^{(\text{MLP}_{\text{vis}})}\) & IUNet (\(\text{MLP}_{\text{vis}}\)) \\ \hline CIFAR10 & 59.266 \(\pm\) 0.050 & 59.668 \(\pm\) 0.171 & 59.349 \(\pm\) 0.174 & **64.847 \(\pm\) 0.121** \\ CIFAR100 & 31.052 \(\pm\) 0.371 & 31.962 \(\pm\) 0.113 & 31.234 \(\pm\) 0.354 & **32.760 \(\pm\) 0.288** \\ SVHN & 84.463 \(\pm\) 0.393 & 85.626 \(\pm\) 0.026 & 84.597 \(\pm\) 0.399 & **89.357 \(\pm\) 0.156** \\ \hline \end{tabular}
\end{table}
Table 1: Comparing different pruning approaches to improve the inductive bias of \(\text{MLP}_{\text{vis}}\) and ResNet on computer vision datasets. Notice, IUNet performs substantially better than existing pruning-based methods by discovering novel architectures that better capture the inductive bias. IUNet flexibly boosts performance of off-the-shelf models.
Invariance Learning Objective: ILOThe goal of supernetwork training is to create a subnetwork, \(f^{P}(\cdot,\theta_{P}^{(0)})\), within the supernetwork, \(f^{M}(\cdot,\theta_{M}^{(T)})\), such that:

1. \(\mathcal{P}(\theta_{M}^{(T)})\) achieves superior performance on the classification task after finetuning.
2. \(\mathcal{P}(\theta_{M}^{(T)})\) captures desirable invariance properties as given by Equation 3.
3. \(\theta_{P}^{(0)}\) has higher weight values than \(\theta_{M}^{(T)}\setminus\theta_{P}^{(0)}\).

Because subnetworks pruned from randomly initialized weights, \(\mathcal{P}(\theta_{M}^{(0)})\), are poor, they include harmful inductive biases that hinders training. Thus, we optimize the trained supernetwork, \(f^{M}(\cdot,\theta_{M}^{(T)})\), on goals (1) and (2) as a surrogate training objective. Goal (3) is handled by PIs, described in the next section.

To achieve (1), we maximize the log likelihood of the data. To achieve (2), we minimize a distance metric, \(\phi(\cdot,\cdot)\), between representations of inputs under invariant perturbations and maximize the metric between different input samples, given by Equation 9. We prove this is equivalent to Supervised Contrastive Learning (SCL) in the Supplementary Material. Hence, (2) can be achieved through SCL.

\[\theta_{M}^{\star}=\underset{\theta_{M}}{argmax}\underset{\begin{subarray}{c}x_ {j}\sim\mathcal{X}\\ g\sim\mathcal{S}\end{subarray}}{\mathbb{E}}\left[\frac{\phi(f_{\mathcal{E}}^{M} (x_{i},\theta_{M}),f_{\mathcal{E}}^{M}(x_{j},\theta_{M}))}{\phi(f_{\mathcal{E}} ^{M}(x_{i},\theta_{M}),f_{\mathcal{E}}^{M}(g(x_{i}),\theta_{M}))}\right] \tag{4}\]

Our final Invariance Learning Objective (ILO) loss function combines these two ideas as shown in Equation 6, where \(\mathcal{L}_{SUP}\) is standard maximum likelihood loss, \(\mathcal{L}_{NCE}\) is a contrastive loss (described in Appendix), \(D_{tr}\) is a labelled training dataset of \((x,y)\) pairs, and \(\lambda\) is a hyperparameter.

\[\mathcal{L}(\theta_{M};\mathcal{S})= \tag{5}\] \[\underset{x,y\sim D_{tr}}{\mathbb{E}}\left[\mathcal{L}_{SUP}(x,y, \theta_{M})+\lambda\mathcal{L}_{NCE}(x,y,\theta_{M};\mathcal{S})\right]\]

Both loss components are crucial to IUNet. With just \(L_{NCE}\), the supernetwork will overfit the contrastive objective [10, 11], causing weights critical for finetuning the supervised objective to be pruned. With just \(L_{SUP}\), the architecture is not explicitly optimized to capture desired invariances.

Proactive Initialization Scheme: PIsDeep neural networks often enter the lazy training regime [13, 12], where the loss steadily decreases while weights barely change. This is particularly harmful to neural networks pruning [12], especially when low-magnitude weights contribute to decreasing the loss and hence should not be pruned.

We propose a simple solution by scaling the weight initialization by a small multiplier, \(\kappa\). We find this alleviates the aforementioned issue by forcing the model to assign large values only to important weights prior to lazy training. Because lazy training is only an issue for pruning, we only apply \(\kappa\)-scaling to the pre-pruning training stage, not the fine-tuning stage. This is done by scaling the initial weights \(\theta_{M}^{(0)}=\kappa\theta_{M^{\dagger}}^{(0)}\), where \(\theta_{M^{\dagger}}^{(0)}\) follows the Kaiming [10] or Glorot [11] initialization.

## Experiment Setup

### Datasets

IUNet is evaluated on _image_ and _tabular_ classification 2:

Footnote 2: More details are provided in the Supplementary.

* **Vision**: Experiments are run on CIFAR10, CIFAR100, and SVHN [13, 12], following baseline work [2]3.
* **Tabular**: Experiments are run on 40 tabular datasets from a benchmark paper [1], covering a diverse range of problems. The datasets were collected from OpenML [10], UCI [1], and Kaggle.

Figure 2: Effect of PIs and ILO on pruned models. The y-axis is the validation accuracy (%) and x-axis is the compression ratio. PIs experiments only alter the supernetwork’s initialization. \(\kappa=1.0\) means normal initialization. ILO experiments only alter the training objective during supernetwork training. After supernetwork training, subnetworks are pruned under different compression ratios, then finetuned. Validation accuracy of trained pruned models are reported.
### Model Setup

IUnet is compared against One-shot Magnitude Pruning (OMP) [1], and \(\beta\)-Lasso pruning [20] on all datasets. We denote the supernetwork used by each pruning method with a superscript. Unless otherwise specified, models are trained via maximum likelihood. In addition, we compare against the following dataset-specific supernetworks (\(\text{MLP}_{\text{VIS}}\), \(\text{MLP}_{\text{TAB}}\), ResNet) and models:

* **Vision**: We consider ResNet[16], \(\text{MLP}_{\text{VIS}}\), a MLP that contains a CNN subnetwork [20], and the aforementioned CNN subnetwork.
* **Tabular**: We consider \(\text{MLP}_{\text{TAB}}\), a 9-layer MLP with hidden dimension 512 [14], XGB [2], TabN[1], a handcrafted tabular deep learning architecture, and \(\text{MLP}_{\text{TAB+C}}\)[15], the state-of-the-art MLP, which was heavily tuned from a cocktail of regularization techniques.

### Considered Invariances

Because contrastive learning is successful on both vision and tabular datasets, our invariant transformations, \(\mathcal{S}\), come from existing works. For computer vision, SimCLR [2] transformations are used: (1) resize crops, (2) horizontal flips, (3) color jitter, and (4) random grayscale. For tabular learning, SCARF [1] transformations are used: (5) randomly corrupting features by drawing the corrupted versions from its empirical marginal distribution.

## Results

### On Inductive Bias

In this section, we compare the effectiveness of the trained subnetwork discovered by IUNet, \(f^{P}(\cdot,\theta_{P}^{(T)})\), against the trained supernetwork, \(f^{M}(\cdot,\theta_{M}^{(T)})\). As seen in Tables 1 and 5, the pruned subnetwork outperforms the original supernetwork, even though the supernetwork has more representational capacity. This supports our claim that IUNet prunes subnetwork architectures with better inductive biases than the supernetwork. Importantly, IUNet substantially improves upon existing pruning baselines by explicitly including invariances via ILO and alleviating the lazy learning issue [11] via PIs.

On _vision_ datasets: As seen in Figure 1, IUNet is a general and flexible framework that improves the inductive bias of not only models like \(\text{MLP}_{\text{VIS}}\) but also specialized architectures like ResNet. As seen in Figure 2, IUNet bridges the gap between MLPs and CNNs. Unlike previous work [16], IUNet does this in an entirely automated procedure. While CNN outperforms IUNet (\(\text{MLP}_{\text{VIS}}\)), we can also apply IUNet to specialized networks, \(\text{IUNet}\) (\(\text{ResNet}\)), which achieves the best overall performance. Figures 1 and 2 show IUNet is a useful tool for injecting inductive bias into arbitrary neural architectures.

On _tabular_ datasets: As seen in Figure 5, the subnetworks derived from MLPs outperform both the original \(\text{MLP}_{\text{TAB}}\) and hand-crafted architectures: TabN and XGB. Unlike vision, how to encode invariances for tabular data is highly nontrivial, making IUNet particularly effective. Note, although IUNet performs competitively against \(\text{MLP}_{\text{TAB+C}}\), they are orthogonal approaches. \(\text{MLP}_{\text{TAB+C}}\) focuses on tuning regularization hyperparameters during finetuning, whereas IUNet improves the model architecture. Note, \(\text{IUNet}\) (\(\text{MLP}_{\text{TAB}}\)) did not use the optimal hyperparameters found by \(\text{MLP}_{\text{TAB+C}}\)[15].

### Ablation Study

To study the effectiveness of (1) pruning, (2) PIs, and (3) ILO, each one is removed from the optimal model. As seen in Table 4, each is crucial to IUNet. Pruning is necessary to encode the inductive bias into the subnetwork's neural architecture. PIs and ILO improves the pruning policy by ensuring weights crucial to finetuning and capturing invariance are not pruned. Notice, without pruning, IUNet No-pune performs worse than the original supernetwork. This highlights an important notion that PIs aims to improve the pruning policy, not the unpruned performance. By sacrificing unpruned performance, PIs ensures important weights are not falsely pruned. PIs is less effective on tabular datasets where the false pruning issue seems less severe. Combining pruning, ILO, and PIs, IUNet most consistently achieves the best performance.
upon existing pruning policies. First, our results support existing findings that (1) OMP does not produce subnetworks that substantially outperform the supernetwork [1] and (2) while unpruned models trained with SCL can outperform supervised ones, pruned models trained with SCL perform substantially worse [14]. PIS flips the trend from (1) by slightly sacrificing unpruned performance, due to poorer initialization, IUNet discovers pruned models with better inductive biases, which improves downstream performance. ILO fixes the poor performance of SCL in (2) by preserving information pathways for both invariance and max likelihood over training. We highlight both these findings are significant among the network pruning community. Finally, Figure 2 confirms IUNet achieves the best performance by combining both PIS and ILO.

In addition to being more effective that the supernetwork, \(f^{M}(\cdot,\theta_{M}^{(T)})\), the pruned network, \(f^{P}(\cdot,\theta_{p}^{(T)})\), is also more efficient. Figure 2 shows IUNet can reach 8-16\(\times\) compression while still keeping superior performance.

### Effect of Proactive Initialization

To further study the role of PIS, the histogram of weight magnitudes is monitored over the course of training. As shown in Figure 5, under the standard OMP pruning setup, the histogram changes little over the course of training, which supports the lazy training hypothesis [13] where performance rapidly improves, while weight magnitudes change very little, decoupling each weight's importance from its magnitude.

With PIS, only important weights grow over the course of training, while most weights remain near zero, barely affecting the output activations of each layer. This phenomenon alleviates the lazy training problem by ensuring (1) pruning safety, as pruned weights are near zero prior which have minimal affect on layer activations, and (2) importance-magnitude coupling, as structurally important connections must grow to affect the output of the layer.

### On Invariance Consistency

To further study whether particular invariances are learned, we compute the consistency metric [20], which measure the percentage of samples whose predicted label would flip when an invariant transformation is applied to the input. As seen in Table 3, the subnetwork found by IUNet, \(f^{P}(\cdot,\theta_{P}^{(0)})\), is able to preserve invariances specified in ILO much better than the supernetwork, \(f^{M}(\cdot,\theta_{M}^{(0)})\). This shows IUNet indeed captures desirable invariances.

### On Weight Visualization

We visualize the supernetwork weights, \(\theta_{M}^{(T)}\), when trained with IUNet compared to standard maximum likelihood (MLP) to determine what structures preserve invariance.

On _vision_ datasets: As seen in Figure 6, IUNet learns more locally connected structures, which improves translation invariance. Prior work [22] found network structure (as opposed to inductive bias) to be the limiting factor for encoding CNN inductive biases into MLPs, which IUNet successfully replicates.

On _tabular_ datasets: As seen in Figure 6, IUNet weights focus more on singular features. This preserves invariance over random feature corruption, as the absence of some tabular features does not greatly alter output activations of most neurons. This structure can also be likened to tree ensembles [11], whose leaves split individual features rather than all features.

## Conclusion

In this work, we study the viability of network pruning for discovering invariant-preserving architectures. Under the computer vision setting, IUNet bridges the gap between deep MLPs and deep CNNs, and reliably boosts ResNet performance. Under the tabular setting, IUNet reliably boosts performance of existing MLPs, comparable to applying the state-of-the-art regularization cocktails. Our proposed novelties, ILO and PIS, flexibly improves existing OMP pruning policies by both successfully integrating contrastive learning and alleviating lazy training. Thus, IUNet effectively uses pruning to tackle invariance learning.

Figure 4: Visualization of weight magnitudes, \(|\theta_{M}^{(T)}|\), trained with different policies. The top row was trained on CIFAR10 and shows the magnitude of each RGB pixel for 6 output logits. The bottom row was trained on arrhythmia and shows the weight matrix of the 1st layer with 280 input and 512 output dimensions. Lighter color means larger magnitude.

Figure 3: Histogram of weight magnitudes, \(|\theta_{M}^{(t)}|\), plotted over each epoch under different \(\kappa\) initializations settings. \(\kappa=1.0\) means normal initialization. Results shown for MLP\({}_{\text{VIS}}\) on the CIFAR10 dataset.
\begin{table}
\begin{tabular}{|c|c|c c c|c c c|} \hline Dataset & MLP\({}_{\text{TAB}}\) & OMP (\({}^{\text{MLP}_{\text{TAB}}}\)) & \(\beta\)-Lasso (\({}^{\text{MLP}_{\text{TAB}}}\)) & IUNet (\({}^{\text{MLP}_{\text{TAB}}}\)) & XGB & TabN & MLP\({}_{\text{TAB+C}}\) \\ \hline credit-g & 70.000 & 70.000 \(\pm\) 0.000 & 67.205 \(\pm\) 0.718 & 63.166 \(\pm\) 0.000 & 68.929 & 61.190 & **74.643** \\ anneal & 99.490 & 99.691 \(\pm\) 0.234 & 99.634 \(\pm\) 0.518 & **99.712 \(\pm\) 0.101** & 85.416 & 84.248 & 89.270 \\ kr-vs-kp & 99.158 & 99.062 \(\pm\) 0.142 & 99.049 \(\pm\) 0.097 & 99.151 \(\pm\) 0.064 & **99.850** & 93.250 & **99.850** \\ arrhythmia & 67.086 & 55.483 \(\pm\) 5.701 & 67.719 \(\pm\) 3.483 & **74.138 \(\pm\) 2.769** & 48.779 & 43.562 & 61.461 \\ mfeat. & 98.169 & 97.959 \(\pm\) 0.278 & 97.204 \(\pm\) 0.620 & **98.176 \(\pm\) 0.121** & 98.000 & 97.250 & 98.000 \\ vehicle & 80.427 & 81.115 \(\pm\) 2.214 & 80.611 \(\pm\) 1.244 & 81.805 \(\pm\) 2.065 & 74.973 & 79.654 & **82.576** \\ kc1 & 80.762 & **84.597 \(\pm\) 0.000** & 83.587 \(\pm\) 1.010 & **84.597 \(\pm\) 0.000** & 66.846 & 52.517 & 74.381 \\ adult & 81.968 & 82.212 \(\pm\) 0.582 & 82.323 \(\pm\) 0.424 & 78.249 \(\pm\) 3.085 & 79.824 & 77.155 & **82.443** \\ walking. & 58.466 & 60.033 \(\pm\) 0.112 & 58.049 \(\pm\) 0.309 & 59.789 \(\pm\) 0.456 & 61.616 & 56.801 & **63.923** \\ phoneme & 84.213 & 86.733 \(\pm\) 0.194 & 84.850 \(\pm\) 1.548 & 87.284 \(\pm\) 0.436 & **87.972** & 86.824 & 86.619 \\ skin-seg. & 99.869 & 99.866 \(\pm\) 0.016 & 99.851 \(\pm\) 0.015 & 99.876 \(\pm\) 0.006 & **99.968** & 99.961 & 99.953 \\ ldpa & 66.590 & 68.458 \(\pm\) 0.140 & 62.362 \(\pm\) 4.605 & 64.816 \(\pm\) 4.535 & **99.008** & 54.815 & 68.107 \\ nomao & 95.776 & 95.682 \(\pm\) 0.046 & 95.756 \(\pm\) 0.074 & 95.703 \(\pm\) 0.110 & **96.872** & 95.425 & 96.826 \\ cnea & 94.080 & 92.742 \(\pm\) 0.404 & 94.808 \(\pm\) 0.254 & **96.075 \(\pm\) 0.242** & 94.907 & 89.352 & 95.833 \\ blood. & 68.965 & 61.841 \(\pm\) 10.012 & 65.126 \(\pm\) 20.792 & **70.375 \(\pm\) 5.255** & 62.281 & 64.327 & 67.617 \\ bank. & **88.300** & **88.300 \(\pm\) 0.000** & 86.923 \(\pm\) 1.948 & **88.300 \(\pm\) 0.000** & 72.658 & 70.639 & 85.993 \\ connect. & 72.111 & 72.016 \(\pm\) 0.112 & 72.400 \(\pm\) 0.214 & 74.475 \(\pm\) 0.445 & 72.374 & 72.045 & **80.073** \\ shuttle & 99.709 & 93.791 \(\pm\) 3.094 & 99.687 \(\pm\) 0.027 & 93.735 \(\pm\) 2.303 & 98.563 & 88.017 & **99.948** \\ higgs & 72.192 & 72.668 \(\pm\) 0.039 & 72.263 \(\pm\) 0.149 & 73.215 \(\pm\) 0.384 & 72.944 & 72.036 & **73.546** \\ australian & 82.153 & 83.942 \(\pm\) 1.578 & 81.667 \(\pm\) 1.572 & 82.562 \(\pm\) 1.927 & **89.717** & 85.278 & 87.088 \\ car & 99.966 & **100.000 \(\pm\) 0.000** & **100.000 \(\pm\) 0.000** & 99.859 \(\pm\) 0.200 & 92.376 & 98.701 & 99.587 \\ segment & 91.504 & 91.603 \(\pm\) 0.508 & 91.317 \(\pm\) 0.074 & 91.563 \(\pm\) 0.000 & **93.723** & 91.775 & **93.723** \\ fashion. & 91.139 & 90.784 \(\pm\) 0.158 & 90.864 \(\pm\) 0.090 & 90.817 \(\pm\) 0.040 & 91.243 & 89.793 & **91.950** \\ jungle. & 86.998 & 92.071 \(\pm\) 0.420 & 87.400 \(\pm\) 0.489 & 95.130 \(\pm\) 0.807 & 87.325 & 73.425 & **97.471** \\ numerai & 51.621 & 51.443 \(\pm\) 0.370 & 51.905 \(\pm\) 0.299 & 51.839 \(\pm\) 0.067 & 52.363 & 51.599 & **52.668** \\ deuragari & 97.550 & 97.573 \(\pm\) 0.031 & 97.549 \(\pm\) 0.014 & 97.517 \(\pm\) 0.014 & 93.310 & 94.179 & **98.370** \\ helena & 29.342 & 28.459 \(\pm\) 0.531 & 29.834 \(\pm\) 0.354 & **29.884 \(\pm\) 0.991** & 21.994 & 19.032 & 27.701 \\ jannis & 68.647 & 66.302 \(\pm\) 3.887 & 69.302 \(\pm\) 0.248 & **69.998 \(\pm\) 1.232** & 55.225 & 56.214 & 65.287 \\ volkert & 70.066 & 68.781 \(\pm\) 0.045 & 69.655 \(\pm\) 0.189 & 70.104 \(\pm\) 0.215 & 64.170 & 59.409 & **71.667** \\ miniboone & 86.539 & 87.575 \(\pm\) 0.855 & 87.751 \(\pm\) 0.398 & 81.226 \(\pm\) 6.569 & **94.024** & 62.173 & 94.015 \\ apsfailure & 97.041 & **98.191 \(\pm\) 0.000** & 98.048 \(\pm\) 0.203 & **98.191 \(\pm\) 0.000** & 88.825 & 51.444 & 92.535 \\ christine & 70.295 & 69.819 \(\
## Additional Related Work

### Tabular Machine Learning

Tabular data is a difficult regime for deep learning, where deep learning models struggle against decision tree approaches. Early methods use forests, ensembling, and boosting Shwartz-Ziv and Armon (2022); Borisov et al. (2022); Chen and Guestrin (2016). Later, researchers handcrafted new deep architectures that mimic trees Popov, Morozov, and Babenko (2019); Arik and Pfister (2021). Yet, when evaluated on large datasets, these approaches are still beaten by XGB Chen and Guestrin (2016); Grinsztajn, Oyallon, and Varoquaux (2022). Recent work found MLPs with heavy regularization tuning Kadra et al. (2021) can outperform decision tree approaches, though this conclusion does not hold on small tabular datasets Joseph and Raj (2022). To specially tackle the small data regime, Bayesian learning and Hopfield networks are combined with MLPs Hollmann et al. (2022); Schafl et al. (2022). There are also work on tabular transformers Huang et al. (2020), though said approaches require much more training data. Without regularization, tree based models still outperform MLPs due to a better inductive bias and resilience to noise Grinsztajn, Oyallon, and Varoquaux (2022). To the best of our knowledge, the state-of-the-art on general tabular datasets remain heavily regularized MLPs (MLP\({}_{\texttt{TAB+C}}\)) Kadra et al. (2021). We aim to further boost regularized MLP performance by discovering model architectures that capture good invariances from tabular data.

### Contrastive Learning

Contrastive learning, initially proposed for metric learning Chopra et al. (2005); Schroff, Kalenichenko, and Philbin (2015); Oh Song et al. (2016), trains a model to learn shared features among images of the same type Jaiswal et al. (2020). It has been widely used in self-supervised pretraining Chen et al. (2020); Chen et al. (2021), where dataset augmentation is crucial. Although contrastive learning was originally proposed for images, it has also shown promising results in graph data Zhu et al. (2019); You et al. (2020), speech data Baewski et al. (2020), and tabular data Bahri et al. (2021). Previous study has showed that speech transformers tend to overfit the contrastive loss in deeper layers, suggesting that removing later layers can be beneficial during finetuning Pasad, Chou, and Livescu (2021). While contrastive learning performs well pretraining unpruned models, its vanilla formulation performs poorly after network pruning Corti et al. (2022). In this work, we establish a connection between contrastive learning and invariance learning and observe that pruned contrastive models fail because of overfitting.

### Neural Architecture Search

Neural Architecture Search (NAS) explores large superarchitectures by leveraging smaller block architectures Wan et al. (2020); Pham et al. (2018); Zoph et al. (2018); Luo et al. (2018); Liu, Simonyan, and Yang (2018). These block architectures are typically small convolutional neural networks (CNNs) or MLPs. The key idea behind NAS is to utilize these blocks Pham et al. (2018); Zoph et al. (2018) to capture desired invariance properties for downstream tasks. Prior works You et al. (2020); Xie et al. (2019) have analyzed randomly selected intra- and inter-block structures and observed performance differences between said structures. However, these work did not propose a method for discovering block architectures directly from data. Our work aims to address this gap by focusing on discovering the architecture within NAS blocks. This approach has the potential to enable NAS in diverse domains, expanding its applicability beyond the current scope.

## Loss Function Details

We provide a more detailed description of our loss function in this section. Following notation from the main paper, we repeat the ILO loss function in Equation 6 below:

\[\mathcal{L}(\theta;\mathcal{S})=\mathbb{E}_{x,y\sim D_{train}}\left[\mathcal{L} _{SUP}(x,y,\theta)+\lambda\mathcal{L}_{NCE}(x,y,\theta;\mathcal{S})\right] \tag{6}\]

To better explain our loss functions, we introduce some new notations. First, we denote the decoder output probability function over classes, \(\mathcal{Y}\), as \(\tilde{p}_{\theta_{\mathcal{D}}}:\mathcal{H}\rightarrow[0,1]^{|\mathcal{Y}|}\), where \(f_{\mathcal{D}}=\text{argmax}\propto\tilde{p}_{\theta_{\mathcal{D}}}\). We denote the model output probability function by combining \(\tilde{p}_{\theta_{\mathcal{D}}}\) with the encoder as follows: \(p_{\theta}=\tilde{p}_{\theta_{\mathcal{D}}}\circ f_{\mathcal{E}}\). We introduce an integer mapping from classes \(\mathcal{Y}\) as \(\mathcal{I}:\mathcal{Y}\rightarrow\{0,1,2,...,|\mathcal{Y}|-1\}\).

We show the maximum likelihood loss, \(\mathcal{L}_{SUP}\), in Equation 7 below.

\[\mathcal{L}_{SUP}(x,y,\theta)=-log(p_{\theta}(x,\theta)_{\mathcal{I}(y)}) \tag{7}\]

We show the supervised contrastive loss, \(\mathcal{L}_{NCE}\), in Equation 8 below. Following SimCLR Chen et al. (2020), we assume that the intermediary representations are \(d\)-dimensional embeddings, \(\mathcal{H}=\mathbb{R}^{d}\), and use the cosine similarity as our similarity function, \(\psi^{(cos)}:\mathbb{R}^{d}\times\mathbb{R}^{d}\rightarrow\mathbb{R}\).

\[\mathcal{L}_{NCE}(x,y,\theta;\mathcal{S})= \tag{8}\]

### Surrogate Objective

We aim to learn invariance-preserving network architectures from the data. In our framework, this involves optimizing our invariance objective, which we repeat in Equation 9. We now prove that by minimizing the supervised contrastive loss in Equation 8 we equivalently maximize the invariance objective, outlined below.

\[\theta^{*}=\underset{\theta}{argmax}\underset{x_{i},x_{j}\sim\mathcal{X}}{ \mathbb{E}}\left[\frac{\phi(f_{\mathcal{E}}^{M}(x_{i},\theta),f_{\mathcal{E}}^{ M}(x_{j},\theta))}{\phi(f_{\mathcal{E}}^{M}(x_{i},\theta),f_{\mathcal{E}}^{M}(g(x_{i} ),\theta))}\right] \tag{9}\]

We convert the distance metric \(\phi\) into similarity metric \(\psi\).
\[\theta^{*}=\underset{\theta}{argmax}\underset{x_{i},x_{j}\sim\mathcal{X}}{ \mathbb{E}}\left[\frac{\psi(f_{\mathcal{E}}^{M}(x_{i},\theta),f_{\mathcal{E}}^{M} (g(x_{i}),\theta))}{\psi(f_{\mathcal{E}}^{M}(x_{i},\theta),f_{\mathcal{E}}^{M} (x_{j},\theta))}\right]\] \[=\underset{\theta}{argmin}\underset{x_{i},x_{j}\sim\mathcal{X}}{ \mathbb{E}}\left[\frac{-\psi(f_{\mathcal{E}}^{M}(x_{i},\theta),f_{\mathcal{E}}^ {M}(g(x_{i}),\theta))}{\psi(f_{\mathcal{E}}^{M}(x_{i},\theta),f_{\mathcal{E}}^ {M}(x_{j},\theta))}\right]\] \[=\underset{\theta}{argmin}\underset{x_{i},x_{j}\sim\mathcal{X}}{ \mathbb{E}}\left[\frac{-\psi(f_{\mathcal{E}}^{M}(x,\theta),f_{\mathcal{E}}^{M} (g(x),\theta))}{\sum\limits_{x^{\prime}\neq x}\psi(f_{\mathcal{E}}^{M}(x, \theta),f_{\mathcal{E}}^{M}(g(x^{\prime}),\theta))}\right]\] \[=\underset{\theta}{argmin}\underset{x_{i},y\sim\mathcal{D}_{tr}}{ \mathbb{E}}\left[\frac{-\psi\left(\frac{f_{\mathcal{E}}^{M}(x,\theta)}{f_{ \mathcal{E}}^{M}(g(x),\theta)}\right)}{\sum\limits_{x^{\prime},y^{\prime} \sim D_{tr}}\psi\left(\frac{f_{\mathcal{E}}^{M}(x,\theta)}{f_{\mathcal{E}}^{M} (g(x^{\prime}),\theta)}\right)}\right]\] \[=\underset{\theta}{argmin}\underset{x_{i},y\sim D_{tr}}{ \mathbb{E}}\left[-log\left(\frac{\psi\left(\frac{f_{\mathcal{E}}^{M}(x, \theta)}{f_{\mathcal{E}}^{M}(g(x),\theta)}\right)}{\sum\limits_{x^{\prime},y^{ \prime}\sim D_{tr}}\psi\left(\frac{f_{\mathcal{E}}^{M}(g,\theta)}{f_{\mathcal{ E}}^{M}(g(x^{\prime}),\theta)}\right)}\right)\right] \tag{10}\]

We set the similarity metric, \(\psi\), to be the same as our contrastive loss: \(\psi(\cdot)=exp(\psi^{(cos)}(\cdot))\).

\[\theta^{*}=\underset{\theta}{argmin}\underset{x,y\sim D_{train}}{ \mathbb{E}}[\mathcal{L}_{NCE}(x,y,\theta;\mathcal{S})] \tag{11}\]

Here, we showed that the vanilla contrastive loss function, Equation 8, serves as a surrogate objective for optimizing our desired invariance objective, Equation 9. By incorporating contrastive learning alongside the maximum likelihood objective in Equation 6, ILO effectively reveals the underlying invariances in the pruned model.

## Additional Discussion on Lazy Training

The lazy training regime [14, 15] is a phenomenon when loss rapidly decreases, while weight values stay relatively constant. This phenomenon occurs on large over-parameterized neural networks [14]. Because the weight values stay relatively constant, the magnitude ordering between weights also changes very little. Therefore, network pruning struggles to preserve such loss decreases in the lazy training regime [14].

Because weights with very small magnitude have minimal effect on the output logits, pruning said weights will not drastically hurt performance. Thus, if the pruning framework can separate very small magnitude weights from normal weights prior to the lazy training regime, we can preserve loss decreases in the lazy training regime. The PIS setting accomplishes this by initializing all weights to be very small so that only important weights will learn large magnitudes. This guarantees that a large percentage of weights will have small magnitudes throughout training, while important larger magnitude weights will emerge over the course of training.

## Additional Experiments

### Effects of Proactive Initialization: Full Results

We provide weight histograms on CIFAR100 and SVHN in Figure 5. As shown, the trends reported in the main text holds on other datasets.

### On Weight Visualization: Full Results

We provide weight visualizations on CIFAR100 and SVHN in Figure 6. As shown, the trends reported in the main text holds on other datasets.

### On Consistency: Full Results

We provide consistency experiments on CIFAR100 and SVHN in Table 6. As shown, the trends reported in the main text holds on other datasets.

## Implementation Details

### Dataset Details

We considered the following _computer vision_ datasets: CIFAR10, CIFAR100 [14], and SVHN [15]. CIFAR10 and CIFAR100 are multi-domain image classification datasets. SVHN is a street sign digit classification dataset. Input images are \(32\times 32\) color images. We split the train set by 80/20 for training and validation. We test on the test set provided separately. We reported dataset statistics in Table 7.

We considered 40 tabular datasets from OpenML [16], UCI [10], and Kaggle, following the \(\text{MLP}_{\text{TAB+C}}\) benchmark [14]. These tabular datasets cover a variety of domains, data types, and class imbalances. We used a 60/20/20 train validation test split, and reported dataset statistics in Table 8. We use a random seed of 11 for the data split, following prior work [14].

\begin{table}
\begin{tabular}{|c|c|c c|} \hline Dataset & \(g(\cdot)\) & \(\text{MLP}_{\text{VIS}}\) & \(\text{IUNet}^{\text{(MLP}_{\text{VIS})}}\) \\ \hline \multirow{4}{*}{CIFAR10} & resize. & 44.096 \(\pm\) 0.434 & **97.349 \(\pm\) 4.500** \\  & horiz. & 80.485 \(\pm\) 0.504 & **99.413 \(\pm\) 1.016** \\  & color. & 56.075 \(\pm\) 0.433 & **98.233 \(\pm\) 3.060** \\  & grayscale. & 81.932 \(\pm\) 0.233 & **99.077 \(\pm\) 1.598** \\ \hline \multirow{4}{*}{CIFAR100} & resize. & 32.990 \(\pm\) 1.065 & **39.936 \(\pm\) 2.786** \\  & horiz. & 70.793 \(\pm\) 0.677 & **77.935 \(\pm\) 1.464** \\  & color. & 31.704 \(\pm\) 0.560 & **51.397 \(\pm\) 2.709** \\  & grayscale. & 71.245 \(\pm\) 0.467 & **76.476 \(\pm\) 1.245** \\ \hline \multirow{4}{*}{SVHN} & resize. & 36.708 \(\pm\) 2.035 & **77.440 \(\pm\) 0.627** \\  & horiz. & 71.400 \(\pm\) 1.651 & **95.082 \(\pm\) 0.166** \\  & color. & 61.341 \(\pm\) 0.946 & **91.097 \(\pm\) 0.395** \\ \cline{1-1}  & grayscale. & 90.344 \(\pm\) 0.233 & **99.259 \(\pm\) 0.073** \\ \hline \end{tabular}
\end{table}
Table 6: Comparing the consistency metric (%) of the untrained supernetwork, \(\text{MLP}_{\text{VIS}}\) and \(\text{MLP}_{\text{TAB}}\), against IUNet’s pruned subnetwork under different invariant transforms, \(g(\cdot)\). IUNet preserves invariances better.
Figure 5: Histogram of weight magnitudes, \(|\theta_{M}^{(t)}|\), plotted over each epoch under different \(\kappa\) initializations settings. \(\kappa=1.0\) means normal initialization. Results shown for \(\text{MLP}_{\text{VIS}}\) on the CIFAR10, CIFAR100, and SVHN datasets.

Figure 6: Visualization of weight magnitudes, \(|\theta_{M}^{(T)}|\), trained with different policies. The models were trained on CIFAR10, CIFAR100, and SVHN. The magnitude of each RGB pixel for 6 output logits are plotted.
### Hyperparameter Settings

All experiments were run 3 times from scratch starting with different random seeds. We report both the mean and standard deviation of all runs. All hyperparameters were chosen based on validation set results.

For all experiments, we used \(\lambda=1\), which was chosen through a grid search over \(\lambda\in\{0.25,0.5,1.0\}\). For all experiments, we used a batch size of 128. For pre-pruning training, we used SGD with Nesterov momentum and a learning rate of 0.001, following past works [1]. For finetuning vision datasets, we used the same optimizer setup except with 16-bit operations except for batch normalization, following \(\beta\)-Lasso [22]. For finetuning tabular datasets, we used AdamW [15], a learning rate of 0.001s, decoupled weight decay, cosine annealing with restart, initial restart budget of 15 epochs, budget multiplier of 2, and snapshot ensembling [16], following prior works [10, 11]. It is important to note we did not tune the dataset and training hyperparameters for each tabular dataset individually like \(\text{MLP}_{\text{TAB+C}}\)[10], rather taking the most effective setting on average.

For tabular datasets, we tuned the compression ratio over the following range of values: \(r\in\{2,4,8\}\) and the PIS multiplier over the following range of values: \(\kappa\in\{0.25,0.125,0.0625\}\). on a subset of 4 tabular datasets. We found that \(r=8\) and \(\kappa=0.25\) performs the most consistently and used this setting for all runs of IUNet in the main paper. It is important to note we did not tune hyperparameters for IUNet on each individual tabular dataset like \(\text{MLP}_{\text{TAB+C}}\)[10], making IUNet a much more efficient model than \(\text{MLP}_{\text{TAB+C}}\). For the tabular baselines [10, 10, 11], we used the same hyperparameter tuning setup as the MLP+C benchmark [10].

For vision datasets, we tuned the compression ratio over the following range of values: \(r\in\{2,4,8,16\}\) on each individual dataset for all network pruning models except \(\beta\)-Lasso4. For \(\beta\)-lasso [22], we tuned the hyperparameters over the range \(\beta=\{50\}\) and L1 regularization in \(l1\in\{10^{-6},2\times 10^{-6},5\times 10^{-6},10^{-5},2\times 10^{-5}\}\) on each individual dataset as done in the original paper. It is important to note that although we tuned both hyperparameters for both IUNet and baselines on each individual datasets, our main and ablation table rankings stay consistent had we chosen a single setting for all datasets, as shown in the detailed pruning experiments in the main paper.

Footnote 4: This is because \(\beta\)-Lasso does not accept a chosen compression ratio as a hyperparameter.

### Supernetwork Architecture

\(\text{MLP}_{\text{VIS}}\) is a deep MLP that contains a CNN subnetwork. Given a scaling factor, \(\alpha\), the CNN architecture consists of 3x3 convolutional layers with the following (out channels, stride) settings: \([(\alpha,1)\), \((2\alpha,2)\), \((2\alpha,1)\), \((4\alpha,2)\), \((4\alpha,1)\), \((8\alpha,2)\), \((8\alpha,1)\), \((16\alpha,2)]\) followed by a hidden layer of dimension \(64\alpha\). It is worth noting that our CNN does not include maxpooling layers for fair comparison with the learned architectures, following the same setup as \(\beta\)-Lasso [22]. To form the MLP Network, we ensured the CNN network structure exists as a subnetwork within the MLP supernetwork by setting the hidden layer sizes to: \([\alpha s^{2},\frac{\alpha s^{2}}{2},\frac{\alpha s^{2}}{2},\frac{\alpha s^{2 }}{4},\frac{\alpha s^{2}}{4},\frac{\alpha s^{2}}{8},\frac{\alpha s^{2}}{8}, \frac{\alpha s^{2}}{16},64\alpha]\). This architecture was also introduced in \(\beta\)-Lasso [22]. All layers are preceded by batch normalization and ReLU activation. We chose \(\alpha=8\) such that our supernetwork can fit onto an Nvidia RTX 3070 GPU.

CNN is the corresponding CNN subnetwork with (out channels, stride) settings: \([(\alpha,1)\), \((2\alpha,2)\), \((2\alpha,1)\), \((4\alpha,2)\), \((4\alpha,1)\), \((8\alpha,2)\), \((8\alpha,1)\), \((16\alpha,2)]\), derived from prior works [22]. Again, we chose \(\alpha=8\) to be consistent with \(\text{MLP}_{\text{VIS}}\).

ResNet[10] is the standard ResNet-18 model used in past benchmarks [1]. Resnet differs from CNN in its inclusion of max-pooling layers and residual connections.

\(\text{MLP}_{\text{TAB}}\) is a 9-layer MLP with hidden dimension 512, batch normalization, and ReLU activation. We did not use dropout or skip connections as it was found to be ineffective on most tabular datasets in MLP+C [10].

### Pruning Implementation Details

Following Shrinkbench [1], we use magnitude-based pruning only on the encoder, \(f_{\mathcal{E}}\), keeping all weights in the decoder, \(f_{\mathcal{D}}\). This is done to prevent pruning a cutset in the decoder architecture, so that all class logits receive input signal. To optimize the performance, we apply magnitude-based pruning globally, instead of layer-wise.

### Hardware

All experiments were conducted on an Nvidia V100 GPU and an AMD EPYC 7402 CPU. The duration of the tabular experiments varied, ranging from a few minutes up to half a day, depending on the specific dataset-model pair and the training phase (pre-pruning training or finetuning). For the vision experiments, a single setting on a single dataset-model pair required a few hours for both pre-pruning training and finetuning.
