Directed Acyclic Graph Neural Networks

1 INTRODUCTION . Graph-structured data is ubiquitous across various disciplines ( Gilmer et al. , 2017 ; Zitnik et al. , 2018 ; Sanchez-Gonzalez et al. , 2020 ) . Graph neural networks ( GNNs ) use both the graph structure and node features to produce a vectorial representation , which can be used for classification , regression ( Hu et al. , 2020 ) , and graph decoding ( Li et al. , 2018 ; Zhang et al. , 2019 ) . Most popular GNNs update node representations through iterative message passing between neighboring nodes , followed by pooling ( either flat or hierarchical ( Lee et al. , 2019 ; Ranjan et al. , 2020 ) ) , to produce a graph representation ( Li et al. , 2016 ; Kipf & Welling , 2017 ; Gilmer et al. , 2017 ; Veličković et al. , 2018 ; Xu et al. , 2019 ) . The relational inductive bias ( Santoro et al. , 2017 ; Battaglia et al. , 2018 ; Xu et al. , 2020 ) —neighborhood aggregation—empowers GNNs to outperform graph-agnostic neural networks . To facilitate subsequent discussions , we formalize a message-passing neural network ( MPNN ) architecture , which computes representations h ` v for all nodes v in a graph G in every layer ` and a final graph representation hG , as ( Gilmer et al. , 2017 ) : h ` v = COMBINE ` ( h ` −1v , AGGREGATE ` ( { h ` −1u | u ∈ N ( v ) } ) ) , ` = 1 , . . . , L , ( 1 ) hG = READOUT ( { hLv , v ∈ V } ) , ( 2 ) where h0v is the input feature of v , N ( v ) denotes a neighborhood of node v ( sometimes including v itself ) , V denotes the node set of G , L is the number of layers , and AGGREGATE ` , COMBINE ` , and READOUT are parameterized neural networks . For notational simplicity , we omit edge attributes ; but they can be straightforwardly incorporated into the framework ( 1 ) – ( 2 ) . Directed acyclic graphs ( DAGs ) are a special type of graphs , yet broadly seen across domains . Examples include parsing results of source code ( Allamanis et al. , 2018 ) , logical formulas ( Crouse et al. , 2019 ) , and natural language sentences , as well as probabilistic graphical models ( Zhang et al. , 2019 ) , neural architectures ( Zhang et al. , 2019 ) , and automated planning problems ( Ma et al. , 2020 ) . ∗To whom correspondence should be addressed . A directed graph is a DAG if and only if the edges define a partial ordering over the nodes . The partial order is an additionally strong inductive bias one naturally desires to incorporate into the neural network . For example , a neural architecture seen as a DAG defines the acyclic dependency of computation , an important piece of information when comparing architectures and predicting their performance . Hence , this information should be incorporated into the architecture representation for higher predictive power . In this work , we propose DAGNNs—directed acyclic graph neural networks—that produce a representation for a DAG driven by the partial order . In particular , the order allows for updating node representations based on those of all their predecessors sequentially , such that nodes without successors digest the information of the entire graph . Such a processing manner substantially differs from that of MPNNs where the information landed on a node is limited by a multi-hop local neighborhood and thus restricted by the depth L of the network . Modulo details to be elaborated in sections that follow , the DAGNN framework reads h ` v = F ` ( h ` −1v , G ` ( { h ` u | u ∈ P ( v ) } , h ` −1v ) ) , ` = 1 , . . . , L , ( 3 ) hG = R ( { h ` v , ` = 0 , 1 , . . . , L , v ∈ T } ) , ( 4 ) where P ( v ) denotes the set of direct predecessors of v , T denotes the set of nodes without ( direct ) successors , and G ` , F ` , and R are parameterized neural networks that play similar roles to AGGREGATE ` , COMBINE ` , and READOUT , respectively . A notable difference between ( 3 ) – ( 4 ) and ( 1 ) – ( 2 ) is that the superscript ` − 1 inside the underlined part of ( 1 ) is advanced to ` in the counterpart in ( 3 ) . In other words , MPNN aggregates neighborhood information from the past layer , whereas DAGNN uses the information in the current layer . An advantage is that DAGNN always uses more recent information to update node representations . Equations ( 3 ) – ( 4 ) outline several other subtle but important differences between DAGNN and MPNNs , such as the use of only direct predecessors for aggregation and the pooling on only nodes without successors . All these differences are unique to the special structure a DAG enjoys . Exploiting this structure properly should yield a more favorable vectorial representation of the graph . In Section 2 , we will elaborate the specifics of ( 3 ) – ( 4 ) . The technical details include ( i ) attention for node aggregation , ( ii ) multiple layers for expressivity , and ( iii ) topological batching for efficient implementation , all of which yield an instantiation of the DAGNN framework that is state of the art . For theoretical contributions , we study topological batching and justify that this technique yields maximal parallel concurrency in processing DAGs . Furthermore , we show that the mapping defined by DAGNN is invariant to node permutation and injective under mild assumptions . This result reassures that the graph representation extracted by DAGNN is discriminative . Because DAGs appear in many different fields , neural architectures for DAGs ( including , notably , D-VAE ( Zhang et al. , 2019 ) ) or special cases ( e.g. , trees ) are scattered around the literature over the years . Generally , they are less explored compared to MPNNs ; and some are rather applicationspecific . In Section 3 , we unify several representative architectures as special cases of the framework ( 3 ) – ( 4 ) . We compare the proposed architecture to them and point out the differences that lead to its superior performance . In Section 4 , we detail our comprehensive , empirical evaluation on datasets from three domains : ( i ) source code parsed to DAGs ( Hu et al. , 2020 ) ; ( ii ) neural architecture search ( Zhang et al. , 2019 ) , where each architecture is a DAG ; and ( iii ) score-based Bayesian network learning ( Zhang et al. , 2019 ) . We show that DAGNN outperforms many representative DAG architectures and MPNNs . Overall , this work contributes a specialized graph neural network , a theoretical study of its properties , an analysis of a topological batching technique for enhancing parallel concurrency , a framework interpretation that encompasses prior DAG architectures , and comprehensive evaluations . Supported code is available at https : //github.com/vthost/DAGNN . 2 THE DAGNN MODEL . A DAG is a directed graph without cycles . Denote by G = ( V , E ) a DAG , where V and E ⊂ V × V are the node set and the edge set , respectively . A ( strong ) partial order over a set S is a binary relation ≤ that is transitive and asymmetric . Some authors use reflexivity versus irreflexivity to distinguish weak partial order over strong partial order . To unify concepts , we forbid self-loops ( which otherwise are considered cycles ) in the DAG and mean strong partial order throughout . A set S with partial order ≤ is called a poset and denoted by a tuple ( S , ≤ ) . A DAG ( V , E ) and a poset ( S , ≤ ) are closely related . For any DAG , one can define a unique partial order≤ on the node set V , such that for all pairs of elements u , v ∈ V , u ≤ v if and only if there is a directed path from u to v. On the other hand , for any poset ( S , ≤ ) , there exists ( possibly more than ) one DAG that uses S as the node set and that admits a directed path from u to v whenever u ≤ v. In a DAG , all nodes without ( direct ) predecessors are called sources and we collect them in the set S. Similarly , all nodes without ( direct ) successors are called targets and we collect them in the set T . Additionally , we let X = { h0v , v ∈ V } be the set of input node features . 2.1 MODEL . The main idea of DAGNN is to process nodes according to the partial order defined by the DAG . Using the language of MPNN , at every node v , we “ aggregate ” information from its neighbors and “ combine ” this aggregated information ( the “ message ” ) with v ’ s information to update the representation of v. The main differences to MPNN are that ( i ) we use the current-layer , rather than the past-layer , information to compute the current-layer representation of v and that ( ii ) we aggregate from the direct-predecessor set P ( v ) only , rather than the entire ( or randomly sampled ) neighborhood N ( v ) . They lead to a straightforward difference in the final “ readout ” also . In the following , we propose an instantiation of Equations ( 3 ) – ( 4 ) . See Figure 1 for an illustration of the architecture . One layer . We use the attention mechanism to instantiate the aggregate operator G ` . For a node v at the ` -th layer , the output message m ` v computed by G ` is a weighted combination of h ` u for all nodes u ∈ P ( v ) at the same layer ` : m ` v︸︷︷︸ message : = G ` ( { h ` u | u ∈ P ( v ) } , h ` −1v ) = ∑ u∈P ( v ) α ` vu ( h ` −1v︸︷︷︸ query , h ` u︸︷︷︸ key ) h ` u︸︷︷︸ value . ( 5 ) The weighting coefficients α ` vu follow the query-key design in usual attention mechanisms , whereby the representation of v in the past layer , h ` −1v , serves as the query . Specifically , we define α ` vu ( h ` −1v , h ` u ) = softmax u∈P ( v ) ( w ` 1 > h ` −1v + w ` 2 > h ` u ) , ( 6 ) where w ` 1 and w ` 2 are model parameters . We use the additive form , as opposed to the usual dotproduct form,1 since it involves fewer parameters . An additional advantage is that it is straightforward to incorporate edge attributes into the model , as will be discussed soon . The combine operator F ` combines the message m ` v with the previous representation of v , h ` −1 v , and produces an updated representation h ` v . We employ a recurrent architecture , which is usually used for processing data in sequential order but similarly suits processing in partial order : h ` v = F ` ( h ` −1v , m ` v ) = GRU ` ( h ` −1v , ︸ ︷︷ ︸ input message︷︸︸︷ m ` v︸︷︷︸ state ) , ( 7 ) where h ` −1v , m ` v , and h ` v are treated as the input , past state , and updated state/output of a GRU , respectively . This design differs from most MPNNs that use simple summation or concatenation to combine the representations . It further differs from GG-NN ( Li et al. , 2016 ) ( which also employs a GRU ) , wherein the roles of the two arguments are switched . In GG-NN , the message is treated as the input and the node representation is treated as the state . In contrast , we start from node features and naturally use them as inputs . The message tracks the processed part of the graph and serves better the role of a hidden state , being recurrently updated . By convention , we define G ` ( ∅ , · ) = 0 for the aggregator so that for nodes with an empty directpredecessor set , the message ( or , equivalently , the initial state of the GRU ) is zero . Bidirectional processing . Just like in sequence models where a sequence may be processed by either the natural order or the reversed order , we optionally invert the directions of the edges in G to create a reverse DAG G̃ . We will use the tilde notation for all terms related to the reverse DAG . For example , the representation of node v in G̃ at the ` -th layer is denoted by h̃ ` v . Readout . After L layers of ( bidirectional ) processing , we use the computed node representations to produce the graph representation . We follow a common practice—concatenate the representations across layers , perform a max-pooling across nodes , and apply a fully-connected layer to produce the output . Different from the usual practice , however , we pull across only the target nodes and concatenate the pooling results from the two directions . Recall that the target nodes contain information of the entire graph following the partial order . Mathematically , the readout R produces hG = FC ( Max-Pool v∈T ( L ‖ ` =0 h ` v ) ‖ Max-Pool u∈S ( L ‖ ` =0 h̃ ` u ) ) . ( 8 ) Note that the target set T̃ of G̃ is the same as the source set S of G. If the processing is unidirectional , the right pooling in ( 8 ) is dropped . Edge attributes . The instantiation of the framework so far has not considered edge attributes . It is in fact simple to incorporate them . Let τ ( u , v ) be the type of an edge ( u , v ) and let yτ be a representation of edges of type τ . We insert this information during message calculation in the aggregator . Specifically , we replace the attention weights α ` vu defined in ( 6 ) by α ` vu ( h ` −1v , h ` u ) = softmax u∈P ( v ) ( w ` 1 > h ` −1v + w ` 2 > h ` u + w ` 3 > yτ ( u , v ) ) . ( 9 ) In practice , we experiment with slightly fewer parameters by setting w ` 3 = w ` 1 and find that the model performs equally well . The edge representations yτ are trainable embeddings of the model . Alternatively , if input edge features are provided , yτ ( u , v ) can be replaced by a neural networktransformed embedding for the edge ( u , v ) .

Summary: This paper introduces a model, Directed Acyclic Graph Neural Network (DAGNN), which processes information according to the flow defined by partial order. DAGNN can be regarded as a special case of previous GNN models, but specific to directed acyclic graph structures. The authors prove that the model satisfies the properties desired by DAG-based graph representation learning.Then they study topology batching on the proposed model to maximize parallel concurrency in processing DAGs. A comprehensive empirical evaluation is conducted on datasets from three domains to verify its effectiveness.