Information Flow in Graph Neural Networks: A Clinical Triage Use Case
Graph Neural Networks (GNNs) have gained popularity in healthcare and other
domains due to their ability to process multi-modal and multi-relational
graphs. However, efficient training of GNNs remains challenging, with several
open research questions. In this paper, we investigate how the flow of
embedding information within GNNs affects the prediction of links in Knowledge
Graphs (KGs). Specifically, we propose a mathematical model that decouples the
GNN connectivity from the connectivity of the graph data and evaluate the
performance of GNNs in a clinical triage use case. Our results demonstrate that
incorporating domain knowledge into the GNN connectivity leads to better
performance than using the same connectivity as the KG or allowing
unconstrained embedding propagation. Moreover, we show that negative edges play
a crucial role in achieving good predictions, and that using too many GNN
layers can degrade performance.
Víctor Valls, Mykhaylo Zayats, Alessandra Pascale
# Information Flow in Graph Neural Networks:

###### Abstract

Graph Neural Networks (GNNs) have gained popularity in healthcare and other domains due to their ability to process multi-modal and multi-relational graphs. However, efficient training of GNNs remains challenging, with several open research questions. In this paper, we investigate how the flow of embedding information within GNNs affects the prediction of links in Knowledge Graphs (KGs). Specifically, we propose a mathematical model that decouples the GNN connectivity from the connectivity of the graph data and evaluate the performance of GNNs in a clinical triage use case. Our results demonstrate that incorporating domain knowledge into the GNN connectivity leads to better performance than using the same connectivity as the KG or allowing unconstrained embedding propagation. Moreover, we show that negative edges play a crucial role in achieving good predictions, and that using too many GNN layers can degrade performance.

## I Introduction

Machine learning algorithms were originally designed to work with data that can be represented as a sequence (e.g., text) or grid (e.g., images). However, these data structures are inadequate for modeling the data of modern applications. For instance, in digital healthcare, a patient's electronic health record (EHR) can include numerous elements, such as demographic information, medical and medication history, laboratory results, etc. One way to model data with arbitrary structure is to use a Knowledge Graph (KG): a graph where nodes represent pieces of information and the edges indicate how the information pieces relate to one another.

