SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Masked Attention for Transformers. As shown in Figure 1, we can implement DAG attention in a very straightforward fashion using a mask that masks out node pairs based on the DAG reachability relation as follows (compare to Equation (1)), with the attention mask  $\mathbf{M}$  being defined as a symmetric matrix over node pairs  $(v, u) \in V \times V$ .
$$
\text {A t t e n t i o n} (\mathbf {X}) = \operatorname {s o f t m a x} \left(\frac {\mathbf {Q K} ^ {T}}{\sqrt {\mathrm {d} _ {K}}} + \mathbf {M}\right) \mathbf {V}, \quad \mathbf {M} (v, u) = \left\{ \begin{array}{l l} 0 & \text {i f} u \in N _ {k} (v) \\ - \infty & \text {o t h e r w i s e .} \end{array} \right.
$$
While this masking represents a simple technique to extend and bias existing transformer models to DAGs, the resulting architecture does not benefit from the restricted node set to be considered, leading to unnecessary, costly matrix operations during attention calculation. Moreover, it consumes  $O(|V|^2)$  of additional memory to store M. Thus, the runtime complexity per layer is still  $O(|V|^2 d)$  for the vanilla transformer - and may be even higher, depending on the underlying transformer.
DAG Attention using Message Passing. Based on the formulation in Equation (3), we propose to follow the message-passing scheme; that is, for a node  $v$ , we compute  $N_{k}(v)$  and only aggregate messages from the nodes in that set to compute the DAG attention. For the latter aggregation, we can use readily available frameworks, such as PyG [Fey and Lenssen, 2019]. We analyze the complexity of this proposal.
Table 1: Statistics of the datasets we used.
ogbg-code2	NA	Self-citation	Cora	Citeseer	Pubmed
# graphs	452,741	19,020	1,000	1	1	1
Avg # nodes	125.2	8.0	59.1	2,708	3,327	19717
Avg n∞	9.78	7.00	4.30	20.88	5.33	60.56
3.6 Computational Complexity
Clearly, the time complexity of the proposed model is lower than that of the standard transformer. We consider two computation steps:
Computing DAG Reachability. To obtain  $N_{k}$ , we compute the transitive closure of  $E$  for each node in the graph using a breadth-first search starting at the node and iteratively expanding  $N_{k}$  based on  $E$ . Hence, the overall complexity of this step is  $O(|E||V|)$ . Observe that, during training, we can consider this step as pre-processing since it only has to be run once, in the very beginning.
DAG Attention. The matrix product  $\mathbf{Q}\mathbf{K}^T$  of self-attention has a cost  $O(|V|^2 d)$  which is quadratic in the number of nodes in  $G$ , and hence especially critical for large graphs. Yet, for our DAG attention, we only aggregate messages from reachable nodes. Therefore, the time complexity of reachability relation attention is  $O(|V| \times n_k \times d)$ , where  $n_k$  is the average size of  $N_k$ . In the worst case, we have  $n_k = |V| - 1$ , but we assume that  $n_k << |V|$  in general. Especially for sparse graphs, the complexity gets significantly lower, i.e.,  $O(|V| \times \overline{deg}^k \times d)$ , where  $\overline{deg}$  is the average node degree.
Finally, observe that directed rooted trees represent a special kind of DAGs broadly seen across domains, such as abstract syntax trees of source code [Allamanis et al., 2018]. For them, the runtime of DAG attention scales almost linearly, as illustrated in the following theorem.
Theorem 1. In a directed rooted tree  $T$ , the runtime of DAG attention is  $O(|V| \times k \times \Delta^{+} \times d)$ , where  $\Delta^{+}$  is the maximal outdegree of  $T$ . When  $k = \infty$ , the runtime of DAG attention is  $O(|V| \times \text{depth}(T) \times \Delta^{+} \times d)$ .
The proof is provided in Appendix A. Note that when  $\Delta^{+}$  is small,  $d$  is a constant and  $\mathrm{depth}(T)$  is  $O(\log |V|)$ , the runtime of DAG attention is almost linear in  $|V|$ . Indeed, we observed this on the ogbg-code2 dataset in our experiments.
We evaluate our proposed architecture on several datasets, comparing it to competitive baselines. Ablation studies also provide more insight into the composition of our model. Primarily, the following questions are investigated:
- Does DAG attention have the expected effects and improves existing graph transformers both in terms of quality and efficiency?
- Is DAG bias encoded through both DAGRA & DAGPE, and how does DAGPE compare to others?
- How does the size of the receptive field - the main difference between regular GNNs and transformers - affect the performance?
Datasets. Table 1 shows the diversity of the datasets we used; see Appendix B for full details.
- ogbg-code2 [Hu et al., 2020]. A large dataset of ASTs derived from Python methods. The node features are syntax tokens. The multi-task classification task is to predict the first 5 tokens of the function name.
- NA [Zhang et al., 2019]. A dataset with much smaller graphs, containing neural architecture DAGs generated by the ENAS software. Each node's features represent a certain neural network component type. The (regression) task is to predict the architecture performance on CIFAR-10.
- Self-citation [ARC, 2021, Luo et al., 2023]. Each DAG in the academic self-citation represents a scholar's academic self-citation network [ARC, 2021]. Each paper has two node attributes: the
Table 2: Code graph classification on ogbg-code2. The baseline results were taken from the OGB leaderboard.
Model	Valid F1 (%)	Test F1 (%)	Time(epoch)
GIN	13.7 ± 0.2	14.9 ± 0.2	181s
GCN	14.0 ± 0.2	15.1 ± 0.2	127s
GIN-Virtual	14.4 ± 0.3	15.8 ± 0.2	155s
GCN-Virtual	14.6 ± 0.1	16.0 ± 0.2	198s
GAT	14.4 ± 0.2	15.7 ± 0.2	134s
PNA	14.5 ± 0.3	15.7 ± 0.3	427s
DAGNN	16.1 ± 0.4	17.5 ± 0.5	6018s
PACE	16.3 ± 0.3	17.8 ± 0.2	2410s
Transformer	15.5 ± 0.2	16.7 ± 0.2	1817s
\( DAG+Transformer \)	17.4 ± 0.1	18.8 ± 0.2	591s
GraphTrans	16.6 ± 0.1	18.3 ± 0.2	1117s
\( DAG+GraphTrans \)	17.0 ± 0.2	18.7 ± 0.2	526s
GraphGPS	17.4 ± 0.2	18.9 ± 0.2	1919s
\( DAG+GraphGPS \)	17.6 ± 0.1	19.3 ± 0.2	608s
SAT (SOTA)	17.7 ± 0.2	19.4 ± 0.3	2437s
\( DAG+SAT \)	18.5 ± 0.1	20.2 ± 0.2	681s
Table 3: Node classification results for the self-citation dataset; AP (%) and ROC-AUC  $(\%)$
Model	AP ↑	ROC-AUC ↑
GIN	57.7 ± 1.8	79.7 ± 0.2
GCN	58.8 ± 0.4	79.9 ± 0.2
GIN-Virtual	57.4 ± 1.2	79.5 ± 0.4
GCN-Virtual	58.9 ± 0.2	80.0 ± 0.1
GAT	55.3 ± 3.7	77.9 ± 1.4
PNA	62.4 ± 0.7	81.0 ± 0.4
DAGNN	61.2 ± 0.6	81.0 ± 0.3
PACE	52.1 ± 1.8	75.9 ± 0.7
Transformer	56.8 ± 1.8	78.7 ± 0.3
DAG+Transformer	63.8 ± 0.8	82.2 ± 0.5
GraphGPS	61.6 ± 2.6	81.3 ± 0.6
DAG+GraphGPS	63.5 ± 1.2	80.8 ± 0.5
SAT	59.8 ± 1.7	79.8 ± 0.7
DAG+SAT	62.7 ± 1.5	80.6 ± 0.7
NodeFormer	39.6 ± 0.6	69.4 ± 0.3
DAG+NodeFormer	64.9 ± 0.8	81.7 ± 0.8
publication year and total citation count (excluding the papers whose citation category is to be inferred). Here we consider the node-level task of predicting whether a paper is highly-cited or not - as a proxy for its impact.
- Cora, Citeseer, Pubmed [Sen et al., 2008]. Established, medium-sized citation graphs. Only for our method, we removed a small number of cyclic citation links to make them DAGs.
Baselines. We chose two basic message-passing GNNs, GCN and GIN; extensions of these models using a virtual node connected to all other graph nodes; the graph attention network GAT, as it showed especially good performance in [Thost and Chen, 2021]; MixHop [Abu-El-Haija et al., 2019], LDS-GNN [Franceschi et al., 2019], IDGL [Chen et al., 2020] and PNA, as more recent GNN proposals. In terms of transformer models, we considered vanilla Transformer (TF), Graph Transformer (GT), GraphTrans, SAT, GraphGPS and NodeFormer which achieved state-of-the-art performance (SOTA). Lastly, we consider neural networks tailored to DAGs: S-VAE, D-VAE, DAGNN and the current SOTA, PACE. For more detailed descriptions, see Section 2.
DAG+ Models. We implemented our DAG attention on top of vanilla Transformer, GraphTrans, SAT, GraphGPS and NodeFormer only (1) modifying their self-attention module by restricting attention in terms of reachability and (2) using DAGPE instead of the original one. We note that there are various alternatives for the latter (e.g., concatenation etc.), therefore we opted for the most simple solution which, at the same time, provides a very direct comparison. Moreover, on ogbg-code2 we did not use any PE since [Chen et al., 2022a] showed that they do not make real impact and we observed the same in preliminary experiments. For fair comparisons, we use the same hyperparameters (including the number of layers, batch size, hidden dimension etc.) and readout as baseline transformers. Given one of the baseline transformers M, we denote the modified model using DAG attention by DAG+M. Unless explicitly specified otherwise, we chose  $k = \infty$  in all experiments. Full details on the experimental setup and hyperparameters are provided in the Appendix B.
Overall Performance, Tables 2, 3, 4 and 5. First, observe that the results are very consistent, although the datasets differ greatly in size, DAG sizes, DAG shape (e.g., in ogbg-code2 we have trees), and nature of data (e.g., node features). The message-passing GNNs represent standardized baselines but do not reach the performance of networks tailored to DAGs, such as DAGNN. Note that the latter is however comparatively bad on the node-level task. This can be explained by its processing. It computes node representations in the order of the partial order, and hence the representations of nodes in the beginning of the order contain only few information about the entire graph. Intuitively, transformers should be able to capture this information, but the transformers we tested do neither perform better, not even the ones tailored to graphs. This shows that they are missing information captured by those message-passing neural networks that were tailored to DAGs. The results clearly show that our DAG attention is successful in providing good improvements, over the best graph
Table 4: Node classification accuracy (%).The baseline results were taken from [Wu et al., 2022].
Model	Cora	Citeseer	Pubmed
GCN	87.06 ± 0.34	75.75 ± 0.37	88.16 ± 0.14
GAT	86.85 ± 0.30	75.92 ± 0.26	86.90 ± 0.22
MixHop	87.59 ± 0.52	73.64 ± 0.73	89.32 ± 0.25
IDGL	87.88 ± 0.34	74.32 ± 0.51	89.22 ± 0.14
LDS-GNN	87.82 ± 0.62	75.22 ± 0.23	OOM
PACE	79.47 ± 0.63	73.65 ± 1.23	OOM
Transformer	75.92 ± 0.86	72.23 ± 1.06	OOM
DAG+Transformer	87.80 ± 0.53	74.42 ± 0.22	89.01 ± 0.13
SAT	75.18 ± 0.62	74.88 ± 0.73	OOM
DAG+SAT	87.48 ± 0.37	76.64 ± 0.26	89.17 ± 0.15
NodeFormer	88.80 ± 0.26	76.33 ± 0.59	89.32 ± 0.25
DAG+NodeFormer	90.49 ± 0.17	78.24 ± 0.33	89.44 ± 0.24
Table 5: Regression. Predictive performance of latent representations over NA.
Model	RMSE ↓	Pearson's r ↑
GCN	0.482 ± 0.003	0.871 ± 0.001
S-VAE	0.521 ± 0.002	0.847 ± 0.001
D-VAE	0.375 ± 0.003	0.924 ± 0.001
DAGNN	0.264 ± 0.004	0.964 ± 0.001
PACE	0.254 ± 0.002	0.964 ± 0.001
Transformer	0.285 ± 0.004	0.957 ± 0.001
GT	0.275 ± 0.003	0.961 ± 0.001
DAG+Transformer	0.253 ± 0.002	0.966 ± 0.001
GraphGPS	0.306 ± 0.004	0.950 ± 0.001
DAG+GraphGPS	0.267 ± 0.005	0.964 ± 0.001
SAT	0.298 ± 0.003	0.952 ± 0.001
DAG+SAT	0.262 ± 0.004	0.964 ± 0.001
Table 6: Ablation results.
Ablation	ogbg-code2	NA	Self-citation
Valid F1	Test F1	RMSE ↓	Pearson's r ↑	AP ↑	ROC-AUC ↑
DAG+TF	0.1731 ± 0.0014	0.1895 ± 0.0014	0.253 ± 0.002	0.966 ± 0.001	0.638 ± 0.008	0.822 ± 0.005
(-) DAGRA	0.1546 ± 0.0018	0.1670 ± 0.0015	0.284 ± 0.003	0.957 ± 0.001	0.573 ± 0.011	0.790 ± 0.003
(-) DAGPE	0.1739 ± 0.0013	0.1879 ± 0.0015	0.263 ± 0.002	0.963 ± 0.001	0.594 ± 0.028	0.782 ± 0.018
(+) RWPE	-	-	0.267 ± 0.003	0.962 ± 0.001	0.628 ± 0.014	0.819 ± 0.010
(+) LapPE	-	-	0.271 ± 0.002	0.961 ± 0.001	0.609 ± 0.017	0.786 ± 0.015
(+) SPD Ying et al. [2021]	0.1749 ± 0.0011	0.1881 ± 0.0017	-	-	0.639 ± 0.006	0.823 ± 0.004
(+) Edge Direction	0.1751 ± 0.0018	0.1870 ± 0.0021	-	-	0.636 ± 0.015	0.817 ± 0.005
TF	0.1546 ± 0.0018	0.1670 ± 0.0015	0.285 ± 0.004	0.957 ± 0.001	0.568 ± 0.018	0.787 ± 0.003
DAG+SAT	0.1821 ± 0.0013	0.1982 ± 0.0010	0.262 ± 0.004	0.964 ± 0.001	0.627 ± 0.015	0.806 ± 0.007
(-) DAGRA	0.1773 ± 0.0023	0.1937 ± 0.0028	0.292 ± 0.003	0.954 ± 0.001	0.598 ± 0.031	0.800 ± 0.012
(-) DAGPE	0.1846 ± 0.0010	0.2018 ± 0.0021	0.282 ± 0.002	0.958 ± 0.001	0.623 ± 0.014	0.806 ± 0.005
(+) SPD Ying et al. [2021]	0.1851 ± 0.0008	0.1991 ± 0.0018	-	-	0.627 ± 0.016	0.810 ± 0.006
(+) Edge Direction	0.1839 ± 0.0014	0.1978 ± 0.0028	-	-	0.623 ± 0.013	0.804 ± 0.007
SAT	0.1773 ± 0.0023	0.1937 ± 0.0028	0.298 ± 0.003	0.952 ± 0.001	0.598 ± 0.017	0.798 ± 0.007
transformer SAT and even better ones over vanilla Transformer. On ogbg-code2, the improvement is smaller for GraphTrans and SAT. However, this benchmark task is heavily dependent on other features (e.g., language understanding) and hence presents a special challenge; in this regard, the NA and self-citation datasets contain "cleaner" graphs. It is interesting to see that our framework lifts the transformers from below the level of DAGNN to above, in terms of all metrics, over these two datasets. Lastly, we observe that the PACE transformer tailored to DAGs is similarly outperformed by our simpler, but more effective technique. Nevertheless, the overall conclusion holds in general: Over all datasets, our DAG attention makes the transformers outperform (1) the original transformers and (2) the neural networks tailored to DAGs. This shows that the DAG-specific bias provided by DAG attention is the right bias. We investigated this in more detail in our ablation experiments.
Ablation Study, Table 6. Recall that our DAG attention is composed of the reachability-based attention (DAGRA) and DAGPE modules. To justify this design, (1) our ablation studies removing DAGRA and DAGPE individually confirm the impact of both modules; (2) we also experimented with replacing DAGPE with LapPE [Dwivedi and Bresson, 2020] and the random-walk-based RWPE [Dwivedi et al., 2021] to show that the DAG-specific nature of our PE is of advantage; (3) we experimented with adding attention bias which captures the graph structure more directly (e.g., shortest-path Ying et al. [2021] and edge directionality), although they are implicit in DAGPE; interestingly, the more direct representation does not make a noticeable difference. For ease of comparison and interpretation (i.e., in terms of magnitude), we also provide the baseline transformer results. Overall, it can be observed that our DAG attention yields highly consistent performance improvements, although occasionally the advantage is less pronounced. This shows that our architecture design provides the right bias on top of (graph) transformers, which makes the improved models better fit for DAGs.
Impact of Size of Receptive Field  $N_{k}$
(1) Average  $n_{\infty}$ , Table 1. Compared to the often only two to three hops considered in message-passing GNNs, our  $k = \infty$  might seem unrealistic. However, as we show in the table, for two of our three and very different datasets, we have that  $n_{\infty}$  is much smaller than the worst case size  $|V|$ . NA
Table 7: Impact of different  $k$  on DAG attention,over ogbg-code2.
k	Avg nk	Test F1 score
DAG+TF	DAG+SAT
1	1.97	0.1724 ± 0.0022	0.1533 ± 0.0108
2	3.91	0.1800 ± 0.0023	0.1918 ± 0.0010
3	5.69	0.1842 ± 0.0006	0.1934 ± 0.0009
4	7.16	0.1849 ± 0.0021	0.1975 ± 0.0019
5	8.29	0.1856 ± 0.0020	0.2000 ± 0.0005
6	9.01	0.1869 ± 0.0020	0.2000 ± 0.0015
7	9.40	0.1869 ± 0.0019	0.2004 ± 0.0014
8	9.60	0.1875 ± 0.0008	0.2005 ± 0.0010
∞	9.78	0.1879 ± 0.0015	0.2018 ± 0.0021
Figure 2: Average training time per epoch for various  $k$  over ogbg-code2, log scale.
represents a special case because the graphs are very small so that we can always find a directed path through them that contains all nodes; however, given that  $|V|$  is generally small, the exception of the rule  $n_{\infty} \ll |V|$  is not critical here.
(2) Performance for Varying  $k$ , Table 7, Figure 2. We ran our model for different  $k$  on ogbg-code2, based upon vanilla Transformer and SAT. We find that incorporating reachable node information leads to small but constant improvements in performance as  $k$  increases. This confirms our intuition about the importance of predecessors and successors in DAGs, and is in line with related works suggesting to iteratively aggregate the DAGs along the partial order.
Figure 2 depicts the training time per epoch. It is easy to see that our framework reduces computation time considerably. For example, even the most time-consuming DAG+SAT requires 10 minutes per epoch, compared to 40 minutes for SAT and 100 minutes for DAGNN on the same GPU type. Since SAT and DAGNN represent the best-performing models among the transformers and, respectively, message-passing GNNs, and DAG attention yields qualitative improvements over both.
Based on the insights from graph neural networks and models for directed acyclic graphs, we have developed a transformer model biased towards DAGs, which allows for incorporating the main characteristic of DAGs - their partial order - into any transformer architecture, by encoding the reachability relation and positions resulting from it. Our architecture is simple and universal as we demonstrated in our experiments. Most importantly, it provides an effective extension improving existing graph transformers over DAGs. While our framework has successfully lifted existing (graph) transformers to the level of the state-of-the-art neural networks tailored to DAGs, our experiments show that there is still room for improvement. There is a variety of challenging tasks over many different kinds of DAGs, which are not fully solved yet; for instance, to tackle the challenge of reasoning over source code graphs, the models will probably need better language understanding. Hence DAG representation learning remains interesting for research, and our work provides an important contribution in closing the gap between transformers and other DAG neural networks.
**Limitations.** Our study is very general, but we found only a limited number of DAG datasets. We tried to resolve this by creating new ones (from the established citation datasets), yet we acknowledge that there is still room for extension. Further, there are likely specific types of DAGs which benefit from more customized modeling. Yet, our study is intended to address a general adaptivity to DAGs.
Broader Impact. Transformers have been popular in artificial intelligence, and it is expected that they also gain importance in graph learning. We hope to provide a tiny piece to the advancement of the area. Our framework is rather abstract, but its simplicity, efficiency, and universality make it practically usable; and we demonstrate good performance on a variety of data.
We extend our gratitude to Yicheng Pan for his invaluable assistance with the computational complexity analysis. We also express our appreciation to all the anonymous reviewers for their insightful and constructive feedback. This work was supported by National Key R&D Program of China (2021YFB3500700), NSFC Grant 62172026, National Social Science Fund of China 22&ZD153, the Fundamental Research Funds for the Central Universities and SKLSDE.
The acl anthology reference corpus. https://aclanthology.org/, Nov. 2021.
Sami Abu-El-Haija, Bryan Perozzi, Amol Kapoor, Nazanin Alipourfard, Kristina Lerman, Hrayr Harutyunyan, Greg Ver Steeg, and Aram Galstyan. Mixhop: Higher-order graph convolutional architectures via sparsified neighborhood mixing. In international conference on machine learning, pages 21–29. PMLR, 2019.
Miliadis Allamanis, Earl T Barr, Premkumar Devanbu, and Charles Sutton. A survey of machine learning for big code and naturalness. ACM Computing Surveys (CSUR), 51(4):1-37, 2018.
Uri Alon, Shaked Brody, Omer Levy, and Eran Yahav. code2seq: Generating sequences from structured representations of code. arXiv preprint arXiv:1808.01400, 2018.
Sergey Brin. The pagerank citation ranking: bringing order to the web. Proceedings of ASIS, 1998, 98:161-172, 1998.
Dexiong Chen, Leslie O'Bray, and Karsten Borgwardt. Structure-aware transformer for graph representation learning. In International Conference on Machine Learning, pages 3469-3489. PMLR, 2022a.
Jinsong Chen, Kaiyuan Gao, Gaichao Li, and Kun He. Nagphormer: Neighborhood aggregation graph transformer for node classification in large graphs. arXiv preprint arXiv:2206.04910, 2022b.
Yu Chen, Lingfei Wu, and Mohammed Zaki. Iterative deep graph learning for graph neural networks: Better and robust node embeddings. Advances in neural information processing systems, 33: 19314-19326, 2020.
Gabriele Corso, Luca Cavalleri, Dominique Beaini, Pietro Lio, and Petar Velickovic. Principal neighbourhood aggregation for graph nets. Advances in Neural Information Processing Systems, 33:13260-13271, 2020.
Maxwell Crouse, Ibrahim Abdelaziz, Cristina Cornelio, Veronika Thost, Lingfei Wu, Kenneth Forbus, and Achille Fokoue. Improving graph neural network representations of logical formulae with subgraph pooling. arXiv preprint arXiv:1911.06904, 2019.
Zehao Dong, Muhan Zhang, Fuhai Li, and Yixin Chen. Pace: A parallelizable computation encoder for directed acyclic graphs. In International Conference on Machine Learning, pages 5360-5377. PMLR, 2022.
Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs. arXiv preprint arXiv:2012.09699, 2020.
Vijay Prakash Dwivedi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Graph neural networks with learnable structural and positional representations. In International Conference on Learning Representations, 2021.
Vijay Prakash Dwivedi, Ladislav Rampasek, Mikhail Galkin, Ali Parviz, Guy Wolf, Anh Tuan Luu, and Dominique Beaini. Long range graph benchmark. arXiv preprint arXiv:2206.08164, 2022.
Matthias Fey and Jan Eric Lenssen. Fast graph representation learning with pytorch geometric. arXiv preprint arXiv:1903.02428, 2019.