[Roadmap] Temporal Graph Support 🚀

Updated on Nov 27th.

This is a roadmap towards better temporal data and temporal GNN support in PyG (targeted for v2.5), which spans across data handling, data loading, models and examples.

Data Handling

Previously, we used TemporalData as our abstraction to represent temporal graphs. However, this data structure is limited, i.e. it can only hold a stream of events, it mixes homogeneous and heterogeneous data, it cannot naturally deal with node features, etc. As such, we want to deprecate TemporalData and move time handling to Data and HeteroData explicitly.

• Deprecate TemporalData
• Bring data.time support to Data/HeteroData:
  • Expose time as a property
  • Support data.up_to(time) -> Data and data.snapshot(start_time, end_time)
  • Support data.sort_by_time (or similar)
• Convert all datasets that currently utilize TemporalData
• Support node-level training labels: This requires a data-structure for which we can query the label efficiently given a node ID and seed timestamp.
• Support Data.update(data)/HeteroData.update(data) in order to add new incoming edges.

Open Questions:

• How do we represent changing node features over time?

Data Loading

With Data and HeteroData having time support, we can utilize PyG's data loaders to take care of temporal subgraph sampling. Importantly, NeighborLoader and LinkNeighborLoader have been already equipped with temporal sampling support.

• Deprecate TemporalDataLoader
• Support node-level/edge-level temporal sampling in NeighborLoader
• Support node-level/edge-level temporal sampling in LinkNeighborLoader

Models

Currently, we only expose TGN as a temporal GNN model, which currently operates on TemporalData and its own sampler implementation.

• Re-factor TGN to be able to operate on Data/HeteroData while utilizing NeighborLoader/LinkNeighborLoader
• Re-factor TGN to support multi-hop sampling (currently capped to num_hops=1)
• Finish our GraphMixer implementation (#8304)
  • Out-source MLPMixer in GraphMixer into an Aggregation module
• Finish our EdgeBank implementation (#7966)
• Add the NAT model

Examples and Tutorials

Afterwards, we need clear and descriptive examples and tutorials on how to leverage PyG for node-level and link-level temporal applications.

• An example of GraphMixer for temporal node-level predictions on TGB dataset
• An example of EdgeBank for temporal link-level predictions on TGB dataset
• A tutorial on PyG's temporal GNN support

Advanced Features

One current limitation of most temporal GNN models is that they learn an embedding per node, which gets updated over time. However, it is not necessarily scalable to keep this embedding on the GPU due to memory constraints. As such, advanced features may include data-structures to query and train these embeddings on CPU, with efficient asynchronous device transfers to avoid host2device/device2host bottlenecks. An example of such a data-structure may be inspired by the GNNAutoScale paper.