Link model

Link Prediction with the Simple API¶

Link prediction Graph Neural Networks (GNNs) can predict the probability for an edge to existe between two nodes.

Part 1: This tutorial shows how to train, analyze, and evaluate a link prediction model on the Arxiv citation graph. The objective is to predict the probability of citation between two papers i.e. train a model "Model(paper1, paper2) -> probability of paper1 citying paper2".

Part 2: Running a model on each pair of papers/nodes is computationally expensive. To improve serving scalability, we can use a decomposable model, i.e., a model decomposable as "Model(paper1, paper2) := Encoder1(paper1) * Encoder2(paper2)". By separating the model into two independent encoders, paper representations can be pre-computed, and only the dot product is required for the final prediction. We show this in this tutorial.

Part 3: If the goal is to retrieve the most likely citation for a specific paper, calculating the probability of all possible other paper is still inefficient. Instead, the output of one encoder can be indexed in a vector database. This allows us for efficient retrieval of the closest paper. We show this in this tutorial.

[Not yet available] Part 4: By default, GNNs use existing edges as both features for message passing and as labels for the training objective. However, some workflows require isolating specific edge sets so the model cannot "see" them during the forward pass. We show how to handle cases where an edgeset acts as a pure label.

[Not yet available] Part 5: Finally, production models must respect temporal constraints, ensuring a model does not access future data to predict past events. Training a model with this restriction, known as temporal-aware modeling, ensures performance metrics remain realistic at serving time. The last section of this tutorial demonstrates how to implement this temporal masking.

Installing GF¶

Make sure your machine has a GPU or TPU, otherwise training is going to take forever. If you are using Google Colab, you can get one for free. Just go to Edit > Notebook settings and select your hardware accelerator as TPU or GPU.

In [ ]:

# Install DGF (Distributed Graph Flow) and OGB (for the toy dataset).
!pip install dgf ogb -U

# Install DGF (Distributed Graph Flow) and OGB (for the toy dataset). !pip install dgf ogb -U

Importing libraries¶

In [ ]:

import os

os.environ["TF_USE_LEGACY_KERAS"] = "1"

import dgf  # Import Graph Flow
import jax
import numpy as np

import os os.environ["TF_USE_LEGACY_KERAS"] = "1" import dgf # Import Graph Flow import jax import numpy as np

Get Graph data¶

We first fetch the Arxiv graph.

In [ ]:

# Download the Mag graph from the OGB repo.
graph, schema = dgf.io.fetch_ogb_graph("arxiv")

# Download the Mag graph from the OGB repo. graph, schema = dgf.io.fetch_ogb_graph("arxiv")

Caching arxiv graph at /tmp/gf_fetch/arxiv.cache
OGB dependency not available. Downloading graph from CNS.

Let's look at the graph structure.

In [ ]:

# Show the schema
dgf.analyse.print_schema(schema)

# Show the schema dgf.analyse.print_schema(schema)

Graph Schema:

Node Sets:
  nodes:
    | Feature   | Format     | Semantic    | Shape   | Num cat. vals   |
    |-----------|------------|-------------|---------|-----------------|
    | #id       | BYTES      | PRIMARY_ID  | None    | None            |
    | #split    | BYTES      | CATEGORICAL | None    | None            |
    | feat      | FLOAT_32   | EMBEDDING   | (128,)  | None            |
    | labels    | INTEGER_64 | CATEGORICAL | None    | 40              |
    | year      | INTEGER_64 | NUMERICAL   | None    | None            |


Edge Sets:
  edges: (Source: nodes, Target: nodes)
    (No features)

This graph has a "label" feature for node prediction. We don't need it for link prediction, so we can mask it in the schema.

In [ ]:

del schema.node_sets["nodes"].features["#split"]
del schema.node_sets["nodes"].features["labels"]
dgf.analyse.print_schema(schema)

del schema.node_sets["nodes"].features["#split"] del schema.node_sets["nodes"].features["labels"] dgf.analyse.print_schema(schema)

Graph Schema:

