Interpretable Neural Architecture Search via Bayesian Optimisation with Weisfeiler-Lehman Kernels

1 INTRODUCTION . Neural architecture search ( NAS ) aims to automate the design of good neural network architectures for a given task and dataset . Although different NAS strategies have led to state-of-the-art neural architectures , outperforming human experts ’ design on a variety of tasks ( Real et al. , 2017 ; Zoph and Le , 2017 ; Cai et al. , 2018 ; Liu et al. , 2018a ; b ; Luo et al. , 2018 ; Pham et al. , 2018 ; Real et al. , 2018 ; Zoph et al. , 2018a ; Xie et al. , 2018 ) , these strategies behave in a black-box fashion , which returns little design insight except for the final architecture for deployment . In this paper , we introduce the idea of interpretable NAS , extending the learning scope from simply the optimal architecture to interpretable features . These features can help explain the performance of networks searched and guide future architecture design . We make the first attempt at interpretable NAS by proposing a new NAS method , NAS-BOWL ; our method combines a Gaussian process ( GP ) surrogate with the Weisfeiler-Lehman ( WL ) subtree graph kernel ( we term this surrogate GPWL ) and applies it within the Bayesian Optimisation ( BO ) framework to efficiently query the search space . During search , we harness the interpretable architecture features extracted by the WL kernel and learn their corresponding effects on the network performance based on the surrogate gradient information . Besides offering a new perspective on interpratability , our method also improves over the existing BO-based NAS approaches . To accommodate the popular cell-based search spaces , which are noncontinuous and graph-like ( Zoph et al. , 2018a ; Ying et al. , 2019 ; Dong and Yang , 2020 ) , current approaches either rely on encoding schemes ( Ying et al. , 2019 ; White et al. , 2019 ) or manually designed similarity metrics ( Kandasamy et al. , 2018 ) , both of which are not scalable to large architectures and ignore the important topological structure of architectures . Another line of work employs graph neural networks ( GNNs ) to construct the BO surrogate ( Ma et al. , 2019 ; Zhang et al. , 2019 ; Shi et al. , 2019 ) ; however , the GNN design introduces additional hyperparameter tuning , and the training of the GNN also requires a large amount of architecture data , which is particularly ∗Equal contribution . Codes are available at https : //github.com/xingchenwan/nasbowl expensive to obtain in NAS . Our method , instead , uses the WL graph kernel to naturally handle the graph-like search spaces and capture the topological structure of architectures . Meanwhile , our surrogate preserves the merits of GPs in data-efficiency , uncertainty computation and automated hyperparameter treatment . In summary , our main contributions are as follows : • We introduce a GP-based BO strategy for NAS , NAS-BOWL , which is highly query-efficient and amenable to the graph-like NAS search spaces . Our proposed surrogate model combines a GP with the WL graph kernel ( GPWL ) to exploit the implicit topological structure of architectures . It is scalable to large architecture cells ( e.g . 32 nodes ) and can achieve better prediction performance than competing methods . • We propose the idea of interpretable NAS based on the graph features extracted by the WL kernel and their corresponding surrogate derivatives . We show that interpretability helps in explaining the performance of the searched neural architectures . As a singular example of concrete application , we propose a simple yet effective motif-based transfer learning baseline to warm-start search on a new image tasks . • We demonstrate that our surrogate model achieves superior performance with much fewer ob- servations in search spaces of different sizes , and that our strategy both achieves state-of-the-art performances on both NAS-Bench datasets and open-domain experiments while being much more efficient than comparable methods . 2 PRELIMINARIES . Graph Representation of Neural Networks Architectures in popular NAS search spaces can be represented as an acyclic directed graph ( Elsken et al. , 2018 ; Zoph et al. , 2018b ; Ying et al. , 2019 ; Dong and Yang , 2020 ; Xie et al. , 2019 ) , where each graph node represents an operation unit or layer ( e.g . a conv3×3-bn-relu in Ying et al . ( 2019 ) ) and each edge defines the information flow from one layer to another . With this representation , NAS can be formulated as an optimisation problem to find the directed graph and its corresponding node operations ( i.e . the directed attributed graph G ) that give the best architecture validation performance y ( G ) : G∗ = arg maxG y ( G ) . Bayesian Optimisation and Gaussian Processes To solve the above optimisation , we adopt BO , which is a query-efficient technique for optimising a black-box , expensive-to-evaluate objective ( Brochu et al. , 2010 ) . BO uses a statistical surrogate to model the objective and builds an acquisition function based on the surrogate . The next query location is recommended by optimising the acquisition function which balances the exploitation and exploration . We use a GP as the surrogate model in this work , as it can achieve competitive modelling performance with small amount of query data ( Williams and Rasmussen , 2006 ) and give analytic predictive posterior mean µ ( Gt|Dt−1 ) and variance k ( Gt , G′t|Dt−1 ) on the heretofore unseen graph Gt given t − 1 observations : µ ( Gt|Dt−1 ) = k ( Gt , G1 : t−1 ) K−11 : t−1y1 : t−1 and k ( Gt , G′t|Dt−1 ) = k ( Gt , G ′ t ) − k ( Gt , G1 : t−1 ) K−11 : t−1k ( G1 : t−1 , G′t ) where G1 : t−1 = { G1 , . . . , Gt−1 } and y1 : t−1 = [ y1 , . . . , yt−1 ] T are the t − 1 observed graphs and objective function values , respectively , and Dt−1 = { G1 : t−1 , y1 : t−1 } . [ K1 : t−1 ] i , j = k ( Gi , Gj ) is the ( i , j ) -th element of Gram matrix induced on the ( i , j ) -th training samples by k ( · , · ) , the graph kernel function . We use Expected Improvement ( Mockus et al. , 1978 ) in this work though our approach is compatible with alternative choices . Graph Kernels Graph kernels are kernel functions defined over graphs to compute their level of similarity . A generic graph kernel may be represented by the function k ( · , · ) over a pair of graphs G and G′ ( Kriege et al. , 2020 ) : k ( G , G′ ) = 〈φ ( G ) , φ ( G′ ) 〉H ( 2.1 ) where φ ( · ) is some feature representation of the graph extracted by the graph kernel and 〈· , ·〉H denotes inner product in the associated reproducing kernel Hilbert space ( RKHS ) ( Nikolentzos et al. , 2019 ; Kriege et al. , 2020 ) . For more detailed reviews on graph kernels , the readers are referred to Nikolentzos et al . ( 2019 ) , Ghosh et al . ( 2018 ) and Kriege et al . ( 2020 ) . Algorithm 1 NAS-BOWL Algorithm . Optional steps of the exemplary use of motifbased warm starting ( Sec 3.2 ) are marked in gray italics . 1 : Input : Maximum BO iterations T , BO batch size b , acquisition function α ( · ) , initial observed data on the target task D0 , Optional : past-task query data Dpast and surrogate Spast 2 : Output : The best architecture G∗T 3 : Initialise the GPWL surrogate S with D0 4 : for t = 1 , . . . , T do 5 : if Pruning based on the past-task motifs then 6 : Compute the motif importance scores ( equation 3.4 ) with Spast/S on Dpast/Dt 7 : while |Gt| < B do 8 : Generate a batch of candidate archi- tectures and reject those which contain none of the top 25 % good motifs ( similar procedure as Fig . 2 ( a ) ) 9 : end while 10 : else 11 : Generate B candidate architectures Gt 12 : end if 13 : { Gt , i } bi=1 = argmaxG∈Gt αt ( G|Dt−1 ) 14 : Evaluate their validation accuracy { yt , i } bi=1 15 : Dt ← Dt−1 ∪ ( { Gt , i } Bi=1 , { yt , i } bi=1 ) 16 : Update the surrogate S with Dt 17 : end for 18 : Return the best architecture seen so far G∗T 3 PROPOSED METHOD . We begin by presenting our proposed algorithm , NAS-BOWL in Algorithm 1 , where there are a few key design features , namely the design of the GP surrogate suitable for architecture search ( we term the surrogate GPWL ) and the method to generate candidate architectures at each BO iteration . We will discuss the first one in Section 3.1 . For architecture generation , we either generate the new candidates via random sampling the adjacency matrices , or use a mutation algorithm similar to those used in a number of previous works ( Kandasamy et al. , 2018 ; Ma et al. , 2019 ; White et al. , 2019 ; Shi et al. , 2019 ) : at each iteration , we generate the architectures by mutating a number of queried architectures that perform the best . Generating candidate architectures in this way enables us to exploit the prior information on the best architectures observed so far to explore the large search space more efficiently . We report NAS-BOWL with both strategies in our experiments . Finally , to give a demonstration of the new possibilities opened by our work , we give an exemplary practical use of intepretable motifs for transfer learning in Algorithm 1 , which is elaborated in Sec 3.2 . 3.1 SURROGATE AND GRAPH KERNEL DESIGN . To enable the GP to work effectively on the graph-like architecture search space , selecting a suitable kernel function is arguably the most important design decision . We propose to use the WeisfeilerLehman ( WL ) graph kernel ( Shervashidze et al. , 2011 ) to enable the direct definition of a GP surrogate on the graph-like search space . The WL kernel compares two directed graphs based on both local and global structures . It starts by comparing the node labels of both graphs via a base kernel kbase ( φ0 ( G ) , φ0 ( G ′ ) ) where φ0 ( G ) denotes the histogram of features at level h = 0 ( i.e . node features ) in the graph , where h is both the index of WL iterations and the depth of the subtree features extracted . For the WL kernel with h > 0 , as shown in Fig . 1 , it then proceeds to collect features at h = 1 by aggregating neighbourhood labels , and compare the two graphs with kbase ( φ1 ( G ) , φ1 ( G ′ ) ) based on the subtree structures of depth 1 ( Shervashidze et al. , 2011 ; Höppner and Jahnke , 2020 ) . The procedure then repeats until the highest iteration level h = H specified and the resulting WL kernel is given by : kHWL ( G , G ′ ) = H∑ h=0 kbase ( φh ( G ) , φh ( G ′ ) ) . ( 3.1 ) In the above equation , kbase is a base kernel ( such as dot product ) over the vector feature embedding . As h increases , the WL kernel captures higher-order features which correspond to increasingly larger neighbourhoods and features at each h are concatenated to form the the final feature vector ( φ ( G ) = [ φ0 ( G ) , ... , φH ( G ) ] ) . The readers are referred to App . A for more detailed algorithmic descriptions of the WL kernel . We argue that WL is desirable for three reasons . First , in contrast to many ad hoc approaches , WL is established with proven successes on labelled and directed graphs , by which networks are represented . Second , the WL representation of graphs is expressive , topology-preserving yet interpretable : Morris et al . ( 2019 ) show that WL is as powerful as standard GNNs in terms of discrimination power . However , GNNs requires relatively large amount of training data and thus is more data-inefficient ( we validate this in Sec . 5 ) . Also , the features extracted by GNNs are harder to interpret compared to those by WL . Note that the WL kernel by itself only measures the similarity between graphs and does not aim to select useful substructures explicitly . It is our novel deployment of the WL procedure ( App . A ) for the NAS application that leads to the extraction of interpretable features while comparing different architectures . We further make smart use of these network features to help explain the architecture performance in Sec 3.2 . Finally , WL is efficient and scalable : denoting { n , m } as the number of nodes and edges respectively , computing the Gram matrix on N training graphs may scale O ( NHm + N2Hn ) ( Shervashidze et al. , 2011 ) . As we show in App . E.3 , in typical cell-based spaces H ≤ 3 suffices , suggesting that the kernel computation cost is likely eclipsed by the O ( N3 ) scaling of GP we incur nonetheless . This is to be contrasted to approaches such as path encoding in White et al . ( 2019 ) , which scales exponentially with n without truncation , and the edit distance kernel in Jin et al . ( 2019 ) , whose exact solution is NP-complete ( Zeng et al. , 2009 ) . With the above-mentioned merits , the incorporation of the WL kernel permits the usage of GP-based BO on various NAS search spaces . This enables the practitioners to harness the rich literature of GP-based BO methods on hyperparameter optimisation and redeploy them on NAS problems . Most prominently , the use of GP surrogate frees us from hand-picking the WL hyperparameter H as we can automatically learn the optimal values by maximising the Bayesian marginal likelihood . As we will justify in Sec . 5 and App . E.3 , this process is extremely effective . This renders a further major advantage of our method as it has no inherent hyperparameters that require manual tuning . This reaffirms with our belief that a practical NAS method itself should require minimum tuning , as it is almost impossible to run traditional hyperparameter search given the vast resources required . Other enhancements , such as improving the expressiveness of the surrogate by combining multiple types of kernels , are briefly investigated in App . C. We find the amount of performance gain depends on the NAS search space and a WL kernel alone suffices for common cell-based spaces .

Summary: The authors propose a new neural architecture search algorithm combining Bayesian optimization with the expressive and popular Weisfeiler-Lehman (WL) Graph Kernel. One advantage of using WL is the interpretable results that stem from the nature of how the kernel is computed, namely a propagation scheme through the graph. Combined the derivative of Eq. 3.2, one can extract subgraphs that are directly responsible for increased performance. In a variety of experiments, the authors show not only increased performance of detected architectures but also find subgraphs that are found by other algorithms as well.