Differentiable Architecture Compression

1 INTRODUCTION . In machine learning , resources at inference time are often much more constrained than at training time . For example , while neural networks for computer vision and natural language processing ( NLP ) are routinely trained using GPUs , trained networks are often deployed on embedded systems or mobile devices with limited memory and computational power . As another example , it is common to train a model that will be applied continuously in production ; while training occurs for a limited time , the machine performing inference may run indefinitely , and so learning a more efficient model can directly reduce costs associated with hardware or energy usage . As a result , many recent papers have studied deep model compression and acceleration . Most of these papers provide resource efficient model components ( Jaderberg et al. , 2014 ; Zhang et al. , 2016 ; Wu et al. , 2017 ) or methods to prune or quantize parameters ( LeCun et al. , 1990 ; Polyak & Wolf , 2015 ; Li et al. , 2016a ; He et al. , 2017 ; Luo & Wu , 2017 ; Luo et al. , 2017 ; Zhuang et al. , 2018 ) . We propose a general paradigm called Differentiable ARchitecture Compression ( DARC ) for learning in the context of constrained resources at inference-time . Rather then suggesting a specific cheap component and using it blindly throughout a neural network , or trying to tune layer hyperparameters as in network pruning or quantization , our approach is inspired by Neural Architecture Search ( NAS ) ( Zoph & Le , 2016 ; Liu et al. , 2018a ; Pham et al. , 2018 ; Kandasamy et al. , 2018 ; Liu et al. , 2018b ) . DARC starts with a resource-intensive network design and uses training data to learn which components can be replaced with efficient alternatives , while maintaining model output quality . The resource requirement of the final model is controlled via a regularization term , which can be flexibly defined depending on the objective , such as minimizing inference time or memory footprint . DARC has a clear intuitive advantage when compared to methods that quantize or prune parameters ; these approaches are inherently restricted in their search space . They can not replace a layer with a structurally different layer , or replace two or more layers with a shallow alternative . As examples , it might be the case that a convolutional layer can not be pruned without hurting performance but can be replaced with a depthwise-separable convolution , or an LSTM layer might not be amenable to weight pruning , but could be replaced with a more efficient self-attention layer . DARC , applied to deep networks , offers a way to search a rich space in the context of model compression . The high-level idea is to partition the network into components , and explore alternatives to these components simultaneously . Replacing all components simultaneously is crucial , as it provides a data-driven way to decide which components can be replaced with cheap alternatives . Indeed , replacing layers blindly may sacrifice too much prediction performance , as can be seen in Mo- bileNets ( Howard et al. , 2017 ; Sandler et al. , 2018 ) , which exclusively use depthwise-separable convolutions ; although computationally efficient , even the most accurate MobileNet model is far less accurate at ImageNet classification than , say , ResNet50 . We frame the problem of learning good replacement components as a sparse ensemble learning problem . This view allows us to draw guidelines from simpler analyzable cases , leading to a simple , gradient-based learning scheme that avoids over-fitting and is fast enough to apply DARC directly to large datasets such as ImageNet . This contrasts from most NAS methods , which first learn an architecture on a small dataset and then fit this architecture on the larger dataset . We present experiments on networks commonly used in computer vision . By applying our techniques on the ResNet Architecture ( He et al. , 2016 ) , we outperform state-of-the-art models on both ImageNet and CIFAR-10 datasets , in terms of accuracy vs. throughput and accuracy vs. model size . A few results from our compression framework : for a WideResNet model achieving 97.2 % accuracy on CIFAR-10 , we improve single-sample inference speed by 2.28× and memory footprint by 5.64× , with no loss in accuracy . For a modified ResNet50 model with 79.15 % Top1 accuracy on ImageNet , we improve inference speed by 1.29× and memory footprint by 3.57× with 1 % loss in accuracy . Both base models are publicly available from the GluonCV Model Zoo ( Mod , 2018 ) . Our experiments empirically demonstrate an intuitive observation that ‘ you get what you optimize for ’ , in that models minimizing model size tend to be quite different from those maximizing throughput . We note that , while our experiments are limited to image classification , DARC is applicable to any deep learning architecture , including models with recurrent cells or transformers , or indeed any sufficiently modular learning algorithm , as described in Section 2 , and any task with a well-defined objective function in which we would like to reduce inference costs . We see this work as a proof-ofconcept for capabilities of the DARC framework , and the results in this paper give a strong indication that DARC can be applied to NLP architectures , or optimized for metrics other than those we try here ( e.g . for latency on devices other than GPU , or energy consumption ) . 2 GENERAL SETTING AND DARC ESTIMATOR . The intuition and motivation for our method starts with a task of model selection . Given a task and J function families corresponding to candidate models , we are interested in finding the best model type for the task . We relax this combinatorial optimization problem to a more tractable differentiable optimization problem ( whence the name “ differentiable architecture compression ” ) , by allowing convex combinations of these candidates . This can be thought of as a constrained form of ensemble learning in which weights are restricted to represent a convex combinations of individual learners . Then , we posit a budget constraint : each model type has an associated cost ( e.g. , memory consumption or latency ) , and the overall cost of the ensemble is the sum of costs of the used models . Our task is then to learn a convex ensemble over a subset of candidates , with total cost within budget . We now formalize this approach , with modifications to address technical challenges as they arise . In the sequel , for any positive integer J , [ J ] = { 1 , 2 , ... , J } denotes the set of positive integers at most J , and ∆J : = { x ∈ [ 0 , 1 ] J : ∑ j∈ [ J ] xj = 1 } denotes the probability simplex over J elements . Consider the conventional supervised learning setting , in which we have an i.i.d . training dataset ( X1 , Y1 ) , ... , ( Xn , Yn ) IID∼ P ( X , Y ) from some joint distribution P ( X , Y ) on X × Y . Fix a loss function L : R × Y → [ 0 , ∞ ] , and a hypothesis class H of R-valued functions . We would like to learn a function h : X → R , h ∈ H that minimizes the risk R ( h ) : = E ( X , Y ) ∼P ( X , Y ) [ L ( h ( X ) , Y ) ] . The usual empirical risk minimization ( ERM ) estimator is ĥERM : = arg minh∈H R̂ ( h ) where , for any hypothesis h ∈ H , R̂ ( h ) : = 1n ∑n i=1 L ( h ( Xi ) , Yi ) denotes the empirical risk . To derive our resource-constrained objective , we impose a few structural assumptions on our hypothesis classH : ( A1 ) H = Conv ⋃ j∈ [ J ] Hj is the convex hull of a union of J classesH1 , ... , HJ . ( A2 ) Each classHj class has a known cost Cj ≥ 0 of using a hypothesis hj ∈ Hj at test-time . ( A3 ) Costs are additive : hypothesis h = ∑J j=1 αjhj ∈ H has cost C ` 0 ( α ) = ∑J j=1 Cj1 { αj > 0 } . ( A4 ) We have a known budget B ≥ 0 for the final model at test-time . As we show in Section 3 , these assumptions arise naturally in architecture compression . Given assumptions equation A1-equation A4 , the constrained ERM estimate is ĝ = ∑ j∈ [ J ] α̂j ĥj , where ( α̂0 , ĥ0 ) : = arg min α∈∆J , ĥj∈Hj R̂ ∑ j∈ [ J ] αjhj , subject to C ` 0 ( α ) ≤ B . ( 1 ) The above estimator is difficult ( NP-hard ) to compute , due to the non-smooth , non-convex budget constraint C ` 0 ( α ) ≤ B . Since this constraint bounds the ` 0 norm of α ( weighted by C ) , the usual remedy would be to relax the constraint to one on the ` 1 norm of α ( weighted by C ) , namely C ` 1 ( α ) : = ∑ j∈ [ J ] Cjαj ≤ B . Unfortunately , due to the constraint that α lies in the probability simplex ∆J ( which implies ∑ j∈ [ J ] αj = 1 ) , the ` 1 constraint is insufficient to induce sparsity on α. Fortunately , sparse optimization on ∆J is well-studied , with many solutions proposed ( Pilanci et al. , 2012 ; Kyrillidis et al. , 2013 ; Li et al. , 2016b ) . Due to ease of implementation , we adopt a simple but effective solution proposed by ( Kyrillidis et al. , 2013 ) , which involves alternating gradient updates with a projection operation P∆J : RJ\RJ− → ∆J , given by P∆J ( α ) = α+ ‖α+‖1 , where α+ = ( max { 0 , α1 } , ... , max { 0 , αj } ) ∈ RJ+ . P∆J is easy to compute , enforces the simplex constraint α ∈ ∆J exactly , and induces sparsity on α . For an intuition of how this works , one can note that ∇‖α‖22/2 = α , so that the update α/‖α‖1 = α − ( 1− 1/‖α‖1 ) α can be viewed as a gradient step for minimizing −‖α‖22/2 with adaptive step size ( 1− 1/‖α‖1 ) . As a technicality , we note that the projection P∆J ( α ) is undefined when α ∈ RJ− has no positive components . However , for realistic gradient step sizes η , this never occurs , since , after each gradient update , ∑ j∈ [ J ] αj ≥ 1−O ( η ) . Finally , a natural initial point for our procedure is one where C ` 1 ( α ) > B , hence we reexpress the constraint C ` 1 ( α ) ≤ B as a penalty λC ` 1 ( α ) . Since the value of λ corresponding to B is not known a priori , we iteratively increase λ until the solution of the optimization problem satisfies the budget constraint . The resulting DARC procedure is shown Algorithm 1 . We note that the “ stopping criterion ” for the inner loop can be as simple as a fixed number of training epochs ( as in our experiments ) , or a more sophisticated early-stopping criterion . Algorithm 1 : DARC algorithm for general hypotheses Data : Training Data { ( Xi , Yi ) } ni=1 , J candidates h1 , w1 , ... , hj , wj with initial parameters w1 , ... , wj and costs c1 , ... , cj ≥ 0 , initial cost penalty parameter λ0 > 0 , budget B . Result : α , w1 , ... , wj such that h = ∑ j∈ [ J ] αjhj , wj has small risk R ( h ) and cost C ` 0 ( α ) ≤ B 1 α← ( 1/J , ... , 1/J ) , λ← λ0 ; 2 while C ` 0 ( α ) > B do 3 while stopping criterion is not met do 4 ( α , w1 , ... , wJ ) ← ( α , w1 , ... , wJ ) − η ( ∇α , w1 , ... , wJ R̂ ( ∑ j∈ [ J ] αjhj , wj ) + C ` 1 ( α ) ) ; 5 α← P∆J ( α ) ; 6 end 7 λ← 2λ 8 end

Summary: The authors proposed a new gradient-based architecture search method that tries to find more efficient alternatives starting from the pre-trained model. The approach is similar to DARTS (Liu et al., 2019) with a budget constraint, such as size and throughput. One major difference is to modify the update of architectural parameters, i.e., mixing weights of all candidate operations, to induce the sparsity rather than to keep the weighted sum of all possible operations. Another difference is that it starts from a well-defined architecture with pre-trained weights. It is simple to apply, but hard to tell. The recognized strengths and concerns are as follows.