End-to-End Neural Network Compression via $\frac{\ell_1}{\ell_2}$
  Regularized Latency Surrogates
Neural network (NN) compression via techniques such as pruning, quantization
requires setting compression hyperparameters (e.g., number of channels to be
pruned, bitwidths for quantization) for each layer either manually or via
neural architecture search (NAS) which can be computationally expensive. We
address this problem by providing an end-to-end technique that optimizes for
model's Floating Point Operations (FLOPs) or for on-device latency via a novel
$\frac{\ell_1}{\ell_2}$ latency surrogate. Our algorithm is versatile and can
be used with many popular compression methods including pruning, low-rank
factorization, and quantization. Crucially, it is fast and runs in almost the
same amount of time as single model training; which is a significant training
speed-up over standard NAS methods. For BERT compression on GLUE fine-tuning
tasks, we achieve $50\%$ reduction in FLOPs with only $1\%$ drop in
performance. For compressing MobileNetV3 on ImageNet-1K, we achieve $15\%$
reduction in FLOPs, and $11\%$ reduction in on-device latency without drop in
accuracy, while still requiring $3\times$ less training compute than SOTA
compression techniques. Finally, for transfer learning on smaller datasets, our
technique identifies $1.2\times$-$1.4\times$ cheaper architectures than
standard MobileNetV3, EfficientNet suite of architectures at almost the same
training cost and accuracy.
Anshul Nasery, Hardik Shah, Arun Sai Suggala, Prateek Jain
End-to-End Neural Network Compression via \(\frac{\ell_{1}}{\ell_{2}}\) Regularized Latency Surrogates

###### Abstract

Neural network (NN) compression via techniques such as pruning, quantization requires setting compression hyperparameters (_e.g.,_ number of channels to be pruned, bitwidths for quantization) for each layer either manually or via neural architecture search (NAS) which can be computationally expensive. We address this problem by providing an end-to-end technique that optimizes for model's Floating Point Operations (FLOPs) or for on-device latency via a novel \(\frac{\ell_{1}}{\ell_{2}}\) latency surrogate. Our algorithm is versatile and can be used with many popular compression methods including pruning, low-rank factorization, and quantization. Crucially, it is fast and runs in almost the same amount of time as _single model training_; which is a significant training speed-up over standard NAS methods. For BERT compression on GLUE fine-tuning tasks, we achieve \(50\%\) reduction in FLOPs with only \(1\%\) drop in performance. For compressing MobileNetV3 on ImageNet-1K, we achieve \(15\%\) reduction in FLOPs, and \(11\%\) reduction in on-device latency _without drop in accuracy_, while still requiring \(3\times\) less training compute than SOTA compression techniques. Finally, for transfer learning on smaller datasets, our technique identifies \(1.2\times\)-\(1.4\times\) cheaper architectures than standard MobileNetV3, EfficientNet suite of architectures at almost the same training cost and accuracy.

## 1 Introduction

Large-scale neural networks consistently provide state-of-the-art performance on complex learning tasks [1; 2; 3]. But they place heavy burden on compute resources such as battery, memory or processor making them hard to deploy on edge devices such as phones, cameras and wearables. Several recent works have designed techniques to compress ML models and make them efficient for inference. However, as detailed below, many of these techniques are hard to use in practice, and often achieve sub-optimal accuracy _vs_ inference time trade-offs.

**Hyperparameter search for compression.** Existing works typically rely on one of the following building blocks to design efficient models: unstructured weights sparsity [4; 5], pruning entire neurons or low-rank factorization [6], quantization [7], distillation [8]. Figuring out an optimal way to combine these building blocks (or to figure out hyper-parameters such as amount of sparsity associated with each block) while satisfying a global latency/FLOPs/resource constraint is difficult and involves a combinatorial search. This problem is further exacerbated when multiple building blocks are used for model compression (_e.g.,_ simultaneous low rank factorization, sparsity/pruning of weights).
Over the past few years, there has been a large body of work that addresses the problem of finding hyperparameters for model compression. Existing literature in this space can be broadly categorized as: (a) methods to find hyperparameters of a _specific_ building block such as unstructured pruning of weights, (b) Neural Architecture Search (NAS) techniques to find hyperparameters of _any_ efficient block. The first set of techniques are naturally limited, but even for the specific blocks considered, their performance on real-world benchmarks is unstable and in fact, can be sub-optimal to the more general approach that we propose in this work (see the left plot in Figure 1). The NAS based techniques are much more generally applicable, but the computational cost of such methods is prohibitive as they typically don't exploit any specific attributes of the popular efficient blocks such as sparsity, pruning.

**Neuron Pruning**: Among the category (a) techniques mentioned above, a prominent line of work has focused on unstructured pruning of weights with non-uniform budget allocation across layers [4, 9, 10, 5]. However, any gain in FLOPs using unstructured pruning is hard to translate to real latency gain as modern hardware - like GPUs, TPUs - are more geared towards dense matrix operations. So it is more fruitful to focus on neuron pruning, which removes entire neurons/channels, and low-rank factorization of weights, which is closely related to neuron pruning. Recent techniques in this line of work add a latency/FLOPs regularizer to the standard cross entropy loss [11, 12] to bias the model towards lower number of neurons. Unfortunately the resulting objective is discrete and difficult to optimize. To alleviate this, existing works have designed continuous surrogates that are more amenable to SGD style optimization. These methods either work in the space of probability distributions over pruned models and optimize the "expected objective" [12, 13, 6] or replace the discontinuous FLOPs regularizer with a continuous surrogate such as \(\ell_{1}\) norms of the weights of the network [11]. However, the former class of techniques are often unstable, hard to implement in practice, and empirical studies indicate that their performance is similar to that of simple magnitude based pruning [14] (also see left plot of Fig. 1). Furthermore, as we show in this work, the latter class of techniques fail to enforce sparsity in the presence of batch, layer normalization (see Section 3).

