ISBNet: Instance-aware Selective Branching Networks

Recent years have witnessed growing interests in designing efficient neural networks and neural architecture search ( NAS ) . Although remarkable efficiency and accuracy have been achieved , existing expert designed and NAS models neglect the fact that input instances are of varying complexity and thus different amounts of computation are required . Inference with a fixed model that processes all instances through the same transformations would incur computational resources unnecessarily . Customizing the model capacity in an instance-aware manner is required to alleviate such a problem . In this paper , we propose a novel Instanceaware Selective Branching Network-ISBNet to support efficient instance-level inference by selectively bypassing transformation branches of insignificant importance weight . These weights are dynamically determined by a lightweight hypernetwork SelectionNet and recalibrated by gumbel-softmax for sparse branch selection . Extensive experiments show that ISBNet achieves extremely efficient inference in terms of parameter size and FLOPs comparing to existing networks . For example , ISBNet takes only 8.70 % parameters and 31.01 % FLOPs of the efficient network MobileNetV2 with comparable accuracy on CIFAR-10 . 1 INTRODUCTION . Deep convolutional neural networks ( CNNs ) ( He et al. , 2016 ; Zoph et al. , 2018 ) have revolutionized computer vision with increasingly larger and more sophisticated architectures . These model architectures have been designed and calibrated by domain experts with rich engineering experience . To achieve good inference results , these models typically comprise hundreds of layers and contain tens of millions of parameters and consequently , consume substantial amounts of computational resources for both training and inference . Recently , there has been a growing interest in efficient network design ( Howard et al. , 2017 ; Iandola et al. , 2016 ; Zhang et al. , 2018 ; Sandler et al. , 2018 ) and neural architecture search ( NAS ) ( Zoph et al. , 2018 ; Real et al. , 2018 ; Liu et al. , 2018b ) , respectively with the objective of devising network architectures that are efficient during inference and automating the architecture design process . Many efficient architectures have indeed been designed in recent years . SqueezeNet ( Iandola et al. , 2016 ) and MobileNet ( Howard et al. , 2017 ) substantially reduce parameter size and computation cost in terms of FLOPs on mobile devices . More recent works such as MobileNetV2 ( Sandler et al. , 2018 ) and ShuffleNetV2 ( Ma et al. , 2018 ) further reduce the FLOPs . It is well recognized that devising these architectures is non-trivial and requires engineering expertise . Automating the architecture design process via neural architecture search ( NAS ) has attracted increasing attention in recent years . Mainstream NAS algorithms ( Zoph & Le , 2016 ; Zoph et al. , 2018 ; Real et al. , 2018 ) search for the network architecture iteratively . In each iteration , an architecture is proposed by a controller , and then trained and evaluated . The evaluation performance is in turn exploited to update the controller . This process is incredibly slow because both the controller and each derived architecture require training . For instance , the reinforcement learning ( RL ) based controller NASNet ( Zoph et al. , 2018 ) takes 1800 GPU days and the evolution algorithm based controller AmoebaNet ( Real et al. , 2018 ) incurs 3150 GPU days to obtain the best architecture . Many acceleration methods ( Baker et al. , 2017 ; Liu et al. , 2018a ; Bender et al. , 2018 ; Pham et al. , 2018 ) have been proposed to accelerate the search process , and more recent works ( Liu et al. , 2018b ; Xie et al. , 2018 ; Wu et al. , 2018 ; Cai et al. , 2018 ) remove the controller and instead optimize the architecture selection and parameters together with gradient-based optimization algorithms . While both expert designed and NAS searched models have produced remarkable efficiency and prediction performance , they have neglected one critical issue that would affect inference efficiency . The architectures of these models are fixed during inference time and thus not adaptive to the varying complexity of input instances . However , in real-world applications , there are only a small fraction of input instances requiring deep representations ( Wang et al. , 2018 ; Huang et al. , 2017a ) . Consequently , expensive computational resources would be wasted if all instances are treated equally . Designing a model with sufficient representational power to cover the hard instances , and meanwhile a finer-grained control to provide just necessary computation dynamically for instance of varying difficulty is therefore essential . In this paper , we propose ISBNet to address the aforementioned issue with its building block Cell as illustrated in Figure 1 . Following the widely adopted strategy in NAS ( Zoph et al. , 2018 ; Pham et al. , 2018 ; Liu et al. , 2018b ; Xie et al. , 2018 ) , the backbone network is a stack of L structurally identical cells , receiving inputs from their two previous cells and each cell contains N inter-connected computational Nodes . The architecture of ISBNet deviates from the conventional wisdom of NAS which painstakingly search for the connection topology and the corresponding transformation operation of each connection . In ISBNet , each node is instead simply connected to its prescribed preceding node ( s ) and each connection transforms via a candidate set of B operations ( branches ) . To allow for instance-aware inference control in the branch level , we integrate L lightweight hypernetworks SelectionNets , one for each cell to determine the importance weight of each branch . Gumbel-softmax ( Jang et al. , 2016 ; Maddison et al. , 2016 ) is further introduced to recalibrate these weights , which enables efficient gradient-based optimization during training , and more importantly , leads to sparse branch selection during inference for efficiency . The contributions of ISBNet can be summarized as follows : • ISBNet is a general architecture framework combining advantages from both efficient network design and NAS , whose components are readily customizable . • ISBNet is a novel architecture supporting the instance-level selective branching mechanism by introducing lightweight SelectionNets , which improves inference efficiency significantly by reducing redundant computation . • ISBNet successfully integrates gumbel-softmax to the branch selection process , which enables direct gradient descent optimization and is more tractable than RL-based method . • ISBNet achieves state-of-the-art inference efficiency in terms of parameter size and FLOPs and inherently supports applications requiring fine-grained instance-level control . Our experiments show that ISBNet is extremely efficient during inference and successfully selects only vital branches on a per-input basis . In particular , with a minor 1.07 % accuracy decrease , ISBNet reduces the parameter size and FLOPs by 10x and 11.31x respectively comparing to the NAS searched high-performance architecture DARTS ( Liu et al. , 2018b ) . Furthermore , with a tiny model of 0.57M parameters , ISBNet achieves much better accuracy while with only 8.03 % and 30.60 % inference time parameter size and FLOPs comparing to the expert-designed efficient network ShuffleNetV2 1.5x ( Ma et al. , 2018 ) . We also conduct ablation studies and visualize the branch selection process to understand the proposed architecture better . The main results and findings are summarized in Sec 4.2 and Sec 4.3 . 2 RELATED WORK . Efficient Network Design . Designing resource-aware networks ( Iandola et al. , 2016 ; Gholami et al. , 2018 ; Ma et al. , 2018 ; Sandler et al. , 2018 ; Gholami et al. , 2018 ) has attracted a great deal of attention in recent years . However , these works mainly focus on reducing parameter size and inference FLOPs in many works . For instances , SqueezeNet ( Hu et al. , 2018 ) reduces parameters and computation with the fire module ; MobileNetV2 ( Sandler et al. , 2018 ) utilize depth-wise and point-wise convolution for more parameter-efficient convolutional neural networks ; ShuffleNetV2 ( Ma et al. , 2018 ) proposes lightweight group convolution with channel shuffle to facilitate the information flowing across the channels . To make inference efficient , many of these transformations are introduced to the candidate operation set in ISBNet . Many recent works explore conditional ( Wang et al. , 2018 ) and resource-constrained prediction ( Huang et al. , 2017a ) for efficiency . SkipNet ( Wang et al. , 2018 ) introduces a gating hypernetwork to determine whether to bypass each residual layer ( He et al. , 2016 ) conditional on the current input instance . Compared with SkipNet , ISBNet provides more efficient and diversified branch selections for the backbone network and the hypernetworks in ISBNet are optimized in an end-to-end training manner instead of generally less tractable policy gradient ( Williams , 1992 ) . MSDNet ( Huang et al. , 2017a ) supports budgeted prediction within prescribed computational resource constraint during inference by inserting multiple classifiers into a 2D multi-scale version of DenseNet ( Huang et al. , 2017b ) . By early-exit into a classifier , MSDNet can provide approximate predictions with minor accuracy decrease . Functionally , ISBNet also supports budgeted prediction by dynamically controlling the number of branches selected , therefore per-input inference cost . Neural Architecture Search . Mainstream NAS ( Zoph et al. , 2018 ; Real et al. , 2018 ) treats architecture search as a stand-alone process whose optimization is severed from candidate architecture optimization . Search algorithms such as RL-based NAS ( Zoph et al. , 2018 ) and evolutionary-based NAS ( Real et al. , 2018 ) obtain state-of-the-art architectures at an unprecedented amount of the GPUtime searching cost . Recently , many works have been proposed to accelerate the search pipeline , e.g. , via performance prediction ( Baker et al. , 2017 ; Liu et al. , 2018a ) , hypernetworks generating initialization weights ( Brock et al. , 2017 ) , weight sharing ( Bender et al. , 2018 ; Pham et al. , 2018 ) . These approaches greatly alleviate the search inefficiency while the scalability issue remains unsolved . A number of proposals ( Liu et al. , 2018b ; Wu et al. , 2018 ; Cai et al. , 2018 ) instead integrate the architecture search process and architecture optimization into the same gradient-based optimization framework . In particular , DARTS ( Liu et al. , 2018b ) relaxes discrete search space to be continuous by introducing operation mixing weights to each connection and optimizes these weights directly with gradient back-propagated from validation loss . Similarly , the discrete search space in SNAS ( Xie et al. , 2018 ) is modeled with sets of one-hot random variables for each connection , which is made differentiable by relaxing the discrete distribution with continuous concrete distribution ( Jang et al. , 2016 ; Maddison et al. , 2016 ) . In terms of architecture optimization , ISBNet also relaxes the discrete branch selection to continuous importance weights optimized by gradient descent ; while instead of direct optimization on the weights , SelectionNets are introduced to dynamically generate these weights which are more effective and meanwhile bring about larger model capacity . Further , SelectionNets enable instance-level architecture customization rather than finding a fixed model . 3 INSTANCE-AWARE SELECTIVE BRANCHING NETWORK . 3.1 THE BACKBONE NETWORK . The backbone network is constructed with a stack of L cells , each of which is a directed acyclic graph consisting of an ordered sequence of N intermediate nodes . As is illustrated in Figure 1 , xl0 and xl1 are the cell input nodes from the two preceding cells ; each intermediate node x l i ( i ≥ 2 ) of the lth cell forms a latent representation and receives n input nodes1 from its preceding nodes : 1n can be larger than 2 for deeper and wider local representation . E.g. , n = i− 1 for each xli leads to dense connection , i.e. , DenseNet ( Huang et al. , 2017b ) . xli = ∑ j∈Sli Fj , i ( xlj ) , Sli ⊂ { 0 , 1 , · · · , i− 1 } ∧ |Sli| = n ( 1 ) Thereby , each cell contains C = n ·N connections in total . The connection passes information from node xlj to x l i after the aggregation of a candidate set of B branches of transformation inspired from widely-adopted transformations in NAS ( Pham et al. , 2018 ; Liu et al. , 2018b ; Xie et al. , 2018 ) and efficient network design ( Iandola et al. , 2016 ; Sandler et al. , 2018 ; Zhang et al. , 2018 ) : Fj , i ( xlj ) = B∑ b=1 wb · Fb ( xlj ) ( 2 ) where wb here represents the importance of the bth branch ( operation ) of the connection and is dynamically generated by the cell hypernetwork rather than a fixed learned parameter as is in existing NAS methods ( Liu et al. , 2018b ; Xie et al. , 2018 ) . We shall introduce the hypernetwork in Section 3.2 . Finally , the output of the cell xlout is aggregated by concatenating the output from all the intermediate nodes . We shall use superscript l , subscript c and b to index the cell , connection , and branch respectively . Recent work ( Xie et al. , 2019 ) reveals that architectures with randomly generated connection achieve surprisingly competitive results comparing to best NAS models , which is confirmed empirically in our experiments on smaller datasets . In this paper , we thus mainly focus on the branch transformation and selection part and their impact on the inference efficiency instead of specifying a detailed connection topology . Under this architecture formulation framework , we can readily adjust the number of candidate branches B and also the specific transformations before training , customizing model capacity and efficiency respectively depending on the difficulty of the task and resource constraints in deployment . 3.2 SelectionNet FOR WEIGHT RECALIBRATION To support instance level inference control , we introduce L lightweight hypernetworks SelectionNet , one for each cell . Each SelectionNet SNetl receives the same input as the lth cell , specifically the two output nodes xl−2out , x l−1 out ( i.e . x l 0 , x l 1 ) from the preceding cells , and concurrently produce C sets of recalibration weights , one for each connection of the cell : W l = SNetl ( xl0 , x l 1 ) ( 3 ) where W l ∈ RC×B is the recalibration weight matrix for the lth cell . The SelectionNet SNetl dynamically generates these weights with a pipeline of m = 2 convolutional blocks , a global average pooling and finally an affine transformation . For the m convolutional transformation , we adopt separable convolution ( Sandler et al. , 2018 ) which contains a point-wise convolution and a depth-wise convolution of stride 2 and kernel size 5 × 5 . The stride reduces the parameter size and computation of SNetl , and the larger kernel size for depth-wise convolution incurs negligible overhead while extracts features for the immediate weight generation with a larger local receptive field . The recalibration weights given by the SelectionNet is reminiscent of convolutional attention mechanism ( Hu et al. , 2018 ; Woo et al. , 2018 ; Newell et al. , 2016 ) , where attention weights are determined dynamically by summarizing information of the immediate input and then exploited to recalibrate the relative importance of different input dimensions , e.g. , channels in SENet ( Hu et al. , 2018 ) . In ISBNet , the recalibration weights are introduced to the branch . Particularly , each candidate operation of the connection is coupled with a rescaling weight . The gumbel-softmax ( Jang et al. , 2016 ; Maddison et al. , 2016 ) technique and the reparameterization trick ( Kingma & Welling , 2013 ) is introduced to further recalibrate these weights generated by the SelectionNet , to enable efficient gradient-based optimization for the whole network during training , and more importantly , ensure a sparse selection of important branches during inference . More specifically , each set of importance weights W lc ∈ RB for the cth connection of the lth cell ( Clc ) after the following recalibration of the gumbel-softmax follows concrete distribution ( Maddison et al. , 2016 ) controlled by a temperature parameter τ : w̃lc , b = exp ( ( wlc , b +G l c , b ) /τ ) ∑B b′=1 exp ( ( w l c , b′ +G l c , b′ ) /τ ) , τ > 0 ( 4 ) where w̃lc , b is then directly used for branch recalibration as in Equation 2 , and G l c , b = − log ( − log ( U lc , b ) ) here is a gumbel random variable coupling with bth branch by sampling U lc , b from Uniform ( 0 , 1 ) ( Jang et al. , 2016 ) . The concrete distribution ( Maddison et al. , 2016 ) suggests that ( 1 ) w̃c , b = 1B , as τ → +∞ , and more importantly ( 2 ) : p ( lim τ→0 w̃lc , b = 1 ) = exp ( w l c , b ) / B∑ b′=1 exp ( wlc , b′ ) ( 5 ) Therefore , high temperature leads to uniform dense branch selection while lower temperature tends to sparsely sample branches following a categorical distribution parameterized by softmax ( W lc ) . 3.3 OPTIMIZATION AND INFERENCE FOR ISBNet With the continuous relaxation of the gumbel-softmax ( Jang et al. , 2016 ; Maddison et al. , 2016 ) and the reparameterization ( Kingma & Welling , 2013 ) , the branch selection process of the SelectionNets is made directly differentiable with respect to the weight wlc , b . In particular , the gradient ∂L ∂w̃lc , b backpropagated from the loss function L to w̃lc , b through the backbone network can be directly backpropagated towlc , b with low variance ( Maddison et al. , 2016 ) , and further to the lth SelectionNet unimpededly . Therefore , parameters of the whole network can be optimized in an end-to-end manner by gradient descent . The temperature τ of Equation 4 regulates the sparsity of the branch selection . A relatively higher temperature forces the weights to distribute more uniformly so that all the branches of each connection are efficiently trained . While a low temperature instead tends to sparsely sample one branch from the categorical distribution parameterized by the importance weights dynamically determined by SelectionNets , thus supporting finer-grained instance-level inference control by bypassing unimportant branches . To leverage both characteristics , we propose a two stage training scheme for ISBNet : ( 1 ) the first stage pretrains the whole network with a fixed relatively high temperature till convergence . ( 2 ) the second stage fine-tunes the parameters with τ steadily annealing to a relatively low temperature . The first stage ensures that branches are sufficiently optimized before the instanceaware selection and the fine-tuning in the second stage helps maintain the performance of ISBNet under sparse branch selection during inference . To further promote inference efficiency and reduce redundancy , a regularization term is explicitly introduced in the fine-tuning stage which takes into account the expectation of the resource consumptionR in the final loss function L for correctly classified instances : L = LCE + λ1||w||22 + λ21ŷ=y logE [ R ] ≈ LCE + λ1||w||22 + λ21ŷ=y log L∑ l=1 C∑ c=1 B∑ b=1 w̃lc , b · R ( F lc , b ( · ) ) ( 6 ) where LCE and λ1||w||22 denotes the cross-entropy loss and the weight decay term , y is the ground truth class label , ŷ the prediction , λ2 controls the regularization strength and R ( · ) calculates the resource consumption of each operation F lc , o ( · ) . The operation importance weight w̃lc , b represents the probability of the corresponding branch F lc , b being selected during inference , and therefore the regularization term E [ R ] corresponds to the expectation of the aggregated resource required for each input instance . The resource regularizer is readily adjustable depending on deployment constraints , which may include the parameter size , FLOPs , and memory access cost ( MAC ) . In this work , we mainly focus on the inference time , specifically FLOPs , which can be calculated beforehand for each branch . R ( F lc , b ( · ) ) is thus a constant here , which means that the regularizerR is also directly differentiable with respect to w̃lc , b . We denote ISBNet trained with regularization strength λ2 as ISBNet-R-λ2 . During inference , the instance-level selective branching is achieved for each connection Clc by selecting branches of the top k largest recalibration weights whose aggregated weight slc just exceeds a threshold T . Denoting W̃ lc sorted in descending order as Ŵ l c , then : slc = min { sk : sk = k∑ b=1 ŵlc , b ∧ sk ≥ T } ( 7 ) After the selection , the recalibration weight w̃lc , b of the selected branch is rescaled by 1 slc to stabilize the scale of the representation . Consequently , the SelectionNet will select only necessary branches2 for each instance depending on the input difficulty and meanwhile the FLOPs of each branch , i.e. , trading off between LCE and R in Equation 6 . Furthermore , the resource consumption of each instance can be precisely regulated in a finer-grained manner by scheduling the threshold dynamically for each connection . In this paper , the same threshold is shared among all connections for simplicity and ISBNet inference with the threshold t is denoted as ISBNet-T-t . Under such an inference scheme , the backbone network comprises up to ( 2B−1 ) L·C possible candidate subnets , corresponding to each unique branch selection of allL·C connections . For a small ISBNet of 10 cells , with 5 candidate operations , 8 connections per cell , there are ( 25 − 1 ) 8·10 ≈ 2·10119 possible candidate architectures of different branch combination , which is orders of magnitudes larger than the search space of conventional NAS ( Pham et al. , 2018 ; Liu et al. , 2018b ; Xie et al. , 2018 ; Cai et al. , 2018 ; Stamoulis et al. , 2019 ) .

Summary: Neural architecture search usually aims to find a single fixed architecture for the task of interest. The paper proposes to condition the architecture on the input instances by introducing a "selection network" that learns to retain a subset of branches in the architecture during each inference pass. The intuition is that easier instances require less compute (hence a shallower/sparser architecture) as compared to the more difficult ones. The authors show improved results on CIFAR-10 and ImageNet in terms of accuracy-latency trade-off over some handcrafted architectures and NAS baselines. The method resembles sparsely gated mixture of experts [1] at a high-level, but has been implemented in a way that better fits the context of architecture search (which is still technically interesting).