Hardware-adaptive Efficient Latency Prediction for NAS via Meta-Learning

1 Introduction . Neural Architecture Search ( NAS ) [ 39 , 1 , 25 , 22 , 21 , 36 , 5 ] , which aims to search for the optimal architecture for a given task , has achieved a huge practical success by overcoming the sub-optimality of manual architecture designs . However , for NAS to be truly practical in real-world scenarios , it should be hardware-aware . That is , we need to search for the neural architectures that satisfy diverse device constraints ( e.g. , memory footprint , inference latency , and energy consumption ) . Due to the practical importance of the problem , many existing works propose to take into account the hardware efficiency constraints ( mostly latency ) in the search process [ 29 , 3 , 34 , 32 , 6 , 2 , 37 , 4 , 33 ] . However , one of the main challenges with such hardware-aware NAS is that collecting training samples ( e.g. , architecture-latency pairs on each target device ) to build reliable prediction models for the efficiency metrics , is computationally costly or requires the knowledge of the hardware devices . Existing hardware-aware NAS methods usually require a large number of samples ( e.g. , 5k ) and train metric predictors for each device from scratch . Additionally , most works [ 29 , 3 , 34 , 32 , 6 ] are task-specific , and thus such collection of performance samples should be done from scratch , for a new task . Thus , the sample collection process is prohibitively costly when we consider real-world ∗These authors contributed equally to this work . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . deployment of a model to a large number of hardware devices , for any tasks . OFA [ 4 ] alleviate the search cost by utilizing a high-performance network trained on ImageNet as the supernet , which does not require training the models from scratch . Yet , they are still sub-optimal since building device-specific predictors for metric still requires a large number of samples to be collected , to build a layer-wise latency predictor for each device . This could take multiple hours depending on the task and the device , and becomes a bottleneck for latency-constrained NAS . BigNAS [ 38 ] considers FLOP as the efficiency metric , but this is a highly inaccurate proxy since latency of an architecture could differ based on its degree of parallelism and memory access cost , for the same FLOP . To overcome such limitations , we propose a novel sample-efficient latency predictor , namely Hardware-adaptive Efficient Latency Predictor ( HELP ) , which supports the latency estimation for multiple devices with a single model , by allowing it to rapidly adapt to unseen devices with only a few latency measurements collected from each device . Specifically , we formulate the latency prediction problem as a few-shot regression task of estimating the latency given an architecture-device pair , and propose a novel hardware embedding that can embed any devices by utilizing the latencies of the reference architectures on each device as its embeddings . Then , we propose a meta-learning framework which combines amortized meta-learning with gradient-based meta-learning , to learn the latency predictor to generalize across multiple devices , utilizing the proposed hardware embeddings . This allows the model to transfer knowledge learned from known devices to a new device , and thus to achieve high sample efficiency when estimating the latency of unseen devices . HELP is highly versatile and general , as it is applicable to any hardware platforms and architecture search spaces , thanks to our device-agnostic hardware embeddings . Also , HELP can be coupled with any NAS frameworks to reduce its computational bottleneck in obtaining latency-constrained architectures . Especially , when coupled with rapid NAS methods such as MetaD2A [ 16 ] , OFA [ 4 ] and HAT [ 33 ] , HELP can perform the entire latency-constrained NAS process for a new device almost instantly . We validate the latency estimation performance of HELP on the NAS-Bench-201 space [ 9 ] with various devices from different hardware platforms , utilizing the latency dataset for an extensive device pool in HW-NAS-Bench dataset [ 18 ] . The results show that our meta-learned predictor successfully generalize to unseen target devices , largely outperforming baseline latency estimation methods using at least 90× less measurements . Then , we combine HELP with existing NAS methods [ 4 , 33 ] and show that our meta-latency predictor largely reduce their total search cost . To summarize , the main contributions of this paper are as follows : • We formulate the latency estimation of a neural architecture for a given device as a few-shot regression problem , which outputs the latency given an architecture-device pair as the input . • To represent heterogeneous hardware devices , we propose a novel device-agnostic hardware embedding , which embeds a device by its latencies on reference architectures . • We propose a novel latency predictor , HELP , which meta-learns a few-shot regression model to generalize across hardware devices , that can estimate the latency of an unseen device using only few measurements . HELP obtains significantly higher latency estimation performance over baselines with minimal total time cost . • We further combine HELP with existing NAS frameworks to show that it leads to find latencyconstrained architectures extremely fast , eliminating the latency-estimation bottleneck in the hardware-constrained NAS . 2 Related Work . Latency Prediction in NAS Hardware-aware NAS [ 29 , 3 , 4 , 34 , 32 , 13 , 33 , 26 ] aims to design neural architectures that achieves a good trade-off between accuracy and latency for efficient deployment to target devices . While they need to consider the actual latencies of the architectures on a target device , evaluating them while searching for architectures is costly . Thus , most existing works replace it with proxies , such as FLOPs [ 38 ] , but they are inaccurate measure of latencies . Earlier hardwareaware NAS methods have been used a layer-wise latency predictor ( lookup table ) [ 29 , 3 , 4 , 34 ] which sums up the latencies measured for each operation in the given neural networks . While this is better than FLOPs , they do not accurately capture the complexities of multiple layer execution on real hardware devices . Thus , recent methods use an end-to-end latency predictor [ 10 , 33 ] that is trained with the latency measurements from the target device , to improve the latency prediction accuracy . BRP-NAS [ 10 ] , which is a GCN-based model , is currently the state-of-the-art end-to-end latency predictor . However , this method is also limited in that it requires a large number of architecture-latency pairs for each device . Since the latency estimator can not generalize to a new device , whenever a new device is given , the NAS system needs to build a new latency estimator , which may take hours to finish . The proposed method significantly reduces the total building cost of the latency predictor for each device , by using a single meta-learned predictor and adapting it to a new device with only a few measurements from it . Meta-learning and Meta-NAS Meta-learning ( learning to learn ) [ 30 ] aims to learn a model that generalizes over a distribution of tasks rather than a single task , such that the model meta-learned over a large number of tasks rapidly adapts to an unseen task . The performance of existing methods [ 31 , 28 , 12 , 24 , 17 ] is usually evaluated on few-shot classification tasks , where the model classifies between instances of unseen classes given only a few training examples . Recently , several works have proposed to utilize meta-learning for NAS . Most of them focus on few-shot classification tasks to search for the architectures and parameters that can generalize well to a new task [ 11 , 20 , 27 ] with gradientbased meta-learning . However , they have limited practical applicability since the computational cost of meta-learning is prohibitively large , for NAS under a standard many-shot setting . A recently proposed meta-NAS framework with amortized meta-learning , MetaD2A [ 16 ] , which utilizes a set encoder to encode a task and uses a graph decoder to generate a task-adaptive architecture , obtained state-of-the-art performance on unseen datasets , with minimal search cost . We also propose a similar amortized meta-learning framework , for hardware-adaptive NAS , based on a novel hardware device embedding . However , after obtaining the device-conditioned initialization parameters , we perform further inner gradient steps for device-specific adaptation , unlike MetaD2A [ 16 ] . 3 Method . Our goal is to design a prediction model that can accurately predict the efficiency metric for a novel architecture-device pair , using a small number of performance samples from the device . While our method is generally applicable to any efficiency metrics that can be measured from the device ( e.g . latency , energy consumption , and memory footprint ) we focus on the latency prediction in this work . 3.1 Problem Definition . Assume that we are given a task specification τ = { hτ , Xτ , Yτ } where hτ ∈ H is a hardware device , Xτ ⊂ X is a set of neural architectures , and Yτ ⊂ Y is a set of latencies of Xτ directly measured on the hardware device hτ . Then , our goal is to learn a regression model f ( x ; θ ) : X → R parameterized by θ that estimates the latency y ∈ Yτ of a neural architecture x ∈ Xτ for a given hardware device hτ by minimizing empirical loss L ( e.g . mean squared error ) on the predicted values f ( Xτ ; θ ) and actual measurements Yτ as follows : min θ L ( fτ ( Xτ ; θ ) , Yτ ) ( 1 ) Learning such a regression model seems like a simple problem , since we can collect any number of measurements from any devices . However , this is a more challenging problem than it seems : 1 . Since we can not generalize across devices , we need to learnN predictors { fτ ( · ; θτ ) } Nτ=1 for N devices , collecting performance samples and training the performance predictor separately for each device , which requires prohibitive computational costs with large number of devices . 2 . Even when assuming that we learn a device-specific predictor , in order not to overfit the regression model , we need to collect a large number of architecture-latency sample pairs for each device ( e.g . 2k [ 33 ] , 5k [ 23 ] samples ) to achieve reliable prediction performance . 3 . With lack of generalization ability across devices and architectures , the NAS framework needs to repeat the time-consuming sample collection process whenever a new device is given , that may take hours to complete . This will become a computational bottleneck even for a rapid meta-NAS framework such as MetaD2A [ 16 ] . To overcome such limitations , we propose a single predictor f ( · ; θ ) which can generalize across devices and architectures , by fast adapting to a new target device and architecture that are unseen during training , by collecting only a small number of architecture-latency pairs from the device ( Xτ Xτ , Y τ Yτ ) . We achieve this goal with a meta-learning framework that can transfer knowledge obtained from the device and architecture pool p ( τ ) . 3.2 Hardware-adaptive Latency Prediction with Device Embeddings . While the measured latency y ← ( x , h ) is dependent on both the device type h and the architecture x , existing latency prediction models takes the form of f ( x ; θ ) , ignoring the device-specific constraints , since the latency predictor is learned for each device separately . This results in poor performance when learning a single latency predictor to perform metric estimation on multiple devices , including unseen ones , which is our main objective . Thus , we propose hardware-conditioned prediction model : f ( x , h ; θ ) : X ×H → R ( 2 ) that can predict the latency differently depending on the device type h , even for the same architecture x . A crucial challenge of our hardware-conditioned prediction model is how to represent the hardware device h , for all devices regardless of their platform types . This is not a trivial problem since the physical architecture of hardware devices could be very different ( e.g . CPU vs FPGA ) . Thus , we simply consider the device as a black box function which outputs the inference latency given an architecture , instead of directly modeling the hardware devices . Then , we obtain the latencies of the device on a fixed set of reference neural architectures as follows : Vh = { y∗1 ( x1 , h ) , y∗2 ( x2 , h ) , ... , y∗d ( xd , h ) } ( 3 ) where E is the set of the reference neural architectures { x1 , x2 , ... , xd } ⊂ X , fixed across all tasks for both meta-training and meta-test , and d is the number of the reference architectures ; in our experiments , we set d = 10 . Further , y∗i ( xi , h ) = { yi ( xi , h ) −min ( V ( 0 ) h ) } / { max ( V ( 0 ) h ) −min ( V ( 0 ) h ) } are standardized latency values ranging from 0 to 1 , where V ( 0 ) h = { y1 ( x1 , h ) , y2 ( x2 , h ) , ... , yd ( xd , h ) } . Since this set of reference devices should be representative , we select them to be diverse and heterogeneous . As for the reference architectures , we randomly sample them from the search space . For more detailed descriptions of the reference devices and architectures , please see the supplementary file . The proposed black-box treatment of hardware devices , and the latency-based hardware embedding allows us to embed a new device without considering its detailed hardware specification .

Summary: Modern neural architecture search methods often require us to train a *latency prediction model* which can predict the latency of any candidate network architecture within a human-defined search space. These are commonly obtained by (i) sampling a set of random network architectures from the search space, (ii) benchmarking each one on the target inference hardware device, and (iii) training a regression model on these <architecture, latency> pairs in order to generalize to new, unseen architectures. The submission explores the idea of (i) training a single model which can generate predictions for many different hardware devices, and (ii) using meta-learning to allow the model to generalize to a new hardware device with little additional data (e.g., 10 samples).