Bi-Level Attention Graph Neural Networks
Recent graph neural networks (GNNs) with the attention mechanism have
historically been limited to small-scale homogeneous graphs (HoGs). However,
GNNs handling heterogeneous graphs (HeGs), which contain several entity and
relation types, all have shortcomings in handling attention. Most GNNs that
learn graph attention for HeGs learn either node-level or relation-level
attention, but not both, limiting their ability to predict both important
entities and relations in the HeG. Even the best existing method that learns
both levels of attention has the limitation of assuming graph relations are
independent and that its learned attention disregards this dependency
association. To effectively model both multi-relational and multi-entity
large-scale HeGs, we present Bi-Level Attention Graph Neural Networks (BA-GNN),
scalable neural networks (NNs) that use a novel bi-level graph attention
mechanism. BA-GNN models both node-node and relation-relation interactions in a
personalized way, by hierarchically attending to both types of information from
local neighborhood contexts instead of the global graph context. Rigorous
experiments on seven real-world HeGs show BA-GNN consistently outperforms all
baselines, and demonstrate quality and transferability of its learned
relation-level attention to improve performance of other GNNs.
Roshni G. Iyer, Wei Wang, Yizhou Sun
# Bi-Level Attention Graph Neural Networks

###### Abstract

Recent graph neural networks (GNNs) with the attention mechanism have historically been limited to small-scale homogeneous graphs (HoGs). However, GNNs handling heterogeneous graphs (HeGs), which contain several entity and relation types, all have shortcomings in handling attention. Most GNNs that learn graph attention for HeGs learn either node-level or relation-level attention, but not both, limiting their ability to predict both important entities and relations in the HeG. Even the best existing method that learns both levels of attention has the limitation of assuming graph relations are independent and that its learned attention disregards this dependency association. To effectively model both multi-relational and multi-entity large-scale HeGs, we present Bi-Level Attention Graph Neural Networks (BA-GNN), scalable neural networks (NNs) that use a novel bi-level graph attention mechanism. BA-GNN models both node-node and relation-relation interactions in a personalized way, by hierarchically attending to both types of information from local neighborhood contexts instead of the global graph context. Rigorous experiments on seven real-world HeGs show BA-GNN consistently outperforms all baselines, and demonstrate quality and transferability of its learned relation-level attention to improve performance of other GNNs.

 graph neural networks, representation learning

## I Introduction

Highly multi-relational data are characteristic of real-world HeGs. Relational data in HeGs are defined as triples of form (_h:_head entity, r:_relation, t:_tail entity_), indicating that two entities are connected by a specific relation type. Figure 1 shows a HeG formed by such triples. However, even comprehensive HeGs [1] remain incomplete. Regarding HeGs completion, despite the recent years' research progress in developing GNNs for representation learning in various domains [7, 8, 14] and adapting the successful attention mechanism [17, 18], most GNNs face several challenges. They either are ill-equipped to handle HeGs [9, 18], or do handle HeGs but do not learn graph attention [4, 6, 13, 23], or learn inaccurate graph attention [5, 12, 19, 22].

Considering the GNNs that learn graph attention, their architectures are limited to only one level of attention, either for nodes or relations, but rarely for both, shown in Table I. This is problematic for modeling HeGs which contain several different entity and relation types. Bi-level attention is more powerful in learning compared to uni-level attention, where only one level of attention is learned by the model. Bi-level attention learns attention at different levels of granularity in HeGs which captures more information about graph components than a uni-level attention mechanism is capable of. h:_nn, one of the few models that attempts to use bi-level attention, unsurprisingly falls short of capturing the associations between the node and relation levels in the HeG. First, h:_nn places unnatural assumptions on the data because it treats graph relations as independent from each other, omitting most relation-relation interactions in HeGs. Second, it requires manually chosen meta paths that force many node-node and node-relation interactions to also be left out, and requires domain specific knowledge to compute. Third, h:_nn lacks a general framework for systematically studying bi-level attention.

To address the above challenges, in this paper, we present Bi-Level **A**ttention **G**raph **Ne**ural **N**etworks (BA-GNN) for HeGs. To summarize, our work makes the following contributions:

1. We design a general framework for bi-level attention, and identify challenges of state-of-art NNs for HeGs.
2. We propose BA-GNN to model both multi-relational and multi-entity large-scale HeGs. BA-GNN avoids manually chosen meta paths, and learns personalized graph properties by integrating graph entity/relation types, graph structure, and graph attention using local graph neighborhoods instead of global graph context. BA-GNN improves accuracy of state-of-art GNNs and scales to million-node/edge graphs, like the AM archaeological dataset.
3. To our knowledge, we are the first to propose efficient bi-level attention GNNs that learn from dependency interactions of both nodes/relations and without meta paths.
4. We rigorously experiment on seven real-world HeGs showing BA-GNN consistently outperforms major state-of-art NN groups, and also demonstrate quality and transferability of BA-GNN's attention-induced change in graph structure to enrich other GNNs.

The remainder of this paper is organized as follows. Section II examines preliminaries and related work. Section III presents a general framework for computing bi-level attention, and describes BA-GNN's architecture. Section IV presents experiment results, ablation studies, and case studies of BA

Fig. 1: Partial HeG of AIFB dataset.
GNN models, and section V concludes.

