Energy-Aware Neural Architecture Optimization with Fast Splitting Steepest Descent

1 INTRODUCTION . Deep neural networks ( DNNs ) have demonstrated remarkable performance in solving various challenge problems such as image classification ( e.g . Simonyan & Zisserman , 2015 ; He et al. , 2016 ; Huang et al. , 2017 ) , object detection ( e.g . He et al. , 2017a ) and language understanding ( e.g . Devlin et al. , 2018 ) . Although large-scale deep networks have good empirical performance , their large sizes cause slow computation and high energy cost in the inference phase . This imposes a great challenge for improving the applicability of deep networks to more real-word domains , especially on mobile and edge devices where the computation and energy budgets are highly limited . It is of urgent demand to develop practical approaches for automatizing the design of small , highly energy-efficient DNN architectures that are still sufficiently accurate for real-world AI systems . Unfortunately , neural architecture optimization has been known to be notoriously difficult . Compared with the easier task of learning the parameters of DNNs , which has been well addressed by back-propagation ( Rumelhart et al. , 1988 ) ; optimizing the network structures casts a much more challenging discrete optimization problem , with excessively large search spaces and high evaluation cost . Furthermore , for neural architecture optimization in energy-efficient settings , extra difficulties arise due to strict constraints on resource usage . Recently , Wu et al . ( 2019 ) investigated similar notations of gradient descent for learning network architectures and framed the architecture optimization problem into a continuous optimization in an infinite-dimensional configuration space , on which novel notions of steepest descent can be derived for incremental update of the neural architectures . In practice , the algorithm optimizes a neural network through a cycle of paramedic updating and splitting phase . In the parametric updating phase , the algorithm performs standard gradient descent to reach a stable local minima ; in the splitting phase , the algorithm expands the network by splitting a subset of exiting neurons into several offsprings in an optimal way . A key observation is that the previous local minima can be turned into a saddle point in the new higher-dimensional space induced by splitting that can be escaped easily ; thus enabling implicitly architecture space exploration and achieving monotonic loss decrease . However , the splitting algorithm in Wu et al . ( 2019 ) treats each neuron equally , without taking into account the different amount of energy consumption incurred by different neurons , thus finding models that may not be applicable in resource-constrained environments . To close the gap between DNNs design via splitting and energy efficiency optimization , we propose an energy-aware splitting procedure that improves over Wu et al . ( 2019 ) by explicitly incorporating energy-related metrics to guild the splitting process . Another practical issue of Wu et al . ( 2019 ) is that it requires eigen-computation of the splitting matrices , which causes a time complexity of O ( nd3 ) and space complexity of O ( nd2 ) when implemented exactly , where n is the number of neurons in the network , and d is the dimension of each neuron . This makes it difficult to implement the algorithm on GPUs for modern neural networks with thousands of neurons , mainly due to the explosion of GPU memory , thus prohibiting efficient parallel calculation on GPUs . In this work , we address this problem by proposing a fast gradientbased approximation of Wu et al . ( 2019 ) , which reduces the time and space complexity to O ( nd2 ) and O ( nd ) , respectively . Critically , our new fast gradient-based approximation can be efficiently implemented on GPUs , hence making it possible to split very large networks , such as these for ImageNet . Our method achieves promising empirical results on challenging benchmarks . Compared with prior art pruning baselines that improve the efficiency by removing the least significant neurons ( e.g . Liu et al. , 2017 ; Li et al. , 2017 ; Gordon et al. , 2018 ) , our method produces a better accuracy-flops trade-off on CIFAR-100 . On the large-scale ImageNet dataset , our method finds more flops-efficient network architectures that achieve 1.0 % and 0.8 % improvements in top-1 accuracy compared with prior expert-designed MobileNet ( Howard et al. , 2017 ) and MobileNetV2 ( Sandler et al. , 2018 ) , respectively . The gain is even more significant on the low-flops regime . 2 SPLITTING STEEPEST DESCENT . Our work builds upon a recently proposed splitting steepest descent approach ( Wu et al. , 2019 ) , which transforms the co-optimization of neural architectures and parameters into a continuous optimization , solved by a generalized steepest descent on a functional space . To illustrate the key idea , assume the neural network structure is specified by a set of size parameters m = { m1 , . . . , mK } , where each mk denotes the number of neurons in the k-th layer , or the number of a certain type of neurons . Denote by Θm the set of possible parameters of networks of size m , then Θ∞ = ∪m∈NK Θm , which we call the configuration space , is the space of all possible neural architectures and parameters . In this view , learning parameters of a fixed network structure is minimizing the loss inside an individual sub-region Θm . In contrast , optimizing in the overall configuration space Θ∞ admits the co-optimization of both architectures and parameters . The key observation is that the optimization in Θ∞ is in fact continuous ( despite being infinite dimensional ) , for which ( generalized ) steepest descent procedures can be introduced to yield efficient and practical algorithms . In particular , Wu et al . ( 2019 ) considered a splitting steepest descent on Θ∞ , which consists of two phases : 1 ) the parametric descent inside each Θm with a fixed network structurem , which reduces to the typical steepest descent on parameters , and 2 ) the architecture descent crossing the boundaries of different sub-regions Θm , which , in the case of Wu et al . ( 2019 ) , corresponds to “ growing ” the network structures by splitting a set of critical neurons to multiple off-springs ( see Figure 1a ) . From the perspective of non-convex optimization , the architecture descent across boundaries of Θm can be viewed as escaping saddle points in the configuration space . As shown in Figure 1b , when the parametric training inside a fixed Θm gets saturated , architecture descent allows us to escape local optima by jumping into a higher dimensional sub-region of a larger network structure . The idea is that the local optima inside Θm can be turned into a saddle point when viewed from the higher dimensional space of larger networks ( Figure 1c ) , which is escaped using splitting descent . Escaping local minima via splitting It requires to fix a proper notion of distance on Θ∞ in order to derive a steepest descent algorithm . In Wu et al . ( 2019 ) , steepest descent with ∞-Wasserstein distance was considered , which is shown to naturally correspond to the practical procedure of splitting neurons . Here we only introduce the intuitive idea from the practical perspective of optimally splitting neurons . The readers are referred to Wu et al . ( 2019 ) for more theoretical discussion . Consider the simplest case of splitting a single neuron . Let σ ( θ , x ) be a neuron inside a neural network that we want to learn from data , where θ is the parameter of the neuron and x its input variable . Assume the loss function of θ has a general form of L ( θ ) : = Ex∼D [ Φ ( σ ( θ , x ) ) ] , ( 1 ) where D is the data distribution , and Φ ( · ) is the map from the output of the neuron to the final loss . In this work , the word “ neuron ” broadly refers to repeatable modules in neural networks , such as the typical hidden neurons , filters in CNNs . Assume we have achieved a stable local optimum of L ( θ ) , that is , ∇L ( θ ) = 0 and∇θθL ( θ ) 0 , so that we can not further decrease the loss by local descent on the parameters . In this case , splitting steepest descent enables further descent by introducing more neurons via splitting . Specifically , we split θ into m off-springs θ : = { θi } mi=1 , and replace the neuron σ ( θ , x ) with a weighted sum of the off-spring neurons ∑m i=1 wiσ ( θi , x ) , where w : = { wi } mi=1 is a set of positive weights assigned on each of the off-springs , and satisfies ∑m i=1 wi = 1 , wi > 0 . See Figure 1a for an illustration . This yields an augmented loss function on θ and w : Lm ( θ , w ) : = Ex∼D [ Φ ( m∑ i=1 wiσ ( θi , x ) ) ] . ( 2 ) It is easy to see that if we set θi = θ for all the off-springs , the network remains unchanged . Therefore , as we change θi in a small neighborhood of θ , it introduces a smooth change on the loss function . Splitting steepest descent is derived by considering the optimal splitting strategies to achieve the steepest descent on loss in a small neighborhood of the original parameters . Deriving splitting steepest descent Derive the optimal splitting strategy involves deciding the number of off-springsm , the values of the weights { wi } and the parameters for the off-springs { θi } . In Wu et al . ( 2019 ) , this is formulated into the following optimization problem : min m , θ , w { Lm ( θ , w ) − L ( θ ) s.t . ||θi − θ|| ≤ , m∑ i=1 wi = 1 , wi > 0 , ∀ i ∈ [ m ] , m ∈ N+ } , ( 3 ) where the parameters θ : = { θi } mi=1 of the off-springs are restricted within an infinitesimal -ball of the original parameter θ , that is , ||θi − θ|| ≤ , with a small positive step size parameter . Note that the number of off-springs m is also optimized , yielding an infinite dimensional optimization . Fortunately , when is very small , the optimum of Equation ( 3 ) is achieved by either m = 1 ( no splitting ) or m = 2 ( two off-springs ) . The property of the optimal solution is characterized ( asymptotically ) by the following key splitting matrix S ( θ ) , S ( θ ) = Ex∼D [ ∇σΦ ( σ ( θ , x ) ) ∇2θθσ ( θ , x ) ] , which is a symmetric d × d matrix ( and d is the dimension of θ ) . The optimum of Equation ( 3 ) , when L ( θ ) reaches a stable local optimum ( i.e. , ∇θL ( θ ) = 0 , ∇θθL ( θ ) 0 ) , is determined by S ( θ ) via min m , θ , w { Lm ( θ , w ) − L ( θ ) } = 2 2 min { λmin ( S ( θ ) ) , 0 } + O ( 3 ) , ( 4 ) where λmin ( S ( θ ) ) denotes the minimum eigenvalue of S ( θ ) , and it is called the splitting index . When λmin ( S ( θ ) ) > 0 , the loss can not be improved by any splitting strategies following ( 4 ) . When λmin ( S ( θ ) ) < 0 , the maximum decrease of loss , which equals 2λmin ( S ( θ ) ) /2 , can be achieved by a simple strategy of splitting the neuron into two copies with equal weights , whose parameters are updated along the minimum eigen-vectors vmin ( S ( θ ) ) of S ( θ ) , that is , m = 2 , θ1 = θ + vmin ( S ( θ ) ) , θ2 = θ − vmin ( S ( θ ) ) , w1 = w2 = 1/2 . ( 5 ) In this case , splitting allows us to escape the parametric local optima to enable further improvement . Splitting deep neural networks As shown in Wu et al . ( 2019 ) , the result above can be naturally extended to more general cases when we need to split multiple neurons in deep neural networks . Consider a neural network with n neurons θ [ 1 : n ] = { θ1 , · · · , θn } . Assume we split a subset A of neurons with the optimal strategy in Equation ( 5 ) following their own splitting matrices , the improvement of the overall loss equals the sum of individual gains : G ( A ) = ∑ ` ∈A λmin ( ` ) , where λmin ( ` ) : = λmin ( S ` ( θ ` ) ) denotes the minimum eigenvalue of the splitting matrix S ` ( θ ` ) associated with neuron ` . Therefore , given a budget of splitting at most a given number of neurons , the optimal subset of neurons to split are the top ranked neurons with the smallest , and negative minimum eigenvalues . Overall , the splitting descent in Wu et al . ( 2019 ) alternates between parametric updates with fixed network architectures , and splitting top ranked neurons to augment the architectures , until a stopping criterion is reached .

Summary: This paper builds on a recently proposed algorithm ("splitting steepest descent", Wu et al 2019) for guiding the growth of a smaller network into a larger one in architecture search. The algorithm in Wu, et al. alternates between two steps, (i) optimization of parameters for a fixed model and (ii) modification of the architecture by identifying a subset of neurons to split into more neurons, based on the "splitting index" of each neuron (amounting to evaluating the smallest eigenvalue of a matrix). This work builds on that in two ways: (i) it incorporates an energy budget into the optimization procedure for choosing which subset of neurons to split, which it approximately solves by a continuous relaxation, and (ii) avoids doing exact eigendecomposition to extract the minimum eigenvalue (splitting index) but instead replaces it with a more efficient SGD on the Rayleigh quotient.