**NAS**: Several works in category (b) formulate model compression as a black-box Neural Architecture Search (NAS) problem and rely on state-of-the-art NAS techniques to search for efficient models [15, 16, 17, 18, 19]. These techniques directly take the latency/FLOPs into account and have the potential to identify the optimal per-layer budget allocation for a wide variety of efficient blocks/compression mechanisms. However, these approaches are often computationally expensive as they take a blackbox view of the problem and perform combinatorial search over the space of architectures. Despite recent advances such as TuNAS [18] and DARTS [20], these techniques can be an order of magnitude slower and less accurate than our proposed method (see Fig 1).

**Our Approach**: In this work, we propose an approach that sits right in the middle of the above two mentioned categories. That is, our approach applies to a large class of efficient building blocks - like unstructured sparsity, neuron pruning, quantization - for which we can write the FLOPs computation with a continuous surrogate (see Table 1). Furthermore, to ensure that our FLOPs, latency regularizers work even in the presence of batchnorm, layernorm, we propose a novel surrogate that is based on \(\frac{\ell_{1}}{\ell_{2}}\) norm. While our surrogates are continuous, they are non-differentiable. In such cases standard

Figure 1: Left plot compares various techniques for BERT compression on GLUE tasks (averaged across tasks). \(x\)-axis is the relative number of FLOPs as compared to BERTBASE. \(y\)-axis is the relative drop in accuracy from the baseline. Pruning SOTA numbers are taken from [21], while distillation baselines are from [22, 23]. Right plot compares various techniques for MobileNetV3 compression on ImageNet-1K dataset. _MobileNetV3_ corresponds to MobileNetV3 models with different width multiplier. _TuNAS, MorphNet_ are SOTA techniques for scalable compression. TuNAS takes a blackbox approach to model compression, whereas MorphNet takes a more direct approach by optimizing FLOPs regularized objective.
optimizers such as SGD, Adam can be quite slow to converge [24]. To overcome this, we propose a projection operation on the mask variables, after each SGD step. Our proposed method speeds up the convergence and also outputs _exact sparse solutions_ thus eradicating need for post-hoc thresholding, while being simple enough to not increase training time significantly.

We implement our algorithm with multiple building blocks including pruning, low-rank factorization, quantization, and apply it on multiple problems in the domain of image classification and NLP. In particular, we demonstrate the effectiveness of our technique for MobileNetV3 compression on ImageNet (see Fig. 1), where our method can learn an architecture with up to 15% (\(11\%\)) lower FLOPs (latency) on Pixel 6 mobile phones, without any drop in accuracy. Here our approach is more accurate than MorphNet, a SOTA technique which focus exclusively on neuron-pruning, as well as, TuNAS, a SOTA NAS technique. Furthermore, in terms of training time, our method is \(3\times\) cheaper than TuNAS. We would like to highlight that MobileNetv3 is a highly optimized architecture found using efficient NAS techniques [25], and our technique is able to compress this architecture further.

One exciting application of our work is that we can apply it to optimize certain "foundational" baseline models for individual fine-tuning tasks. For example, for compression of BERT on GLUE benchmarks, our method achieved \(40-50\%\) reduction in FLOPs with only \(1\%\) drop in accuracy (see Fig 1). Moreover, our technique dominates standard model compression baselines. Similarly for smaller vision classification tasks, our technique compresses MobileNetV3, EfficientNet suite of architectures and identifies \(1.2\times\)-\(1.4\times\) cheaper architectures without significant loss in accuracy (see Figure 3). We would like to note that all these results are obtained at almost the same cost as that of training a single model for the task. Finally, we also demonstrate the versatility of our method by using it to quantize a CNN on CIFAR-10, and learning the bit-widths (\(2,4,8,16\)) for each of its layers. Our technique found a model that is 55% smaller than the baseline float-16 model, while achieving the same accuracy (see Figure 5). While low-bit quantization is not usually exploited by general purpose accelerators to speed up computation, it can still lead to reduction in inference times of large language models such as GPT as these models are memory bandwidth bound [26]. Here is a summary of our contributions:

**(1).** We provide an end-to-end neural network compression technique that directly optimizes the FLOPs/latency regularized objective during leading to compression during training. Our algorithm can be used with many popular efficient building blocks including pruning, low-rank factorization, quantization, and can optimize for on-device inference latency.

**(2).** We design a novel \(\frac{\ell_{1}}{\ell_{2}}\) regularized surrogate for latency that works even in the presence of batchnorm, layernorm. Our algorithm is fast and runs in the same amount of time as single model training, and doesn't require any post-processing steps.

**(3).** We demonstrate the performance of our technique on both language and vision tasks. Moreover, for transfer learning settings where the goal is to take a baseline architecture and optimize it for individual tasks, our techniques outperform SOTA techniques in the broad-domain of automated neural compression.

