SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Veronika Thost*
MIT-IBM Watson AI Lab, IBM Research
Lei Shi*
Transformer models have recently gained popularity in graph representation learning as they have the potential to learn complex relationships beyond the ones captured by regular graph neural networks. The main research question is how to inject the structural bias of graphs into the transformer architecture, and several proposals have been made for undirected molecular graphs and, recently, also for larger network graphs. In this paper, we study transformers over directed acyclic graphs (DAGs) and propose architecture adaptations tailored to DAGs: (1) An attention mechanism that is considerably more efficient than the regular quadratic complexity of transformers and at the same time faithfully captures the DAG structure, and (2) a positional encoding of the DAG's partial order, complementing the former. We rigorously evaluate our approach over various types of tasks, ranging from classifying source code graphs to nodes in citation networks, and show that it is effective in two important aspects: in making graph transformers generally outperform graph neural networks tailored to DAGs and in improving SOTA graph transformer performance in terms of both quality and efficiency.
Graph-structured data is ubiquitous in various disciplines [Gilmer et al., 2017, Zitnik et al., 2018, Sanchez-Gonzalez et al., 2020] and hence graph representation learning has the potential to provide huge impact. There are various types of graph neural networks (GNNs), the majority of which is based on a so-called message-passing scheme where node representations are computed iteratively by aggregating the embeddings of neighbor nodes [Gilmer et al., 2017]. Yet, this mechanism in its basic form has limited expressivity [Xu et al., 2018] and research is focusing on extensions.
Transformer models have recently gained popularity in graph learning as they have the potential to learn complex relationships beyond the ones captured by regular GNNs, in a different way [Dwivedi and Bresson, 2020, Kreuzer et al., 2021]. Technically, they can be considered as graph neural networks operating on fully-connected computational graphs, decoupled from the input graphs. The main research question in this context is how to inject the structural bias of the given input graphs (i.e., which nodes are actually connected) into the transformer architecture by adapting their attention and positional encoding appropriately. Several promising proposals have been made to encode undirected molecular graphs [Ying et al., 2021, Ross et al., 2022] and recent works take into account the scalability challenge [Dwivedi et al., 2022], also over larger network graphs [Rampášek et al., 2022, Chen et al., 2022b].
Figure 1: Overview of our DAG attention. A mask matrix restricts the receptive field of the node under consideration to nodes that are directly reachable in the DAG. Our positional encoding (right), additionally captures its position as its depth in the DAG.
We focus on directed acyclic graphs (DAGs), which are of special interest across various domains; examples include parsing results of source code [Allamanis et al., 2018], logical formulas [Crouse et al., 2019], and conversational emotion recognition [Shen et al., 2021], as well as probabilistic graphical models and neural architectures [Zhang et al., 2018, 2019, Knyazev et al., 2021]. In a DAG, the edges define a partial order over the nodes. This partial order represents an additional strong inductive bias, which offers itself to be integrated into the models.
In fact, various kinds of neural networks tailored to DAG structures have been proposed over the years; from early descriptions of recursive neural networks over DAGs [Sperduti and Starita, 1997, Frasconi et al., 1998], which have pointed out a similarity to sequence learning, to more recent works, such as DAG-RNN [Shuai et al., 2016], DAG-LSTM [Crouse et al., 2019], D-VAE [Zhang et al., 2019], and DAGNN [Thost and Chen, 2021]. The latter models focus on encoding the entire DAG; in a nutshell, they compute the embedding in a message-passing fashion by iterating over the DAG nodes in an asynchronous way, given by the partial order, and thereafter aggregating the final node representations into one for the DAG. This procedure yields state-of-the-art performance in terms of quality, but its asynchronous nature leads to comparatively (to regular GNNs) very slow runtime performance. The more parallel computation of transformers would seem to make them more suitable models and [Dong et al., 2022] have recently proposed one such model, PACE, which shows convincing performance in terms of both quality and efficiency.
In this paper, we focus on transformers over DAGs more generally, motivated by the highly flexible, parallel, and expressive nature of their computation. In particular, their success in sequence learning opens up the question how they can be tailored to DAGs, which are essentially sequential graphs. Based on the above-mentioned insights in DAG learning, we propose a straightforward and efficient DAG attention framework, which effectively biases any transformer towards DAGs.
- As outlined in Figure 1, we (1) adapt the attention mechanism by restricting the receptive field of each node to its predecessor and successor nodes so that it faithfully captures the DAG structure - in terms of its reachability relation - and at the same time gets considerably more efficient than the regular quadratic complexity; and (2) employ the positional encoding in a complementary way, to further bias the computation towards the DAG topology, by explicitly encoding the node depth.
- We show that the attention is not restricted effectively, even if it seems so technically. Moreover, we draw a connection to the random walk theory.
- We rigorously evaluate our proposal in ablation studies and show that it successfully improves different kinds of baseline transformers, from vanilla transformers [Vaswani et al., 2017] to state-of-the-art graph transformers [Wu et al., 2021, Chen et al., 2022a, Rampasek et al., 2022, Wu et al., 2022], over various types of DAG data. Our experiments range from classifying source code graphs to nodes in citation networks, and even go far beyond related works' problem scope. Most importantly, our proposal is proven effective in two important aspects: in making graph
transformers generally outperform graph neural networks tailored to DAGs and in improving SOTA graph transformer performance in terms of both quality and efficiency.
- Finally, our DAG attention framework can be implemented in two possible and rather straightforward ways, either on top of existing transformer models, or based on message-passing GNNs. Our implementation is available at https://github.com/LUOyk1999/DAGformer.
Directed Acyclic Graphs. We refer to a graph as tuple  $G = (V, E, \mathbf{X})$ , with node set  $V$ , edge set  $E \subseteq V \times V$ , and node features  $\mathbf{X} \in \mathbb{R}^{n \times d}$ , with each row representing the feature vector of one node, with number of nodes  $n$  and feature dimension  $d$ , the features of node  $v$  are denoted by  $x_v \in \mathbb{R}^d$ . A directed acyclic graph (DAG) is a directed graph without directed cycles. For a DAG  $G$ , we can define a unique strong partial order  $\leqslant$  on the node set  $V$ , such that, for all pairs of nodes  $u, v \in V$ ,  $u \leqslant v$  if and only if there is a directed path from  $u$  to  $v$ . We define the reachability relation  $\leqslant$  over a DAG based on that partial order. That is,  $u \leqslant v$  if and only if  $v$  is reachable from  $u$ ; further, if  $u \leqslant v$ , then  $u$  is called a predecessor of  $v$ , and  $v$  is a successor of  $u$ . All nodes without predecessors are source nodes, and all nodes without successors are sink nodes. We further define a restricted form of reachability  $\leqslant_K$  with  $u \leqslant_k v$  if and only if there is a directed path of length at most  $k$  from  $u$  to  $v$ . The depth of a node  $v \in V$  is defined below. The depth of a DAG is the maximum depth over its nodes.
$$
d e p t h (v) = \left\{ \begin{array}{l l} 0 & \text {i f v i s a s o u r c e n o d e} \\ 1 + \max  _ {(u, v) \in E} d e p t h (u) & \text {o t h e r w i s e} \end{array} \right.
$$
Message-Passing Graph Neural Networks. Most modern graph neural networks (GNNs) are or can be formulated in terms of the message-passing architecture, a framework proposed in [Gilmer et al., 2017]. In that architecture, node representations are computed iteratively by aggregating the embeddings of their neighbor nodes (i.e., the messages), and a final graph representation can be obtained by aggregating the node embeddings. Yet, in its basic form, this mechanism has limited expressivity [Xu et al., 2018] and graph neural networks remain an active research area. Notable works we compare to include the graph convolutional network (GCN) [Kipf and Welling, 2017], a very early, basic architecture; the graph attention network (GAT) [Veličković et al., 2018], aggregating the neighbor node embeddings using attention; the graph isomorphism network (GIN) [Xu et al., 2018], which integrates extra MLP layers into message passing for increased expressivity; and the principal neighbourhood aggregation (PNA) model [Corso et al., 2020], a recent proposal focusing on adapting the aggregation mechanism to the node under consideration (e.g., scaling based on its degree). For a broader overview of the field, we refer the reader to [Wu et al., 2020].
Transformers on Graphs. Transformer models [Vaswani et al., 2017] have gained popularity in graph learning as they have the potential to learn complex relationships beyond the ones captured by regular GNNs and in a different way. The architecture is composed of two main parts: a self-attention module and a feed-forward neural network, each followed by a residual connection with a normalization layer. Formally, the self-attention projects the input node features  $\mathbf{X}$  using three matrices  $\mathrm{W}_{\mathbf{Q}} \in \mathbb{R}^{d \times d_K}$ ,  $\mathrm{W}_{\mathbf{K}} \in \mathbb{R}^{d \times d_K}$  and  $\mathrm{W}_{\mathbf{V}} \in \mathbb{R}^{d \times d_K}$  to the corresponding representations for query (Q), key (K), and value (V), and is described as follows :
$$
\mathbf {Q} = \mathbf {X W} _ {\mathbf {Q}}, \mathbf {K} = \mathbf {X W} _ {\mathbf {K}}, \mathbf {V} = \mathbf {X W} _ {\mathbf {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}}}\right) \mathbf {V}. \tag {1}
$$
Over graphs, we focus on computing node instead of token embeddings (recall Figure 1). Technically, transformers can be considered as message-passing GNNs operating on fully-connected computational graphs, decoupled from the input graphs. The main research question in the context of graphs is how to inject the structural bias of the given input graphs by adapting their attention and by adding extensions to the original features, via so-called positional encodings (PEs). Graph Transformer [Dwivedi and Bresson, 2020] represents an early work using Laplacian eigenvectors as positional encodings, and various extensions and other models have been proposed since then [Min et al., 2022]. For instance, [Mialon et al., 2021] proposed a relative PE [Shaw et al., 2018] by means of kernels on graphs to bias the self-attention calculation. Notably, GraphTrans [Wu et al., 2021]
was the first hybrid architecture, using a stack of message-passing GNN layers before the regular transformer layers. Finally, [Chen et al., 2022a] have reformulated the self-attention mechanism as a kernel smoother as shown below and incorporated structure information into their structure-aware transformer (SAT) architecture by extracting a subgraph representation rooted at each node before attention computation:
$$
\text {A t t e n t i o n} \left(x _ {v}\right) = \sum_ {u \in V} \frac {\kappa \left(x _ {v} , x _ {u}\right)}{\sum_ {w \in V} \kappa \left(x _ {v} , x _ {w}\right)} f \left(x _ {u}\right), \forall v \in V, \tag {2}
$$
with  $f(x) = \mathbf{W}_{\mathbf{V}}x$  , and non-symmetric exp. kernel  $\kappa$  ..
$$
\kappa \left(x, x ^ {\prime}\right) = \exp \left(\frac {\langle \varphi (x) \mathbf {W} _ {\mathbf {Q}} , \varphi \left(x ^ {\prime}\right) \mathbf {W} _ {\mathbf {K}} \rangle}{\sqrt {d _ {K}}}\right), \quad \varphi (x) = \left\{ \begin{array}{l l} x & \text {i n v a n i l l a t r a n s f o r m e r} \\ \operatorname {G N N} _ {G} (x) & \text {i n S A T} \end{array} \right.
$$
where  $\langle \cdot, \cdot \rangle$  is the dot product on  $\mathbb{R}^d$  and  $\mathrm{GNN}_G(x)$  is an arbitrary GNN model. Most of the aforementioned works focus on the classification of smaller graphs, such as molecules; yet, recently, GraphGPS [Rampášek et al., 2022] is also considering larger graphs and the area is focusing on the development of scalable models; for instance, Nodeformer [Wu et al., 2022] is designed to address the issue of scalability and expressivity in node classification. Altogether, the transformer architecture opens new and promising avenues for graph representation learning, beyond message-passing GNNs.
Neural Networks over DAGs. The additional strong inductive bias present in DAGs has motivated researchers to formulate special neural architectures tailored to this kind of graphs [Sperduti and Starita, 1997, Frasconi et al., 1998]. Various types of models have been proposed over the years in many different application areas. There are works in the context of syntactic graphs, [Tai et al., 2015, Shuai et al., 2016] logical formulas [Crouse et al., 2019], source code representations [Allamanis et al., 2018] and neural architectures [Zhang et al., 2018, 2019, Knyazev et al., 2021]. We particularly compare to the most recent proposals S-VAE and D-VAE [Zhang et al., 2019]; and to DAGNN from [Thost and Chen, 2021], who also showed that using attention for neighbor node aggregation is beneficial in several DAG tasks.
While the models can be formulated in terms of the message-passing framework [Thost and Chen, 2021], their processing is special for GNNs: they compute a graph embedding by iterating over the DAG nodes in an asynchronous way, given by the partial order, and starting from the source nodes (or the other way around, from the sink nodes). That is, the messages are passed along the partial order of the DAG and, at any point during computation, capture the information of the subgraph consisting of all the predecessors of the node currently under consideration. Thereafter, a final graph representations can be obtained by aggregating the embeddings of all sink nodes. In contrast to regular message-passing GNNs, these customized works for DAGs usually focus on encoding graphs (i.e., instead of nodes or edges, which are other common GNN targets) and pass messages over the entire graph, while regular GNNs consider a fixed radius - usually rather small - around each node and essentially obtain a neighborhood embedding. Furthermore, the nature of the proposed architectures shows that recurrent techniques have proven especially useful, analogy to learning over sequences. In summary, these DAG models are elegant and effective, yielding state-of-the-art results, but their asynchronous nature leads to very slow and practically inhibitive performance compared to regular GNNs.
Transformers over DAGs. We are aware of only few proposals for transformers tailored to DAGs. [Huang et al., 2022] developed a Directed Acyclic Transformer in the context of machine translation in order to simultaneously capture multiple translations within the decoded DAG, but the model's encoder is a vanilla transformer. [Kotnis et al., 2021] proposed BIQE, which leverages the depth of nodes in DAG paths as positional encodings, with a specific focus on answering complex queries in knowledge graphs. [Gagrani et al., 2022] proposed Topoformer, which is also an encoder-decoder model but has been developed for finding optimal topological orders in DAGs. Since the study entirely focuses on the latter goal (e.g., in terms of training objective and evaluation) it is very specific and different from our more general study. The DAG encoder itself is also rather different in that it uses a Laplacian PE, as it is used with regular graph transformers; and a more complex attention mechanism, which does not only model the reachability but also several other relations over the graph. Closest to our method is PACE [Dong et al., 2022], which similarly focuses on modeling the sequential nature of DAGs inside a transformer model and independently of a specific application in mind. However, (1) it applies a rather complex, node-specific positional encoding, whereas our PE only distinguishes
node depth; (2) the attention is based on the directed transitive closure, while we use reachability and show that this provides better quality; and (3) it's implemented using regular transformers, with runtime complexity  $O(|V|^2 d)$ , while we propose a much more efficient and scalable implementation based on message passing. Altogether, our proposal is simpler and considerably more efficient. In addition, we do not propose a single model but a framework which can be flexibly applied on top of existing graph transformers and thus complement any custom, graph-specific transformer and adapt it to DAGs.
3 Transformers for Directed Acyclic Graphs
Transformers have revolutionized deep learning, in particular sequence learning, and yield promising performance over graphs. Furthermore, we posit that their benefits actually match well the above-mentioned shortcomings of DAG neural networks. Therefore we developed a graph transformer framework tailored to DAGs. See Figure 1 for an overview.
3.1 Attention based on DAG Reachability
In contrast to regular graphs, the partial order in DAGs creates particular relations between connected nodes, in the sense that a given node's predecessors and successors are most likely more important to it than other graph nodes. Note that this intuition is also captured in the processing of DAG neural networks. Hence the reachability relation suggests itself to be exploited in our architecture. We apply it to restrict the receptive field of nodes during attention to their predecessors and successors in the graph. [Vaswani et al., 2017] already mentioned the possibility to use restricted attention in order to reduce complexity and, indeed, our proposal does not only yield an architecture which is biased towards the DAG structure, but additionally a considerably more efficient model.
While graphs to classify are usually of manageable size, graph learning in general may face much larger ones. For this reason, we formulate our model in a more general way, based on a bounded reachability relation, representing the receptive field of each node:  $N_{k}(v) = \{(u,v)\in \leqslant_{k}\} \cup \{(v,u)\in \leqslant_{k}\}$ . We adapt Equation (2) for our reachability-based attention (DAGRA):
$$
\mathrm {A t t e n t i o n} _ {\mathrm {D A G}} \left(x _ {v}\right) = \sum_ {u \in N _ {k} (v)} \frac {\kappa \left(x _ {v} , x _ {u}\right)}{\sum_ {w \in N _ {k} (v)} \kappa \left(x _ {v} , x _ {w}\right)} f \left(x _ {u}\right), \forall v \in V.
$$
The number  $k$  represents a hyperparameter and can be chosen according to the data. In our ablation study (see Section 4),  $k = \infty$  yielded best performance consistently. Observe that this choice of  $k = \infty$  is still very different from both regular GNNs, which usually considerably restrict  $k$ , and regular transformers (Eq. 1), whose receptive field is not restricted at all (i.e., in terms of reachability).
3.2 Positional Encodings based on DAG Depth
As outlined in Section 2, positional encodings have been recognized as important and proven effective for incorporating graph structure into transformers. Observe that the sequential nature of DAGs makes them possess special position information, the depth of a node within the DAG. We propose to include this knowledge in the form of directed acyclic graph positional encodings (DAGPE) as follows, exactly as suggested for the original transformer architecture [Vaswani et al., 2017]:
$$
P E _ {(v, 2 i)} = \sin \left(\frac {\operatorname {d e p t h} (v)}{1 0 0 0 0 ^ {\frac {2 i}{d}}}\right), \quad P E _ {(v, 2 i + 1)} = \cos \left(\frac {\operatorname {d e p t h} (v)}{1 0 0 0 0 ^ {\frac {2 i}{d}}}\right),
$$
where  $d$  is the node feature dimension and  $i$  the index of the dimension under consideration.
DAG Attention. We obtain the following model, combining DAGRA and DAGPE, for  $v \in V$ :
$$
\text {A t t e n t i o n} \left(x _ {v}\right) = \sum_ {u \in N _ {k} (v)} \frac {\kappa \left(x _ {v} + P E _ {v} , x _ {u} + P E _ {u}\right)}{\sum_ {w \in N _ {k} (v)} \kappa \left(x _ {v} + P E _ {v} , x _ {w} + P E _ {w}\right)} f \left(x _ {u}\right). \tag {3}
$$
Our attention incorporates both similarity of node features and of node depths. We argue that these are the most critical aspects of DAGs and our evaluation will show their impact and general effectiveness. In particular, note that most kinds of DAG data (e.g., citation networks) do not require us distinguishing between predecessor or successor nodes of same depth.
3.3 Expressive Power of DAG Attention
Technically, we restrict the attention to reachable nodes and, in this way, obtain considerable efficiency gains. Yet, our architecture is tailored to the special DAG structure, and we can show that this design offers similar expressivity to regular transformers.
It is important to note that in our framework, all nodes directly communicate with at least one source node (i.e., node without predecessors) by the DAG structure. This is specifically the case because we do not restrict the radius  $k$  of the receptive field, but consider  $k = \infty$ . Hence, 2 layers are always enough to establish communication beyond any two nodes that have a common source node. Observe that this is often the case in practice; especially in DAG classification, many kinds of DAGs contain only a single source node (e.g., ogbg-code2 Hu et al. [2020] and NA Zhang et al. [2019]).
For DAGs with  $m$  source nodes we need  $2m$  layers for full communication, if we assume the DAG to be connected. In the latter case, every pair of source nodes has a common successor through which communication can happen. Further, connectedness is a reasonable assumption, otherwise communication is likely not needed in most scenarios.
3.4 Theoretical Intuition
We have shown that our architecture's bias emphasizes DAG relationships while re-directing the remaining relationships in the regular transformer's full attention matrix though the source nodes. This can also be shown to be in line with random walk theory, we draw a connection to PageRank [Brin, 1998, Gasteiger et al.]. Specifically, we consider a PageRank variant that takes the root node into account - personalized PageRank. We define the root node  $x$  via the teleport vector  $i_x$ , which is a one-hot indicator vector. The personalized PageRank can be obtained for node  $x$  using the recurrent equation  $\pi_G(i_x) = (1 - \alpha)A_{rw}\pi_G(i_x) + \alpha i_x$ , with teleport (or restart) probability  $\alpha \in (0,1]$ . Solving this equation, we obtain:  $\pi_G(i_x) = \alpha (I_n - (1 - \alpha)A_{rw})^{-1}i_x$ , where  $A_{rw} = AD^{-1}$ , with  $A$  and  $D$  being the adjacency and the degree matrix, respectively [Gasteiger et al.]. We invert the directions of the edges in  $G$  to create a reverse DAG  $\tilde{G}$ . We can show that, for every node  $x$ , only nodes  $y$  not reachable from  $x$  (i.e.,  $y \notin N_{\infty}(x)$ ) will satisfy that  $\pi_G(i_x)[y] + \pi_{\tilde{G}}(i_x)[y]$  (the  $y$ -th element of  $\pi_{\tilde{G}}(i_x)$ ) equals 0. The proof is provided in Appendix A. This means that for the random walk's limit distribution, the probability of nodes that are not reachable from  $x$  is 0.
We describe two ways of implementing our model, especially DAGRA, based upon transformers and message-passing GNNs, respectively.