Training Strategies

Training GNNs on graph databases introduces unique challenges. We benchmarked several strategies to understand the tradeoffs between speed, memory, and accuracy.

Baseline: PyG with Neo4j Sampling

The standard approach: use Neo4j as the data source, but perform training in Python with PyTorch Geometric.

How it works:

• Sample k-hop neighborhoods from Neo4j using Cypher queries
• Fetch node features and labels
• Run training step in PyG

Pros:

• Simple, uses standard PyG training loops
• Flexible sampling strategies

Cons:

• High data transfer between Neo4j and Python
• Can be slow for large graphs

Profiling the Cypher sampler revealed that operators like Expand and Projection dominate the per-batch time:

Java UDP Sampling

Use the custom Java plugin for sampling instead of Cypher queries.

How it works:

• Call the Java sampling procedure from Python
• Fetch compact byte-array features
• Train in PyG

Pros:

• Faster sampling (optimized Java code)
• Reduced data transfer (byte arrays vs. floats)
• Statistically identical accuracy to baseline

Cons:

• Requires building and deploying the Java plugin

Critically, the Java UDP sampler achieves the same validation accuracy as PyG’s native sampler and the Cypher-based approach:

Pre-aggregation

Pre-compute aggregations to speed up training.

How it works:

• Pre-compute aggregated messages for each node
• Store pre-aggregated features in Neo4j
• Training uses pre-computed values

Pros:

• Faster training (less computation per step)

Cons:

• Requires pre-processing step
• Less flexible (aggregation is fixed)

Pre-aggregation shows the most benefit on high-dimensional feature vectors (Coauthor Physics, 4,993 dims), where feature fetch latency dominates:

GraphSAINT

GraphSAINT is a sampling strategy that samples subgraphs rather than individual neighborhoods. We compared Neo4j and PyG implementations.

How it works:

• Sample entire subgraphs from the graph
• Train on subgraphs instead of individual neighborhoods
• Reduces redundancy in sampling

Pros:

• More efficient sampling for some graph structures
• Can reduce training time

Cons:

• More complex implementation
• May not suit all graph types

Key Insights

1. Java UDP sampling was significantly faster than Cypher sampling while maintaining identical accuracy. The reduced data transfer (byte arrays) was the main factor.

2. Feature format matters: Using embedding_bytes (byte arrays) was ~2x faster than embedding_bytes_floats for training. However, in-database Cypher inference requires float features, creating a tradeoff.

3. Pre-aggregation can speed up training but reduces flexibility. It’s a good option when the aggregation strategy is fixed and the graph is static.

4. GraphSAINT showed promise but required careful tuning of subgraph size and sampling strategy.

Scalability

The Java UDP approach scaled best to large graphs. On ogbn-papers100M (111M nodes), the Cypher-based approach was prohibitively slow, while the Java plugin made training feasible on modest hardware:

The chart shows that the baseline Neo4j approach spends 153ms per batch (mostly on feature fetching), while the UDP approach cuts this to 45ms, achieving a 3.4x speedup.