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: The paper present a method for learning a compressible subspace of neural networks that contains a fine-grained spectrum of models that range from highly efficient to highly accurate. The proposed method allows choosing the proper point of trade-off between accuracy and efficiency at inference time, according to the available resources. There are also efforts to reduce the runtime tweaking overhead like replacing BatchNorm with GroupNorm.