Temporal Graph Neural Networks: TGN and TGATConv in TGraphX

Most graph neural network research assumes that graph structure is static: edges exist or they do not, and node features do not change over time. Real-world graphs rarely satisfy this assumption. Social networks gain and lose connections by the minute. Transaction networks record events at precise timestamps. Citation networks grow continuously as new papers appear. Temporal graph neural networks (TGNNs) were developed to handle these evolving graphs, where both topology and features change as a function of time.

TGraphX provides two temporal GNN modules: tgraphx.temporal.tgn for the Temporal Graph Network (TGN) memory-based architecture, and tgraphx.temporal.tgat for the Temporal Graph Attention (TGAT) convolution. This article explains the principles behind each, their architectural differences, and how to use them.

What Makes a Graph Temporal

A temporal graph is defined by a sequence of timestamped events. Each event typically records a source node, a destination node, a timestamp, and optionally an edge feature vector. Formally:

Events = { (s_i, d_i, t_i, e_i) }  where t_1 ≤ t_2 ≤ ... ≤ t_N
        

Three distinct representations arise from this event stream:

• Continuous-time dynamic graphs (CTDG): Events arrive at arbitrary timestamps. The graph has no fixed snapshot structure. TGN and TGAT operate in this setting.
• Discrete-time dynamic graphs (DTDG): The timeline is divided into equal-length windows, producing a sequence of static graph snapshots.
• Temporal edge streams: A simpler representation where edges carry a timestamp attribute.

TGN and TGATConv are designed for the CTDG setting. They process events as they arrive and maintain representations that track node history.

Prerequisites

Before working with temporal GNNs, you should understand:

• Standard GNN message passing (see the GIN architecture guide for GNN fundamentals)
• PyTorch sequence models (LSTM/GRU) at a conceptual level
• How attention mechanisms work in transformers and GAT (the TGAT architecture reuses this)

Install TGraphX with:

pip install tgraphx
        

TGN: Memory-Based Temporal Graph Networks

TGN (Rossi et al., 2020) introduces a memory module — a persistent state per node that summarizes all interactions the node has participated in, up to the current time. This memory is updated after each interaction using a recurrent neural network (typically a GRU or LSTM).

The TGN pipeline processes each event (s, d, t, e) in four steps:

1. Memory lookup: Retrieve current memory states m_s and m_d for source and destination nodes.
2. Message computation: Compute messages using the node memories, the edge feature, and the time difference since the last interaction.
3. Memory update: Update m_s and m_d using the GRU with the new messages.
4. Embedding computation: Produce temporal node embeddings by combining current memory with neighborhood aggregation.

The time encoding is a key innovation. Rather than directly using raw timestamps, TGN encodes time differences using a learnable feature:

Φ(Δt) = [cos(ω_1 · Δt + φ_1), ..., cos(ω_d · Δt + φ_d)]
        

This time2vec encoding maps time differences into a feature vector that a downstream attention or MLP can consume.

Using TGN in TGraphX

import torch
        from tgraphx.temporal.tgn import TGN
        
        # Instantiate TGN
        model = TGN(
            num_nodes=1000,
            node_feat_dim=64,
            edge_feat_dim=16,
            memory_dim=128,
            time_dim=32,
            embedding_dim=128,
            num_neighbors=10,
        )
        
        # Simulate a batch of temporal events
        src = torch.randint(0, 1000, (64,))       # source node indices
        dst = torch.randint(0, 1000, (64,))       # destination node indices
        timestamps = torch.linspace(0.0, 100.0, 64)  # event timestamps
        edge_feats = torch.randn(64, 16)           # per-event edge features
        node_feats = torch.randn(1000, 64)         # static node features
        
        # Forward: returns embeddings for src and dst nodes after processing events
        src_emb, dst_emb = model(
            src, dst, timestamps, edge_feats, node_feats
        )
        print(src_emb.shape)  # [64, 128]
        print(dst_emb.shape)  # [64, 128]
        

The memory module is updated in-place during the forward pass. For evaluation, you typically run through the event stream chronologically and then query embeddings for the nodes of interest.

TGATConv: Temporal Graph Attention Convolution

