Directed Acyclic Graph Neural Networks

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 $\mathcal{P}(v)$ only, rather than the entire (or randomly sampled) neighborhood $\mathcal{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^{\ell}$. For a node $v$ at the $\ell$-th layer, the output message $m_{v}^{\ell}$ computed by $G^{\ell}$ is a weighted combination of $h_{u}^{\ell}$ for all nodes $u\in\mathcal{P}(v)$ at the same layer $\ell$:

$$\underbrace{m_{v}^{\ell}}_{\text{message}}:=G^{\ell}\Big{(}\{h_{u}^{\ell}\mid u\in\mathcal{P}(v)\},\,h_{v}^{\ell-1}\Big{)}=\sum_{u\in\mathcal{P}(v)}\alpha_{vu}^{\ell}\Big{(}\underbrace{h_{v}^{\ell-1}}_{\text{query}},\,\underbrace{h_{u}^{\ell}}_{\text{key}}\Big{)}\underbrace{h_{u}^{\ell}}_{\text{value}}.$$

(5)

The weighting coefficients $\alpha_{vu}^{\ell}$ follow the query-key design in usual attention mechanisms, whereby the representation of $v$ in the past layer, $h_{v}^{\ell-1}$, serves as the query. Specifically, we define

$$\alpha_{vu}^{\ell}\Big{(}h_{v}^{\ell-1},\,h_{u}^{\ell}\Big{)}=\operatorname*{softmax}_{u\in\mathcal{P}(v)}\Big{(}{w_{1}^{\ell}}^{\top}h_{v}^{\ell-1}+{w_{2}^{\ell}}^{\top}h_{u}^{\ell}\Big{)},$$

(6)

where $w_{1}^{\ell}$ and $w_{2}^{\ell}$ are model parameters. We use the additive form, as opposed to the usual dot-product form,¹¹1The usual dot-product form reads $\alpha_{vu}^{\ell}(h_{v}^{\ell-1},\,h_{u}^{\ell})=\operatorname*{softmax}(\langle{W_{1}^{\ell}}^{\top}h_{v}^{\ell-1},\,{W_{2}^{\ell}}^{\top}h_{u}^{\ell}\rangle)$. We find that in practice the dot-product form and the additive form perform rather similarly, but the former requires substantially more parameters. We are indebted to Hyoungjin Lim who pointed out that, however, in the additive form, the query term will be canceled out inside the softmax computation. 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^{\ell}$ combines the message $m_{v}^{\ell}$ with the previous representation of $v$, $h_{v}^{\ell-1}$, and produces an updated representation $h_{v}^{\ell}$. We employ a recurrent architecture, which is usually used for processing data in sequential order but similarly suits processing in partial order:

$$h_{v}^{\ell}=F^{\ell}\Big{(}h_{v}^{\ell-1},\,m_{v}^{\ell}\Big{)}=\textsc{Gru}^{\ell}\Big{(}\underbrace{h_{v}^{\ell-1},}_{\text{input}}\,\overbrace{\underbrace{m_{v}^{\ell}}_{\text{state}}}^{\text{message}}\Big{)},$$

(7)

where $h_{v}^{\ell-1}$, $m_{v}^{\ell}$, and $h_{v}^{\ell}$ 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^{\ell}(\emptyset,\cdot)=0$ for the aggregator so that for nodes with an empty direct-predecessor 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 $\mathcal{G}$ to create a reverse DAG $\widetilde{\mathcal{G}}$. We will use the tilde notation for all terms related to the reverse DAG. For example, the representation of node $v$ in $\widetilde{\mathcal{G}}$ at the $\ell$-th layer is denoted by $\widetilde{h}_{v}^{\ell}$.

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

$$h_{\mathcal{G}}=\textsc{FC}\Big{(}\operatorname*{Max-Pool}_{v\in\mathcal{T}}\big{(}\operatorname*{\parallel}_{\ell=0}^{L}h_{v}^{\ell}\big{)}\,\,\operatorname*{\parallel}\,\,\operatorname*{Max-Pool}_{u\in\mathcal{S}}\big{(}\operatorname*{\parallel}_{\ell=0}^{L}\widetilde{h}_{u}^{\ell}\big{)}\Big{)}.$$

(8)

Note that the target set $\widetilde{\mathcal{T}}$ of $\widetilde{\mathcal{G}}$ is the same as the source set $\mathcal{S}$ of $\mathcal{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 $\tau(u,v)$ be the type of an edge $(u,v)$ and let $y_{\tau}$ be a representation of edges of type $\tau$. We insert this information during message calculation in the aggregator. Specifically, we replace the attention weights $\alpha_{vu}^{\ell}$ defined in (6) by

$$\alpha_{vu}^{\ell}\Big{(}h_{v}^{\ell-1},\,h_{u}^{\ell}\Big{)}=\operatorname*{softmax}_{u\in\mathcal{P}(v)}\Big{(}{w_{1}^{\ell}}^{\top}h_{v}^{\ell-1}+{w_{2}^{\ell}}^{\top}h_{u}^{\ell}+{w_{3}^{\ell}}^{\top}y_{\tau(u,v)}\Big{)}.$$

(9)

In practice, we experiment with slightly fewer parameters by setting $w_{3}^{\ell}=w_{1}^{\ell}$ and find that the model performs equally well. The edge representations $y_{\tau}$ are trainable embeddings of the model. Alternatively, if input edge features are provided, $y_{\tau(u,v)}$ can be replaced by a neural network-transformed embedding for the edge $(u,v)$.

A key difference to MPNN is that DAGNN processes nodes sequentially owing to the nature of the aggregator $G^{\ell}$, obeying the partial order. Thus, for computational efficiency, it is important to maximally exploit concurrency so as to better leverage parallel computing resources (e.g., GPUs). One observation is that nodes without dependency may be grouped together and processed concurrently, if their predecessors have all been processed. See Figure 2 for an illustration.

To materialize this idea, we consider topological batching, which partitions the node set $\mathcal{V}$ into ordered batches $\{\mathcal{B}_{i}\}_{i\geq 0}$ so that (i) the $\mathcal{B}_{i}$’s are disjoint and their union is $\mathcal{V}$; (ii) for every pair of nodes $u,v\in\mathcal{B}_{i}$ for some $i$, there is not a directed path from $u$ to $v$ or from $v$ to $u$; (iii) for every $i>0$, there exists one node in $\mathcal{B}_{i}$ such that it is the tail of an edge whose head is in $\mathcal{B}_{i-1}$. The concept was propsoed by Crouse et al. (2019);²²2See also an earlier implementation in https://github.com/unbounce/pytorch-tree-lstm in what follows, we derive several properties that legitimizes its use in our setting. First, topological batching produces the minimum number of sequential batches such that all nodes in each batch can be processed in parallel.

The partitioning procedure may be as follows. All nodes without direct predecessors, $\mathcal{S}$, form the initial batch. Iteratively, remove the batch just formed from the graph, as well as the edges emitting from these nodes. The nodes without direct predecessors in the remaining graph form the next batch.

In the following, we summarize properties of the DAGNN model; they are consistent with the corresponding results for MPNNs. To formalize these results, we let $\mathcal{M}:\mathcal{V}\times\mathcal{E}\times\mathcal{X}\to h_{\mathcal{G}}$ denote the mapping defined by Equations (3)–(4). For notational consistency, we omit bidirectional processing, and thus ignore the tilde term in (8). The first results state that DAGNN produces the same graph representation invariant to node permutation.

The next result states that the framework will not produce the same graph representation for different graphs (i.e., non-isomorphic graphs), under a common condition.

The condition required by Theorem 4 is not restrictive. There exist (infinitely many) injective multiset functions $G^{\ell}$, $F^{\ell}$, and $R$, although the ones instantiated by (5)–(8) are not necessarily injective. The modification to injection can be done by using the $\epsilon$-trick applied in GIN (Xu et al., 2019), but, similar to the referenced work, the $\epsilon$ that ensures injection is unknown. In practice, it is either set as zero or treated as a tunable hyperparameter.

The OGBG-CODE dataset (Hu et al., 2020) contains 452,741 Python functions parsed into DAGs. We consider the TOK task, predicting the tokens that form the function name; it is included in the Open Graph Benchmark (OGB). Additionally, we introduce the LP task, predicting the length of the longest path of the DAG. The metric for TOK is the F1 score and that for LP is accuracy. Because of the vast size, we also create a 15% training subset, OGBG-CODE-15, for similar experiments.

For this dataset, we consider three basic baselines and several GNN models for comparison. For the TOK task, the Node2Token baseline predicts tokens from the attributes of the second graph node, while the TargetInGraph baseline predicts tokens that appear in both the ground truth and in the attributes of some graph node. These baselines exploit the fact that the tokens form node attributes and that the second node’s attribute contains the function name if it is part of the vocabulary. For the LP task, the MajorityInValid baseline constantly predicts the majority length seen from the validation set. The considered GNN models include four from OGB: GCN (Kipf & Welling, 2017), GIN (Xu et al., 2019), GCN-VN, GIN-VN (where -VN means adding a virtual node connecting all existing nodes); two using attention/gated-sum mechanisms: GAT (Veličković et al., 2018), GG-NN (Li et al., 2016); two hierarchical pooling approaches using attention: SAGPool (Lee et al., 2019), ASAP (Ranjan et al., 2020); and the D-VAE encoder.

The NA dataset (Zhang et al., 2019) contains 19,020 neural architectures generated by the ENAS software. The task is to predict the architecture performance on CIFAR-10 under the weight-sharing scheme. Since it is a regression task, the metrics are RMSE and Pearson’s $r$. To gauge performance with Zhang et al. (2019), we similarly train (unsupervised) autoencoders and use sparse Gaussian process regression on the latent representation to predict the architecture performance. DAGNN serves as the encoder and we pair it with an adaptation of the D-VAE decoder (see Appendix D). We compare to D-VAE and all the autoencoders compared therein: S-VAE (Bowman et al., 2016), GraphRNN (You et al., 2018), GCN (Zhang et al., 2019), and DeepGMG (Li et al., 2018).

The BN dataset (Zhang et al., 2019) contains 200,000 Bayesian networks generated by using the R package bnlearn. The task is to predict the BIC score that measures how well a BN fits the Asia dataset (Lauritzen & Spiegelhalter, 1988). We use the same metrics and baselines as for NA.

Prediction performance, token prediction (TOK), Table 1. The general trend is the same across the full dataset and the 15% subset. DAGNN performs the best. GAT achieves the second best result, surprisingly outperforming D-VAE (the third best). Hence, using attention as aggregator during message passing benefits this task. On the 15% subset, only DAGNN, GAT, and D-VAE match or surpass the TargetInGraph baseline. Note that not all ground-truth tokens are in the vocabulary and thus the best achievable F1 is 90.99. Even so, all methods are far from reaching this ceiling performance. Furthermore, although most of the MPNN models (middle section of the table) use as many as five layers for message passing, the generally good performance of DAGNN and D-VAE indicates that DAG architectures not restricted by the network depth benefit from the inductive bias.

Prediction performance, length of longest path (LP), Table 1. This analytical task interestingly reveals that many of the findings for the TOK task do not directly carry over. DAGNN still performs the best, but the second place is achieved by D-VAE while GAT lags far behind. The unsatisfactory performance of GAT indicates that attention alone is insufficient for DAG representation learning. The hierarchical pooling methods also perform disappointingly, showing that ignoring nodes may modify important properties of the graph (in this case, the longest path). It is worth noting that DAGNN and D-VAE achieve nearly perfect accuracy. This result corroborates the theory of Xu et al. (2020), who state that when the inductive bias is aligned with the reasoning algorithm (in this case, path tracing), the model learns to reason more easily and achieves better sample efficiency.

Prediction performance, scoring the DAG, Table 2. On NA and BN, DAGNN also outperforms D-VAE, which in turn outperforms the other four baselines (among them, DeepGMG works the best on NA and S-VAE works the best on BN, consistent with the findings of Zhang et al. (2019).) While D-VAE demonstrates the benefit of incorporating the DAG bias, DAGNN proves the superiority of its architectural components, as will be further verified in the subsequent ablation study.

Time cost, Figure 3. The added expressivity of DAGNN comes with a tradeoff: the sequential processing of the topological batches requires more time than does the concurrent processing of all graph nodes, as in MPNNs. Figure 3 shows that such a trade-off is innate to DAG architectures, including the D-VAE encoder. Moreover, the figure shows that, when used as a component of a larger architecture (autoencoder), the overhead of DAGNN may not be essential. For example, in this particular experiment, DeepGMG (paired with the S-VAE encoder) takes an order of magnitude more time than does DAGNN (paired with the D-VAE decoder). Most importantly, not reflected in the figure is that DAGNN learns better and faster at larger learning rates, leading to fewer learning epochs. For example, DAGNN reaches the best performance at epoch 45, while D-VAE at around 200.

Ablation study, Table 3. While the D-VAE encoder performs competitively owing similarly to the incorporation of the DAG bias, what distinguishes our proposal are several architecture components that gain further performance improvement. In Table 3, we summarize results under the following cases: replacing attention in the aggregator by gated sum; reducing the multiple layers to one; replacing the GRUs by fully connected layers; modifying the readout by pooling over all nodes; and removing the edge attributes. One observes that replacing attention generally leads to the highest degradation in performance, while modifying other components yields losses too. There are two exceptions. One occurs on LP-15, where gated-sum aggregation surprisingly outperforms attention by a tight margin, considering the standard deviation. The other occurs on the modification of the readout for the BN dataset. In this case, a Bayesian network factorizes the joint distribution of all variables (nodes) it includes. Even though the DAG structure characterizes the conditional independence of the variables, they play equal roles to the BIC score and thus it is possible that emphasis of the target nodes adversely affects the predictive performance. In this case, pooling over all nodes appears to correct the overemphasis.

Sensitivity analysis, Table 4 and Figure 5. It is well known that MPNNs often achieve best performance with a small number of layers, a curious behavior distinct from other neural networks. It is important to see if such a behavior extends to DAGNN. In Table 4, we list the results for up to four layers. One observes that indeed the best performance occurs at either two or three layers. In other words, one layer is insufficient (as already demonstrated in the ablation study) and more than three layers offer no advantage. We further extend the experimentation on TOK-15 with additional layers and plot the results in Figure 5. The trend corroborates that the most significant improvement occurs when going beyond a single layer. It is also interesting to see that a single layer yields the highest variance subject to randomization.

Structure learning, Figure 5. For an application of DAGNN, we extend the use of the BN dataset to learn the Bayesian network for the Asia data. In particular, we take the Bayesian optimization approach and optimize the BIC score over the latent space of DAGs. We use the graphs in BN as pivots and encode every graph by using DAGNN. The optimization yields a DAG with BIC score $-11107.29$ (see Figure 5). This DAG is almost the same as the ground truth (see Figure 2 of Lauritzen & Spiegelhalter (1988)), except that it does not include the edge from “visit to Asia?” to “Tuberculosis?”. It is interesting to note that the identified DAG has a higher BIC score than that of the ground truth, $-11109.74$. Furthermore, the BIC score is also much higher than that found by using the D-VAE encoder, $-11125.75$ (Zhang et al., 2019). This encouraging result corroborates the superior encoding quality of DAGNN and the effective use of it in downstream tasks.