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: This paper proposes a new architecture search method called "DARC" that utilizes a differentiable objective function. Since a naive formulation of architecture search is reduced to a combinatorial optimization which is not differentiable, the optimization requires much computational cost. To overcome this difficulty, this paper proposes a L1-norm relaxation and apply such relation in a layer-wise manner. The method shares a similar spirit with NAS, but the proposed model is more like "model selection" from a fixed candidates, and thus there is a Rademacher complexity guarantee. The effectiveness of DARC is justified by thorough numerical experiments.