Implementing Generative and Analytical Models to Create and Enrich Knowledge Graphs for RAGs
Retrieval-augmented generation (RAG) systems are an advanced combination of generative AI and retrieval-based technologies. The purpose of RAGs is to improve the quality of generated text by incorporating externally retrieved data into the generative procedure. The foundation of this integration is the application of knowledge graphs, which substantially enhance the functionalities of RAGs in many ways. Let us understand the advantages of using Knowledge Graphs within RAGs:
- Reliable domain corpus: Knowledge graphs are structured databases that store and arrange facts, relationships, and semantic information about different domain entities. They provide a large, domain-specific corpus to support RAG systems so that relevant data can be retrieved.
- Mitigated hallucinations in text generation: A significant drawback in adopting generative AI is the prevalence of ‘hallucinations' – which denote the production of inaccurate or deceptive information masquerading as factual. With Knowledge Graphs, RAG systems access the facts before incorporating them into the generated content as a validation step. This can also manifest as "Explainability" if needed.
- Response performance optimization: Knowledge graphs facilitate context-aware retrievals to augment the performance of RAG systems. They enable the system to understand the individual facts and the interrelationships between functional components.
- Dynamic knowledge synthesis: Knowledge graphs can hold timestamps with data and ensure the latest and updated information is used especially in the fields whre the changes and evolution is drastic, ensuring that the content produced accurately reflects the most up-to-date and precise information available.
Knowledge graphs are essential to retrieval-augmented generation systems, not just supplements. These infrastructure components improve data retrieval, reduce content generation errors, and update informational outputs. RAGs become more reliable and effective in AI-driven content generation.
Creating and Enhancing Knowledge Graphs is no longer a manual task. Multiple generative and analytical models enable the creation of new knowledge graphs and enhance existing knowledge by embedding newer facts inferred. Fact embeddings convert facts, which usually include entities and their relationships, into a lower-dimensional space of continuous vectors. Facts in knowledge graphs are nothing but triplets (entity1, relation, entity2). These are embedded where the related geometric relationships in the vector space accurately represent the semantic relationships within the knowledge graph. In this article(series), we explore generative, analytical, and hybrid models** to create and enrich these Knowledge Graphs.
**Disclaimer: I am inspired by some great work done in focused research, please check the references at the bottom of this article. Please note that there is ongoing research on benchmarking the measuring accuracy of RAG efficiencies
Part 1 of this article has the following sections:
- Section 1: Understanding all available models [Generative, Analytical, and Hybrid] that can be leveraged in creating and enriching Graphs
- Section 2: Compare and Contrast the Advantages, Limitations, Suitable Tasks, Technology Stack, Complexity and Implementation approach
- Section 3: For each Generative Model, consider a real-world data set and implement it end-to-end with hands-on coding and guidance given. [This article will focus only on the three generative models]
Finally, evaluate the best output for the graph and use that to enhance your RAG. Refer to my previous articles on how to use Knowledge Graphs in RAGs and how to combine text and knowledge embeddings to achieve better performing RAG.
Section 1
Models to create and enrich Knowledge Graphs
To enhance or create new Knowledge Graphs, many generative and analytical models can be used. In this section, we will list all the models and understand the tasks that each of these models help achieve when applied to Graph Data.
Generative Models
Generative models are specifically engineered to generate new graph data beyond what was learned in the training set. Understanding the structure's underlying distribution is necessary to accomplish this. Typically, they are implemented in situations that necessitate the creation of unique graph structures, as in the case of synthetic graph generation or data augmentation. Generative Models covered in this article include GraphRNN, VAE, and Diffusion Models.
Analytical Models
Analytical models are applied on graphs to analyze the data and generate insights . These models do not generate new graphs but help understand and explore the structural and feature information present in existing graphs for tasks like classification, prediction, and clustering. Some of the Analytical Models covered in this article include GNN, GAT, and Knowledge Graph Embedding Models.
Hybrid Models (Both Generative and Analytical Models)
Models such as GAE possess the ability to both analyze and generate graphs. They acquire analytically valuable representations that can be utilized for the reconstruction or generation of graph data, rendering them flexible instruments in the graph Machine learning arsenal.
Section 2
Model Comparison
Each model possesses distinct strengths and applications. The selection of the appropriate model is contingent upon the particular task at hand and the inherent attributes of the graph data that necessitate processing. By possessing a comprehensive comprehension of the domain problem, we can deploy the most efficient model tailored to specific requirements and employ the suitable technology stack for implementation.
In the following sections, we will implement each of the generative models and evaluate the graph output. The examples taken below deliberately use simpler cases to ease the explanation of the concept.
Drug-Protein Dataset [SNAP]
Dataset under consideration: *BioSNAP Drug-Protein Data: https://snap.stanford.edu/biodata/datasets/10015/10015-ChG-TargetDecagon.html
The dataset provided comprises data about drugs and their corresponding target proteins, which have been validated through biological experiments or pharmacological studies. This dataset is comprehensive and has data about drug and gene nodes.
*SNAP Datasets are BSD licensed and can be used for academic and commercial uses.
Section 3 – Generative Models
Generative Model 3.1
GraphRNN [Graph Recurrent Neural Networks] for Graphs
GraphRNN is an innovative method that utilizes recurrent neural networks (RNNs) to create intricate and lifelike graphs. Conventional approaches to graph generation often depend on pre-established models or heuristics, which may fail to capture the complex patterns observed in real-world graphs.
GraphRNN is a generative model that has been specifically developed to create graphs by leveraging the knowledge gained from a dataset of pre-existing graphs. GraphRNN utilizes RNNs to effectively model the sequential process of graph generation, enabling it to capture both the local and global structural properties of graphs. This approach stands apart from traditional methods.
From the comparison table above, we see that GraphRNN is highly effective in generating new graph structures by leveraging a dataset of pre-existing graphs. Now, I'll guide you through an example of utilizing GraphRNN to generate a basic graph.
Here, I present a detailed strategy for moving forward with this dataset:
Pre-requisite: Install required libraries: PyTorch, Networkx, and matplotlibp
Step 1: Download and Process Data
Use the above link to download the drug-protein dataset from the above bioSMAP dataset link. I have taken a subset of data from the complete dataset and loaded to create the initial Knowledge Graph.
import pandas as pd
file_path = '/content/sample_data/Drug-Protein-Small.csv'
data = pd.read_csv(file_path)
print(data.head())The output of the loaded dataset looks like below:
Drug Protein Interaction
0 CID000003488 1559.0 1.0
1 CID000003488 8647.0 0.0
2 CID000077992 3351.0 0.0
3 CID000077992 3350.0 1.0
4 CID000077992 3352.0 1.0Step 2: Build the Knowledge Graph
After understanding the data's structure, we will proceed to construct a graph that visually represents the drugs and targets, along with their interactions. We will use the NetworkX framework to build the Graph. We will use Python libraries such as NetworkX to transform this data into a graph structure. In this representation, nodes are the drugs and proteins and edges are the interactions.
import networkx as nx
import matplotlib.pyplot as plt
G = nx.from_pandas_edgelist(data, 'stitch_id', 'protein_id', create_using=nx.Graph())
def visualize_graph(graph, title="Graph Visualization"):
plt.figure(figsize=(10, 10))
pos = nx.spring_layout(graph)
nx.draw(graph, pos, with_labels=True, node_color='lightblue', edge_color='#3b0a45', node_size=100, font_size=6)
plt.title(title)
plt.show()
visualize_graph(G, "Initial Drug-Protein Interaction Graph")The output of the initial graph from the dataset:
Step 3: Implement GraphRNN
3.1 Convert the graph into an adjacency matrix for the GraphRNN model and store the labels to print the new Graphs:
# Assuming 'Drug' and 'Protein' columns exist with a binary interaction column 'Interaction'
# If you pivot this data to create the graph (optional depending on your approach)
interaction_matrix = data.pivot_table(index='Drug', columns='Protein', values='Interaction', fill_value=0)
# Create index-to-label mapping
index_to_label = {i: drug for i, drug in enumerate(interaction_matrix.index)}
# You might also need to map proteins if they are columns and involved in graph generation
protein_index_to_label= {i + l: protein for i, protein in enumerate(interaction_matrix.columns)}3.2 Define and Train GraphRNN Model
Model instantiation: The input, hidden, and output layers are sized. Input_size and output_size are 99, likely reflecting the adjacency matrix (the graph has 99 nodes).
Optimizer, Loss Function: The Adam optimizer is a popular neural network trainer due to its efficiency and adaptive learning rate. For binary classification, the BCELoss (Binary Cross-Entropy Loss) loss function is used. This function measures the difference between actual values and predicted probabilities, ideal for graph link prediction.
import torch.nn as nn
class GraphRNN(nn.Module):
def __init__(self, input_size, hidden_size, output_size):
super(GraphRNN, self).__init__()
self.rnn = nn.GRU(input_size=input_size, hidden_size=hidden_size, batch_first=True)
self.fc = nn.Linear(hidden_size, output_size)
def forward(self, x):
x, _ = self.rnn(x) # Ignore hidden state
x = self.fc(x)
return torch.sigmoid(x) # Use sigmoid to output probabilities
# Instantiate the model
model = GraphRNN(input_size=99, hidden_size=64, output_size=99)Training Loop: The network is trained over 200 epochs using an adj_tensor with adjacency matrices for different graph structures. Forward passes are calculated for each adjacency matrix, and the loss is compared to the actual matrix. This loss adjusts model weights. The loss is accumulated and printed every 10 epochs to track training progress and assess how well the model is reconstructing input graphs. This setup is typical for graph data tasks where each graph is represented as a sequence or set of sequences, and the goal is to accurately predict or recreate.
# Train the model (simplified example)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = nn.BCELoss() # Use Binary Cross-Entropy loss for binary classification
# Assuming batch size of 1 for simplicity
for epoch in range(200): # Training for 200 epochs
total_loss = 0
for adj_t in adj_tensor:
optimizer.zero_grad()
output = model(adj_t.unsqueeze(0)) # Add batch dimension
loss = criterion(output, adj_t.unsqueeze(0))
loss.backward()
optimizer.step()
total_loss += loss.item()
if (epoch + 1) % 10 == 0:
print(f'Epoch {epoch+1}, Loss: {total_loss/len(adj_tensor)}')The output from the collab notebook is:
Epoch 10, Loss: 0.00032655306752374857
Epoch 20, Loss: 4.582361790851038e-05
Epoch 30, Loss: 1.8102605722423036e-05
Epoch 40, Loss: 8.829843779099662e-06
Epoch 50, Loss: 4.680652554353318e-06
...
Epoch 200, Loss: 3.369176377562367e-09Step 5: Integrating output for new discoveries
Finally, let's create a new graph using the trained model
with torch.no_grad():
new_adj = model(torch.zeros(1, 99, 99)) # Start with a zero matrix
# Generate the graph (assuming you already have 'new_adj' from your VAE model)
new_graph = nx.from_numpy_array(new_adj.squeeze().numpy() > 0.5)
# Apply the label mapping to the nodes
new_graph = nx.relabel_nodes(new_graph, protein_index_to_label)
new_graph = nx.relabel_nodes(new_graph, index_to_label)
visualize_graph(new_graph, "Newly created Graph using GraphRNN")![Newly created Graph using GraphRN by author
This entire process involves loading a real-world dataset, creating a graph, visualizing it, training a GraphRNN model, and utilizing the model to generate new potential drug-protein interactions. This approach utilizes real-world data to effectively apply a generative graph model, making it well-suited for investigating new interactions in the field of computational drug discovery.
Generative Model 3.2
[VAE] Variational Autoencoders for Graphs
Variational Autoencoders are a type of generative model that operates within the framework of Bayesian inference. They employ deep learning techniques to acquire knowledge about the underlying probability distribution of training data. A VAE consists of two primary components:
- The encoder is responsible for compressing input data into a set of latent variables, typically using a series of neural network layers. The encoder offers parameters for a probability distribution that represents the data.
- The decoder component utilizes neural networks to reconstruct the input data from the compressed form, known as the latent representation.
During the training process, the parameters are fine-tuned to minimize both the reconstruction loss, which evaluates how accurately the decoder reproduces the input data from the latent representation, and the Kullback-Leibler divergence, which measures the divergence between two probability distributions.
Benefits of using VAEs in graphs
- Creating Unique Graph Structures: VAEs have the ability to generate fresh graph data by acquiring knowledge of the intricate distributions found in pre-existing graphs. This is especially valuable in fields where the exploration of potential new structures can yield valuable insights, such as in the realms of drug discovery or social network analysis.
- Using data augmentation techniques, new instances of graphs can be generated to enhance training datasets and boost the performance of graph-based machine learning models.
- Feature Extraction: Through the compression of graphs into a latent space, VAEs are able to encode graph properties into a condensed representation. This can prove valuable for various downstream tasks, including graph classification or clustering.
**Step-by-step guide to implement VAEs for Knowledge Graphs*****
***IMP: Pre-requisites, Step 1 and Step 2, follow the same instructions as explained in GraphRNN section.
Step 3: Implement the VAE Model
In this step, we establish the architecture of the Variational Autoencoder (VAE), which plays a key role in acquiring the latent representations of the graph's adjacency matrix.
3.1 Define VAE Model
In this step, we define the architecture of the Variational Autoencoder (VAE) that is required for latent representations of the adjacency matrix of the graph. Training the VAE includes optimizing the parameters of both the encoder and the decoder to minimize the loss between the input adjacency matrices and their reconstructions. The VAE is composed of two primary components: the encoder and the decoder.
Encoder The encoder's main task is to condense the input data, such as the adjacency matrix that represents drug-protein interactions, into a more compact and concentrated form known as the latent space. This is accomplished by utilizing a sequence of linear (fully connected) layers and applying nonlinear activations (ReLU in this instance) to gradually decrease the dimensionality of the input data. The encoder's final output consists of latent variables that capture the necessary information for reconstructing the input data.
The code defines an encoder with the following layers:
- The input layer accepts the input data, which represents the dimensions of the adjacency matrix.
- Hidden Layers: A sequence of layers that decrease the dimensionality, employing ReLU activation functions for non-linearity and to enhance the complexity of the model.
- The latent space layer is the last layer of the encoder that produces the latent representation of the data.
Decoder The decoder's function is to reconstruct the original input data as accurately as possible using the latent representation. This process involves reversing the encoding process by progressively increasing the dimensionality of the latent representation to match the original input size. The decoder's output is a reconstructed adjacency matrix, utilizing a sigmoid activation in the final layer to guarantee that the output values fall within the range of 0 and 1, which is ideal for binary data.
The decoder has a structure that mirrors the encoder, working in reverse to reconstruct the adjacency matrix using the latent variables.
- Beginning with the latent variables, the latent input layer is initiated.
- Hidden Layers: Expands the dimensionality back to its original size.
- The output layer generates the reconstructed adjacency matrix, utilizing a sigmoid activation function to ensure that the outputs are limited within the range of 0 to 1.
This framework enables the VAE to acquire a concise depiction of the graph structure and produce new graphs that exhibit statistical resemblance to the ones encountered during training.
import torch
from torch import nn, optim
class VAE(nn.Module):
def __init__(self, input_dim):
super(VAE, self).__init__()
self.encoder = nn.Sequential(
nn.Linear(input_dim, 100),
nn.ReLU(),
nn.Linear(100, 50),
nn.ReLU(),
nn.Linear(50, 20) # This will output the latent dimension
)
self.decoder = nn.Sequential(
nn.Linear(20, 50),
nn.ReLU(),
nn.Linear(50, 100),
nn.ReLU(),
nn.Linear(100, input_dim),
nn.Sigmoid() # Use sigmoid to ensure outputs are between 0 and 1
)
def forward(self, x):
z = self.encoder(x)
return self.decoder(z)
# Initialize the VAE
input_dim = interaction_matrix.shape[1] # number of proteins
vae = VAE(input_dim=input_dim)Step 3.2 Training VAE
Training the VAE requires optimizing the parameters of both the encoder and the decoder to minimize the loss between the input adjacency matrices and their reconstructions. The loss function utilized is the binary cross-entropy loss, which is well-suited for binary data commonly found in adjacency matrices in graphs.
The training process utilizes an Adam optimizer, known for its efficiency in handling sparse gradients and its adaptiveness in tuning the learning rate for different weights. The Binary Cross-Entropy Loss quantifies the disparity between the observed values in the adjacency matrix and the forecasted values generated by the decoder
The training runs for a specified number of epochs, representing one complete pass through the entire dataset. During each epoch, the optimizer's gradients are set to zero using optimizer.zero_grad(). Passing the input tensor through the VAE yields the model's reconstructed adjacency matrix.
- Compare the reconstructed output to the original input to calculate loss.
- Loss.backward() updates model weights through backpropagation.
- Optimizers update model parameters using gradients (optimizer.step()).
- The loss is displayed throughout training to track progress and make adjustments like changing the learning rate or epochs.
optimizer = torch.optim.Adam(vae.parameters(), lr=0.001)
criterion = nn.BCELoss()
epochs = 100
input_tensor = torch.tensor(adjacency_matrix, dtype=torch.float32)
for epoch in range(epochs):
optimizer.zero_grad()
outputs = vae(input_tensor)
loss = criterion(outputs, input_tensor)
loss.backward()
optimizer.step()
if epoch % 10 == 0:
print(f'Epoch {epoch}, Loss: {loss.item()}')The output with oss function / epoch is as below:
Epoch 0, Loss: 0.6963541507720947
Epoch 10, Loss: 0.6716635227203369
Epoch 20, Loss: 0.6221835613250732
Epoch 30, Loss: 0.48909738659858704
Epoch 40, Loss: 0.23997312784194946The training process allows the VAE to efficiently encode the structural information of the graph into the latent space and then decode it to reconstruct the graph. The quality of the reconstruction can serve as an indicator of the model's performance and its capacity to generate realistic graphs using learned distributions.
Step 5: Generate new Graphs
Use the trained model to generate new graphs.
with torch.no_grad():
sample = torch.randn(1, 64) # Generate a random latent vector
generated_matrix = vae.decoder(sample)
print(generated_matrix)The output of the newly generated adjacency matrix is as follows. Using this matrix a new graph is created.
tensor([[0.4930, 0.4707, 0.4603, 0.4887, 0.5281, 0.5275, 0.4841, 0.4661, 0.5138,
0.4350, 0.4646, 0.4732, 0.4951, 0.5150, 0.4886, 0.4639, 0.4906, 0.4774,
0.5030, 0.4478, 0.5161, 0.4951, 0.4475, 0.5236, 0.5367, 0.4726, 0.5264,
0.4875, 0.5230, 0.5047, 0.4568, 0.4803, 0.4753, 0.4704, 0.4897, 0.5096,
0.5130, 0.4968, 0.5014, 0.4583, 0.4765, 0.5401, 0.5193, 0.5270, 0.5005,
0.4927, 0.5322, 0.5051, 0.5398, 0.5179, 0.4905, 0.5042, 0.5142, 0.5199,
0.5199, 0.5213, 0.5105, 0.4526, 0.5220, 0.5568, 0.5225, 0.4491, 0.5013,
0.5148]])Generative Model 3.3
Diffusion Models for Graphs
Diffusion models are computational frameworks that simulate the spread of information, behaviors, or states through a network. They are extensively utilized in a wide range of fields, such as epidemiology, social network analysis, marketing, and many others. When it comes to graphs, diffusion models provide valuable insights into the interplay between entities (nodes) and the spread of relationships (edges) within the network.
There are different ways to categorize diffusion models based on how they simulate the spread of information, behaviors, or states through a network. These are some of the most common types. In this article, we will implement the Independent Cascade Model on Graphs.
The Independent Cascade (IC) model is a widely used diffusion model that is commonly employed to simulate the propagation of information, behaviors, or diseases across a network. This model, along with the Linear Threshold model, is widely leveraged in social network analysis and epidemiology to gain insights into the propagation of influence or contagion among individuals or nodes. Implementing Diffusion Models in Graphs
Here is a brief overview of the methodology for utilizing diffusion models in graphs:
Step 1: Initial Graph Representation
Implementing the Independent Cascade Model in Python using NetworkX to simulate the diffusion process in a citation network. Considering AI Papers and Citations data. [!0 Papers are considered]
Step 2: Diffusion Process
We define the function for the Independent Cascade Model to simulate the diffusion process.
- Activation: For a specified probability (
prob=0.3). Every active node can activate its inactive neighbors - Propagation: Recently activated nodes strive to activate their dormant neighbors in the subsequent round.
- Adding New Nodes: From time to time, additional nodes are incorporated into the graph to simulate the inclusion of fresh research papers
def independent_cascade_with_new_nodes(G, seeds, prob=0.3):
new_edges = []
new_nodes = []
activated_nodes = set(seeds)
newly_activated_nodes = set(seeds)
while newly_activated_nodes:
next_activated_nodes = set()
for node in newly_activated_nodes:
neighbors = set(G.successors(node)) - activated_nodes
for neighbor in neighbors:
if random.random() < prob:
next_activated_nodes.add(neighbor)
new_edges.append((node, neighbor))
# Adding new nodes with a certain probability
if random.random() < prob:
new_node = f"Paper {chr(75 + len(new_nodes))}" # Generate a new paper name starting from 'K'
new_nodes.append(new_node)
G.add_node(new_node)
new_edge = (node, new_node)
new_edges.append(new_edge)
next_activated_nodes.add(new_node)
newly_activated_nodes = next_activated_nodes
activated_nodes.update(newly_activated_nodes) return new_edges, new_nodeStep 3: Improving Graphs and refinement through iterations
New edges discovered during the diffusion process are incorporated into the graph, enriching it with additional connections. This we see is across two iterations.
# First iteration of the diffusion process
new_edges, new_nodes = independent_cascade_with_new_nodes(G, seeds)
print("New edges identified through diffusion:", new_edges)
print("New nodes added through diffusion:", new_nodes)
G.add_edges_from(new_edges)
# Update labels to include new nodes
for new_node in new_nodes:
papers[new_node] = new_node
# Visualize the enhanced graph
visualize_graph(G, "Enhanced Citation Network after First Iteration", 'lightgreen')
# Second iteration of the diffusion process
new_edges, new_nodes = independent_cascade_with_new_nodes(G, seeds)
print("New edges identified in the second iteration:", new_edges)
print("New nodes added in the second iteration:", new_nodes)
G.add_edges_from(new_edges)
# Update labels to include new nodes
for new_node in new_nodes:
papers[new_node] = new_node
# Visualize the further enhanced graph
visualize_graph(G, "Enhanced Citation Network after Second Iteration", 'lightcoral')The output for iteration 1:
New edges identified through diffusion: [('Paper A', 'Paper K'), ('Paper B', 'Paper D'), ('Paper D', 'Paper F'), ('Paper K', 'Paper L'), ('Paper F', 'Paper H'), ('Paper H', 'Paper J'), ('Paper H', 'Paper M'), ('Paper J', 'Paper N')]
New nodes added through diffusion: ['Paper K', 'Paper L', 'Paper M', 'Paper N']
The output of iteration 2:
New edges identified in the second iteration: [('Paper A', 'Paper C'), ('Paper A', 'Paper K'), ('Paper K', 'Paper L'), ('Paper K', 'Paper K')]
New nodes added in the second iteration: ['Paper K']
Diffusion models play a crucial role in gaining insights into the spread of information, behaviors, or states across networks. By representing your data as a graph and applying these models, you can discover previously unknown connections, enrich your knowledge graph, and gain a more profound understanding of the patterns within your data. For those interested in modeling the spread of diseases, analyzing social networks, or studying the adoption of new technologies, diffusion models offer a solid framework to explore and gain insights into the diffusion process.
Summary
This article explores how Knowledge Graphs can enhance RAGs. We explored the importance of having well-informed Knowledge Graphs and various models that can be utilized to create or improve them. This is the first part of the series, where we delve into the practical application of three generative models: GraphRNN, VAE, and Diffusion Models. In the upcoming article, we will demonstrate the implementation of the three analytical models GNN, GAT, and Knowledge Graph Embeddings.
References and Citations
- https://arxiv.org/abs/1802.08773 – GraphRNN Paper
- https://academic.oup.com/bioinformatics/article/36/2/603/5542390
- https://arxiv.org/abs/2312.16890 / https://github.com/HKUDS/DiffKG
- https://link.springer.com/chapter/10.1007/978-3-031-00126-0_15
Citation for the drug dataset used in this article
- Zitnik, M.; Rok Sosič, S.M.; Leskovec, J. BioSNAP Datasets: Stanford Biomedical Network Dataset Collection. 2018. Available online: http://snap.stanford.edu/biodata (accessed on 22 July 2022).