## 2 Related Work

### Neural Architecture Search

Early works on NAS treated the problem as a purely blackbox optimization (BO) problem. These works relied on BO techniques such as random search [27], Gaussian process optimization [17], and zeroth-order gradient descent [15; 16], evolutionary algorithms to optimize the NAS objective and identify a good architecture. Several works have improved upon these algorithms using heuristics such as early stopping [27]. Nonetheless, these techniques are computationally expensive, as evaluating the optimization objective at any point requires training a neural network from scratch. Moreover, due to computational complexity, these techniques perform a very coarse grained search and are not suited for fine-grained search over sparsity or low-rank structures.

Recent works have tried to open the blackbox a bit. In these techniques, the search space is first transformed to the space of probability distributions over architectures. Next, a surrogate model (which takes an architecture as input and tries to output the optimal set of weights for the architecture) is trained to quickly evaluate the optimization objective at any input [18; 20; 28; 29; 12]. While these techniques are fast, they involve joint training of the surrogate model during the search process. This joint training often makes the optimization process unstable [30].
**NAS for Efficient ML.** Several recent works at the intersection of efficient ML and NAS have realized the importance of explicitly accounting for the hardware in the search process [15; 31; 32; 33; 34; 35]. These works incorporate the actual inference time in their search objectives, instead of surrogates such as FLOPs. The inference time maybe estimated using another neural network, or through latency tables for basic arithmetic operations on the target platform [19]. Many of these works rely on greedy, random search heuristics to solve the resulting objective [32; 33]. However, these heuristics either take a lot of time to find the optimal architecture or are not guaranteed to converge to an optimal solution. There are some works that rely on the NAS algorithms described above [15; 31; 18]. However, these techniques face the same issues as previously mentioned.

**Hardware, Neural Architecture codesign.** Certain hardware level parameters such as tiling configurations of tensors significantly affect the inference time of a model. Recent hardware-aware NAS techniques expose these hardware level parameters to the NAS algorithm and simultaneously search over neural architectures and hardware configurations [35]. These techniques have the potential to achieve better performance than vanilla NAS techniques which do not search over hardware configurations.

### Model Compression

The field of model compression is vast. Here, we focus on techniques that perform training-time compression (as opposed to post-training compression) using the following building blocks: unstructured sparsity, pruning and low-rank factorization. Early works in unstructured sparsity and pruning relied on magnitude, gradient based pruning [4; 36; 14]. Several works have explored more sophisticated scoring metrics for pruning [37; 38; 39; 40; 41]. Other techniques include adding sparsity inducing norms such as \(\ell_{0},\ell_{1}\)to the training objective [13; 5]. A number of works have also explored low-rank factorization for model compression [42; 43; 44]. Some of these techniques again rely on sparsity inducing regularizers to induce the low-rank structure [6]. Others rely on SVD based pruning. Some recent works try and optimize FLOPs regularized objective to perform pruning, low-rank factorization [11; 12]. However, as we discussed in the introduction, the resulting optimization techniques are often unstable and difficult to use in practice.

## 3 Method

In this section, we describe our approach for model compression. For simplicity of presentation, we illustrate our technique on feed-forward networks and restrict ourselves to pruning. The ideas here can be extended to other architectures (_e.g.,_ 1x1 convolutions in CNNs), and other efficient building blocks (_e.g.,_ unstructured sparsity, low-rank factorization, quantization) in a straightforward manner (see Table 1 for details). Consider the following problem: we are given a pre-trained feed forward neural network (FFN) \(f^{*}(x)=\sigma(W_{D}^{T}\sigma(W_{D-1}^{*}\sigma(\ldots\sigma(W_{1}^{*}x))))\), where \(W_{i}^{*}\in\mathbb{R}^{d_{i+1}\times d_{i}}\) for all \(i\in[D]\), and a dataset \(\{(x_{i},y_{i})\}_{i=1}^{n}\). Our goal is to compress \(f^{*}\) while simultaneously performing well on the learning task. This problem can be formulated as the following optimization problem

\[\min_{\mathcal{W}}\frac{1}{n}\sum_{i=1}^{n}\ell(x_{i},y_{i};\mathcal{W})+ \lambda\times\text{Latency}(\mathcal{W}). \tag{1}\]

Here \(\mathcal{W}=\{W_{i}\}_{i=1}^{D},\) with \(W_{i}\in\mathbb{R}^{d_{i+1}^{\prime}\times d_{i}^{\prime}}\) being the weight matrix at layer \(i\), \(\lambda\) is the regularization parameter which trades-off latency with accuracy and \(\ell\) is the supervised loss.1. Directly optimizing the above objective is intractable because \(\text{Latency}(\mathcal{W})\) is a discrete function of the dimensions of weight matrices, and is hardware specific.

Footnote 1: In this objective, we search over \(d_{i}^{\prime}\) such that \(d_{i}^{\prime}\leq d_{i}\)

We now present our technique for solving Equation (1). To begin with, we substitute \(\text{Latency}(\mathcal{W})\) with \(\text{FLOPs}(\mathcal{W})\)2. Later, we extend it to actual latency. The objective in this case is given by

Footnote 2: FLOPs is also a discrete function of dimensions of \(W_{i}\), and the resulting optimization problem is still intractable