## II Preliminary and Related Work

Here, we introduce HeG concepts and discuss the achievement of various state-of-art NNs, summarized in Table I.

**Definition 1**: _Heterogeneous Graph: We define HeGs as \(\mathcal{G}=(\mathcal{V},\mathcal{E})\) with nodes \(v_{i}\in\mathcal{V}\), and edges \(e_{i,j}\in\mathcal{E}\) connecting source \(v_{i}\) and target \(v_{j}\). Nodes in the HeG are associated with entity types through an entity mapping function \(\Lambda(v_{i}):\mathcal{V}\rightarrow\mathcal{B},\mathcal{B}=\{b|b\in\mathcal{ B}\}\), for entity type \(b\). Edges in the HeG are associated with relation types through a relation mapping function \(\Gamma(e_{i,j}):\mathcal{E}\rightarrow\mathcal{R},\mathcal{R}=\{r|r\in \mathcal{R}\}\) for relation type \(r\). For efficiency in our model, we compute entity and relation mapping functions for node and relation neighborhoods rather than globally such that \(\Lambda(v_{i}):\mathcal{V}_{i}\rightarrow\mathcal{B}_{i},\mathcal{B}_{i} \subset\mathcal{B}\) and \(\Gamma(e_{i,j}):\mathcal{E}_{i}\rightarrow\mathcal{R}_{i},\mathcal{R}_{i} \subset\mathcal{R}\). In this paper, we consider local HeG neighborhoods of a node \(v_{i}\in\mathcal{V}\) consisting of both one-hop nodes, \(\{v_{j}|v_{j}\in\mathcal{V}_{i}\}\) and one-hop relations, \(\{r|r\in\mathcal{R}_{i}\}\)._

**Definition 2**: _Meta Relation: The meta relation for \(e_{i,j}\) between source \(v_{i}\) and target \(v_{j}\) is \((\Lambda(v_{i}),\Gamma(e_{i,j}),\Lambda(v_{j}))\), and \(\Gamma(e_{i,j})^{-1}=\Gamma(e_{j,i})\) is the inverse of \(\Gamma(e_{i,j})\). In this paper, we loosely use the term relation to denote meta relation. Traditional meta paths are a sequence of such meta relations._

**Definition 3**: _Graph Attention: Graph attention enables NNs to learn useful graph representations by selectively attending to different nodes and relations. Multiplicative and additive attention are state-of-art attention mechanisms used in NNs [17, 18], both of which operate on encoder states. Multiplicative attention uses an inner product or cosine similarity of encoder states while additive attention is a linear combination or concatenation of encoder states._

### _GNNs for Homogeneous Graphs_

Successful models in this category, like GAT [18], use attention-based neural architectures for learning representations.

**Graph Attention Networks**GAT [18] are additive attention-based GNNs that effectively leverage graph structure and sparsity to compute a node's attention. GAT models, however, are limited to HoGs and cannot handle HeGs which contain different relations that may have varying levels of importance for different nodes.

### _GNNs for Heterogeneous Graphs_

Successful models (1) leverage different graph relations, like R-GCN, (2) learn bi-level attention, like HAN, and (3) learn multiplicative attention, like transformer-based NNs.

**Relational Graph Convolutional Networks**R-GCNs [13] extend GCNs and GAT, which operate on local graph neighborhoods of HoGs, to operate on multi-relational graphs by distinguishing nodes by relation type. R-GCNs, however, treat all relation-specific nodes as equally important. Further, R-GCNs do not utilize graph attention as they are limited to directly learning from weight parameters.

**Heterogeneous Graph Attention Networks**To address limitations of the above models, HAN integrates bi-level attention, which learns node- and relation-level attention, with GNNs to learn node embeddings. However, HAN uses a global learnable weight vector lacking local inter-relation comparison. Besides, HAN uses pre-defined metapaths which are computationally expensive to design and compute, and result in sub-optimal graph components learned by the model.

**Transformer**transformer models [17], although successful in natural language processing for small text sequences, have limitations for multi-relational and multi-entity HeGs. This is because transformer attends to all other tokens in the sequence, making it infeasible for large-scale input. While recent works extend transformer-like attention to other graph domains, they have limitations, shown in Table I.

## III BA-GNN Architecture

We design a general bi-level attention framework for computing hierarchical attention and then discuss BA-GNN's architecture. Source code and data are at: [https://github.com/roshnigyier/BA-GNN](https://github.com/roshnigyier/BA-GNN). \(\mathrm{READ.md}\) details dataset properties, data splits, and hyperparameters of BA-GNN models.

### _General Bi-Level Attention Framework_

Bi-level attention in HeGs incorporates interactions between relation-specific nodes for learning lower level attention which informs the higher level attention that captures inter-relation interactions. In this way, bi-level attention jointly attends to node-node, relation-relation and node-relation interactions to collectively produce a representative node embedding. Uni-level attention models omit these critical graph interactions and ability for the two-levels of attention to jointly inform each other. Eq. 1 describes the general bi-level attention framework to compute embeddings for each \(v_{i}\) in the HeG:

