+1 (415) 649-9454

Link Prediction in Neo4j GDS: Building an ML Pipeline That Survives Production

Node similarity and kNN answer "who looks like whom right now?" Link prediction answers a different, harder question: which relationships are missing from this graph, and which of them are about to exist? That is the question behind "who should this customer buy from next", "which two identities belong to the same person", "which supplier is quietly a single point of failure", and "which parts co-fail in the field".

Neo4j Graph Data Science ships a supervised machine-learning pipeline for exactly this. It is not just an algorithm call — it is a train/test split with negative sampling, a feature-engineering step, model selection across candidate models, and a stored model you can call from Cypher at inference time. Most teams we work with either do not know the pipeline API exists or misuse the split, which quietly inflates their AUCPR to a number they cannot reproduce in production. This tutorial walks the whole thing end to end.

Prerequisites

  • Neo4j 2025.x or 2026.x with the Graph Data Science plugin (Enterprise or Aura with GDS; GDS Community supports pipelines but not model persistence to disk)
  • Python 3.10+ and the graphdatascience client, which is far more pleasant than raw Cypher procedure calls for this workflow
  • A graph with one relationship type worth predicting. We will use a co-purchase graph: (:Customer)-[:BOUGHT]->(:Product) projected into (:Product)-[:CO_PURCHASED]-(:Product)
pip install "graphdatascience>=1.12" "neo4j>=6"
from graphdatascience import GraphDataScience
import os

gds = GraphDataScience(
    os.environ["NEO4J_URI"],
    auth=(os.environ["NEO4J_USER"], os.environ["NEO4J_PASSWORD"]),
)
print(gds.version())

Step 0 — the one modelling rule that matters

Link prediction in GDS only works on an undirected, homogeneous relationship type. One node label (or one projected label), one relationship type, undirected. If your real question spans labels — "which Customer will buy which Product" — you have two options: project a monopartite graph (customers connected to customers through shared products), or reframe it as a node-classification / recommendation problem instead.

Fix that before you write any pipeline code. Half the failed link-prediction projects we review failed here, not in the ML.

Step 1 — build the co-purchase projection

// Materialise the monopartite graph. Threshold aggressively: without a
// minimum support you will create a near-complete graph and every
// downstream algorithm will crawl.
MATCH (p1:Product)<-[:BOUGHT]-(:Customer)-[:BOUGHT]->(p2:Product)
WHERE elementId(p1) < elementId(p2)
WITH p1, p2, count(*) AS shoppers
WHERE shoppers >= 3
MERGE (p1)-[r:CO_PURCHASED]-(p2)
SET r.weight = shoppers;

Then project it into memory, undirected:

G, project_stats = gds.graph.project(
    "copurchase",
    {"Product": {"properties": ["price", "review_count"]}},
    {"CO_PURCHASED": {"orientation": "UNDIRECTED", "properties": ["weight"]}},
)
print(G.node_count(), G.relationship_count(), G.density())

Sanity-check the density. Below roughly 0.001 you have a sparse graph and negative sampling will be easy; above 0.1 your "negatives" are largely links that simply have not happened yet, and your metrics will be pessimistic and noisy.

Step 2 — create the pipeline and add features

pipe, _ = gds.beta.pipeline.linkPrediction.create("copurchase-lp")