Node Sets:
  nodes:
    | Feature   | Format     | Semantic   | Shape   | Num cat. vals   |
    |-----------|------------|------------|---------|-----------------|
    | #id       | BYTES      | PRIMARY_ID | None    | None            |
    | feat      | FLOAT_32   | EMBEDDING  | (128,)  | None            |
    | year      | INTEGER_64 | NUMERICAL  | None    | None            |


Edge Sets:
  edges: (Source: nodes, Target: nodes)
    (No features)

Part 1: Training model¶

We train a model to predict edges between papers.

In [ ]:

model = dgf.learning.train_link_model(
    graph=graph,
    schema=schema,
    target_edgeset="edges",
    # Reduce the number of hops and train steps, for the demo to be faster.
    num_sampling_hops=1,
    num_train_steps=5000,
)

model = dgf.learning.train_link_model( graph=graph, schema=schema, target_edgeset="edges", # Reduce the number of hops and train steps, for the demo to be faster. num_sampling_hops=1, num_train_steps=5000, )

Using gpu JAX backend
Graph input schema:
Node Sets:
  nodes:
    | Feature   | Format     | Semantic   | Shape   | Num cat. vals   |
    |-----------|------------|------------|---------|-----------------|
    | #id       | BYTES      | PRIMARY_ID | None    | None            |
    | feat      | FLOAT_32   | EMBEDDING  | (128,)  | None            |
    | year      | INTEGER_64 | NUMERICAL  | None    | None            |


Edge Sets:
  edges: (Source: nodes, Target: nodes)
    (No features)

Preparing dataset
Num. training seed edges: 1158243, Num. validation seed edges: 8000
Create graph sampler
Compute feature statistics
  Source stats:
GraphFeatureStatistics:
  Node Sets (1):
    'nodes':
      '#id': count=19307, min=nan, max=nan
      'feat': count=19307, min=nan, max=nan
      'year': count=19307, min=1990.0000, max=2020.0000, quantiles=(100)[1990.0000, 2007.0000, 2009.0000, ..., 2020.0000, 2020.0000, 2020.0000]

  Target stats:
GraphFeatureStatistics:
  Node Sets (1):
    'nodes':
      '#id': count=1544, min=nan, max=nan
      'feat': count=1544, min=nan, max=nan
      'year': count=1544, min=1971.0000, max=2020.0000, quantiles=(100)[1971.0000, 2006.0000, 2009.0000, ..., 2020.0000, 2020.0000, 2020.0000]

[Warning] No normalizer created for node set 'nodes', feature '#id'.
[Warning] No normalizer created for node set 'nodes', feature '#id'.

Compute graph statistics for padding
  positive source padding: Node Sets:
  nodes: 219 nodes

Edge Sets:
  edges: 213 edges
  positive target padding: Node Sets:
  nodes: 270 nodes

Edge Sets:
  edges: 261 edges
  negative target padding: Node Sets:
  nodes: 906 nodes

Edge Sets:
  edges: 848 edges
Preparing dataset finished in 2.17 seconds
Source normalizer:
Graph Normalizer:

Node Sets:
  nodes:
    - feat: IdentityNormalizer
    - year: SoftQuantileNormalizer

Edge Sets:
  edges: (Source: nodes, Target: nodes)
    (No normalizers)

Target normalizer:
Graph Normalizer:

Node Sets:
  nodes:
    - feat: IdentityNormalizer
    - year: SoftQuantileNormalizer

Edge Sets:
  edges: (Source: nodes, Target: nodes)
    (No normalizers)

Source normalized graph schema:
Node Sets:
  nodes:
    | Feature            | Format   | Semantic   | Shape   | Num cat. vals   |
    |--------------------|----------|------------|---------|-----------------|
    | feat               | FLOAT_32 | EMBEDDING  | (128,)  | None            |
    | year_SOFT_QUANTILE | FLOAT_32 | EMBEDDING  | None    | None            |


Edge Sets:
  edges: (Source: nodes, Target: nodes)
    (No features)

Target normalized graph schema:
Node Sets:
  nodes:
    | Feature            | Format   | Semantic   | Shape   | Num cat. vals   |
    |--------------------|----------|------------|---------|-----------------|
    | feat               | FLOAT_32 | EMBEDDING  | (128,)  | None            |
    | year_SOFT_QUANTILE | FLOAT_32 | EMBEDDING  | None    | None            |


Edge Sets:
  edges: (Source: nodes, Target: nodes)
    (No features)

Caching validation dataset

Caching validation dataset: 100%|██████████| 1000/1000 [00:07<00:00, 129.79it/s]

Caching validation dataset finished in 7.71 seconds
Number of cache validation batches: 1000
Training model
Generate first batch to initialize model
Create model variables

...Tracing model
Create model variables finished in 25.23 seconds
Will validate model every 1000 step(s)
Will checkpoint model every 1000 step(s)
Start training. The first two steps are generally slow.

Training:   0%|          | 0/5000 [00:00<?, ?it/s]

...Tracing model

Training:  20%|█▉        | 995/5000 [00:40<01:15, 52.93it/s, step=1000, train-auc=0.8841, train-hit_at_1=0.5487, train-hit_at_5=0.9663, train-loss=0.2153, train-mrr=0.7232]

...Tracing model

Training:  20%|██        | 1008/5000 [00:43<06:39,  9.98it/s, step=1000, train-auc=0.8841, train-hit_at_1=0.5487, train-hit_at_5=0.9663, train-loss=0.2153, train-mrr=0.7232]

Validation loop took 2.43s (only printed once)
step:1000 train-auc:0.8841 train-hit_at_1:0.5487 train-hit_at_5:0.9663 train-loss:0.2153 train-mrr:0.7232 valid-auc:0.9068 valid-hit_at_1:0.6295 valid-hit_at_5:0.9683 valid-loss:0.1815 valid-mrr:0.7741

Training:  40%|████      | 2012/5000 [01:03<02:16, 21.90it/s, step=2000, train-auc=0.9191, train-hit_at_1=0.6887, train-hit_at_5=0.9688, train-loss=0.1151, train-mrr=0.8094, valid-auc=0.9068, valid-hit_at_1=0.6295, valid-hit_at_5=0.9683, valid-loss=0.1815, valid-mrr=0.7741]

step:2000 train-auc:0.9191 train-hit_at_1:0.6887 train-hit_at_5:0.9688 train-loss:0.1151 train-mrr:0.8094 valid-auc:0.9285 valid-hit_at_1:0.7046 valid-hit_at_5:0.9782 valid-loss:0.0949 valid-mrr:0.8219

Training:  60%|██████    | 3011/5000 [01:23<01:38, 20.14it/s, step=3000, train-auc=0.9409, train-hit_at_1=0.7300, train-hit_at_5=0.9825, train-loss=0.0834, train-mrr=0.8429, valid-auc=0.9285, valid-hit_at_1=0.7046, valid-hit_at_5=0.9782, valid-loss=0.0949, valid-mrr=0.8219]

step:3000 train-auc:0.9409 train-hit_at_1:0.7300 train-hit_at_5:0.9825 train-loss:0.0834 train-mrr:0.8429 valid-auc:0.9460 valid-hit_at_1:0.7655 valid-hit_at_5:0.9849 valid-loss:0.0767 valid-mrr:0.8608

Training:  80%|████████  | 4009/5000 [01:43<00:47, 20.66it/s, step=4000, train-auc=0.9536, train-hit_at_1=0.8000, train-hit_at_5=0.9825, train-loss=0.0711, train-mrr=0.8818, valid-auc=0.9460, valid-hit_at_1=0.7655, valid-hit_at_5=0.9849, valid-loss=0.0767, valid-mrr=0.8608]

step:4000 train-auc:0.9536 train-hit_at_1:0.8000 train-hit_at_5:0.9825 train-loss:0.0711 train-mrr:0.8818 valid-auc:0.9568 valid-hit_at_1:0.7964 valid-hit_at_5:0.9906 valid-loss:0.0650 valid-mrr:0.8819

Training: 100%|██████████| 5000/5000 [02:02<00:00, 40.67it/s, step=4900, train-auc=0.9584, train-hit_at_1=0.7963, train-hit_at_5=0.9938, train-loss=0.0637, train-mrr=0.8817, valid-auc=0.9568, valid-hit_at_1=0.7964, valid-hit_at_5=0.9906, valid-loss=0.0650, valid-mrr=0.8819]

step:5000 train-auc:0.9584 train-hit_at_1:0.7963 train-hit_at_5:0.9938 train-loss:0.0637 train-mrr:0.8817 valid-auc:0.9593 valid-hit_at_1:0.8030 valid-hit_at_5:0.9919 valid-loss:0.0611 valid-mrr:0.8863
Final metrics: {'step': '4900', 'train-auc': '0.9584', 'train-hit_at_1': '0.7963', 'train-hit_at_5': '0.9938', 'train-loss': '0.0637', 'train-mrr': '0.8817', 'valid-auc': '0.9568', 'valid-hit_at_1': '0.7964', 'valid-hit_at_5': '0.9906', 'valid-loss': '0.0650', 'valid-mrr': '0.8819'}
Training model finished in 149.38 seconds

Once trained, it is important to look at the model:

In [ ]:

model.describe()

model.describe()

Out[ ]:

Objective Train logs Hyper-parameters Schemas Feature statistics Graph sampling Architecture Padding

Node prediction model: Predict the value of a node feature.

• Target edgeset: edges

• Number of training seed edges: 1158243
• Number of validation seed edges: 8000
• Training duration: 2m 39s

num_sampling_hops=1
sampling_width=15
num_layers=2
batch_size=8
max_training_time_seconds=None
num_train_steps=5000
random_seed=42
node_embedding_dim=128
learning_rate=0.001
opt_weight_decay=0.0001
dropout=0.1
message_pooling='sum'
architecture=
num_negative_nodes=8
message_passing_on_target_edgeset=True
negative_edges='random'
random_walk_num_walks_per_negative=10

Raw schema

Node Sets:
  nodes:
    | Feature   | Format     | Semantic   | Shape   | Num cat. vals   |
    |-----------|------------|------------|---------|-----------------|
    | #id       | BYTES      | PRIMARY_ID | None    | None            |
    | feat      | FLOAT_32   | EMBEDDING  | (128,)  | None            |
    | year      | INTEGER_64 | NUMERICAL  | None    | None            |


Edge Sets:
  edges: (Source: nodes, Target: nodes)
    (No features)

Source Normalized schema

Node Sets:
  nodes:
    | Feature            | Format   | Semantic   | Shape   | Num cat. vals   |
    |--------------------|----------|------------|---------|-----------------|
    | feat               | FLOAT_32 | EMBEDDING  | (128,)  | None            |
    | year_SOFT_QUANTILE | FLOAT_32 | EMBEDDING  | None    | None            |


Edge Sets:
  edges: (Source: nodes, Target: nodes)
    (No features)

Target Normalized schema

Node Sets:
  nodes:
    | Feature            | Format   | Semantic   | Shape   | Num cat. vals   |
    |--------------------|----------|------------|---------|-----------------|
    | feat               | FLOAT_32 | EMBEDDING  | (128,)  | None            |
    | year_SOFT_QUANTILE | FLOAT_32 | EMBEDDING  | None    | None            |


Edge Sets:
  edges: (Source: nodes, Target: nodes)
    (No features)

Source feature statistics

GraphFeatureStatistics:
  Node Sets (1):
    'nodes':
      '#id': count=19307, min=nan, max=nan
      'feat': count=19307, min=nan, max=nan
      'year': count=19307, min=1990.0000, max=2020.0000, quantiles=(100)[1990.0000, 2007.0000, 2009.0000, ..., 2020.0000, 2020.0000, 2020.0000]

Target feature statistics

GraphFeatureStatistics:
  Node Sets (1):
    'nodes':
      '#id': count=1544, min=nan, max=nan
      'feat': count=1544, min=nan, max=nan
      'year': count=1544, min=1971.0000, max=2020.0000, quantiles=(100)[1971.0000, 2006.0000, 2009.0000, ..., 2020.0000, 2020.0000, 2020.0000]

Source sampling plan

Root: nodes
├── edges [width=15] ➔ nodes
└── edges (reversed) [width=15] ➔ nodes

Target sampling plan

Root: nodes
├── edges [width=15] ➔ nodes
└── edges (reversed) [width=15] ➔ nodes

Model Structure

EmbedGraph(cat-embedding=64)
Dense(128)
Activation(silu)
Norm(layer_norm)
Graph Convolution Block x2:
    X = ...
    MPNN:
      Message:
        Dense(128)
        Activation(silu)
        Dense(128)
      Update:
        Dense(128)
        Activation(silu)
        Dropout(0.1)
        Dense(128)
    Residual(X)
    # Post MPNN
    X = ...
    Norm(rms_norm)
    Dense(512)
    Activation(silu)
    Dense(128)
    Dropout(0.1)
    Residual(X)
Identity

Model Weights

{'float32': 1154048}

Positive Source padding

Node Sets:
  nodes: 219 nodes

Edge Sets:
  edges: 213 edges

Positive Target padding

Node Sets:
  nodes: 270 nodes

Edge Sets:
  edges: 261 edges

Negative Target padding

Node Sets:
  nodes: 906 nodes

Edge Sets:
  edges: 848 edges

After training, a model is generally evaluated on a test dataset/graph.

Note: We don't have a test graph, so we use our training dataset here. In a real pipeline, evaluating a model on a training dataset make little sense.

In [ ]:

model.evaluate(graph)

model.evaluate(graph)

Evaluating model on 10000 edges

Inference:   0%|          | 0/11250 [00:00<?, ?it/s]

...Tracing inference model

Inference:   0%|          | 1/11250 [00:01<3:40:03,  1.17s/it]

...Tracing inference model

Inference: 100%|██████████| 11250/11250 [00:36<00:00, 309.33it/s]

Out[ ]:

Evaluation

• Num Examples: 10000
• MRR: 0.8867838492063469
• AUC: 0.960125
• Hit@N:
  • @1: 0.803
  • @5: 0.9937

We can make predictions for individual nodes:

In [ ]:

# Predict the probability of an edge between nodes 0 and 1.
predictions = model.predict(graph, source_node_idxs=[0], target_node_idxs=[1])
print()
print("Source node id:", graph.node_sets["nodes"].features["#id"][0])
print("Target node id:", graph.node_sets["nodes"].features["#id"][1])
print("Probability of an edge between the two:", predictions)

# Predict the probability of an edge between nodes 0 and 1. predictions = model.predict(graph, source_node_idxs=[0], target_node_idxs=[1]) print() print("Source node id:", graph.node_sets["nodes"].features["#id"][0]) print("Target node id:", graph.node_sets["nodes"].features["#id"][1]) print("Probability of an edge between the two:", predictions)

Inference:   0%|          | 0/1 [00:00<?, ?it/s]

...Tracing inference model

Inference: 100%|██████████| 1/1 [00:01<00:00,  1.01s/it]

Source node id: b'n0'
Target node id: b'n1'
Probability of an edge between the two: [0.26800677]

You can make predictions for groups of nodes (more efficient):

In [ ]:

predictions = model.predict(
    graph, source_node_idxs=[0], target_node_idxs=[1, 2, 3, 4, 5]
)
predictions

predictions = model.predict( graph, source_node_idxs=[0], target_node_idxs=[1, 2, 3, 4, 5] ) predictions

Inference:   0%|          | 0/1 [00:00<?, ?it/s]

...Tracing inference model

Inference: 100%|██████████| 1/1 [00:01<00:00,  1.01s/it]

Out[ ]:

array([0.26814568, 0.48160702, 0.34378505, 0.28163916, 0.06047892],
      dtype=float32)

Part 2: Decomposable model embedding a.k.a. Two tower model.¶

Models trained with train_link_model are decomposable by default. You can compute the embeddings with predict_embedding.

In [ ]:

source_embedding = model.predict_embedding(
    graph, node_idxs=[0], encoder="source"
)
target_embedding = model.predict_embedding(
    graph, node_idxs=[1, 2, 3, 4, 5], encoder="target"
)
print()
print("source_embedding:", source_embedding.shape)
print("target_embedding:", target_embedding.shape)

source_embedding = model.predict_embedding( graph, node_idxs=[0], encoder="source" ) target_embedding = model.predict_embedding( graph, node_idxs=[1, 2, 3, 4, 5], encoder="target" ) print() print("source_embedding:", source_embedding.shape) print("target_embedding:", target_embedding.shape)

Embedding (source): 100%|██████████| 1/1 [00:00<00:00,  1.76it/s]
Embedding (target): 100%|██████████| 1/1 [00:00<00:00,  1.74it/s]

source_embedding: (1, 128)
target_embedding: (5, 128)

With the embedding, we can re-compute the probabilities predicted by predict.

Note: GNN inference is stocastic. The values will not exactly match.

In [ ]:

jax.nn.sigmoid(np.sum(source_embedding * target_embedding, axis=1))

jax.nn.sigmoid(np.sum(source_embedding * target_embedding, axis=1))

Out[ ]:

Array([0.26822478, 0.4816262 , 0.34422588, 0.28125998, 0.06113339],      dtype=float32)

Part 3: Indexing embeddings for fast retrieval¶

First, let's compute target embeddings and index them using a vector database library. We will use Faiss.

In [ ]:

# Compute all the embeddings.
num_target_nodes = graph.node_sets["nodes"].num_nodes
target_embeddings = model.predict_embedding(
    graph, node_idxs=range(num_target_nodes), encoder="target"
)

# Compute all the embeddings. num_target_nodes = graph.node_sets["nodes"].num_nodes target_embeddings = model.predict_embedding( graph, node_idxs=range(num_target_nodes), encoder="target" )

Embedding (target): 100%|██████████| 21168/21168 [00:29<00:00, 720.49it/s]

Now, let's index them.

Note: As a temporary solution, we index values with a simple (but inefficient) NumPy-based solution.

In [ ]:

# As a temporary solution, we index values with a simple (but inefficient) NumPy-based solution.
# TODO: Update code when Faiss is available.
def index_embeddings(embeddings: np.ndarray):

  def sigmoid(x):
    return 1 / (1 + np.exp(-x))

  def find_closest_nodes(query: np.ndarray, n_neighbors: int = 5) -> np.ndarray:
    # Compute similarities
    similarities = np.dot(embeddings, query.T).squeeze()
    # Get top k indices (highest similarity)
    top_idxs = np.argsort(similarities)[-n_neighbors:][::-1]
    # Convert similarity to probabilities
    top_probas = sigmoid(similarities[top_idxs])
    return top_idxs, top_probas

  return find_closest_nodes


# Query the index
index = index_embeddings(target_embeddings)

# As a temporary solution, we index values with a simple (but inefficient) NumPy-based solution. # TODO: Update code when Faiss is available. def index_embeddings(embeddings: np.ndarray): def sigmoid(x): return 1 / (1 + np.exp(-x)) def find_closest_nodes(query: np.ndarray, n_neighbors: int = 5) -> np.ndarray: # Compute similarities similarities = np.dot(embeddings, query.T).squeeze() # Get top k indices (highest similarity) top_idxs = np.argsort(similarities)[-n_neighbors:][::-1] # Convert similarity to probabilities top_probas = sigmoid(similarities[top_idxs]) return top_idxs, top_probas return find_closest_nodes # Query the index index = index_embeddings(target_embeddings)

Finally, we can take the source_embedding we computed in Part 2 and query the index to retrieve the most likely connections instantly, without having to run the model on all pairs.

In [ ]:

# Query the index using our source embedding
idxs, probas = index(source_embedding[0])

print("5 most likely connections to query node:")
for idx, proba in zip(idxs, probas):
  id = graph.node_sets["nodes"].features["#id"][idx]
  print(f"\tid={id} idx={idx} proba={proba}")

# Query the index using our source embedding idxs, probas = index(source_embedding[0]) print("5 most likely connections to query node:") for idx, proba in zip(idxs, probas): id = graph.node_sets["nodes"].features["#id"][idx] print(f"\tid={id} idx={idx} proba={proba}")

5 most likely connections to query node:
	id=b'n157579' idx=157579 proba=0.9983875751495361
	id=b'n112608' idx=112608 proba=0.9982813596725464
	id=b'n121245' idx=121245 proba=0.9981396198272705
	id=b'n4697' idx=4697 proba=0.9980858564376831
	id=b'n50028' idx=50028 proba=0.9974484443664551