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: This paper proposes an instance-aware dynamic network, ISBNet, for efficient image classification. The network consists of layers of cell structures with multiple branches within. During the inference, the network uses SelectionNet to compute a "calibration weight matrix", which essentially controls which branches within the cell should be used to compute the output. Similar to previous works in NAS, this paper uses Gumbel Softmax to compute the branch selection probability. The network is trained to minimize a loss function that considers both the accuracy and the inference cost. Training of the network is divided into two stages: First, a high temperature is used to ensure all the branches are sufficiently optimized, and at the second stage, the authors aneal the temperature. During the inference, branches are selected if their probability computed by Gumbel Softmax is larger than a certain threshold.