LCS: Learning Compressible Subspaces for Adaptive Network Compression at Inference Time

When deploying deep learning models to a device , it is traditionally assumed that available computational resources ( compute , memory , and power ) remain static . However , real-world computing systems do not always provide stable resource guarantees . Computational resources need to be conserved when load from other processes is high or battery power is low . Inspired by recent works on neural network subspaces , we propose a method for training a compressible subspace of neural networks that contains a fine-grained spectrum of models that range from highly efficient to highly accurate . Our models require no retraining , thus our subspace of models can be deployed entirely on-device to allow adaptive network compression at inference time . We present results for achieving arbitrarily finegrained accuracy-efficiency trade-offs at inference time for structured and unstructured sparsity . We achieve accuracies on-par with standard models when testing our uncompressed models , and maintain high accuracy for sparsity rates above 90 % when testing our compressed models . We also demonstrate that our algorithm extends to quantization at variable bit widths , achieving accuracy on par with individually trained networks . 1 INTRODUCTION . Deep neural networks are deployed to a variety of computing platforms , including phones , tablets , and watches ( Dhar et al. , 2019 ) . Networks are generally designed to consume a fixed budget of resources , but the compute resources available on a device can vary while the device is in use . Computational burden from other processes , as well as battery life , may influence the availability of resources to a neural network . Adaptively adjusting inference-time load is beyond the capabilities of traditional neural networks , which are designed with a fixed architecture and a fixed resource usage . A simple approach to the problem of providing an accuracy-efficiency trade-off is to train multiple networks of different sizes . In this approach , multiple networks are stored on the device and loaded into memory when needed . There is a breadth of research in the design of efficient architectures that can be trained with different capacities , then deployed on a device ( Howard et al. , 2017 ; Sandler et al. , 2018 ; Howard et al. , 2019 ) . However , there are a few drawbacks with using multiple networks to provide an accuracy-efficiency trade-off : ( 1 ) it requires training and deploying multiple networks ( which induces training-time computational burden and on-device storage burden ) , ( 2 ) it requires all compression levels to be specified before deployment , and ( 3 ) it requires new networks to be loaded into memory when changing the compression level , which prohibits real-time adaptive compression . Previous methods address the first issue in the setting of structured sparsity by training a single network conditioned to perform well when varying numbers of channels are pruned ( Yu et al. , 2018 ; Yu & Huang , 2019 ) . However , these methods require BatchNorm ( Ioffe & Szegedy , 2015 ) statistics to be recalibrated for every accuracy-efficiency configuration before deployment . This requires users to know every compression level in advance and requires storage of extra parameters ( which can be problematic for low-compute devices ) . Another method trains a larger network that can be queried for efficient subnetworks ( Cai et al. , 2020 ) , but requires users to reload network weights , and to know every compression level before deployment . We address the issue of deploying a model that can be compressed in real-time without specifying the compression levels beforehand . Inspired by Wortsman et al . ( 2021 ) , we formulate this problem as computing an adaptively compressible network subspace ( Figure 1a ) . We train a linear subspace using two sets of weights , conditioning them so that a convex combination of them produces an accurate network . We expand on Wortsman et al . ( 2021 ) by ( 1 ) specializing our linear subspace to contain an arbitrarily fine-grained efficiency-accuracy trade-off ( meaning , any compression level within predefined limits can be chosen ) , and ( 2 ) replacing BatchNorm ( Ioffe & Szegedy , 2015 ) with GroupNorm ( Wu & He , 2018 ) to avoid BatchNorm recalibration . Contributions : Our contributions are as follows . ( 1 ) We introduce our method , Learning Compressible Subspaces ( LCS ) , for training models that can be compressed with arbitrarily fine-grained compression levels in real-time after deployment . To our knowledge this has not been done before . ( 2 ) We demonstrate a new use for neural network subspaces , specializing at one end for high-accuracy and at the other end for high-efficiency . ( 3 ) We provide an empirical evaluation of our method using structured sparsity , unstructured sparsity , and quantization . ( 4 ) We open source our code at https : //withheld.for.review . 2 RELATED WORK . Architectures Demonstrating Efficiency-Accuracy Trade-Offs : Several recent network architectures include a hyperparameter to control the number of filters in each layer ( Howard et al. , 2017 ; Sandler et al. , 2018 ; Howard et al. , 2019 ) or the number of blocks in the network ( Tan & Le , 2019 ) . In Once For All ( Cai et al. , 2020 ) , the need for individually training networks of different sizes is circumvented . Instead , a single large network is trained , then queried for efficient subnetworks . We differ from these methods in two ways . First , we compress on-device without specifying the compression levels before deployment . Previous works require training separate networks ( or in the case of Once For All , querying a larger model for a compressed network ) before deployment . Second , we don ’ t require deploying a set of weights for every compression level . A method for post-training quantization to variable bit widths appears in Shkolnik et al . ( 2020 ) , but it relies on calibration data and can not be run on low-compute edge devices . Thus , compression levels must be specified before deployment . They also present a method for quantization-aware training to variable bit widths , but most of their results keep activation bit widths fixed , whereas we vary it . Adaptive Compression : Recent works train a single neural network which can be adaptively configured at inference time to execute at different compression levels . These methods are the closest to our work . In Learning Efficient Convolutions ( LEC ) ( Liu et al. , 2017 ) , the authors train a single network , then fine-tune it at different structured sparsity rates ( though they also present results without fine-tuning , which is closer to our use case ) . Other methods train a single set of weights conditioned to perform well when channels are pruned , but require recalibration ( or preemptive storage ) of BatchNorm ( Ioffe & Szegedy , 2015 ) statistics at each sparsity level . These methods include Network Slimming ( NS ) ( Yu et al. , 2018 ) and Universal Slimming ( US ) ( Yu & Huang , 2019 ) . Similar methods train a single network to perform well at various levels of quantization by storing extra copies of BatchNorm statistics ( Guerra et al. , 2020 ) . We differ from these methods by avoiding the need to recalibrate or store BatchNorm statistics ( Section 3.3 ) and by allowing for fine-grained selection of compression levels at inference time , neither of which have been done before ( Table 1 ) . Additionally , our method is broadly applicable to a variety of compression methods , whereas previous works focus on a single compression method . Other Post-Training Compression Methods : Other works have investigated post-training compression . In Nagel et al . ( 2019 ) , a method is presented for compressing 32 bit models to 8 bits , though it has not been evaluated in the low-bit regime . It involves running an equalization step and assuming a Conv-BatchNorm-ReLU network structure . A related post-training compression method is shown in Horton et al . ( 2020 ) , which shows results for both quantization and sparsity . However , their sparsity method requires a lightweight training phase . We differ from these methods in providing real-time post-deployment compression for both sparsity and quantization without making assumptions about the network structure . Sparse Networks : Sparsity is a well-studied approach to reduce neural network compute . Recent works have achieved state-of-the-art accuracy using simple magnitude-based pruning ( Han et al. , 2015 ; Zhu & Gupta , 2017 ; See et al. , 2016 ; Narang et al. , 2017 ) . When experimenting with structured sparsity , we simply remove filters from the network , as described in Yu et al . ( 2018 ) . When investigating unstructured sparsity , we use TopK pruning ( Zhu & Gupta , 2017 ) . Quantized Networks : Network quantization represents the weights and/or activations of neural networks with lower-bit integral representations , rather than 32-bit floating point . This reduces model size and memory bandwidth and can accelerate inference . Our work follows the affine quantization scheme presented in Jacob et al . ( 2018 ) , in which a scale and offset are applied to an integer tensor to approximate the original floating-point tensor . However , our method is general enough to utilize any other quantization scheme . Neural Network Subspaces : The idea of learning a neural network subspace is introduced in Wortsman et al . ( 2021 ) ( though another formulation was introduced concurrently in Benton et al . ( 2021 ) ) . Multiple sets of network weights are treated as the corners of a simplex , and an optimization procedure updates these corners to find a region in weight space in which points inside the simplex correspond to accurate networks . We note that Wortsman et al . ( 2021 ) recalibrate BatchNorm statistics before evaluation , which we avoid . 3 COMPRESSIBLE SUBSPACES . 3.1 COMPRESSIBLE LINES . In our method , we train a neural network subspace ( Wortsman et al. , 2021 ) that contains a spectrum of networks that each have a different efficiency-accuracy trade-off . In Wortsman et al . ( 2021 ) , an architecture is chosen , and its collection of weights is denoted by ω . A linear subspace is parameterized by two sets of network weights , ω1 and ω2 . At training time , a network is sampled from the line defined by this pair of weights , ω∗ ( α ) = αω1 + ( 1−α ) ω2 , where α ∈ [ 0 , 1 ] . The forward and backward pass of gradient descent are computed using ω∗ ( α ) , and gradient updates are backpropagated to ω1 and ω2 . A cosine distance regularization term is used to ensure that ω1 and ω2 don ’ t converge to a single point . We propose training a linear subspace of networks in a similar manner , but we also bias the subspace to contain high-accuracy solutions at one end and high-efficiency solutions at the other end . To accomplish this , we introduce a parameter γ used to control the compression level and a compression function f ( ω , γ ) used to compress a model ’ s weights ω using a compression level γ . The parameter γ induces different compression levels at different positions in the linear subspace , and therefore can be considered a function γ ≡ γ ( α ) . When training , we first sample a network from the subspace , then compress it , obtaining a network with weights f ( ω∗ ( α ) , γ ( α ) ) . We then perform a standard forward and backward pass of gradient descent with it , backpropagating gradients to ω1 and ω2 . In summary , our training method adjusts the standard subspace training method to include the application of a compression function f at a compression level γ ( α ) . Note that we apply the compression after sampling the model ω∗ ( α ) from the subspace , since we want the sampled model to be compressed ( rather than ω1 and ω2 ) for efficient inference . We experiment with formulating our compression function as structured sparsity , unstructured sparsity , and quantization ( Section 4 ) . In all cases , the cost of computing f is small compared to the cost of a network forward pass ( Section A.1 ) .

Summary: This paper introduces learnable compressible subspaces, which attempts to learn a set of models that can be switched at inference time to adapt to different resource requirements. This work is motivated by previous work in neural subspaces and slimmable networks. It is evaluated on CIFAR10 and ImageNet and compared against other recent works for adaptable inference models. These results show under certain conditions LCS can maintain higher accuracies at larger sparsities compared to other works.