\[\widetilde{\mathbf{h}}_{i}^{(l+1)} =\mathbf{HigherAtt}\big{(}\mathbf{LowerAtt}(\cdot),\mathcal{R}, \widetilde{\mathbf{h}}^{(l)}\big{)} \tag{1}\] \[=\mathrm{AGG}\Bigg{(}\bigg{\{}\mathbf{f}_{v}\Big{(}\big{(} \mathbf{LowerAtt}(\cdot)\big{|}r\in\mathcal{R},\widetilde{\mathbf{h}}^{(l)} \big{)}\Big{)}\bigg{|}r\in\mathcal{R}\bigg{\}}\Bigg{)},\]

\begin{table}
\begin{tabular}{|c|c|c|c|c|c|c|c|c|} \hline Type & Model & [A] & [B] & [C] & [D] & [T] & [E] & [F] & [G] \\ \hline – & Transformer [17] & ✓ & ✓ & ✓ & ✗ & ✓ & ✓ & – \\ \hline \multirow{4}{*}{(1)} & TransEe [2] & ✓ & ✓ & ✗ & ✗ & ✗ & ✓ & ✗ \\ \cline{2-10}  & HQE [10] & ✓ & ✓ & ✗ & ✗ & ✗ & ✓ & ✗ \\ \cline{2-10}  & DistMIT [21] & ✓ & ✓ & ✗ & ✗ & ✗ & ✓ & ✗ \\ \cline{2-10}  & Complex [16] & ✓ & ✓ & ✗ & ✗ & ✗ & ✓ & ✗ \\ \hline \multirow{2}{*}{(2)} & GCN [9] & ✗ & ✓ & ✗ & ✗ & ✗ & ✓ & ✗ \\ \cline{2-10}  & GAT [18] & ✗ & ✓ & ✓ & ✗ & ✓ & ✓ & ✗ \\ \hline \multirow{4}{*}{(3B)} & metaframe\({}_{\mathcal{V}}\) & ✓ & ✓ & ✗ & ✗ & ✗ & ✗ \\ \cline{2-10}  & HEEc [14] & ✓ & ✓ & ✗ & ✗ & ✗ & ✗ \\ \cline{2-10}  & HDEV [6] & ✓ & ✓ & ✗ & ✗ & ✗ & ✗ \\ \cline{2-10}  & HGEAN [7] & ✓ & ✓ & ✗ & ✗ & ✓ & ✗ \\ \cline{2-10}  & TemporalGAT [5] & ✗ & ✓ & ✓ & ✓ & ✓ & ✓ & ✗ \\ \cline{2-10}  & HGTNN [23] & ✓ & ✓ & ✗ & ✗ & ✗ & ✓ & ✗ \\ \hline \multirow{4}{*}{(3B)} & H-GCN [13] & ✓ & ✓ & ✓ & ✗ & ✗ & ✓ & ✗ \\ \cline{2-10}  & HAN [19] & ✓ & ✓ & ✓ & ✓ & ✗ & ✗ & ✗ \\ \hline \multirow{4}{*}{(3B)} & DySAT [12] & ✓ & ✗ & ✓ & ✗ & ✓ & ✓ & ✓ \\ \cline{2-10}  & TGM [20] & ✗ & ✗ & ✓ & ✗ & ✓ & ✓ & ✓ \\ \cline{2-10}  & HGT [8] & ✓ & ✓ & ✓ & ✗ & ✓ & ✓ & ✓ \\ \cline{2-10}  & GTN [22] & ✓ & ✓ & ✓ & ✗ & ✓ & ✓ & ✓ \\ \hline \multirow{4}{*}{(3B)} & BA-GNN (ours) & ✓ & ✓ & ✓ & ✓ & ✓ & ✓ & ✓ \\ \cline{2-10}  & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} \\ \cline{2-10}  & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} \\ \cline{2-10}  & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} \\ \cline{2-10}  & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} \\ \cline{2-10}  & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} \\ \cline{2-10}  & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} \\ \cline{2-10}  & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} & \multicolumn{10}{c}{} \\ \hline \multirow{4}{*}{(3B)} & HGT [20] & ✗ & ✗ & ✓ & ✗ & ✓ & ✓
\[\mathbf{LowerAtt}(\cdot) =\mathrm{AGG}\big{(}\big{\{}\mathbf{g}_{\gamma}(e_{i,j}|r,\widetilde{ \mathbf{h}}^{(l)})\big{|}v_{j}\in N_{i}^{r}\big{\}}\big{)}, \tag{2}\] \[\mathcal{R} =\{\Gamma(e_{i,j})|v_{j}\in\mathcal{V}\}, \tag{3}\]

where \(\mathbf{h}_{i}\) and \(\widetilde{\mathbf{h}}_{i}\) are the initial and projected node features respectively, \(\mathcal{R}_{i}\) is the relation set on the edge of \(v_{i}\), and \(\mathbf{g}_{\gamma}(\cdot)\) is a vector-output function of the node-level attention, \(\gamma\), that provides a relation-specific embedding summary using learned node representations from the previous layer, \(\widetilde{\mathbf{h}}^{(l)}\), which is aggregated, \(\mathrm{AGG}(\cdot)\), over edges \(e_{i,j}\) in the relation-specific neighborhood context of \(v_{j}\in N_{i}^{r}\). \(\mathbf{f}_{\psi}(\cdot)\) is a vector-output function of the relation-level attention, \(\psi\), that are attended relation-specific local context embeddings, \(\mathbf{LowerAtt}(\cdot)\), which are aggregated over relations in the neighborhood context to form the layer's final node embedding, \(\widetilde{\mathbf{h}}_{i}^{(l+1)}\).

