Triple-Search: Differentiable Joint-Search of Networks, Precision, and Accelerators

1 INTRODUCTION . The powerful performance and prohibitive complexity of deep neural networks ( DNNs ) have fueled a tremendous demand for efficient DNN accelerators which could boost DNN acceleration efficiency by orders-of-magnitude ( Chen et al. , 2016 ) . In response , extensive research efforts have been devoted to developing DNN accelerators . Early works decouple the design of efficient DNN algorithms and their accelerators . On the algorithms level , pruning , quantization , or neural architecture search ( NAS ) are adopted to trim down the model complexity ; On the hardware level , various FPGA-/ASIC-based accelerators have been developed to customize the micro-architectures ( e.g. , processing elements dimension , memory sizes , and network-on-chip design ) and algorithm-to-hardware mapping methods ( e.g. , loop tiling strategies and loop orders ) in order to optimize the acceleration efficiency for a given DNN . Later , hardware-aware NAS ( HA-NAS ) has been developed to further improve DNNs ’ acceleration efficiency for different applications ( Tan et al. , 2019 ) . More recently , it has been recognized that ( 1 ) optimal DNN accelerators require a joint consideration/search for all the following different yet coupled aspects , including DNNs ’ network structure , the adopted precision , and their accelerators ’ micro-architecture and mapping methods , and ( 2 ) merely exploring a subset of these aspects will lead to sub-optimal designs in terms of hardware efficiency or task accuracy . For example , the optimal accelerators for networks with different structures ( e.g. , width , depth , and kernel size ) can be very different ; while the optimal networks and their bitwidths for different accelerators can differ a lot ( Wu et al. , 2019 ) . However , the direction of jointly designing or searching for all the three aspects has only been slightly touched on previously . For example , ( Chen et al. , 2018 ; Gong et al. , 2019 ; Wang et al. , 2020 ) proposed to jointly search for the structure and precision of DNNs for a fixed target hardware ; ( Abdelfattah et al. , 2020 ; Yang et al. , 2020 ; Jiang et al. , 2020a ; b ) made the first attempt to jointly search for the networks and their accelerators , yet either their network or accelerator choices are limited due to the prohibitive time cost required by their adopted reinforcement learning ( RL ) based methods ; and EDD ( Li et al. , 2020 ) contributed a pioneering effort towards this direction by formulating a differentiable joint search framework , which however only consider one single accelerator parameter ( i.e. , parallel factor ) and more importantly , has not yet fully solved the challenges of such joint search . Although differentiable search is one of the most promising ways in terms of search efficiency to explore the huge joint search space as discussed in Sec . 4.2 , plethora of challenges exist to achieve an effective generic joint search for the aforementioned three aspects . First , Challenge 1 : to jointly search for a network and its precision via differentiable search , there exists a dilemma whether to activate all the paths during search . On one hand , the required memory consumption can easily explode and thus constrain the search ’ s scalability to more complex tasks if all paths are activated ; on the other hand , partially activating a subset of the paths can lead to a sequential training of different precision on the same weights , which might result in inaccurate accuracy ranking among different precision as discussed in ( Jin et al. , 2020 ) . Second , Challenge 2 : the accelerators ’ parameters are not differentiable , and it is non-trivial to derive the operation-wise hardware-cost penalty in order to perform differentiable search ( in considering search efficiency ) . This is because the optimal accelerator is often determined by the whole network instead of one specific operation/layer due to the fact that some accelerator parameters ( e.g. , the loop order ) need to be optimized for the whole network . In this paper , we aim to address the aforementioned challenges towards scalable generic joint search for the network , precision , and accelerator . Specifically , we make the following contributions : • We propose a Triple-Search ( TRIPS ) framework to jointly search for the network , precision , and accelerator in a differentiable manner to efficiently explore the huge joint search space which can not be afforded by previous RL-based methods due to their required prohibitive search cost . TRIPS identifies and tackles the aforementioned challenges towards scalable generic joint search of the three for maximizing both the accuracy and acceleration efficiency . • We develop a heterogeneous sampling strategy for simultaneous updating the weights and network structures to ( 1 ) avoid the need to sequentially train different precision and ( 2 ) achieve unbiased search with constant memory consumption , i.e. , solve the above Challenge 1 . In addition , we develop a novel co-search pipeline that integrates a differentiable hardware search engine to address the above Challenge 2 . • Extensive experiments and ablation studies validate the effectiveness of our proposed TRIPS framework in terms of the resulting search time , task accuracy , and accelerator efficiency , when benchmarked over state-of-the-art ( SOTA ) co-search/exploration techniques , HA-NAS methods , and DNN accelerators . Furthermore , we visualize the searched accelerators by TRIPS to provide insights towards efficient DNN accelerator design in Appendix . 2 RELATED WORKS . Hardware-aware NAS . Hardware-aware NAS has been proposed to automate the design of efficient DNNs . Early works ( Tan et al. , 2019 ; Howard et al. , 2019 ; Tan & Le , 2019 ) utilize RL-based NAS that requires a massive search time/cost , while recent works ( Wu et al. , 2019 ; Wan et al. , 2020 ; Cai et al. , 2018 ; Stamoulis et al. , 2019 ) explore the design space in a differentiable way ( Liu et al. , 2018 ) with much improved searching efficiency . Along another direction , one-shot NAS methods ( Cai et al. , 2019 ; Guo et al. , 2020 ; Yu et al. , 2020 ) pretrain the supernet and directly evaluate the performances of the sub-networks in a weight sharing manner as a proxy of their independently trained performances at the cost of a longer pretrain time . In addition , NAS has been adopted to search for quantization strategies ( Wang et al. , 2019 ; Wu et al. , 2018 ; Cai & Vasconcelos , 2020 ; Elthakeb et al. , 2020 ) for trimming down the complexity of a given DNN . However , these works leave unexplored the hardware design space , which is a crucial enabler for DNN ’ s acceleration efficiency , thus can lead to sub-optimal solutions . DNN accelerators . Motivated by customized accelerators ’ large potential gains , SOTA accelerators ( Du et al. , 2015 ; Chen et al. , 2017 ) innovate micro-architectures and algorithm-to-hardware mapping methods to optimize the acceleration efficiency , given a DNN and the hardware specifications . However , it is non-trivial to design an optimal accelerator as it requires cross-disciplinary knowledge in algorithm , micro-architecture , and circuit design . SOTA accelerator design relies on either experts ’ manual design , which is very time consuming or design flow ( Chen et al. , 2005 ; 2009 ; Rupnow et al. , 2011 ) and DNN accelerator design automation ( Wang et al. , 2016 ; Zhang et al. , 2018a ; Guan et al. , 2017 ; Venkatesan et al. , 2019 ; Wang et al. , 2018a ; Gao et al. , 2017 ) . As they merely explore the accelerator design space , they can result in sub-optimal solutions as compared to SOTA co-search/exploration methods and our TRIPS framework . Co-exploration/search techniques . Pioneering efforts have been made towards jointly searching of DNNs and their accelerators to some extent . For joint searching of DNNs and their precision , ( Chen et al. , 2018 ; Gong et al. , 2019 ; Wang et al. , 2020 ) adopt either differentiable or evolutionary algorithms yet without exploring their hardware accelerators . For joint searching of DNNs and their accelerators , ( Abdelfattah et al. , 2020 ; Yang et al. , 2020 ; Jiang et al. , 2020a ; b ) conduct RL-based search for the networks and some accelerator parameters/templates , where they strictly constrain the search space of the network or accelerator to achieve a practical RL search time , limiting their scalability and achievable efficiency . ( Lin et al . ) is another pioneering work which co-designs the newtork and accelerator in a sequential manner based on the fact that the accelerator ’ s design cycle is longer than the networks . EDD ( Li et al. , 2020 ) extends differentiable NAS to search for layer-wise precision and the accelerators ’ parallel factor , which is most relevant to our TRIPS . EDD has not yet fully solved the joint search challenges . First , it does not discuss or address the potentially explosive memory consumption issue of such joint search ; second , EDD ’ s accelerator search space only includes the parallel factor , which can be strictly limited to their accelerator template and can not generalize to include common accelerator parameters such as the tiling strategies . Built upon prior art , our TRIPS targets a scalable generic joint search framework to optimally search for the network , its precision , and adopted accelerator in a differentiable manner for improving efficiency . 3 THE PROPOSED TECHNIQUES . In this section , we describe our proposed techniques for enabling TRIPS , where Sec . 3.1 provides TRIPS ’ s formulation , Sec . 3.2 and Sec . 3.3 introduce TRIPS ’ s enablers that address the key challenges of scalable generic joint search for networks , precision , and accelerators , and Sec . 3.4 unifies the enablers to build a comprehensive co-search framework . 3.1 TRIPS : FORMULATION . Fig . 1 shows an overview of TRIPS , which jointly searches for the networks ( e.g. , kernel size , channel expansion , and group number ) , precision ( e.g. , 4-/6-/8-/12-/16-bit ) , and the accelerators ( e.g. , PE array type , buffer size , and tiling strategies of each memory hierarchy ) in a differentiable manner . TRIPS targets a scalable yet generic joint search framework , which we formulate as a bi-level optimization problem : min α , β Lval ( ω∗ , net ( α ) , prec ( β ) ) + λLcost ( hw ( γ∗ ) , net ( α ) , prec ( β ) ) ( 1 ) s.t . ω ∗ = arg min ω Ltrain ( ω , net ( α ) , prec ( β ) ) , ( 2 ) s.t . γ ∗ = arg min γ Lcost ( hw ( γ ) , net ( α ) , prec ( β ) ) ( 3 ) Where α , β , and γ are the continuous variables parameterizing the probability of different choices for the network operators , precision bitwidths , and accelerator parameters as in ( Liu et al. , 2018 ) , ω is the supernet weights , Ltrain , Lval , and Lcost are the loss during training and validation and the hardware-cost loss , and net ( α ) , prec ( β ) , and hw ( γ ) denote the network , precision , and accelerator characterized by α , β , γ , respectively .

Summary: This paper proposes Triple-Search (TRIPS), a differentiable framework of jointly searching for network architecture, quantization precision, and accelerator parameters. To address the dilemma between exploding training memory and biased search, the proposed framework leverages heterogeneous sampling where soft Gumbel Softmax is used for weight update and hard Gumbel Softmax is used for probabilities \beta. To integrate accelerator search, hard Gumbel Softmax is used on hardware design choices and the overall hardware cost is used for penalization. Experiments are conducted on the FPGA platform for CIFAR and ImageNet dataset to show the superiority of TRIPS over NAS-only methods.