\[\min_{\mathcal{W}}\frac{1}{n}\sum_{i=1}^{n}\ell(x_{i},y_{i};\mathcal{W})+ \lambda\sum_{i=1}^{D}d_{i}^{\prime}d_{i+1}^{\prime}. \tag{2}\]
To solve this objective, we associate masks with each neuron in the network. In particular, we parameterize the weight matrix in the \(i^{th}\) layer as \(W_{i}\times\text{diag}(\alpha_{i})\). Here \(\alpha_{i}\in\{0,1\}^{d_{i}}\) are the mask variables of layer \(i\). If \(\alpha_{i,j}\) is set to \(0\), then the \(j^{th}\) neuron in the \((i-1)^{th}\) layer will be pruned. The FLOPs regularizer can now be written in terms of masks as \(\sum_{i=1}^{D}\|\alpha_{i}\|_{0}\|\alpha_{i+1}\|_{0}\), where \(\alpha_{D+1}\) is the static vector of all \(1\)'s. The resulting objective though is not continuous. To make it continuous and amenable to gradient based optimization, one class of techniques place a Bernoulli distribution \(\text{Bern}(p_{i,j})\) over each of the masks \(\alpha_{i,j}\) and solve the following smoothed objective [12; 13; 6]

\[\min_{\mathcal{W},p}\mathbb{E}\left[\frac{1}{n}\sum_{i=1}^{n}\ell(x_{i},y_{i}; p,\mathcal{W})+\lambda\sum_{i=1}^{D}\|\alpha_{i}\|_{0}\|\alpha_{i+1}\|_{0} \right].\]

The expectation above is taken w.r.t the random masks \(\alpha_{i}\)'s. It is easy to see that the above objective is equivalent to Equation (2), and is consequently as hard as solving the latter. In fact, the above problem can be shown to be NP-hard using the observation that sparse linear regression is a special case of it [45]. Furthermore, the discrete nature of \(\alpha_{i}\)'s makes the optimization process unstable [13]. To overcome this, [12; 13; 6] rely on a heuristic which involves relaxing Bernoulli distribution to a continuous distribution such as LogisticSigmoid. However, the main drawback of the resulting algorithm is that it is hard to implement in practice and requires very careful annealing of the parameters of LogisticSigmoid distribution. Another drawback of this class of techniques is that their performance is not well understood theoretically, even for simple and fundamental problems such as sparse linear regression.

Another approach to convert the discrete objective in Equation (2) into a continuous function is to replace the \(\ell_{0}\) norm on \(\alpha_{i}\)'s with \(\ell_{1}\) norm

\[\min_{\mathcal{W},\alpha_{i}\in\mathbb{R}^{d_{i}}}\frac{1}{n}\sum_{i=1}^{n} \ell(x_{i},y_{i};\alpha,\mathcal{W})+\lambda\sum_{i=1}^{D}\|\alpha_{i}\|_{1} \|\alpha_{i+1}\|_{1}. \tag{3}\]

This approach is much more attractive than the previous approach as it is known to recover the optimal sparse solutions for a variety of statistical problems including sparse linear regression, low-rank matrix completion [46; 47]. Furthermore, it is much simpler to implement in practice, with numerous algorithms being proposed for fast convergence to the stationary points of the objective [24; 48]. Consequently, recent SOTA compression techniques relied on \(\ell_{1}\) norm surrogates to compute the FLOPs regularizer [11]. A major drawback of \(\ell_{1}\) norm though is that it does not promote sparsity in the presence of batch normalization and layer normalization [49; 50]. To see this, consider the following \(1\)-hidden layer network: \(\sigma(\text{BN}(W_{2}\text{diag}(\alpha_{2})\sigma(\text{BN}(W_{1}\text{ diag}(\alpha_{1})x))))\). One can scale down all entries of \(\alpha_{1}\) and scale up the weights \(W_{1}\) without affecting the output of the network. Doing this reduces the objective value in Equation (3), but doesn't induce any sparsity in the network. In practice, we in fact notice this behaviour during optimization of Equation (3), which leads to sub-optimal solutions (see Section 3.2). Note that adding \(\ell_{2}\) penalty on the weights (_i.e.,_ weight decay) doesn't mitigate this issue as any scaling of \(\alpha\)'s can be absorbed by the batch norm parameters without changing the output of the network.

### Inducing sparsity through \(\frac{\ell_{1}}{\ell_{2}}\) regularizer

We now introduce our approach for making the objective in Equation (2) continuous. We replace \(\ell_{0}\) norm over masks (\(\|\alpha_{i}\|_{0}\)) with \(\frac{\ell_{1}}{\ell_{2}}\) penalty (\(\sqrt{d_{i}}\|\alpha_{i}\|_{1}/\|\alpha_{i}\|_{2}\)) and solve the following optimization problem

\[\min_{\mathcal{W},\alpha_{i}\in\mathbb{R}^{d_{i}}}\frac{1}{n}\sum_{i=1}^{n} \ell(x_{i},y_{i};\alpha,\mathcal{W})+\lambda\sum_{i=1}^{D}\frac{\sqrt{d_{i}} \|\alpha_{i}\|_{1}}{\|\alpha_{i}\|_{2}}\frac{\sqrt{d_{i+1}}\|\alpha_{i+1}\|_{1 }}{\|\alpha_{i+1}\|_{2}}. \tag{4}\]

