Temporal Graph Networks: Learning from Events
Published:
TGN’s Design Philosophy
TGN separates two concerns:
- Memory: long-term history of a node’s interactions, stored as a fixed-size vector \(s_v\)
- Embedding: current structural context, computed by aggregating recent neighbours
This separation allows efficient online updates (only memory changes on each event) while preserving rich structural context when embeddings are needed.
The TGN Components
1. Memory Module
Each node \(v\) has a memory state \(s_v \in \mathbb{R}^{d_s}\). A node that has never been seen starts from \(s_v = 0\).
When an interaction \((u, v, t, e_{uv})\) occurs — node \(u\) interacting with node \(v\) at time \(t\) with edge features \(e_{uv}\) — a raw message is computed for each of the two endpoints:
Here \(s_u(t^-)\) is \(u\)’s memory just before the event, and \(\Delta t_u = t - t_u^-\) is the time elapsed since \(u\)’s last memory update. Source and destination get separate message functions \(\mathrm{msg}_{\text{s}}\) and \(\mathrm{msg}_{\text{d}}\), because an interaction is not symmetric (a user clicking an item is not the same event as an item being clicked by a user).
Memory update (a GRU cell, with the message as input and the old memory as hidden state):
The GRU’s gating mechanism naturally handles the trade-off between old memory and new information.
If several events touching \(u\) arrive in the same batch, TGN first collapses their raw messages into one with an aggregator (the paper’s default is simply “keep the most recent”), then applies a single memory update.
2. Temporal Graph Attention (Embedding Module)
When we need node \(u\)’s embedding at time \(t\) (for inference), we aggregate from temporal neighbours:
Time encoding: encode elapsed time as a feature using learnable frequencies — the Bochner-theorem construction of TGAT:
This gives the model a sense of recency — events further in the past have lower encoded similarity to the current time.
Temporal attention over neighbours:
where \(\mathcal{N}_k(u, t) = \{(w, t_{uw}) : t_{uw} < t\}\) restricted to the \(k\) most recent such neighbours. The strict inequality is the point: the neighbourhood is defined by a temporal cut, so an embedding at time \(t\) can never see an edge that has not happened yet.
3. Link Prediction Decoder
Given node embeddings \(h_u(t)\) and \(h_v(t)\), compute interaction probability:
Trained with binary cross-entropy + negative sampling (sample random non-interacting pairs as negatives).
Worked Example: One TGN Memory Update
Suppose user \(u\) has memory \(s_u = [0.5, -0.2, 0.8]\) (so \(d_s = 3\)). At time \(t = 100\), \(u\) interacts with item \(v\) (memory \(s_v = [0.1, 0.9, 0.3]\)) via edge features \(e_{uv} = [1]\) (e.g. a “click”). The time since \(u\)’s last memory update is \(\Delta t_u = 100 - 85 = 15\).
Step 1: Raw message for u
m_u = MSG_s(s_u, s_v, delta_t, e_uv)
= Linear([s_u || s_v || delta_t || e_uv])
= Linear([0.5, -0.2, 0.8, 0.1, 0.9, 0.3, 15, 1])
→ suppose m_u = [0.3, 0.7, -0.1] (after linear + activation)
Step 2: GRU memory update
s_u(t=100) = GRU(m_u, s_u(t^-))
= GRU([0.3, 0.7, -0.1], [0.5, -0.2, 0.8])
→ suppose s_u = [0.45, 0.6, 0.3] (gate blends old + new)
The GRU’s forget gate suppresses the old memory dimension 3 (0.8 → 0.3) because the new message strongly activates it differently. Dimension 2 rises (−0.2 → 0.6) reflecting the new interaction. This is all differentiable — gradients flow back through the GRU to learn what to remember.
The Full TGN Loop
The ordering matters more than it looks. The memory used to predict an event must be the memory before that event:
For each batch of events, in chronological order:
1. Flush the raw-message store: update memories using the
messages generated by EARLIER batches only
2. Compute temporal embeddings h_u(t), h_v(t) from the
updated memories + neighbours with timestamp < t
3. Predict p(u, v, t) and compute the loss
4. Compute the raw messages for THIS batch's events and
push them to the store — to be consumed at step 1
of the next batch
Steps 1 and 4 are what keep the model causal. If you updated the memory with an event and then predicted that same event, the target would already be encoded in the input and the reported accuracy would be meaningless.
Variants and Ablations
Several prior architectures fall out of the TGN template as particular choices of memory updater, message function and embedding module:
| Architecture | Memory | Embedding module |
|---|---|---|
| JODIE | RNN | Time projection of the memory |
| DyRep | RNN | Identity (embedding = memory); attention enters via the message function |
| TGAT | None | Temporal graph attention |
| TGN-attn | GRU | Temporal graph attention |
Note that DyRep does carry a memory — what distinguishes it from TGN-attn is that its embedding module is the identity, so the memory is used directly as the node representation, and the graph attention is folded into how its messages are built.
In the paper’s ablations, TGN-attn is the strongest of these variants on the Wikipedia and Reddit link-prediction benchmarks, and it is also the most expensive: memory plus attention is strictly more computation than either alone.
Inductive Capability
TGN naturally handles new nodes: when a previously unseen node \(v\) appears, its memory is initialised to \(0\). After its first few interactions, the memory accumulates history. This is the key advantage over transductive methods that require all nodes at training time. The caveat is that a node’s first prediction is made from an empty memory, so inductive performance leans on the embedding module and edge features rather than the memory.
Batch Processing and Training
Online, one-event-at-a-time processing is not directly compatible with batched GPU training, and there is a subtler problem: if the memory is updated with an interaction and the loss is then computed on that same interaction, the memory-related modules receive no useful gradient at all — the answer is already in the input.
TGN’s fix is the raw-message store described above. Within a batch, memories are brought up to date using messages from previous batches, the loss is computed, and only then are this batch’s messages stored. The memory updater therefore sits on the path from an earlier batch’s events to the current batch’s prediction, and does receive gradient. Backpropagation is truncated at the batch boundary rather than run over the entire event history.
A consequence worth remembering: batch size is not a free hyperparameter in TGN. Larger batches mean staler memory at prediction time, which changes the model, not just the optimisation.
Summary
| Component | Purpose |
|---|---|
| Memory \(s_v\) | Long-term history (fixed-size, GRU-updated) |
| Raw messages | Per-event information for memory update |
| Raw-message store | Defers updates by one batch — keeps the model causal and trainable |
| Time encoding | Recency signal for temporal attention |
| Temporal attention | Structural context from neighbours strictly before \(t\) |
| Link decoder | Interaction probability from embeddings |
TGN is the standard baseline for continuous-time dynamic graph link prediction. Its modular design (memory + embedding + decoder) allows ablation studies and component swapping — making it a useful research framework as well as a practical model.
References
- Rossi, E., Chamberlain, B., Frasca, F., Eynard, D., Monti, F., & Bronstein, M. (2020). Temporal Graph Networks for Deep Learning on Dynamic Graphs. ICML GRL+ Workshop 2020 (TGN: memory modules + temporal graph attention for continuous-time dynamic graphs).
- Xu, D., Ruan, C., Körpeoglu, E., Kumar, S., & Achan, K. (2020). Inductive Representation Learning on Temporal Graphs. ICLR 2020 (TGAT: temporal attention without memory; time-encoding via Bochner’s theorem).
- Kumar, S., Zhang, X., & Leskovec, J. (2019). Predicting Dynamic Embedding Trajectory in Temporal Interaction Networks. KDD 2019 (JODIE: bipartite temporal interaction model using RNN projections).