Source-linked AI summary
PyG 2.0: Scalable Learning on Real World Graphs
Matthias Fey, Jinu Sunil, Akihiro Nitta, Rishi Puri, Manan Shah, Blaž Stojanovič, Ramona Bendias, Alexandria Barghi, Vid Kocijan, Zecheng Zhang, Xinwei He, Jan Eric Lenssen, Jure Leskovec
TL;DR
Graph learning has expanded to massive, heterogeneous, temporal, and application-specific graphs, creating infrastructure and usability demands beyond small benchmark settings. This paper presents PyG 2.0 and its subsequent evolution as a modular framework addressing these demands through scalable graph infrastructure, neural support, and explainability. In one highlighted application, GNNs improved accuracy from 16% for an LLM-based Agentic RAG baseline to 32% for GNN+LLM-based Graph RAG.
Problem
Graph learning increasingly involves massive, heterogeneous, temporal, and multi-modal graphs, requiring infrastructure that supports efficient training and practical deployment beyond small benchmarks.
Method
The paper traces PyG 2.0 and later architectural developments across heterogeneous neural processing, scalable graph infrastructure, explainability, and application integrations.
Results
16% to 32% accuracy: GNN+LLM-based Graph RAG versus an LLM-based Agentic RAG baseline.
Takeaways & Limitations
PyG provides a modular framework for scalable real-world graph learning, including heterogeneous and temporal graphs, efficient data workflows, and explainable GNNs.
Takeaways & Limitations
Temporal sampling applies constraints only to node and edge types with timestamps; types without timestamps are sampled without temporal constraints.
Abstract
from arXiv · showhide
PyG (PyTorch Geometric) has evolved significantly since its initial release, establishing itself as a leading framework for Graph Neural Networks. In this paper, we present Pyg 2.0 (and its subsequent minor versions), a comprehensive update that introduces substantial improvements in scalability and real-world application capabilities. We detail the framework's enhanced architecture, including support for heterogeneous and temporal graphs, scalable feature/graph stores, and various optimizations, enabling researchers and practitioners to tackle large-scale graph learning problems efficiently. Over the recent years, PyG has been supporting graph learning in a large variety of application areas, which we will summarize, while providing a deep dive into the important areas of relational deep learning and large language modeling.
1 Introduction
PyG has evolved from a flexible GNN message-passing framework into infrastructure for scalable, heterogeneous, temporal, and practical graph learning. PyG 2.0 and its ongoing development focus on heterogeneity, scaling and efficiency, explainability, and broad applications.
- Motivation: PyG helped move graph learning beyond small benchmark tasks toward efficient training on web-scale, heterogeneous, temporal, and multi-modal graphs.The framework also supports explainability and deployment across practical applications.
- Contribution: PyG 2.0 introduced architectural changes whose design principles and decisions are traced through the library’s subsequent evolution.The paper covers improvements integrated through the most recent version.
- Heterogeneity: Heterogeneity support covers diverse node and edge types, heterogeneous message passing, and learning on temporal graphs.PyG natively supports heterogeneous data types and message passing, together with temporal graph functionality.
- Scaling and Efficiency: Scaling and efficiency address massive graphs of approximately 10 billion nodes through distributed processing, efficient formats, loaders, samplers, accelerated message passing, and compilation.These capabilities target large-scale graph-learning use cases requiring optimized loading and training APIs.
- Explainability: Explainability is addressed through a plug-and-play method intended to make any PyG GNN explainable out-of-the-box.The paper discusses explainability particularly in heterogeneous graph learning settings.
- Applications: The paper surveys applications in chemistry, material design, computer vision, weather, and traffic forecasting, alongside deeper application-focused discussions.This application overview is used to showcase PyG’s generality.
PyG 2.0: End-to-End Graph Learning
PyG 2.0 provides a modular end-to-end framework for scalable graph learning, combining heterogeneous processing, optimized neural operators, external storage interfaces, and explainability. Its standardized components support experimentation across graph sizes, data types, and pipeline stages.
- Framework architecture: PyG organizes graph infrastructure, neural modeling, and post-processing through modular components that can be independently replaced.The design supports graph storage, sampling, model architectures, explanations, evaluation metrics, and related routines within one pipeline.
- Neural optimization: PyG 2.0 introduces EdgeIndex and cached sparse-format conversions to reduce memory bottlenecks and select efficient message-passing computation paths.The unified approach preserves backward compatibility with COO-based workflows and supports optimized sparse matrix multiplication and segmented aggregation.
- Neural framework: Aggregation is exposed as a first-class, composable concept spanning simple, advanced, learnable, and unconventional aggregation functions.The unified abstraction enables centralized optimization and specialized implementations for message passing and global readouts.
- Heterogeneous learning: PyG 2.0 natively supports heterogeneous graphs and provides data types, transformations, samplers, and message-passing mechanisms for multiple node and edge types.Grouped and segmented matrix multiplications implement parallel projections across varying node and edge types.
- Neural optimization: 2–3× speedup in runtime is observed with torch.compile while maintaining predictive accuracy across GNN architectures.Kernel fusion combines sparse computations and feature transformations while reducing memory-access and kernel-launch overhead.
- Scalable infrastructure: FeatureStore and GraphStore remote backends let PyG train on externally stored large-scale graphs while retaining the familiar training loop and PyTorch abstractions.WholeGraph-based workflows provide 2x-8x data-loading speedups with minimal code change and can extend to multi-node multi-GPU setups.
- Explainability: PyG’s universal Explainer interface generates node, edge, and feature attributions while handling data flow, visualization, and evaluation protocols.Researchers can focus on developing explainer algorithms while the framework manages surrounding explanation workflows.
PyG 2.0 in Action
PyG supports end-to-end relational deep learning and graph-enhanced language-model workflows, alongside applications spanning chemistry, weather, traffic, optimization, social networks, computer vision, and an expanding ecosystem.
- Relational Deep Learning: PyG supports the full relational deep learning blueprint, including multi-modal data, historical subgraph queries, and recommender systems.Its FeatureStore integrates PyTorch Frame for semantic multimodal features, while graph learning supports cross-table information exchange and recommendation retrieval.
- Relational Deep Learning: Each relational-table row can be encoded with tabular deep learning before message-passing GNNs exchange information across tables.This combines table-encoding algorithms with GNN message passing in an end-to-end workflow.
- Large Language Models: GraphRAG retrieves contextual knowledge-graph subgraphs for a natural-language query, encodes them with a GNN, and projects node embeddings into an LLM.PyG exposes this workflow through customizable FeatureStore and GraphStore extensions.
- Large Language Models: Adding GNNs increased accuracy from 16% for LLM-based Agentic RAG to 32% for GNN+LLM-based Graph RAG.The G-Retriever model supports combining PyG GNNs with HuggingFace LLMs.
- Applications and Ecosystem: PyG also provides TXT2KG, an interface for converting unstructured text datasets into knowledge graphs through parsing and prompt engineering.
- Applications and Ecosystem: PyG has been applied across chemistry, weather forecasting, traffic analysis, combinatorial optimization, social-network analysis, and computer vision.Its surrounding ecosystem adds distributed training, AutoML, temporal graphs, anomaly detection, federated learning, tabular encoders, and relational benchmarks.
4 Conclusion
PyG 2.0 advances graph learning through scalable infrastructure, a modular neural framework, and explainability, while supporting applications including relational deep learning and LLM-based RAG systems.
- Conclusion: PyG 2.0 offers scalable real-world graph learning while maintaining ease of use and flexibility.The paper organizes its advances into scalable graph infrastructure, the neural framework, and post-processing techniques such as explainability.
- Conclusion: The framework’s modular design supports applications from relational deep learning to RAG systems in large language models.