The \(\sqrt{d_{i}}\) term in the numerator normalizes the penalty to lie between \([0,d_{i}]\). When \(\alpha_{i}\)'s are all \(1\)'s, the regularizer evaluates to FLOPs. Observe that this regularizer is invariant to scaling of \(\alpha\)'s. Consequently, the value of the regularizer cannot simply be reduced by scaling down \(\alpha_{i}\)'s. In our experiments in sec 3.2 and Appendix C.2, we show that this handles batch, layer normalizations better than \(\ell_{1}\) regularizer. Several works have studied this regularizer in the context of sparse linear regression and showed that is recovers the underlying sparse signal under mild conditions on the
data [51; 52; 53]. [54] used a similar \(\frac{\ell_{1}}{\ell_{2}}\) regularizer for network pruning, but their technique doesn't optimize latency or FLOPs, and relies on post-training thresholding to get sparsity.

For certain technical reasons described later, we add a positivity constraint on \(\alpha_{i}\)'s and solve the following objective

\[\min_{\mathcal{W},\alpha_{i}\in\mathbb{R}^{d_{i}}_{+}}\frac{1}{n}\sum_{i=1}^{n }\ell(x_{i},y_{i};\alpha,\mathcal{W})+\lambda\sum_{i=1}^{D}\frac{\sqrt{d_{i}} \sum_{j=1}^{d_{i}}\alpha_{i,j}}{\|\alpha_{i}\|_{2}}\frac{\sqrt{d_{i+1}}\sum_{ j=1}^{d_{i+1}}\alpha_{i+1,j}}{\|\alpha_{i+1}\|_{2}}. \tag{5}\]

Note that we consider \(\alpha\in\mathbb{R}^{d_{i}}_{+}\) rather that discrete or bounded values. We would like to highlight that this change doesn't reduce the representational power of our model. It is mainly done for computational reasons. In the sequel, we use the shorthand \(\|\alpha_{i}\|_{1p}\) (\(p\) for positive) to denote \(\sum_{j=1}^{d_{i}}\alpha_{i,j}\).

**Importance of positivity constraints.** The objective in Equation (4) is continuous, but not smooth. For such losses, standard optimization techniques such as SGD, Adam are slow to converge to stationary points [55]. Furthermore, these algorithms don't output exact sparse solutions. This forces additional post-processing steps to be introduced into the compression pipeline. For example, [11; 54] rely on Adam optimizer and add a pruning step at the end, where masks that are close to \(0\) are pruned away. This is quite cumbersome in practice as one needs to choose appropriate thresholds for pruning, which introduces an additional tunable hyper-parameter, and needs re-training after pruning. To overcome this, we add a positivity constraint to the mask variables and modify the objective to Equation (5). This makes the regularizer smooth (except at all 0's vector), and easy to optimize using SGD, Adam. After each SGD/Adam update, we simply project the masks back to the space of positive real numbers. The overall update looks as follows

\[\mathcal{W}\leftarrow\mathcal{W}-\eta\nabla_{\mathcal{W}}(\mathcal{L}(\alpha, \mathcal{W})+\lambda\mathcal{R}(\alpha)),\quad\alpha\leftarrow\max(0,\alpha- \eta\nabla_{\alpha}(\mathcal{L}(\alpha,\mathcal{W})+\lambda\mathcal{R}( \alpha))).\]

Here \(\mathcal{L}(\alpha,\mathcal{W})\) is the empirical risk and \(\mathcal{R}(\alpha)\) is the regularizer. Notice, the only additional step compared to traditional optimization, is the clipping of \(\alpha\)'s. In our ablation studies in Sec 3.2 and App C.2, we validate the importance of this projection step, together with \(\frac{\ell_{1}}{\ell_{2}}\) norm, in encouraging sparse solutions.

### Verification of design choices

To empirically demonstrate the drawbacks of using \(\ell_{1}\) penalty for model compression, we perform experiments on the FashionMNIST dataset with a single hidden layer fully connected network which has a batch norm layer after the first linear layer. We prune out the input to the network using a mask \(\alpha\) on the input. We compare the performance of networks compressed using FLOPs regularizer induced by \(\ell_{1}\) and \(\frac{\ell_{1}}{\ell_{2}}\) norms. We use SGD for optimization of both the objectives. Furthermore, we pre-train the network using standard CE loss, and initialize \(\alpha=\mathbf{1}\). We track the variance of the absolute values of the entries of \(\alpha\), i.e. \(\frac{\sum_{i=1}^{d}(|\alpha_{i}|-\mu_{\alpha})^{2}}{d}\), where \(\mu_{\alpha}=\frac{\sum_{i=1}^{d}|\alpha_{i}|}{d}\). We also track the mean \(\mu_{\alpha}\) of the absolute values of the entries of \(\alpha\). Finally, we plot out the curve between FLOPs and the considered norm of \(\alpha\) (_i.e., \(\ell_{1}\)_, \(\frac{\ell_{1}}{\ell_{2}}\)). Figure 2 presents the results from these experiments. We can see that the \(\ell_{1}\) objective is mis-aligned with the actual value of FLOPs, while the regularizer computed using \(\frac{\ell_{1}}{\ell_{2}}\) is a better proxy. We also find that the mean and variance of \(\alpha\)'s sharply decreases when \(\ell_{1}\) induced FLOPs regularizer is used for compression. This indicates that all entries of \(\alpha\) are uniformly