Many learning problems on KGs can be cast as predicting links between nodes. Fig. 1 shows an example of a chronic disease prediction problem on a KG. The patient (IDXA98) is connected to its EHR (with the patient's information such as name, dob, medical conditions, etc.), and the goal is to predict to which chronic disease nodes the patient is connected (colored arrow in Fig. 1). Making such a prediction is possible by analyzing the EHR of other patients with _known_ chronic diseases. Another example of a link prediction problem on a KG is when a patient is already diagnosed with a disease (e.g., SARS-CoV-2), and the goal is to find the most effective drug/treatment to help the patient recover [1].

While there exist several methods for predicting edges on graphs [2, 3, 4], Graph Neural Networks (GNNs) have emerged as one of the most widely used techniques. In brief, GNNs were developed in parallel by two communities: _geometric deep learning_ and _graph representation learning_. The first community focused on applying neural networks for prediction tasks on graph data, while the latter community concentrated on learning low-dimensional vector representations of the nodes and edges in the graph [5]. Current GNNs approaches combine the efforts of both communities and include important extensions such as the ability to handle multi-modal and multi-relational data [6, 7].

GNNs' ability to process multi-modal and multi-relational graphs boosted their popularity in various domains, including healthcare. Some applications include the prediction of hospital readmission [8, 9], chronic diseases [10], and ICU mortality [11, 12]. However, despite their popularity, efficient training of GNNs remains challenging. Previous work has primarily focused on designing new architectures for embedding aggregation [6], with little emphasis on how embedding information should be exchanged in the network. For instance, the works in [6, 13] suggest that the GNNs' connectivity--which determines how nodes receive information from their neighbors--should align with the connectivity of the KG. However, there are cases where it could be advantageous to explore more complex GNN connectivities that are tailored to the specific task at hand. For example, exchanging embeddings based on the KG connectivity depicted in Fig. 1 precludes medical conditions from influencing patient embeddings. Yet, incorporating such interactions can be beneficial in tasks such as chronic disease prediction, where it is essential to capture the patients' existing medical conditions (e.g., hyperglycemia for predicting diabetes) in their embeddings.

In this paper, we investigate how the flow of embeddings within a GNN affects the prediction of links in a clinical triage use case. The paper makes the following contributions:

1. We present a mathematical model for predicting links

Fig. 1: Example of a Knowledge Graph (KG) representing the medical record of a patient (Orla). The gray boxes represent the nodes, and the arrows the edges. The dashed and colored arrow with the question mark is the link we would like to predict.
on KGs with GNNs, where we cast the prediction task as an optimization problem and leverage GNNs as an algorithmic tool to solve it (Sec. II). This model emphasizes that the GNN design parameters, such as the GNN connectivity, can be decoupled from the underlying structure of the graph data (i.e., the KG).
2. We show how to map the link prediction optimization to a program in PyG (Sec. III) and study how the GNN parameters affect the link prediction accuracy in a clinical triage use case (Sec. IV). Our findings suggest that a GNN connectivity that considers domain knowledge is more effective than just using the connectivity of the graph data, and that allowing embeddings to flow in any direction may result in poor performance (Sec. IV-C1). Additionally, we demonstrate that negative edges play a crucial role in achieving good predictions (Sec. IV-C3), and that using too many GNN layers can degrade performance (Sec. IV-C2).

## II Link Prediction Model

### _Multi-relational Knowledge Graph (KG)_

A multi-relational Knowledge Graph (KG) is a graph with \(n\) nodes and \(m\)_directed_ links, where each link is associated with a relation \(r\) that represents the type of connection between the nodes. For instance, in the semantic triple _patient_ (node) _suffers from_ (relation) _anemia_ (node), the relation _suffers from_ indicates the type of connection between the node _patient_ and the node _anemia_. The nodes in the graph are also associated with a type or class, e.g., the node _anemia_ can be of the type _medical condition_.

Besides the links' relation, a link can be _positive_, _negative_, or _unknown_. A positive link indicates the two nodes are connected, while a negative link implies no connection. For example, if a patient has tested positive for diabetes, there will be a (positive) link between the patient's node and the diabetes node in the KG. Conversely, if the patient has tested negative, a (negative) link will indicate that such a connection does not exist. Unknown links, as the name suggests, are links whose existence is unknown from the data. This is the type of link that we would like to predict.

### _Link prediction as an optimization problem_

We can model a link prediction problem in a KG as follows. Every node \(i\in\{1,\ldots,n\}\) is associated with a feature vector \(e_{i}\in\mathbf{R}^{d}\), which is also known as the _node's embedding_, or _embedding_ for short. Similarly, every link is associated with a relation matrix \(W_{r}\in\mathbf{R}^{d\times d}\), \(r\in\{1,\ldots,R\}\). Next, for every pair of nodes \((i,j)\) and relation \(r\), we define the links' "score" as

\[x_{ij}^{(r)}:=f(e_{i},W_{r},e_{j}) \tag{1}\]

where \(f\) is a function that takes \(e_{i}\), \(W_{r}\), and \(e_{j}\) as inputs and returns a real number in the interval \([0,1]\).1 Similarly, for every link connecting nodes \((i,j)\) with a relation \(r\), we define the labels

Footnote 1: For example, \(\mathcal{L}\) can be \(\|\mathbf{x}-\mathbf{y}\|_{2}\), i.e., the \(\ell_{2}\)-norm.

\[y_{ij}^{(r)}=\begin{cases}1&\text{there is a relation $r$ from node $i$ to $j$},\\ 0&\text{there is not a relation $r$ from node $i$ to $j$}.\end{cases}\]

With the above model, we can formulate the optimization problem

\[\underset{e_{i},W_{r}}{\text{minimize}}\quad\mathcal{L}(\mathbf{x},\mathbf{y}) \tag{2}\]

where \(\mathcal{L}:\mathbf{R}^{m}\times\mathbf{R}^{m}\rightarrow\mathbf{R}\), \(\mathbf{x}=(x_{ij}^{(r)})\in\mathbf{R}^{m}\), \(\mathbf{y}\in\{0,1\}^{m}\). The role of the loss function \(\mathcal{L}\) is to penalize vector \(\mathbf{x}\) being different from vector \(\mathbf{y}\) component-wise.2 Namely, by minimizing \(\mathcal{L}\) in (2), we are finding the nodes' embeddings \(e_{i}\) and the matrices \(W_{r}\) such that the score \(x_{ij}^{(r)}\) is equal to (or, close to) the label \(y_{ij}^{(r)}\), which indicates the presence of a positive/negative link.

Footnote 2: For example, \(\mathcal{L}\) can be \(\|\mathbf{x}-\mathbf{y}\|_{2}\), i.e., the \(\ell_{2}\)-norm.

### _Solving the link prediction problem with a GNN_

GNNs tackle the optimization problem (2) by computing nodes' embeddings based on their connectivity patterns with neighboring nodes. To illustrate the concept, we show in Figure 2 a toy example where a "patient" feature vector is built with the patient's medical conditions embeddings. In particular, the feature vectors of nodes \(i\) and \(j\) (medical conditions) are combined _linearly_ to obtain the embedding of node \(k\) (a patient).

GNNs combine the neighbors' embeddings by using multiple _non-linear_ functions. Fig. 3 shows an example of how

Fig. 3: Example of a GNN with three nodes and two NN layers per node. Vectors \(e_{i}^{(0)}\), \(e_{j}^{(0)}\), \(e_{k}^{(0)}\) are the initial embeddings of nodes \(i\), \(j\), \(k\), i.e., the input in the first NN layer. Vectors \(e_{i}^{(2)}=e_{i}\), \(e_{j}^{(2)}=e_{j}\), \(e_{k}^{(2)}=e_{k}\) are the embedding outputs of layer 2.

Fig. 2: Toy example of how the embedding of a patient is the linear combination of two medical conditions embeddings.
a GNN uses multiple layers (i.e., functions) to combine the embeddings. Each layer \(l\in\{1,\ldots,L\}\) fuses the feature vector in the \((l-1)\)-th layer of the node with embeddings of its neighbors, and the output is passed to the next layer where the process is repeated.

The GNN connectivity depends on how nodes are connected in the KG (but is not necessarily the same), and it determines how the embedding information propagates. Designing a GNN connectivity that enables efficient learning is use-case dependent as it requires knowing how nodes should interact. In Sec. IV-C1, we will show how different GNN connectivities affect the link-prediction performance for a clinical triage use case.

### _Link prediction_

Predicting an (unknown) link/relation in a KG consists of evaluating (1) with the embeddings and relation weights learned during the training. For example, suppose we have a patient and want to predict whether the patient may be _infected_ by COVID-19. Then, we use \(e_{\text{patient}}\), \(W_{\text{infected}}\), and \(e_{\text{COVID-19}}\) in (1), and if the score is larger than a confidence threshold (e.g., \(0.9\)), we can determine a positive link exists.

Often, we need to predict links that connect nodes not seen during the training. In that case, we need to first compute the embeddings of the new nodes by combining their initial embeddings3 with the embeddings of their neighbors seen in the training. For example, a new patient may be connected to nodes that appeared during the training (e.g., fever, headache, cough). Then, the embedding of the unseen node (i.e., the patient) is calculated with the embeddings of the neighboring nodes.

Footnote 3: The input in the first NN layer. See Fig. 3.

## III Link Prediction in PyG

This section presents how to implement the link prediction optimization in Sec. II with PyG [14]--a python library for GNNs built upon PyTorch. We follow a similar approach as in the PyG tutorial for node classification [15].

### _Creating the KG and tensors_

The first step is constructing a KG with the format in Table I. Each row in the table corresponds to a subject-relation-object "triple" indicating how nodes are connected. Recall the edges in the KG are directed, where the _subject_ and _object_ are the _source_ and _target_ nodes, respectively.4 The column _link type_ indicates whether such a link is positive (True) or negative (False), and the columns _sub. type_ and _obj. type_ indicate--as the names suggest--the types of nodes. Having different node types is useful, for example, to control how the embedding information flows in the GNN. In Sec. IV-C1, we will show how different GNN connectivities affect the link prediction performance.

Footnote 4: The names in the subject and object columns identify a single node in the KG. That is, there cannot be multiple nodes called _London_ referring to different cities, e.g., England (UK), Ontario (Canada), Texas (USA), etc.

The second step is to map the KG to a PyG data object that contains the nodes' initial embeddings (data.x), the network connectivity (data.edge_index), the types of relations (data.edge_type), and the labels (data.y) that indicate whether the links are positive or negative. Fig. 4 shows how to map the KG in Table I to a data object, where the nodes and relations have unique IDs (integers).

### _GNN model_

The GNN model consists of two core parts: (i) the generation of the embeddings (i.e., encoder) and (ii) the scoring function (i.e., decoder). Fig. 5 shows an example of a GNN model with two RGCNConv layers [6].5 The first part is to initialize the scoring function6 and the NN functions that will generate the embeddings. The initial embeddings (data.x) in the encoder and can be set manually--e.g., using a pre-trained natural language model that maps the node's name (i.e., a string) to a vector of a fixed size (e.g., as in [16])--or they can be variables in the optimization.

Footnote 5: RGCNConv is a type of layer/function to compute the embeddings.

Footnote 6: The weights \(W_{r}\) are part/defined in the scoring function.

The forward function computes the nodes' embeddings and scores for every link in the KG. The embedding information is obtained with a communication mask7 that controls how the nodes propagate their embeddings to their neighbors in the GNN, i.e., the GNN connectivity.

Footnote 7: A mask is a vector of booleans that selects which edges (i.e., the tensor’s rows) to use.

### _Training the model_

The training of the GNN is shown in Fig. 6. The process follows the steps in the tutorial for node classification in [15] with two differences. The first one is that we can control the embedding communication in the GNN, which affects how the initial embeddings and the embedding in intermediate layers are combined. The second difference is that we can filter the edges for which we want to evaluate the loss function, which

Fig. 4: Example of a torch tensor that maps the KG in Table I with unique IDs. Each row corresponds to a row in Table I where nodes’ IDs are assigned sequentially: Orla (0), Paul (1), London (2), cholesterol (3), New York (4), diabetes (5). The relations’ IDs are also assigned sequentially: born (0), has (1).
is useful to specialize the model to the types of links we would like to predict.8

Footnote 8: This is similar to the mask used in [15] to filter the training data.

### _Link prediction_

The link prediction task consists of calling the function model(data) where data includes the links we would like to predict. For instance, if we want to predict if the node _Orla_ is connected to the node _diabetes_ with the relation type _has_ (with the mapping used in the example in Fig. 4), we need to add the tensor([0, 5]) to data.edge_index and tensor([1]) to data.edge_type.

To predict links connecting nodes not seen in the training, we must first assign the new nodes unique IDs and generate their embeddings using the same function used in the training. However, the GNN communication mask employed to create the embeddings must prevent the unseen/new nodes from affecting the embeddings of the nodes seen in the training.

## IV Use case: Clinical triage with Synthea

This section presents a numerical evaluation of the GNN for clinical triage with the Synthea dataset generator [17]. The experiments' goal is to illustrate how (i) the GNN parameters and (ii) the domain knowledge affect the link prediction accuracy. In particular, we study different GNN connectivities (Sec. IV-C1), embedding sizes, and number of GNN layers (Sec. IV-C2), and the importance of negative edges in the construction of the KG (Sec. IV-C3).

### _Use case and dataset overview_

#### Iv-A1 Use case

The _clinical triage_ problem involves determining the appropriate course of care when a patient presents with symptoms or medical conditions at the first point of contact. This includes deciding whether the patient requires immediate attention from a healthcare professional, such as in an emergency situation (e.g., a heart attack).

#### Iv-A2 Dataset

Synthea is a _synthetic_healthcare dataset generator that simulates realistic patient medical records. The generated patient records include a variety of information such as demographics, medical history, medications, allergies, and encounters with healthcare providers. The resulting data is designed to be representative of the United States population in terms of age, gender, and ethnicity, and it includes data on over 10 million synthetic patients.

For the clinical triage problem, we access the patient's medical records generated by Synthea and extract the patient's medical conditions and encounters. Each encounter is associated with conditions (e.g., diabetes) and observations (e.g., fever), and belongs to a class that corresponds to one of the following care actions: _wellness_, _inpatient_, _outpatient_, _ambulatory_, and _emergency_. The goal of clinical triage problem is: Given a patient's medical encounter with some medical conditions and observations, determine the type of care action the patient should receive, i.e., to which care action node should the encounter be connected.

### _Experiment setup_

#### Iv-B1 Kg

We generate a KG for the clinical triage problem with Synthea as shown in Fig. 7. There are five types of nodes (_encounter_, _observation_, _condition_, _patient_, and _care action_) connected by four different types of relations (_encounter-careaction_, _encounter-observation_, _encounter-condition_, and _patient-encounter_). As a remark, Synthea does not provide information about negative edges. Still, since an encounter can only be connected to one care action, we add negative links9 between an encounter and the other care actions. For example, suppose an encounter has a positive link with the

Fig. 6: Example of the optimization procedure using the steps in [15].

Fig. 7: (a) KG schema generated with Synthea for the clinical triage problem. (b) Example of a graph with 5 patients. The edges in (b) are shown as undirected due to the figure size.

Fig. 5: Example of a torch module to implement a GNN with two RGCNConv layers. The layer architecture was proposed in [6] and is available directly in PyG.
care action _inpatient_. In that case, we add negative links in the KG between the encounter and the care actions _wellness_, _outpatient_, _ambulatory_, and _emergency_.

#### V-B2 Gnn

We will introduce the GNN connectivities in the experiment in Sec. IV-C1. The initial embeddings of the node types _observation_, _condition_, and _care action_ are variables in the optimization problem, while the initial embeddings of the node types _encounter_ and _patient_ are set to zero. This choice of initial embeddings is because, in the testing, we do not want to infer the embeddings of observations, conditions, or care actions not seen in the training. However, we allow new patients and new encounters.

#### V-B3 Training parameters

We conducted the numerical experiments with 50/10 patients in the training/testing sets--sampled uniformly at random. All the experiments are run for 50 realizations, where each realization involves a random sample of 50/10 patients from the Synthea generated data.10 The scoring function is DistMult [6] with weights initialized at uniformly at random, the loss function is logistic regression, and the GNN architecture uses RGCNConv layers [6]. The training is carried out with Adam [18] for 1000 epochs with variable learning rate11 and weight decay 0.0005.

Footnote 10: A training sample has, on average, 35k edges and 1.9k nodes (50 patients, 1603 encounters, 153 observations, 107 conditions, and 5 care actions).

Footnote 11: The learning rate is equal to 0.1 for the first 100 epochs, 0.01 for the following 600 epochs, and 0.001 for the last 300 epochs.

### _Experiments_

#### V-C1 GNN connectivity

This experiment studies how the GNN connectivity for the propagation of the embedding information affects the link prediction performance. We consider the four GNN connectivities shown in Figure 8. In short, the C1 connectivity corresponds to the connectivity of the positive links in the KG. The C2 connectivity is obtained by adding "reverse" links in C1.12 Thus, the embedding information flows in any direction. The C3 connectivity is as C2 but without the edge from the node care action to encounter. Finally, the C4 connectivity allows only embedding information to flow from observation and condition nodes to encounter nodes. The rationale behind C4 is that the node types observation and condition can be regarded as the "attributes" or "properties" of an encounter, and therefore the embedding of an encounter should not affect the embeddings of the observation and condition nodes.13

Footnote 12: The reverse links have relation type {_sub. type_}-{_obj. type_}.

Fig. 9 shows the number of correct care action predictions for the four GNN connectivities, where _total_ (blue bar) indicates the number of care actions of that type in the testing data (ground truth). Observe from the figure that the frequency of care actions is skewed, with _wellness_ being the most common and _emergency_ being the least common. Regarding correct predictions, the C4 connectivity has the best overall performance, closely followed by the C3 despite C3 having more than twice as many edges as C4 (see also Table II). The C1 and C2 connectivities did not perform well, but for two different reasons. First, the C1 connectivity does not allow the encounter nodes to receive embeddings from the observation and condition nodes, which are the "characteristics that define an encounter." The C2 connectivity fails because we allow the encounter nodes to access information that is not available in the testing. Specifically, the links between care action and encounter nodes do not exist when creating the embeddings in the testing--since they are the links we would like to predict.14

Footnote 13: i.e., there should be no edge from an encounter node to an observation/condition node.

Footnote 14: We just want to predict the link from care action to encounter, but the reverse edge from care action to encounter will not exist either in the testing.

**Conclusions:** GNN connectivities that may appear intuitive (C1 and C2) do not perform well because (i) the associated KG connectivity does not capture how the nodes' embeddings should interact, and (ii) the training uses links for computing

Fig. 8: The four GNN connectivities used in the experiments in Sec. IV-C1. Accronyms: condition (C), observation (O), encounter (E), patient (P), and care action (CA).

Fig. 9: Illustrating how the GNNs connectivities in Fig. 8 affect the prediction of care actions. The results are the average of 50 random samples from the Synthea generated dataset. Each sample consists of 50/10 patients in the training/testing sets. The nodes’ embeddings have size 5 and the GNN has 2 layers. The total number of predictions per care action is indicated in blue.
the nodes' embeddings that are not present in the testing. Connecting nodes in every direction (C3) obtains a good performance, but it is slightly outperformed by a bespoke GNN connectivity (C4) that considers only essential connections.

#### Iv-B2 Embedding size and number of GNN layers

This experiment investigates the impact of two basic GNN design parameters: The number of layers and the size of the nodes' embedding. The GNN connectivity used here corresponds to the C4 connectivity described in Sec. IV-C1.

Fig. 9(a) shows the average prediction accuracy of care action as a function of the embedding size for GNNs with 2 and 3 layers. Observe from the figure that, in both cases, the accuracy improves rapidly for embedding sizes ranging from 1 to 3, but beyond that point, the increase in accuracy becomes more gradual.

Fig. 9(b) shows the prediction accuracy as a function of the number of GNN layers when the embedding sizes are fixed to 5 and 10. Observe from the figure that adding more layers decreases the GNN performance, which is in stark contrast to _deep_ CNNs, which use many layers. We reckon this behavior is because of the "over-smoothing" phenomenon also noted in the literature [19, 20], where adding more layers makes the GNN "too well connected," and therefore, the nodes embeddings become "too similar."

**Conclusions:** The link prediction performance improves with the embeddings' size, but the improvement gains diminish once the embeddings are large enough. Using a large number of GNN layers has a negative impact on performance.

#### Iv-B3 Negative edges

This experiment studies the impact of removing negative edges in the KG. Recall that the negative edges in the KG are from the encounter to the care actions with no positive edge (see also Sec. IV-B1). We use the C4 GNN connectivity introduced in Sec. IV-C1 and obtain the results shown in Figure 11. The figure shows that not using negative edges considerably drops the link prediction performance. This behavior is because the link prediction can be thought of as a binary classification of an edge, and negative edges are a source of negative samples of such "classification." Notably, negative edges are not readily available in the Synthea generated data, and we had to use domain knowledge (i.e., understanding of the data) to add those.

**Conclusions:** Negative edges are crucial for making good link predictions in this use case. The negative links were not available in Synthea directly, and we had to use domain knowledge (i.e., understanding of the graph data) to include them in the KG.

## V Conclusions

This paper studied the flow of embedding information within GNNs and its impact on performance, specifically in a clinical triage use case. We proposed a mathematical model that decouples the GNN connectivity from the connectivity of graph data and found that incorporating domain knowledge in the GNN connectivity is more effective than relying solely on graph data connectivity. Our results also show that negative edges play a crucial role in achieving good performance, while using many GNN layers can lead to performance degradation.

A future research direction is to evaluate how the approach performs in other datasets, and how to automatize the learning of the "domain knowledge." Specifically, the identification of key GNN links for transporting embedding information, and the identification of negative edges in the KG that may not be explicitly present in the data.