In Sections III-B and III-C, we propose a novel semi-supervised attention-based GCN model, BA-gnN, for multi-relational and multi-entity HeGs. BA-gnN performs attention aggregation on node and _relation_ levels, rather than nodes and edges. For node-level attention, attention is placed on the _edges_ of neighbor nodes. For relation-level attention, attention is placed on the relations, which are formed by _grouping_ edges by relation type. In this way, our model uses a hierarchical attention mechanism. The higher-order graph considers relation-type-specific edge groups, and the lower-order graph considers nodes in their local contexts. Figure 2 summarizes BA-gnN's attention mechanism. BA-gnN models use \(L\) stacked layers, each of which is defined through Eq. 1, and for further efficiency, all nodes and relations are restricted to their neighborhoods. Model input can be chosen as pre-defined features or as a unique one-hot vector for each node.

### _Node-level Attention_

Node-level attention distinguishes different roles of nodes in the neighborhood context for learning relation-specific node embeddings. As node-level attentions are target-node-specific, they are different for different target nodes. Our model learns node embeddings such that it intrinsically captures graph attributes and structure in its neighborhood. In HeGs, neighbor nodes may belong to different feature spaces due to different node types, so we utilize an entity-specific type transformation matrix, \(\boldsymbol{\mathcal{T}}_{\Lambda(v_{i})}\), to project all node features to the same space through \(\widetilde{\mathbf{h}}_{i}=\boldsymbol{\mathcal{T}}_{\Lambda(v_{i})}\cdot \mathbf{h}_{i}\). \(\boldsymbol{\mathcal{T}}_{\Lambda(v_{i})}\in\mathbb{R}^{d\times|\mathcal{R}_{ i}|},\mathbf{h}_{i}\in\mathbb{R}^{|\mathcal{R}_{i}|}\) if the initial features are chosen to be a one-hot vector of dimension \(d\), where \(\widetilde{\mathbf{h}}_{i}\) is continuously updated to learn the final embedding.

BA-gnN's node-level attention uses additive attention inspired by GAT, discussed in Section II, but overcomes GAT's limitation by extending the attention to HeGs. GAT performs projections that do not consider the different relation types in the HeG. We address this by using a learnable relation-specific attention vector, \(\mathbf{a}_{r}^{(l)}\in\mathbb{R}^{2d}\). For a specific relation \(r\), the attention is shared for all node pairs, so that each node is influenced by its neighborhood context. The attention is also asymmetric since the importance of \(v_{j}\) to \(v_{i}\) may be different from the importance of \(v_{i}\) to \(v_{j}\). We compute relation-specific node-level attention at layer \(l\) as follows, with a \(\mathrm{softmax}(\cdot)\) activation applied to normalize each node-pair attention weight and where \(v_{j},v_{k}\in N_{i}^{r}\) and \(\mathbf{x}^{T(l)}\) is transpose of \(\mathbf{x}\) at layer