Figure 2: **Comparison of \(\ell_{1}\), \(\frac{\ell_{1}}{\ell_{2}}\) induced FLOPs regularizer for pruning on FashionMNIST: Figures (a) and (b) depict the evolution of the statistics of the mask variables (\(\alpha\)) as training progresses. Figure (c) shows the relation between the actual FLOPs of the model and the value of the proxy computed by Equations 3, 4. Figure (d) shows the evolution of the Frobenius norm of the weight matrix.**
scaled down to a small, non-zero value, reducing the value of the regularizer, while not providing any sparsity. As seen from the figure, \(\frac{\ell_{1}}{\ell_{2}}\) does not suffer from this drawback. Finally, we note that the frobenius norm of the weight matrix \(W\) increases when \(\ell_{1}\) regularization is used on \(\alpha\), suggesting that the network is simply scaling down \(\alpha^{\prime}\)s and scaling up the weights to evade the regularizer.

### Hardware aware model compression

In this section, we extend the FLOPs regularizer to take the latency on the target hardware into account. The resulting regularizer is especially useful for performing hardware aware network compression. Our key observation is that the inference on a neural network can be broken down into a series of matrix multiplication operations. For example, inference on a depth \(D\) FFN involves \(D\) matrix-vector multiplications, which take-up majority of the time. So, getting a good estimate of the inference time of the overall network boils down to having a good estimate of the latency of matrix-vector multiplication. To this end, we rely on lookup tables. Before the start of the pruning phase, we construct a \(2\)-dimensional lookup-table \(T\) whose \((d_{1},d_{2})^{th}\) entry is the on-device latency of multiplying a matrix of size \(d_{1}\times d_{2}\) with a vector of size \(d_{2}\). Such a table is easy to construct, given access to the target device. Next, to incorporate the look-up table \(T\) into our pruning algorithm, we convert it into a continuous function by performing linear interpolation on the entries in the table [56]. To be precise, for any \((x,y)\in[d_{1},d_{1}+1]\times[d_{2},d_{2}+1]\), where \(d_{1},d_{2}\in\mathbb{N}\cup\{0\}\), we define \(T(x,y)\) as: \(T(x,y)=t_{1}+(t_{2}-t_{1})(y-d_{2})\), where \(t_{1}=T(d_{1},d_{2})+(T(d_{1}+1,d_{2})-T(d_{1},d_{2}))(x-d_{1}),\) and \(t_{2}=T(d_{1},d_{2}+1)+(T(d_{1}+1,d_{2}+1)-T(d_{1},d_{2}+1))(x-d_{1}).\) Note that in contrast to black-box NAS techniques like [19] which search over a discrete space of number of filters for each block, our approach needs the latency surrogate to be differentiable, and hence we need interpolated latency tables. See the appendix for details on how we construct the tables.

We use this interpolated lookup table to construct our _latency_ regularizer as follows

\[\sum_{i=1}^{D}T\left(\frac{\sqrt{d_{i}}\|\alpha_{i}\|_{1p}}{\|\alpha_{i}\|_{ 2}},\frac{\sqrt{d_{i+1}}\|\alpha_{i+1}\|_{1p}}{\|\alpha_{i+1}\|_{2}}\right). \tag{6}\]

In the above expression, our differentiable surrogate for \(\|\alpha_{i}\|_{0}\) (_i.e.,_\(\sqrt{d_{i}}\|\alpha_{i}\|_{1p}/\|\alpha_{i}\|_{2}\)), is used to index the lookup table. We note that \(\frac{\ell_{1}}{\ell_{2}}\) norm is very crucial for this technique to be successful. This is because \(\frac{\sqrt{d_{i}}\|\alpha_{i}\|_{1p}}{\|\alpha_{i}\|_{2}}\) is normalized and always lies between \([0,d_{i}]\). In contrast, using \(\ell_{1}\) norm surrogate in the regularizer gives us \(T(\|\alpha_{i}\|_{1},\|\alpha_{i+1}\|_{1})\). Scaling \(\alpha_{i}\) by a constant can drastically change this regularizer, and makes the optimization unstable.