A pipeline has two halves: node property steps (executed inside the pipeline, on the training graph only — this is the part that prevents leakage) and link features (functions that combine two nodes' properties into a feature vector for the candidate pair).

# Node property steps: structural signal the raw graph does not carry.
pipe.addNodeProperty(
    "fastRP", mutateProperty="embedding", embeddingDimension=128,
    randomSeed=7, relationshipWeightProperty="weight",
)
pipe.addNodeProperty("degree", mutateProperty="degree")
pipe.addNodeProperty(
    "wcc", mutateProperty="component", relationshipTypes=["*"],
)

# Link features: how a pair of nodes becomes a row.
pipe.addFeature("hadamard", nodeProperties=["embedding"])
pipe.addFeature("l2", nodeProperties=["embedding"])
pipe.addFeature("cosine", nodeProperties=["embedding"])
pipe.addFeature("hadamard", nodeProperties=["degree", "price", "review_count"])

Why these: hadamard on an embedding is the workhorse (elementwise product, high signal for "these two live in the same neighbourhood"), l2 captures distance, cosine captures direction independent of magnitude, and degree/price features let the model learn popularity bias explicitly instead of smuggling it in.

Crucially, FastRP is added as a pipeline step, not computed beforehand with gds.fastRP.mutate. Computing embeddings on the full graph before splitting leaks test-set topology into your features and can add ten or more points of phantom AUCPR. If you remember one thing from this tutorial, remember that.

Step 3 — configure the split

pipe.configureSplit(
    testFraction=0.2,       # held-out positive links
    trainFraction=0.2,      # of the remainder, used for model selection
    validationFolds=5,
    negativeSamplingRatio=20.0,
)

negativeSamplingRatio deserves a paragraph. Real graphs are overwhelmingly non-edges, so a 1:1 sample teaches the model a world that does not exist and its scores will be wildly overconfident in production. Set the ratio to reflect roughly the class imbalance you will actually score against — 20:1 is a reasonable starting point for a sparse graph — and read AUCPR, never accuracy, as your metric. Accuracy on a 20:1 imbalance is 95% for a model that says "no" to everything.

Step 4 — candidate models and model selection

pipe.addLogisticRegression(penalty=(0.001, 1.0), patience=3, tolerance=0.001)
pipe.addRandomForest(maxDepth=(5, 20), numberOfDecisionTrees=(50, 200))
pipe.addMLP(hiddenLayerSizes=[64, 16], penalty=(0.0001, 0.1))

pipe.configureAutoTuning(maxTrials=15)

Ranges (tuples) are auto-tuned; scalars are fixed. Start with logistic regression only — it trains in seconds, it is interpretable, and it is a genuine baseline. Add the forest and MLP once you know what score you have to beat.

Step 5 — train

model, train_result = pipe.train(
    G,
    modelName="copurchase-lp-v1",
    targetRelationshipType="CO_PURCHASED",
    randomSeed=7,
    metrics=["AUCPR", "OUT_OF_BAG_ERROR"],
)

print(train_result["modelSelectionStats"]["bestParameters"])
print(model.metrics()["AUCPR"])

Read the three AUCPR numbers the result gives you — outer train, validation, and test — together:

PatternDiagnosisAction
train ≫ validation ≈ testOverfittingRaise penalty, lower maxDepth, shrink embedding dimension
train ≈ validation ≈ test, all lowUnderfitting / no signalAdd node properties, richer features, check your projection
test ≫ validationLucky splitIncrease validationFolds, re-run with a different randomSeed
AUCPR > 0.95 on first tryAlmost always leakageCheck that features come from pipeline steps, not precomputed

That last row is the common one. A suspiciously good link-prediction model is a bug report.

Step 6 — predict

Two modes. Exhaustive scoring of all candidate pairs is O(n²) and fine for tens of thousands of nodes; approximate kNN-style search scales to millions.

# Approximate: each node gets its topN best candidate links
predictions = model.predict_stream(
    G, topN=10, threshold=0.6, sampleRate=0.5, topK=20, randomSeed=7,
)
print(predictions.sort_values("probability", ascending=False).head())

To persist them as real relationships you can query from your application:

model.predict_mutate(
    G, topN=10, threshold=0.6, mutateRelationshipType="PREDICTED_CO_PURCHASE",
    sampleRate=0.5, topK=20,
)
gds.graph.writeRelationship(G, "PREDICTED_CO_PURCHASE", relationshipProperty="probability")

Serving them is then ordinary Cypher, fast enough for a request path:

MATCH (p:Product {sku: $sku})-[r:PREDICTED_CO_PURCHASE]-(rec:Product)
WHERE r.probability >= 0.7
RETURN rec.sku, rec.name, r.probability
ORDER BY r.probability DESC LIMIT 8;

Step 7 — make it repeatable

A model in memory is a demo. Production needs four more things:

  1. Persist and publish the model. gds.model.store("copurchase-lp-v1") survives restarts (Enterprise); gds.model.publish shares it across users. Re-load it on startup, and fail loudly if it is missing rather than silently falling back to no recommendations.
  2. Retrain on a schedule, not on vibes. Co-purchase topology drifts. Retrain weekly or monthly, and gate promotion on the new model's test AUCPR beating the incumbent's on the same seed and split configuration.
  3. Monitor score distribution, not just metrics. If the median predicted probability shifts by 0.1 between runs, something upstream changed. This slots naturally into the dashboards from our Neo4j observability tutorial.
  4. Budget the memory. gds.graph.project.estimate and pipe.train_estimate before you run anything on the real graph; a 128-dimension FastRP over a dense projection is the usual cause of a heap-exhausted training job. See sizing Neo4j memory for the surrounding configuration.

Then drop the in-memory graph when you are done — it holds heap until you do:

G.drop()

Where link prediction earns its keep

Beyond recommendations, the three engagements where we reach for this pipeline most often:

  • Entity resolution. Predicted links between identity nodes, fed into the blocking-and-clustering flow from our entity resolution tutorial, catch duplicates that string similarity alone misses.
  • Fraud rings. A high-probability predicted link between two accounts that share no explicit attribute is a strong investigation signal — it complements the graph features described in real-time graph features for ML models.
  • Knowledge-graph completion. LLM extraction produces sparse graphs with obviously missing edges; link prediction proposes them, and a human or a second model confirms. That pairing is increasingly how the GraphRAG systems we build stay useful as corpora grow.

Next steps

If your first honest AUCPR is disappointing, the fix is almost always in the graph rather than the hyperparameters: a better projection, a sensible support threshold, node properties that carry real domain signal. That is where we start too. If you want a senior Neo4j data scientist to review your pipeline configuration or build the first production version with your team, see our services or get in touch.