\[l\text{: }\gamma_{i,j}^{(l),r}\ =\ \frac{\exp\!\left(\mathrm{LeakyReLU}\big{(} \mathbf{a}_{r}^{r(l)}\big{[}\widetilde{\mathbf{h}}_{i}^{(l)}\big{|}\big{|} \widetilde{\mathbf{h}}_{i}^{(l)}\big{|}\big{)}\right)}{\sum_{v_{k}\in N_{i}^{r}} \exp\!\left(\mathrm{LeakyReLU}\big{(}\mathbf{a}_{r}^{r(l)}\big{[}\widetilde{ \mathbf{h}}_{i}^{(l)}\big{|}\big{|}\widetilde{\mathbf{h}}_{i}^{(l)}\big{|} \big{|}\widetilde{\mathbf{h}}_{i}^{(l)}\big{]}\big{)}\right)},\]

where \(\mathbf{a}_{r}^{(l)}\) attends over the concatenated, \(||\), node features of \(v_{i}\) and \(v_{j}\) with applied \(\mathrm{LeakyReLU}(\cdot)\) and \(\mathrm{softmax}(\cdot)\) activations. By restricting the attention to within the relation-specific local context of nodes, sparsity structural information is injected into the model through adjacency-masked attention layers.

Node \(v_{i}\)'s relation-specific embedding, \(\mathbf{z}_{i}^{(l),r}\), can then be learned with \(\mathrm{AGG}(\cdot)\) from Eq. 1 being a weighted summation of the neighbor's projected features as follows:

\[\mathbf{z}_{i}^{(l),r} =\mathbf{LowerAtt}(\cdot) \tag{4}\] \[=\mathrm{AGG}\big{(}\big{\{}\mathbf{g}_{\gamma}(e_{i,j}|r, \widetilde{\mathbf{h}}^{(l)})\big{|}v_{j}\in N_{i}^{r}\big{\}}\big{)}\] \[=\sum_{v_{j}\in N_{i}^{r}}\Big{[}\mathbf{g}_{\gamma}(e_{i,j}|r, \widetilde{\mathbf{h}}^{(l)})\Big{]}\] \[=\sum_{v_{j}\in N_{i}^{r}}\Big{[}\gamma_{i,j}^{(l),r}\widetilde{ \mathbf{h}}_{j}^{(l)}\Big{]},\]

where \(\mathbf{z}_{i}^{(l),r}\) provides a summary of relation \(r\) for \(v_{i}\) at layer \(l\). We also add skip-connections for corresponding \(\mathbf{LowerAtt}(\cdot)\) from the previous layer, \(l-1\), to preserve learned node-level representations as the depth of the NN is extended.

### _Relation-level Attention_

Relation-level attention distinguishes roles of different relations in the neighborhood context for learning more comprehensive node embeddings. In HeGs, different relations may play different roles of importance for \(v_{i}\), in addition to \(v_{i}\)'s relation-specific neighbor nodes. So, we learn relation-level attention to better fuse \(v_{i}\)'s relation-specific node embeddings. One could design a simple node-relation attention mechanism to encode the effect of relations between nodes, but this would fail to capture relation-relation dependencies hidden in HeGs.

We address transformer's inefficiency for large-scale HeGs through an approximation technique by sampling relation-specific node embeddings from the local graph context. Further, instead of using the same single set of projections for all words, we enable each relation-specific embedding to learn a distinct set of personalized projection weights, while maximizing parameter sharing. This technique captures unique relation-dependent characteristics such that each relation-specific node embedding is also influenced by its local relation context.

Fig. 2: Bi-level attention visualization. (a) Node-level aggregating: A node’s features is a weighted combination of its prior layer’s relation-specific embeddings, \(\mathbf{z}_{i}^{r}\). (b) Relation-level aggregating: Relation-level attention is learned via multiplicative attention using neighborhood similarity to determine relative relation importance.
Node \(v_{i}\)'s relation-specific transformer-based query \(\mathbf{q}_{r,i}^{(l)}\), key \(\mathbf{k}_{r,i}^{(l)}\), and value \(\mathbf{v}_{r,i}^{(l)}\) vectors are computed as follows: \(\mathbf{q}_{r,i}^{(l)};\mathbf{k}_{r,i}^{(l)};\mathbf{v}_{r,i}^{(l)}=\mathbf{W} _{1,r}\mathbf{z}_{i}^{(l),r};\mathbf{W}_{2,r}\mathbf{z}_{i}^{(l),r};\mathbf{W}_{ 3,r}\mathbf{z}_{i}^{(l),r}\), such that \(\mathbf{z}_{i}^{(l),r}\) is projected onto the learnable weight matrices \(\mathbf{W}_{1,r},\mathbf{W}_{2,r},\mathbf{W}_{3,r}\in\mathbb{R}^{d\times d}\). The relation-level attention for relations \((r,r^{\prime})\) are computed by iterating over all possible relation pairs in the neighborhood context, \(r,r^{\prime}\in\mathcal{R}_{i}\), where \(\mathcal{R}_{i}=\{\Gamma(e_{i,j})|v_{j}\in\mathcal{V}_{i}\}\). The importance of relation \(r^{\prime}\) of node \(v_{i}\) is as follows, with relation similarity being captured through \(\psi_{i}^{(l),r,r^{\prime}}=\mathrm{softmax}(\mathbf{q}_{r,i}^{T^{(l)}}\mathbf{ k}_{\mathcal{V}_{r,i}}^{(l)})\), where the more similar \(r^{\prime}\) is to \(r\), the greater the attention weights of \(r^{\prime}\), which results in more contribution of \(r^{\prime}\)'s embedding to \(v_{i}\)'s final embedding. A \(\mathrm{softmax}(\cdot)\) activation is then applied to normalize each relation pair's attention weight.

A node's relation-specific embedding is then informed by a weighted summation of its similarity to other local context relations, \(\psi_{i}^{(l),r,r^{\prime}}\). To reduce information loss, we add a self-connection of a special relation type per node, which is projected onto \(\mathbf{W}_{i}\), and aggregated to the attended relation-specific embedding, \(\psi_{i}^{(l),r,r^{\prime}}\mathbf{v}_{r^{\prime},i}^{(l)}\). Lastly, a \(\mathrm{ReLU}(\cdot)\) activation is applied to get the overall relation-specific embedding, \(\mathbf{\delta}_{i}^{(l),r}=\mathrm{ReLU}(\sum_{r^{\prime}\in\mathcal{R}_{i}}\psi_ {i}^{(l),r,r^{\prime}}\mathbf{v}_{r^{\prime},i}^{(l)}+\mathbf{W}_{i}\widetilde {\mathbf{h}}_{i}^{(l)})\).

Node \(v_{i}\)'s final embedding is learned with \(\mathrm{AGG}(\cdot)\) in Eq. 1 being a summation of all attended relation-specific embeddings through iteration of neighborhood relations, \(r\in R_{i}\). We also apply multi-head attention of \(S=\{1,...,K\}\) heads to allow BA-GNN to jointly attend to different representation subspaces of nodes/relations, with aggregation, \(\mathrm{AGG}(\cdot)\), via averaging:

\[\widetilde{\mathbf{h}}_{i}^{(l+1)} =\mathbf{HigherAtt}(\mathbf{LowerAtt}(\cdot),\mathcal{R}, \widetilde{\mathbf{h}}^{(l)}) \tag{5}\] \[=\frac{1}{K}\sum_{k=1}^{K}\sum_{r\in\mathcal{R}_{i}}\Big{[} \mathbf{f}_{\mathbf{v}}\Big{(}\big{\{}\mathbf{LowerAtt}(\cdot)|r\in\mathcal{R },\widetilde{\mathbf{h}}^{(l)}\big{\}}\Big{)}\Big{]}\] \[=\mathrm{AGG}\Big{(}\big{\{}\sum_{r\in\mathcal{R}_{i}}\big{[} \mathbf{\delta}_{i}^{(l),r}\big{]}|k\in S\big{\}}\Big{)}.\]

We also add skip-connections for corresponding \(\mathbf{HigherAtt}(\cdot)\) from the previous layer, \(l-1\), to preserve learned higher-order relation-level representations as depth of the NN is extended.

The final representation of a node at layer \((l+1)\) is:

\[\widetilde{\mathbf{h}}_{i}^{(l+1)}=\mathrm{AGG}\bigg{(}\Big{\{} \sum_{r\in\mathcal{R}_{i}}\mathrm{ReLU}(\] \[\sum_{r^{\prime}\in\mathcal{R}_{i}}\mathrm{softmax}(\mathbf{q}_{r,i}^{T^{(l)}}\mathbf{k}_{r^{\prime},i}^{(l)})\mathbf{v}_{r^{\prime},i}^{(l)}+ \mathbf{W}_{i}\widetilde{\mathbf{h}}_{i}^{(l)})\Big{|}k\in S\Big{\}}\bigg{)}. \tag{6}\]

### _Analysis of Proposed Attention_

We use multiplicative attention at the relation level, instead of additive attention, because learning attention through a concatenation of features does not compute feature similarity which inner product operations capture. Since relation features are characterized by single attribute relation types, a relation's scaling can be directly determined by its feature similarity to other relations in its neighborhood. This is unlike node-level attention, where node features may have several attributes, making it more difficult to learn latent similarity of nodes through direct feature comparison such as through inner product computation. Rigorous evaluation of mixed combinations of additive/multiplicative attention shown in experiments further support our bi-level attention combination choice.

## IV Experiments

In this section, we evaluate BA-GNN on seven large-scale heterogeneous datasets (HDs). We conduct experiments on node classification and link prediction using Pytorch Geometric and Deep Graph Library frameworks on an Nvidia Tesla V100 GPU cluster, and report model test accuracies.

DatasetsWe evaluate on benchmark Resource Description Framework (RDF) format datasets [11] for node classification: AIFB, MUTAG, BGS, and AM. For link prediction, we evaluate on FB15k [15], WN18 [3], and FB15k-237 [15].

### _Node Classification_

Node classification is the semi-supervised classification of nodes to entity types. For evaluation consistency against primary baseline models, we implement BA-GNN with \(L=2\) and where the output of the final layer uses a \(\mathrm{softmax}(\cdot)\) activation per node. Our model follows the same node classification evaluation procedure as [13], using cross-entropy loss with parameters learned from the Adam Optimizer.

BaselinesTable I summarizes our baselines. To adapt the models to our problem setting of multi-relational, static HeGs, we made the following modifications. For GAT[18] and GCN[9], we omit HeG relations. For TemporalGAT[5], we omit the temporal convolutional network used for temporal interactions. For HetGNN[23], we consider the neighbors to be the entire set of neighbor nodes and relations. For DySAT[12], we omit temporal attention. For TGAT[20], we omit functional time encoding. For HGT[8], we omit relative temporal encoding.

ResultsExperiment results are in Table II. Results show BA-GNN significantly and consistently outperforms all baselines for all tasks on all datasets. For example, on AIFB, MUTAG, BGS, and AM, against the most competitive NNs per category, BA-GNN achieves relative performance gains of up to 22%, 25%, 20%, and 24% respectively, and overall performance gains of up to 32%, 33%, 36%, and 35% respectively. Further, the Welch t-test of unequal variance shows that BA-GNN's relative performance compared to each model per dataset is statistically significant to be greater, with \(\mathrm{p}\)-\(\mathrm{value}<0.001\).
indicates that BA-GNN's relation-level attention is more effective across the different data domains and that its personalized attention to local graph contexts yields performance gain.

### _Link Prediction_

Link prediction involves assigning confidence scores to HeG triples to determine how likely predicted edges belong to true relations. Our models follow the same evaluation framework as [2] and [13] using negative sampling and cross-entropy loss with parameters learned from the Adam Optimizer. We use evaluation metrics of mean reciprocal rank (MRR) and Hits @ n, in raw and filtered settings. The same number of negative samples, \(w=1\), are used to make datasets comparable.

_Baselines_: We evaluate standalone GNNs (BA-GNN, R-GCN), KGE models, and GNN+KGE autoencoder models using the same setup procedure as [2] and [13]. The autoencoder models include: BA-GNN\({}_{x}\) and R-GCN\({}_{y}\) where \(x\), \(y\) are Transe (\(T\)), Hole (\(H\)), DistMult (\(D\)), and ComplEx (\(C\)).

_Results and Ablation Studies_: Experiment results are in Table III. Results show that the best BA-GNN models outperform R-GCN models on all datasets for all metrics of both tasks of MRR and Hits @ n = 1, 3, 10. We observe that BA-GNN outperforms R-GCN when comparing standalone models. Further, results show that autoencoder models outperform each of GNN and KGE standalone models, showing that GNNs

\begin{table}
\begin{tabular}{c|c|c c c c c c|c c c c} \hline \hline \multirow{2}{*}{Type} & \multicolumn{2}{c}{Model} & \multicolumn{2}{c}{AlfB} & \multicolumn{2}{c}{MUTAG} & \multicolumn{2}{c}{BGS} & \multicolumn{2}{c}{AM} & BA-GNN (\% improve) \\ \hline \multirow{4}{*}{(1) Non-GNN-based KGE models for HeGs} & \multicolumn{2}{c}{Transe [2]} & 66.89 \(\pm\) 1.26 & 54.41 \(\pm\) 0.73 & 58.46 \(\pm\) 0.38 & 61.23 \(\pm\) 0.57 & +\{32.05, 36.29\} \\  & \multicolumn{2}{c}{Hole [10]} & 67.98 \(\pm\) 0.42 & **63.17**\(\pm\) 0.26 & 72.74 \(\pm\) 0.22 & 65.25 \(\pm\) 0.38 & +\{22.01, 31.43\} \\  & \multicolumn{2}{c}{DistMult [21]} & 73.64 \(\pm\) 0.10 & 62.06 \(\pm\) 0.18 & 68.19 \(\pm\) 0.13 & 69.12 \(\pm\) 0.24 & +\{25.3, 27.56\} \\  & ComplEx [16] & **77.24**\(\pm\) 0.15 & 62.38 \(\pm\) 0.22 & **74.72**\(\pm\) 0.18 & **72.39**\(\pm\) 0.16 & +\{20.03, 25.43\} \\ \hline \multirow{4}{*}{(2) GNNs for HGGs} & GCN [9] & 91.99 \(\pm\) 0.21 & **67.02**\(\pm\) 0.08 & **78.74**\(\pm\) 0.16 & 86.82 \(\pm\) 0.60 & +\{69.5, 20.79\} \\  & GAT [18] & **92.50**\(\pm\) 0.29 & 66.18 \(\pm\) 0.10 & 77.93 \(\pm\) 0.17 & **88.52**\(\pm\) 1.65 & +\{64.4, 21.63\} \\ \hline \multirow{4}{*}{(3) Non-transformer-based GNNs for HeGs} & \multicolumn{2}{c}{metapatra2vec [4]} & 89.52 \(\pm\) 0.12 & 66.04 \(\pm\) 0.27 & 78.34 \(\pm\) 0.13 & 85.48 \(\pm\) 0.11 & +\{9.42, 21.77\} \\  & HERec [14] & 91.03 \(\pm\) 0.15 & 66.96 \(\pm\) 0.18 & 79.36 \(\pm\) 0.25 & 85.98 \(\pm\) 0.07 & +\{7.91, 20.85\} \\  & HN2vec [6] & 91.63 \(\pm\) 0.17 & 66.29 \(\pm\) 0.14 & 79.01 \(\pm\) 0.12 & 86.22 \(\pm\) 0.21 & +\{73.1, 21.52\} \\  & HEGAN [7] & 92.33 \(\pm\) 0.13 & 68.07 \(\pm\) 0.08 & 81.60 \(\pm\) 0.27 & 86.79 \(\pm\) 0.14 & +\{66.61, 19.74\} \\  & TemporalGAT [5] & 93.42 \(\pm\) 0.11 & 66.88 \(\pm\) 0.24 & 79.14 \(\pm\) 0.13 & 89.10 \(\pm\) 0.13 & +\{55.2, 20.93\} \\  & HetGNN [23] & 95.18 \(\pm\) 0.16 & 75.64 \(\pm\) 0.09 & 82.05 \(\pm\) 0.25 & 89.67 \(\pm\) 0.05 & +\{3.76, 12.70\} \\  & R-GCN [13] & 95.31 \(\pm\) 0.62 & 73.23 \(\pm\) 0.48 & 83.10 \(\pm\) 0.80 & 89.29 \(\pm\) 0.35 & +\{3.63, 14.58\} \\  & HAN [19] & **96.25**\(\pm\) 0.12 & **76.46**\(\pm\) 0.07 & **86.84**\(\pm\) 0.21 & **90.68**\(\pm\) 0.23 & +\{2.69, 11.35\} \\ \hline \multirow{4}{*}{(3) Transformer-based GNNs for HeGs} & \multicolumn{2}{c}{DYSAT [12]} & 92.64 \(\pm\) 0.21 & 66.57 \(\pm\) 0.05 & 78.02 \(\pm\) 0.19 & 88.80 \(\pm\) 0.15 & +\{63.0, 21.24\} \\  & TGAT [20] & 92.84 \(\pm\) 0.14 & 67.19 \(\pm\) 0.21 & 78.35 \(\pm\) 0.15 & 89.43 \(\pm\) 0.28 & +\{6.10, 20.62\} \\  & HGT [8] & 95.97 \(\pm\) 0.15 & **76.84**\(\pm\) 0.12 & **86.01**\(\pm\) 0.18 & 90.33 \(\pm\) 0.13 & +\{2.97, 10.97\} \\  & GTN [22] & **96.04**\(\pm\) 0.17 & 76.32 \(\pm\) 0.12 & 85.38 \(\pm\) 0.24 & **90.56**\(\pm\) 0.10 & +\{2.90, 11.49\} \\ \hline \multirow{4}{*}{BA-GNN variants} & BA-GNN-node & 95.46 \(\pm\) 0.13 & 73.19 \(\pm\) 0.25 & 84.23 \(\pm\) 0.22 & 89.45 \(\pm\) 0.02 & - \\  & BA-GNN-relation & 95.28 \(\pm\) 0.23 & 76.17 \(\pm\) 0.22 & 85.43 \(\pm\) 0.34 & 90.52 \(\pm\) 0.18 & - \\  & BA-GNN(node)+HAN(rel.) & 96.30 \(\pm\) 0.09 & 76.48 \(\pm\) 0.04 & 86.80 \(\pm\) 0.24 & 90.67 \(\pm\) 0.35 & - \\  & BA-GNN(\(\text{M}_{\text{\tiny{\text{score}}}}/\text{A}_{\text{\tiny{\text{all. }}}}\)) & 96.02 \(\pm\) 0.13 & 76.20 \(\pm\) 0.09 & 86.90 \(\pm\) 0.17 & 90.54 \(\pm\) 0.11 & - \\  & BA-GNN(\(\text{M}_{\text{\tiny{\text{score}}}}/\text{A}_{\text{\tiny{\text{all. }}}}\)) & 96.38 \(\pm\) 0.14 & 76.61 \(\pm\) 0.22 & 87.00 \(\pm\) 0.09 & 90.73 \(\pm\) 0.05 & - \\  & BA-GNN(M\({}_{\text{\tiny{\text{score}}}}/\text{A}_{\text{\tiny{\text{all. }}}}\)) & 96.44 \(\pm\) 0.16 & 77.01 \(\pm\) 0.07 & 86.92 \(\pm\) 0.12 & 90.78 \(\pm\) 0.08 & - \\  & BA-GNN(ours) & **98.94**\(\pm\) 0.13 & **87.81**\(\pm\) 0.11 & **94.75**\(\pm\) 0.08 & **96.68**\(\pm\) 0.14 & - \\ \hline \multirow{4}{*
and KGE models can each be benefited by their joint learning. Results also show that BA-GNN autoencoders outperform R-GCN autoencoders on all datasets for all tasks and metrics.

### _Case Study_

We conduct experiments to determine the quality of relation-level attention and graph-structure of BA-GNN. We modify the AM dataset to contain the following relation types, each with cummulative 10% splits: (1) relations randomly selected, (2) relations with the highest relation-level attention weights from BA-GNN, and (3) relations with the lowest relation-level attention weights from BA-GNN. Experiment figures on node classification for HAN and BA-GNN models are in Figure 3(a).

(2)'s graph structure yields the highest test accuracy on all splits of AM compared to (1) or (3), while (3) yields the lowest test accuracy. (1) is as expected in between test accuracies of (2) and (3). Models that do not learn relation-level attention (BA-GNN-node) still benefit from the graph structure identified by (2). This suggests that BA-GNN's relation-level attention can selectively identify important graph components and that its learning of graph structure can enhance other leading GNNs.

### _Attention Visualization_

We randomly sample two nodes belonging to entity _Persons_ on AIFB and plot its learned relation-level attention weights from layer \(l=L\) using heat maps, seen in Figures 3(b) and (c). The corresponding partial graphs of _person 1_ and _person 2_ are in Figure 1. In Figure 3(b), _author\({}^{-1}\)_ and _member\({}^{-1}\)_ have high attention to _is_worked_on_by_ because a person is likely to have publications and research affiliations in their research area. _name_of_ has high attention to _homepage_of_, observed in both Figures 3(b) and (c), because a homepage may directly contain personal identifying information. In Figure 3(c), _head\({}^{-1}\)_ and _member\({}^{-1}\)_ have high attention to each other, since head of a research group is a member. Further, members of the research group are likely to work on the group's projects and focus on a particular research domain, explaining why _works_at_project\({}^{-1}\)_ has higher attention to _head\({}^{-1}\)_ and _member\({}^{-1}\)_, and _works_at_project\({}^{-1}\)_ and members of the research group also have higher attention to _is_worked_on_by_.

## V Conclusion

We propose Bi-Level Attention Graph Neural Networks (BA-GNN) for modeling multi-entity and multi-relational large-scale heterogeneous graphs (HeGs), via entity type and meta relation information to learn graph structure and properties. Further, BA-GNN distinguishes nodes and relations using a novel smart-sampling bi-level attention mechanism to guide the model when aggregating features in graph neighborhoods. We conduct extensive experiments on seven real-world heterogeneous datasets, and show BA-GNN learns effective and efficient embeddings. We observe that BA-GNN outperforms all state-of-art GNN baselines on various information recovery tasks.

## VI Acknowledgements

This work was supported by NSF 1705169, 1829071, 1937599, 2031187, 2106859; NIH R35-HL135772, NIBIB R01-EB027650; DARPA HR00112090027; Okawa Foundation Grant; and Amazon Research Awards. We also thank Kai-Wei Chang and Yunsheng Bai for helpful discussions.