\begin{table}
\begin{tabular}{|c||c|c|c|} \hline
**Efficient** & **Parameterization of \(W_{i}\)** & **FLOPs** & **Regularizer (FLOPs surrogate)** \\
**Building Block** & \(W_{i}\times\text{diag}(\alpha_{i})\) & \(\|\alpha_{i}\|_{0}\|\alpha_{i+1}\|_{0}\) & \(\frac{\sqrt{d_{i}}[\|\alpha_{i}\|_{1p}}{\|\alpha_{i}\|_{2}}\frac{\sqrt{d_{i+1 }}\|\alpha_{i+1}\|_{p}}{\|\alpha_{i+1}\|_{2}}\) \\ \hline Pruning & \(W_{i}\odot\alpha_{i}\), where \(\alpha_{i}\in\mathbb{R}^{d_{i+1}\times d_{i+1}}\), \(\odot\) is the elementwise & \(\|\text{Vec}(\alpha_{i})\|_{0}\) & \(\frac{\sqrt{d_{i}}[\|\alpha_{i}\|_{1p}}{\|\text{Vec}(\alpha_{i})\|_{2}}\) \\  & multiplication operator & & \\ \hline Low-rank Factorization & \(U\),diag\((\beta_{i})\),\(\nabla_{\alpha_{i}}\) & \((d_{i}+d_{i+1})\|_{\beta_{i}}\|_{0}\) & \((d_{i}+d_{i+1})\frac{\sqrt{d_{i+1}}[\beta_{i}]_{1p}}{\|\beta_{i}\|_{2}}\) \\  & \(d_{i+1}=\min\{d_{i},d_{i+1}\}\) & & \((d_{i}+d_{i+1})\frac{\sqrt{d_{i+1}}[\beta_{i}]_{1p}}{\|\beta_{i}\|_{2}}\) \\ \hline \multirow{4}{*}{\begin{tabular}{c} Quantization \\ \((1,2,4\) bit quantization) \\ \end{tabular} } & \(W_{i,+1}+\alpha_{i}\alpha_{i}(\Delta_{i}+2,d_{i}(\Delta_{i},d_{i}))\), where \(\alpha_{i,2,4,4,4}\in\mathbb{I},1\), are & \(|(1-\alpha_{i,2})|_{0}d_{i,4,1}+\) & \(\frac{\ell_{1}}{\ell_{2}}\) norm over \\  & mask variables, \(W_{i,\pm}\) is the & \(2|\alpha_{i}(1-\alpha_{i,2})|_{0}d_{i,4,1}+\) & \(2|\alpha_{i,2}|(1-\alpha_{i,4})|_{0}d_{i,4,1}+\) & \(2\alpha_{i,2}|(1-\alpha_{i,4})\), \(4\alpha_{i,2}\alpha_{i,4}|d_{i,4}|\) \\  & \(\Delta_{i,4}=W_{i,4}-W_{i,2}\) & & \(\lambda_{i,2}=W_{i,2}-W_{i,1}\) & \(\lambda_{i,2}=W_{i,2}-W_{i,1}\) & \(\lambda_{i,2}=W_{i,2}-W_{i,1}\) \\ \hline \multirow{4}{*}{
\begin{tabular}{c} Pruning + \\ Low-rank Factorization \\ \end{tabular} } & \(U\);diag\((\beta_{i})\);\(V\);diag\((\alpha_{i})\), where \(U_{i}\in\mathbb{R}^{d_{i+1}\times d_{i+1}}\), & \((\|\alpha_{i}\|_{0}+\|\alpha_{i+1}\|_{0})\|\beta_{i}\|_{0}\) & \(\left(\begin{array}{c}\frac{\sqrt{d_{i}}[\|\alpha_{i}\|_{1p}}{\|\alpha_{i}\|_{ 1p}}\frac{\sqrt{d_{i+1}}\|\alpha_{i+1}\|_{p}}{\|\alpha_{i+1}\|_{2}}\\ \frac{\sqrt{d_{i+1}}[\beta_{i}]_{1p}}{\|\alpha_{i}\|_{1p}}\end{array}\right)\) \\  & \(d_{i,\pm}=\min\{d_{i},d_{i+1}\}\) & & \(\times\frac{\sqrt{d_{i+1}}[\beta_{i}]_{1p}}{\|\alpha_{i}\|_{2}}\) \\ \hline \end{tabular}
\end{table}
Table 1: Table describing regularizers used by our technique for various efficient building blocks (refer to Appendix for details on quantization). One can easily design regularizers for searching over a combination of building blocks. For example, last row presents regularizer for low-rank + pruning, which we use in our large-scale experiments.
## 4 Experiments

In this section, we apply our framework to large scale pre-training and transfer learning tasks on standard language and vision benchmarks. To demonstrate the versatility of our technique, we perform experiments on multiple model families (MobileNet, EfficientNet [2], BERT), and multiple building blocks (pruning, low-rank factorization, quantization). We also present a case study using the actual on-device latency instead of FLOPs. See Appendix C.2 for other ablation studies.

### ImageNet Pre-training

We begin by comparing the performance of our technique with baselines on MobileNetV3 compression, for ImageNet classification. We rely on low-rank factorization + pruning for the compression. The results from this experiment are presented in Figure 1. By varying the strength of our regularization, we obtain models with different MACs and accuracies. We find that models produced by our method significantly outperform MobileNetV3 and TuNAS in the high and mid-MACs regime. In particular, for the same accuracy as MobileNetV3Large, our approach finds a model with \(15\%\) fewer MACs. In comparison with TuNAS, we achieve 30% reduction in MACs at the same level of accuracy. We however find that our model is at par with MobileNetV3Small in the low MACs regime, indicating that the former is already well-tuned for this task. In terms of compute needed for training, TuNAS is the most expensive among all the techniques we tried; it took 2 days to train with our hardware setup. In contrast, our method took 13 hours (\(3-4\times\) faster than TuNAS), and MorphNet took 10 hours.

### Transfer Learning

A common paradigm in deploying machine learning models today is to first pre-train them on a large scale dataset such as ImageNet, and then fine-tune them for the desired target task. However, deploying large models is not feasible on edge devices. Our technique provides a light-weight modification to the standard fine-tuning procedure by producing a compressed model with comparable transfer learning performance on the specific task. We demonstrate this on vision and language tasks.

**Vision tasks.** We consider the task of fine-tuning an ImageNet pre-trained model for a smaller dataset. We consider Cars196 [57] and Food101 [58] as the target datasets, and compare against the MobileNetV3 and EfficientNet families of models. We use ImageNet pre-trained models for initialization. We plot the FLOP-accuracy curves in Fig 3. We compress MobileNetv3Large and EfficientNet-B4 and EfficientNet-B2 architectures while transferring them to the target target task. We find that our method consistently improves over baseline architectures across various FLOPs regimes. This is because our technique is able to adaptively prune the model based on the difficulty of the classification task. On both the tasks, we see 1% accuracy gains over MobileNetV3 small. The accuracy gains persist at the latency footprint of MobileNetV3Large-0.75, where we see over 1.5% accuracy gains on both datasets. On EfficientNet, we see upto 40% reduction in FLOPs without any drop in accuracy on Food101, and around 20% reduction in FLOPs on the Cars196 dataset for the largest models (B4). We also see around 30% FLOP reduction while maintaining the transfer learning performance of the B1 and B0 variants. This demonstrates that our learnt models can scale better than the heuristic scaling described in [2]. See the appendix for additional results.

**Fine-tuning BERT on GLUE.** We consider 5 datasets of the GLUE benchmark [59] that are commonly used in the literature, and fine-tune a pre-trained BERT-Base model with our FLOPs regularizer. We re-parameterize the weight matrices of the feed forward network of each transformer block with our low-rank+sparse parameterization. We compare our approach against model pruning,

Figure 3: **Accuracy-FLOPs trade-off on transfer learning tasks: Figures (a) and (b) depict the the fine-tuning performance of models found by our method while compressing MobileNetv3Large and baseline MobileNetV3 on Cars-196 and Food-101 datasets. Figures (c) and (d) show the performance on the EfficientNet family of architectures, where baselines are EfficientNetB0-B4, while our method compresses EfficientNet B4 and B2.**
where SOTA numbers are taken from Fig. 6 of [21], reporting the maximum accuracy among [60; 61; 62; 63; 64; 65]. We also report the performance of widely-used distillation based baselines [22; 23]. Figure 1 presents the average performance on the 5 datasets, and Figure 6 in appendix presents the individual performance. In both these figures, we plot the relative FLOPs of the compressed model w.r.t BERT-base against the drop in accuracy w.r.t BERT-base (similar to [21]). We find that on 4 of the 5 datasets considered, our technique provides a higher accuracy for the same number of FLOPs, indicating the efficacy of our method. On MRPC, a dataset with very few samples, our method is worse off on higher FLOPs, but outperforms the baselines in the low FLOP regime.

### Additional Experiments

**Using the latency regularizer.** In Eq 6, we propose a latency surrogate for optimizing the actual on-device inference latency. In this section, we provide empirical evidence of the effectiveness of this approach for MobileNetv3 on Pixel 6. We compare the accuracy-latency curves of models produced using FLOPs, latency regularizers (see Fig 4). Observe that using the latency regularizer leads to models with smaller latencies and consequently better latency-accuracy tradeoff compared to using the FLOP regularizer. We also find these models to have better performance than MobileNetV3 (\(0.5-2\%\) improvement in accuracy for similar latency), despite MobileNetv3 being hand-crafted for faster inference on mobile devices.

**Quantization.** In this set of experiments, we consider CIFAR-10 classification and compress a 3 layer CNN using quantization. We use the quantization formulation presented in Table 1 and search over \(\{2,4,8,16\}\) bit quantizations for each layer. We compare with a baseline which uses the same level of quantization at each layer. Fig 5 presents the results from this experiments. The details of the implementation can be found in the appendix. We find that our technique compresses the model size by almost 55% without drop in accuracy (as compared to a model with 16-bit weights). Our technique also outputs a model which is 1.4% more accurate than a 2-bit quantized model with only 4% more FLOPs. In the plot on the right in Fig 5, we visualize the learned bit-widths of our models. We find that later layers are assigned a smaller bit width, indicating the importance of learning expressive filters early in the network. The different models in our plots were found by varying the the value of the regularizer coefficient, and hence no combinatorial search over bit-widths is required.

## 5 Conclusion and Future Work

In this work, we presented an end-to-end technique for neural network compression. Our approach applies to a wide variety of efficient blocks including pruning, unstructured sparsity, quantization. At

Figure 4: Left plot shows the accuracy-latency curves of models obtained using FLOPs, latency regularizers. Right table compares the performance of our latency regularized models with MobileNetV3 baseline.

Figure 5: **Quantization on CIFAR-10: Figure (a) compares the performance of our technique for dynamic quantization against fixed-bit quantization for a 4 layer CNN on CIFAR-10. The baselines have weights quantized to 2,4,8, 16 bits. Fig (b) depicts the learnt bitwidths for different layers of the models found by our technique, with the labels denoting the number of MACs (in Bn) of the models.**
the core of our algorithm is a novel surrogate for FLOPs, latency that relies on \(\frac{\ell_{1}}{\ell_{2}}\) norms, and works with batchnorm, layernorm. Our algorithm is computationally efficient and runs in same amount of time as needed for training a single model. We demonstrated the efficacy of our approach on various pre-training and transfer learning tasks on standard language and vision benchmarks. As a future work, it will useful to incorporate more efficient building blocks such as block diagonal matrices into our framework. Another interesting direction would be to make our technique more hardware aware by incorporating hardware level parameters such as tiling into our search process.
