Generalizing and Tensorizing Subgraph Search in the Supernet

1 INTRODUCTION . Deep learning ( Goodfellow et al. , 2017 ) has been successfully applied in many applications , such as image classification for computer vision ( CV ) ( LeCun et al. , 1998 ; Krizhevsky et al. , 2012 ; He et al. , 2016 ; Huang et al. , 2017 ) and language modeling for natural language processing ( NLP ) ( Mikolov et al. , 2013 ; Devlin et al. , 2018 ) . While the architecture design is of great importance to deep learning , manually designing proper architectures for a certain task is hard and requires lots of human efforts or sometimes even impossible ( Zoph & Le , 2017 ; Baker et al. , 2016 ) . Recently , neural architecture search ( NAS ) techniques ( Elsken et al. , 2019 ) have been developed to alleviate this issue , which mainly focuses on CV and NLP tasks . Behind existing NAS methods , a multi-graph ( Skiena. , 1992 ) structure , i.e. , supernet ( Zoph et al. , 2017 ; Pham et al. , 2018 ; Liu et al. , 2018 ) , where nodes are connected by edges with multiple choices , has played a central role . In such context , the choices on each edge are different operations , and the subgraphs correspond to different neural architectures . The objective here is to find a suitable subgraph in this supernet , i.e . better neural architectures for a given task . However , the supernet does not only arise in CV/NLP field and we find it also emerge in many other deep learning areas ( see Table 1 ) . An example is logical chain inference on knowledge graphs ( Yang et al. , 2017 ; Sadeghian et al. , 2019 ; Qu & Tang , 2019 ) , where the construction logical rules can be modeled by a supernet . Another example is meta-path discovery in heterogeneous information networks ( Yun et al. , 2019 ; Wan et al. , 2020 ) , where the discovery of meta-paths can also be modeled by a supernet . Therefore , we propose to broaden the horizon of NAS , i.e. , generalize it to many deep learning fields and solve the new NAS problem under a unified framework . Since subgraphs are discrete objects ( choices on each edge are discrete ) , it has been a common approach ( Liu et al. , 2018 ; Sadeghian et al. , 2019 ; Yun et al. , 2019 ) to transform it into a continuous optimization problem . Previous methods often introduce continuous parameters separately for each edge . However , this formulation can not generalize to different supernets as the topological structures of supernets are highly task-dependent and diverse . Therefore , it will fail to capture the supernet ’ s topology and hence be ineffective . In this paper , we propose a novel method TRACE to introduce a continuous parameter for each subgraph ( all these parameters will form a tensor ) . Then , we propose to construct a tensor network ( TN ) ( andrzej . et al. , 2016 ; 2017 ) based on the topological structure of supernet . For different tensor networks , we introduce an efficient algorithm for optimization on supernets . Extensive experiments are conducted on diverse deep learning tasks . Empirical results demonstrate that TRACE performs better than the state-of-the-art methods in each domain . As a summary , our contributions are as follows : • We broaden the horizon of existing supernet-based NAS methods . Specifically , we generalize the concept of subgraph search in supernet from NAS to other deep learning tasks that have graph-like structures and propose to solve them in a unified framework by tensorizing the supernet . • While existing supernet-based NAS methods ignore the topological structure of the supernet , we encode the supernet in a topology-aware manner based on the tensor network and propose an efficient algorithm to solve the search problem . • We conduct extensive experiments on various learning tasks , i.e. , architecture design for CV , logical inference for KG , and meta-path discovery for HIN . Empirical results demonstrate that our method can find better architectures , which lead to state-of-the-art performance on various applications . 2 RELATED WORKS . 2.1 SUPERNET IN NEURAL ARCHITECTURE SEARCH ( NAS ) . There have been numerous algorithms proposed to solve the NAS problem . The first NAS work , NASRL ( Zoph & Le , 2017 ) , models the NAS problem as a multiple decision making problem and proposes to use reinforcement learning ( RL ) ( Sutton & Barto , 2018 ) to solve this problem . However , this formulation does not consider the repetitively stacked nature of neural architectures and is very inefficient as it has to train many different networks to converge . To alleviate this issue , NASNet ( Zoph et al. , 2017 ) first models NAS as an optimization problem on supernet . The supernet formulation enables searching for transferrable architectures across different datasets and improves the searching efficiency . Later , based on the supernet formulation , ENAS ( Pham et al. , 2018 ) proposes weight-sharing techniques , which shares the weight of each subgraph in a supernet . This technique further improves searching efficiency and many different methods have been proposed under this framework ( see Table 1 ) , including DARTS ( Liu et al. , 2018 ) , SNAS ( Xie et al. , 2018 ) , and NASP ( Yao et al. , 2020 ) . DARTS is the first to introduce deterministic formulation to NAS field , and SNAS uses a similar parametrized method with DARTS under stochastic formulation . NASP improves upon DARTS by using proximal operator ( Parikh & Boyd , 2014 ) and activates only one subgraphs in each iteration to avoid co-adaptation between subgraphs . 2.2 TENSOR METHODS IN MACHINE LEARNING . A tensor ( Kolda & Bader , 2009 ) is a multi-dimensional array as an extension to a vector or matrix . Tensor methods have found wide applications in machine learning , including network compression ( Novikov et al. , 2015 ; Wang et al. , 2018 ) and knowledge graph completion ( Liu et al. , 2020 ) . Recently , andrzej . et al . ( 2016 ; 2017 ) have proposed a unified framework called tensor network ( TN ) , which uses an undirected graph to represent tensor decomposition methods . By using different graphs , TN can cover many tensor decomposition methods as a special case , e.g. , CP ( Bro , 1997 ) , Tucker ( Tucker , 1966 ) , tensor train ( Oseledets , 2011 ) , and tensor ring decomposition ( Zhao et al. , 2016 ) . However , it is not an easy task to construct a tensor network for a given problem , as the topological structure of tensor network is hard to design ( Li & Sun , 2020 ) . 3 PROPOSED METHOD . Here , we describe our method for the supernet search problem . Section 3.1-3.2 introduce how we tensorize the supernets . Section 3.3 propose an optimization algorithm for the search problem which can be utilized for both stochastic and deterministic objectives . Finally , how supernet appears beyond existing NAS works and can be generalized to these new tasks are presented in Section 3.4 . Notations . In this paper , we will use S to denote a supernet , and P to denote a subgraph in supernet . For a supernet S with T edges , we let all edges be indexed by e1 , . . . , eT , and define Ct to be the number of choices on edge et with t ∈ { 1 , . . . , T } . With subscript i1 , . . . , iT denoted by i− for short , we use Si− to denote the subgraph with choices i1 ∈ { 1 , . . . , C1 } , . . . , iT ∈ { 1 , . . . , CT } . And we use softmax ( oi ) = exp ( oi ) / ∑n j=1 exp ( oj ) to denote the softmax operation over a vector oi ∈ Rn . 3.1 A TENSOR FORMULATION FOR SUPERNET . While existing works ( Liu et al. , 2018 ; Zoph et al. , 2017 ; Pham et al. , 2018 ) introduce parameters separately for each edge et , since edges may correlate with each other , a more general and natural approach is to introduce a continuous parameter directly for each subgraph P ∈ S. Note that a subgraph P can be distinguished by its choices on each edge , we propose to encode all possible choices into a tensor T ∈ RC1×···×CT and take these choices as indices , i.e. , i− , to the tensor T . As a consequence , the architecture of the subgraph P is indexed as Si− , and Ti− ∈ [ 0 , 1 ] represent how “ good ” P can be . Let f ( w , P ) stand for a learning model with the parameter w and subgraph P ∈ T , L be the loss of f on the training dataset Dtr , and J measure the performance of f on the validation set Dval . Under above tensor formulation , the search objective here is : max P∈T J ( f ( w∗ ( P ) , P ) ; Dval ) , s.t . { w∗ ( P ) = arg minw L ( f ( w , P ) ; Dtr ) ∑ i− Ti− = 1 , ( 1 ) Same as existing supernet search works , the subgraph P is searched in the upper-level , while the network weight w is trained in the lower-level . Dtr is the training set and Dval is the validation set . However , we have the extra constraint ∑ i− Ti−=1 here , which is to ensure that the sum of probabilities for all subgraphs is one . Next , we show how P and T can be parameterized with topological information from the supernet in Section 3.2 . Then , a gradient-based algorithm which can effectively handle the constrained bi-level optimization problem ( 1 ) is proposed in Section 3.3 . 3.2 ENCODING SUPERNET TOPOLOGY BY TENSOR NETWORK ( TN ) . Existing methods consider each edge separately and can be seen as rank-1 factorization of the full tensor T ( see Table 1 ) under ( 1 ) , i.e . Ti− = θ1i1 . . . θ T iT where θtit ∈ R is the continuous parameter for choice it on edge t. However , this formulation ignores the topological structure of different supernets , as it uses the same decomposition method for all supernets . Motivated by such limitation , we propose to introduce tensor network ( TN ) to better encode the topological structure of supernet . Our encoding process is described in Algorithm 1 , where N ( S ) to denotes the set of nodes in this supernet and N ′ ( S ) ⊆ N ( S ) denotes the set of nodes that are connected to more than one edge . Specifically , we introduce a third-order tensor αt for each edge , which is based on previous methods ( e.g. , DARTS and SNAS ) , but uses tensors instead of vectors to allow more flexibility . And we introduce hyper-parametersRN1 ( t ) andRN2 ( t ) , which are corresponding to the rank of tensor network . Then , we use index summation to reflect the topological structure ( common nodes ) between different edges , i.e. , Ti− = ∑Rn rn , n∈N ′ ( S ) ∏T t=1 αtrN1 ( t ) , it , rN2 ( t ) , ( 2 ) We also give two examples in 1Figure 2 to illustrate our tensorizing process for two specific supernets . Algorithm 1 Supernet encoding process ( a step-by-step and graphical illustration is in Appendix.G ) . Input : Supernet S ; 1 : Introduce αt ∈ RRN1 ( t ) ×Ct×RN2 ( t ) for edge et which connects nodes N1 ( t ) and N2 ( t ) ; 2 : Remove isolated nodes and obtain N ′ ( S ) ; 3 : Compute Ti− by ( 2 ) ; 4 : return encoded supernet Ti− ; The reason of using N ′ ( S ) instead of N ( S ) is 2Proposition 1 , which shows that using N ′ ( S ) will not restrict the expressive power but allows the usage of less parameters . Proposition 1 . Any tensors that can be expressed by Ti− = ∑Rn rn , n∈N ( S ) ∏T t=1 α t rN1 ( t ) , it , rN2 ( t ) can also be expressed by Ti− = ∑Rn rn , n∈N ′ ( S ) ∏T t=1 α t rN1 ( t ) , it , rN2 ( t ) .

Summary: + This paper generalizes the supernet search problem on a broader horizon. Specifically, some of the current NAS methods use supernet to co-training different neural architectures for further architecture search. This paper does not just consider supernet as a tool for NAS, but also consider supernet as a graphical model and extend supernet to several general tasks in the form of graph data. (+)