TGAT (Xu et al., 2020) takes a different approach. Rather than maintaining a per-node memory state, TGAT uses attention over temporally encoded neighborhoods. For a target node v at time t, TGAT aggregates over neighbors u that interacted with v before time t, weighting each neighbor by an attention score that incorporates both the neighbor's features and the time difference t - t_uv.

The attention score is:

α(v, u) = softmax( q(h_v || Φ(t - t_uv)) · k(h_u || Φ(t - t_uv)) )
        

where || denotes concatenation, q and k are learned query and key projections, and Φ is the time encoding.

TGAT is simpler than TGN in that it has no memory module — it reads historical interactions directly from a neighbor cache. This makes it stateless and easier to deploy but means it can only look as far back as the cache retains.

import torch
        from tgraphx.temporal.tgat import TGATConv
        
        # TGAT convolution layer
        tgat = TGATConv(
            in_channels=64,
            out_channels=64,
            heads=4,
            time_dim=16,
            edge_feat_dim=8,
            dropout=0.1,
        )
        
        # Node features and temporal edge index
        N = 500
        x = torch.randn(N, 64)
        edge_index = torch.randint(0, N, (2, 2000), dtype=torch.long)
        edge_times = torch.rand(2000) * 100.0   # per-edge timestamps
        edge_feats = torch.randn(2000, 8)       # per-edge features
        query_time = torch.full((N,), 100.0)   # query time for each node
        
        out = tgat(x, edge_index, edge_times, query_time, edge_feats)
        print(out.shape)  # [500, 64]
        

TGN vs TGATConv: Key Differences

Choose TGN when your application requires tracking long-term user or entity behavior across thousands of events — for example, predicting future transactions based on months of history. Choose TGATConv when you need simpler deployment, faster inference, and the relevant history is shallow (recent neighbors suffice).

A Complete Temporal Link Prediction Training Loop

Link prediction is the most common task for temporal GNNs: given events up to time t, predict whether nodes s and d will interact at time t + δ.

import torch
        import torch.nn as nn
        import torch.nn.functional as F
        from tgraphx.temporal.tgn import TGN
        
        class TemporalLinkPredictor(nn.Module):
            def __init__(self, num_nodes, node_dim, edge_dim, memory_dim):
                super().__init__()
                self.tgn = TGN(
                    num_nodes=num_nodes,
                    node_feat_dim=node_dim,
                    edge_feat_dim=edge_dim,
                    memory_dim=memory_dim,
                    time_dim=32,
                    embedding_dim=memory_dim,
                    num_neighbors=10,
                )
                self.link_predictor = nn.Sequential(
                    nn.Linear(memory_dim * 2, 128),
                    nn.ReLU(),
                    nn.Linear(128, 1),
                )
        
            def forward(self, src, dst, neg_dst, timestamps, edge_feats, node_feats):
                src_emb, dst_emb = self.tgn(src, dst, timestamps, edge_feats, node_feats)
                # Negative sample embeddings (use dst embeddings of negative pairs)
                _, neg_emb = self.tgn(src, neg_dst, timestamps, edge_feats, node_feats)
        
                pos_score = self.link_predictor(torch.cat([src_emb, dst_emb], dim=-1))
                neg_score = self.link_predictor(torch.cat([src_emb, neg_emb], dim=-1))
                return pos_score.squeeze(), neg_score.squeeze()
        
        # Training step
        num_nodes = 10000
        model = TemporalLinkPredictor(num_nodes, node_dim=64, edge_dim=16, memory_dim=128)
        optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
        node_feats = torch.randn(num_nodes, 64)
        
        for epoch in range(5):
            # In practice, iterate over chronologically sorted event batches
            src = torch.randint(0, num_nodes, (128,))
            dst = torch.randint(0, num_nodes, (128,))
            neg_dst = torch.randint(0, num_nodes, (128,))
            timestamps = torch.linspace(0, 1000, 128)
            edge_feats = torch.randn(128, 16)
        
            optimizer.zero_grad()
            pos_score, neg_score = model(src, dst, neg_dst, timestamps, edge_feats, node_feats)
            loss = F.binary_cross_entropy_with_logits(
                torch.cat([pos_score, neg_score]),
                torch.cat([torch.ones(128), torch.zeros(128)])
            )
            loss.backward()
            optimizer.step()
            print(f"Epoch {epoch}: loss={loss.item():.4f}")
        

Temporal Sampling Considerations

Temporal GNNs require strict chronological ordering of events. A common source of bugs is feeding events out of order, which causes data leakage — the model sees future interactions when computing embeddings for past queries.

Always sort your event stream by timestamp before passing it to TGN or TGATConv. When batching, ensure that no batch contains events from multiple timestamps in a shuffled order. For details on avoiding data leakage in GNN tasks, see the GNN research reproducibility guide.

Limitations and Honest Notes

Memory-based models are harder to parallelize. TGN's memory update is sequential — you cannot process events for the same node in parallel without race conditions. This limits GPU utilization compared to static GNNs.

TGAT's expressiveness is bounded by neighborhood depth. If the relevant history for a prediction is more than a few hops or many time steps back, TGAT cannot capture it. TGN handles long-range history through its memory module, but at training cost.

Time encoding sensitivity. Both TGN and TGAT are sensitive to the time encoding dimensionality and initialization. Very small time dimensions can fail to distinguish events separated by small time differences; very large dimensions increase model parameters without improving performance.

No support for DTDG (snapshot-based) graphs. If your temporal data is naturally structured as graph snapshots (e.g., monthly network states), TGN and TGATConv are not the right tools. Consider applying static GNNs to each snapshot and merging representations separately.

Evaluation must respect causality. Test-set links must come from strictly later times than training-set links. Cross-temporal validation splits — where a future link appears in training — overestimate model performance significantly.

Frequently Asked Questions

Can TGN handle node feature updates over time?
The memory module captures changes in a node's interaction history but does not directly update static node features. If node attributes change over time, you must supply the updated features at the time of each event.

How does TGraphX handle the cold-start problem for new nodes?
New nodes with no interaction history start with a zero-initialized memory. TGN can still produce embeddings using the current static features and any available edge features from the new node's first interaction.

What datasets are commonly used to benchmark temporal GNNs?
JODIE (Reddit, Wikipedia, MOOC, LastFM), SNAP temporal graph datasets, and UCI Message dataset are standard benchmarks. The TGraphX GitHub repository documents supported data formats.

Is TGAT related to standard GAT?
Yes — TGAT extends the GAT attention mechanism with time encodings injected into both the query and key projections. See the TGraphX vs PyTorch Geometric comparison for discussion of static GAT.

How should I set the time_dim parameter?
A time dimension of 16–64 is typical. Larger values can represent more complex temporal patterns but increase model size. Start with 32 and reduce if memory is constrained.

Can temporal GNNs be combined with node feature GNN layers like GraphSAGE?
Yes. A common pattern is to use TGN's memory module to produce a temporal node embedding, then pass it as the node feature input to a static GNN layer for further message passing over a snapshot of the graph. The two components can be trained jointly.

Reproducibility in Temporal GNN Experiments

Temporal GNN experiments have unique reproducibility challenges beyond standard static GNN experiments:

• Event ordering matters. Shuffle the training set by shuffling batches of chronologically ordered events, not individual events.
• Memory state must be reset between independent runs. TGN's memory is stateful; if you run validation immediately after training without resetting memory, you are evaluating with memory states from the training run. Always reset memory before computing validation or test metrics.
• Negative sampling at evaluation must respect causality. Negative samples for test-set events should be drawn from nodes that have not interacted with the source node up to the test time, not from the full node set.

Use TGraphX's reproducibility context and document your event ordering and negative sampling protocol when reporting results:

from tgraphx.reproducibility import set_reproducible
        
        with set_reproducible(seed=42):
            # Sort events chronologically before creating any batches
            sort_idx = timestamps.argsort()
            src = src[sort_idx]
            dst = dst[sort_idx]
            timestamps = timestamps[sort_idx]
            edge_feats = edge_feats[sort_idx]
            # ... training and evaluation ...
        

For broader GNN reproducibility practices, see the GNN research reproducibility guide.