Temporal Graph Networks for Deep Learning on Dynamic Graphs

Introduction

In the past few years, graph representation learning (Bronstein et al., 2017; Hamilton et al., 2017b; Battaglia et al., 2018) has produced a sequence of successes, gaining increasing popularity in machine learning. Graphs are ubiquitously used as models for systems of relations and interactions in many fields (Battaglia et al., 2016; Qi et al., 2018; Monti et al., 2016; Choma et al., 2018; Duvenaud et al., 2015; Gilmer et al., 2017; Parisot et al., 2018; Rossi et al., 2019), in particular, social sciences (Ying et al., 2018; Monti et al., 2019; Rossi et al., 2020) and biology (Zitnik et al., 2018; Veselkov et al., 2019; Gainza et al., 2019). Learning on such data is possible using graph neural networks (GNNs) (Hamilton et al., 2017a) that typically operate by a message passing mechanism (Battaglia et al., 2018) aggregating information in a neighborhood of a node and create node embeddings that are then used for node classification (Monti et al., 2016; Velickovic et al., 2018; Kipf & Welling, 2017), graph classification (Gilmer et al., 2017), or edge prediction (Zhang & Chen, 2018) tasks.

The majority of methods for deep learning on graphs assume that the underlying graph is static. However, most real-life systems of interactions such as social networks or biological interactomes are dynamic. While it is often possible to apply static graph deep learning models (Liben-Nowell & Kleinberg, 2007) to dynamic graphs by ignoring the temporal evolution, this has been shown to be sub-optimal (Xu et al., 2020), and in some cases, it is the dynamic structure that contains crucial insights about the system. Learning on dynamic graphs is relatively recent, and most works are limited to the setting of discrete-time dynamic graphs represented as a sequence of snapshots of the graph (Liben-Nowell & Kleinberg, 2007; Dunlavy et al., 2011; Yu et al., 2019; Sankar et al., 2020; Pareja et al., 2019; Yu et al., 2018). Such approaches are unsuitable for interesting real world settings such as social networks, where dynamic graphs are continuous (i.e. edges can appear at any time) and evolving (i.e. new nodes join the graph continuously). Only recently, several approaches have been proposed that support the continuous-time scenario (Xu et al., 2020; Trivedi et al., 2019; Kumar et al., 2019; Ma et al., 2018; Nguyen et al., 2018; Bastas et al., 2019).

In this paper, we first propose the generic inductive framework of Temporal Graph Networks (TGNs) operating on continuous-time dynamic graphs represented as a sequence of events, and show that many previous methods are specific instances of TGNs. Second, we propose a novel training strategy allowing the model to learn from the sequentiality of the data while maintaining highly efficient parallel processing. Third, we perform a detailed ablation study of different components of our framework and analyze the tradeoff between speed and accuracy. Finally, we show state-of-the-art performance on multiple tasks and datasets in both transductive and inductive settings, while being much faster than previous methods.

Background

A static graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$ comprises nodes $\mathcal{V}=\{1,\ldots,n\}$ and edges $\mathcal{E}\subseteq\mathcal{V}\times\mathcal{V}$, which are endowed with features, denoted by $\mathbf{v}_{i}$ and $\mathbf{e}_{ij}$ for all $i,j=1,\ldots,n$, respectively. A typical graph neural network (GNN) creates an embedding $\mathbf{z}_{i}$ of the nodes by learning a local aggregation rule of the form

There exist two main models for dynamic graphs. Discrete-time dynamic graphs (DTDG) are sequences of static graph snapshots taken at intervals in time. Continuos-time dynamic graphs (CTDG) are more general and can be represented as timed lists of events, which may include edge addition or deletion, node addition or deletion and node or edge feature transformations.

Our temporal (multi-)graph is modeled as a sequence of time-stamped events $\mathcal{G}=\{x(t_{1}),x(t_{2}),\ldots\}$, representing addition or change of a node or interaction between a pair of nodes at times $0\leq t_{1}\leq t_{2}\leq\ldots$. An event $x(t)$ can be of two types: 1) A node-wise event is represented by $\mathbf{v}_{i}(t)$, where $i$ denotes the index of the node and $\mathbf{v}$ is the vector attribute associated with the event. If the index $i$ has not been seen before, the event creates node $i$ (with the given features), otherwise it updates the features. 2) An interaction event between nodes $i$ and $j$ is represented by a (directed) temporal edge $\mathbf{e}_{ij}(t)$ (there might be more than one edge between a pair of nodes, so technically this is a multigraph). We denote by $\mathcal{V}(T)=\{i\,:\,\exists\mathbf{v}_{i}(t)\in\mathcal{G},t\in T\}$ and $\mathcal{E}(T)=\{(i,j)\,:\,\exists\mathbf{e}_{ij}(t)\in\mathcal{G},t\in T\}$ the temporal set of vertices and edges, respectively, and by $\mathcal{N}_{i}(T)=\{j\,:\,(i,j)\in\mathcal{E}(T)\}$ the neighborhood of node $i$ in time interval $T$. $\mathcal{N}^{k}_{i}(T)$ denotes the $k$-hop neighborhood. A snapshot of the temporal graph $\mathcal{G}$ at time $t$ is the (multi-)graph $\mathcal{G}(t)=(\mathcal{V}[0,t],\mathcal{E}[0,t])$ with $n(t)$ nodes. Deletion events are discussed in Appendix A.1.

Temporal Graph Networks

Following the terminology in (Kazemi et al., 2020), a neural model for dynamic graphs can be regarded as an encoder-decoder pair, where an encoder is a function that maps from a dynamic graph to node embeddings, and a decoder takes as input one or more node embeddings and makes a task-specific prediction e.g. node classification or edge prediction. The key contribution of this paper is a novel Temporal Graph Network (TGN) encoder applied on a continuous-time dynamic graph represented as a sequence of time-stamped events and producing, for each time $t$, the embedding of the graph nodes $\mathbf{Z}\left(t)=(\mathbf{z}_{1}(t),\ldots,\mathbf{z}_{n(t)}(t)\right)$.

The memory (state) of the model at time $t$ consists of a vector $\mathbf{s}_{i}(t)$ for each node $i$ the model has seen so far. The memory of a node is updated after an event (e.g. interaction with another node or node-wise change), and its purpose is to represent the node’s history in a compressed format. Thanks to this specific module, TGNs have the capability to memorize long term dependencies for each node in the graph. When a new node is encountered, its memory is initialized as the zero vector, and it is then updated for each event involving the node, even after the model has finished training. While a global (graph-wise) memory can also be added to the model to track the evolution of the entire network, we leave this as future work.

For each event involving node $i$, a message is computed to update $i$’s memory. In the case of an interaction event $\mathbf{e}_{ij}(t)$ between source node $i$ and target node $j$ at time $t$, two messages can be computed:

Similarly, in case of a node-wise event $\mathbf{v}_{i}(t)$, a single message can be computed for the node involved in the event:

Resorting to batch processing for efficiency reasons may lead to multiple events involving the same node $i$ in the same batch. As each event generates a message in our formulation, we use a mechanism to aggregate messages $\mathbf{m}_{i}(t_{1}),\ldots,\mathbf{m}_{i}(t_{b})$ for $t_{1},\ldots,t_{b}\leq t$,

As previously mentioned, the memory of a node is updated upon each event involving the node itself:

The embedding module is used to generate the temporal embedding $\mathbf{z}_{i}(t)$ of node $i$ at any time $t$. The main goal of the embedding module is to avoid the so-called memory staleness problem (Kazemi et al., 2020). Since the memory of a node $i$ is updated only when the node is involved in an event, it might happen that, in the absence of events for a long time (e.g. a social network user who stops using the platform for some time before becoming active again), $i$’s memory becomes stale. While multiple implementations of the embedding module are possible, we use the form:

where $h$ is a learnable function. This includes many different formulations as particular cases:

Temporal Graph Attention (attn): A series of $L$ graph attention layers compute $i$’s embedding by aggregating information from its $L$-hop temporal neighborhood.

The input to the $l$-th layer is $i$’s representation $\mathbf{h}_{i}^{(l-1)}(t)$, the current timestamp $t$, $i$’s neighborhood representation $\{\mathbf{h}_{1}^{(l-1)}(t),\ldots,\mathbf{h}_{N}^{(l-1)}(t)\}$ together with timestamps $t_{1},\ldots,t_{N}$ and features $\mathbf{e}_{i1}(t_{1}),\ldots,\mathbf{e}_{iN}(t_{N})$ for each of the considered interactions which form an edge in $i$’s temporal neighborhood:

Temporal Graph Sum (sum): A simpler and faster aggregation over the graph:

The graph embedding modules mitigate the staleness problem by aggregating information from a node’s neighbors memory. When a node has been inactive for a while, it is likely that some of its neighbours have been recently active, and by aggregating their memories, TGN can compute an up-to-date embedding for the node. The temporal graph attention is additionally able to select which neighbors are more important based on both features and timing information.

2 Training

TGN can be trained for a variety of tasks such as edge prediction (self-supervised) or node classification (semi-supervised). We use link prediction as an example: provided a list of time ordered interactions, the goal is to predict future interactions from those observed in the past. Figure 1 shows the computations performed by TGN on a batch of training data.

The complexity in our training strategy relates to the memory-related modules (Message function, Message aggregator, and Memory updater) because they do not directly influence the loss and therefore do not receive a gradient. To solve this problem, the memory must be updated before predicting the batch interactions. However, updating the memory with an interaction $\mathbf{e}_{ij}(t)$ before using the model to predict that same interaction, causes information leakage. To avoid the issue, when processing a batch, we update the memory with messages coming from previous batches (which are stored in the Raw Message Store), and then predict the interactions. Figure 2 shows the training flow for the memory-related modules. Pseudocode for the training procedure is presented in Appendix A.2.

More formally, at any time $t$, the Raw Message Store contains (at most) one raw message $rm_{i}$ for each node $i$The Raw Message Store does not contain a message for $i$ only if $i$ has never been involved in an event in the past., generated from the last interaction involving $i$ before time $t$. When the model processes the next interactions involving $i$, its memory is updated using $rm_{i}$ (arrows 1, 2, 3 in Figure 2), then the updated memory is used to compute the node’s embedding and the batch loss (arrows 4, 5, 6). Finally, the raw messages for the new interaction are stored in the raw message store (arrows 7). It is also worth noticing that all predictions in a given batch have access to the same state of the memory. While from the perspective of the first interaction in the batch the memory is up-to-date (since it contains information about all previous interactions in the graph), from the perspective of the last interaction in the batch the same memory is out-of-date, since it lacks information about previous interactions in the same batch. This disincentives the use of a big batch size (in the extreme case where the batch size is a big as the dataset, all predictions would be made using the initial zero memory). We found a batch size of $200$ to be a good trade-off between speed and update granularity.

Related Work

Early models for learning on dynamic graphs focused on DTDGs. Such approaches either aggregate graph snapshots and then apply static methods (Liben-Nowell & Kleinberg, 2007; Hisano, 2018; Sharan & Neville, 2008; Ibrahim & Chen, 2015; Ahmed & Chen, 2016; Ahmed et al., 2016), assemble snapshots into tensors and factorize (Dunlavy et al., 2011; Yu et al., 2017; Ma et al., 2019), or encode each snapshot to produce a series of embeddings. In the latter case, the embeddings are either aggregated by taking a weighted sum (Yao et al., 2016; Zhu et al., 2012), fit to time series models (Huang & Lin, 2009; Güneş et al., 2016; da Silva Soares & Prudêncio, 2012; Moradabadi & Meybodi, 2017), used as components in RNNs (Seo et al., 2018; Narayan & Roe, 2018; Manessi et al., 2020; Yu et al., 2019; Chen et al., 2018; Sankar et al., 2020; Pareja et al., 2019), or learned by imposing a smoothness constraint over time (Kim & Han, 2009; Gupta et al., 2011; Yao et al., 2016; Zhu et al., 2017; Zhou et al., 2018; Singer et al., 2019; Goyal et al., 2018; Fard et al., 2019; Pei et al., 2016). Another line of work encodes DTDGs by first performing random walks on an initial snapshot and then modifying the walk behaviour for subsequent snapshots (Mahdavi et al., 2018; Du et al., 2018; Xin et al., 2016; De Winter et al., 2018; Yu et al., 2018). Spatio-temporal graphs (considered by Zhang et al. (2018); Li et al. (2018) for traffic forecasting) are specific cases of dynamic graphs where the topology of the graph is fixed.

CTDGs have been addressed only recently. Several approaches use random walk models (Nguyen et al., 2018; 2018; Bastas et al., 2019) incorporating continuous time through constraints on transition probabilities. Sequence-based approaches for CTDGs (Kumar et al., 2019; Trivedi et al., 2017; 2019; Ma et al., 2018) use RNNs to update representations of the source and destination node each time a new edge appears. Other recent works have focused on dynamic knowledge graphs (Goel et al., 2019; Xu et al., 2019; Dasgupta et al., 2018; García-Durán et al., 2018). Many architectures for continuous-time dynamic graphs are based on a node-wise memory updated by an RNN when new interactions appear. Yet, they lack a GNN-like aggregation from a node’s neighbors when computing its embedding, which makes them susceptible to the staleness problem (i.e. a node embedding becoming out of date) while at the same time also limiting their expressive power.

Experiments

We use three datasets in our experiments: Wikipedia, Reddit (Kumar et al., 2019), and Twitter, which are described in detail in Appendix A.3. Our experimental setup closely follows (Xu et al., 2020) and focuses on the tasks of future edge (‘link’) prediction and dynamic node classification. In future edge prediction, the goal is to predict the probability of an edge occurring between two nodes at a given time. Our encoder is combined with a simple MLP decoder mapping from the concatenation of two node embeddings to the probability of the edge. We study both the transductive and inductive settings. In the transductive task, we predict future links of the nodes observed during training, whereas in the inductive tasks we predict future links of nodes never observed before. For node classification, the transductive setting is used. For all tasks and datasets we perform the same 70%-15%-15% chronological split as in Xu et al. (2020). All the results were averaged over 10 runs. Hyperparameters and additional details can be found in Appendix A.4.

Our strong baselines are state-of-the-art approaches for continuous time dynamic graphs (CTDNE (Nguyen et al., 2018), Jodie (Kumar et al., 2019), DyRep (Trivedi et al., 2019) and TGAT (Xu et al., 2020)) as well as state-of-the-art models for static graphs (GAE (Kipf & Welling, 2016), VGAE (Kipf & Welling, 2016), DeepWalk (Perozzi et al., 2014), Node2Vec (Grover & Leskovec, 2016), GAT (Velickovic et al., 2018) and GraphSAGE (Hamilton et al., 2017b)).

1 Performance

Table 2 presents the results on future edge prediction. Our model clearly outperforms the baselines by a large margin in both transductive and inductive settings on all datasets. The gap is particularly large on the Twitter dataset, where we outperfom the second-best method (DyRep) by over 4% and 10% in the transductive and inductive case respectively. Table 3 shows the results on dynamic node classification, where again our model obtains state-of-the-art results, with a large improvement over all other methods.

Due to the efficient parallel processing and the need for only one graph attention layer (see Section 5.2 for the ablation study on the number of layers), our model is up to $30\times$ faster than TGAT per epoch (Figure 3(a)), while requiring a similar number of epochs to converge.

2 Choice of Modules

We perform a detailed ablation study comparing different instances of our TGN framework, focusing on the speed vs accuracy tradeoff resulting from the choice of modules and their combination. The variants we experiment with are reported in Table 1 and their results are depicted in Figure 3(a).

We compare a model that does not make use of a memory (TGN-no-mem), with a model which uses memory (TGN-attn) but is otherwise identical. While TGN-att is about $3\times$ slower, it has nearly 4% higher precision than TGN-no-mem, confirming the importance of memory for learning on dynamic graphs, due to its ability to store long-term information about a node which is otherwise hard to capture. This finding is confirmed in Figure 3(b) where we compare different models when increasing the number of sampled neighbors: the models with memory consistently outperform the models without memory. Moreover, using the memory in conjunction with sampling the most recent neighbors reduces the number of neighbors needed to achieve the best performance when used.

Figure 3(a) compare models with different embedding modules (TGN-id, TGN-time, TGN-attn, TGN-sum). The first interesting insight is that projecting the embedding in time seems to slightly hurt, as TGN-time underperforms TGN-id. Moreover, the ability to exploit the graph is crucial for performance: we note that all graph-based projections (TGN-attn, TGN-sum) outperform the graph-less TGN-id model by a large margin, with TGN-attn being the top performer at the expense of being only slightly slower than the simpler TGN-sum. This indicates that the ability to obtain more recent information through the graph, and to select which neighbors are the most important are critical factors for the performance of the model.

We compare two further models, one using the last message aggregator (TGN-attn) and another a mean aggregator (TGN-mean) but otherwise the same. While TGN-mean performs slightly better, it is more than $3\times$ slower.

While in TGAT having 2 layers is of fundamental importance for obtaining good performances (TGAT vs TGAT-1l has over 10% difference in average precision), in TGN the presence of the memory makes it enough to use just one layer to obtain very high performance (TGN-attn vs TGN-2l). This is because when accessing the memory of the 1-hop neighbors, we are indirectly accessing information from hops further away. Moreover, being able to use only one layer of graph attention speeds up the model dramatically.

Conclusion

We introduce TGN, a generic framework for learning on continuous-time dynamic graphs. We obtain state-of-the-art results on several tasks and datasets while being faster than previous methods. Detailed ablation studies show the importance of the memory and its related modules to store long-term information, as well as the importance of the graph-based embedding module to generate up-to-date node embeddings. We envision interesting applications of TGN in the fields of social sciences, recommender systems, and biological interaction networks, opening up a future research direction of exploring more advanced settings of our model and understanding the most appropriate domain-specific choices.

References

Appendix A Appendix

The TGN frameworks also support edge and node deletions events.

In the case of an edge deletion event $(i,j,t^{\prime},t)$ where an edge between nodes $i$ and $j$ which was created at time $t^{\prime}$ is deleted at time $t$, two messages can be computed for the source and target nodes that respectively started and received the interaction:

In case of a node deletion event, we simply remove the node (and its incoming and outgoing edges) from the temporal graph so that when computing other nodes embedding this node is not used during the temporal graph attention. Additionally, it would be possible to compute a message from the node’s feature and memory and use it to update the memories of all its neighbors.

A.2 TGN Training

When parallelizing the training of TGN, it is important to mantain the temporal dependencies between interactions. If the events on each node were independent, we could treat them as a sequence and train an RNN on each node independently using Back Propagation Through Time (BTPP). However, the graph structure introduces dependencies between the events (the update of a node depends on the current memory of other nodes) which prevent us from processing nodes in parallel. Previous methods either process the interactions one at a time, or use the t-batch (Kumar et al., 2019) training algorithm, which however does not satisfy temporal consistency when aggregating from the graph as in the case of TGN (since the update does not only depend on the memory of the other node involved the interaction, but also on the neighbors of the two nodes).

This issues motivate our training algorithm, which processes all interactions in batches following the chronological order. It stores the last message for each node in a message store, to process it before predicting the next interaction for the node. This allows the memory-related modules to receive a gradient. Algorithm 1 presents the pseudocode for TGN training, while Figure 4 shows a schematic diagram of TGN.

A.3 Datasets

Reddit and Wikipedia are bipartite interaction graphs. In the Reddit dataset, users and sub-reddits are nodes, and an interaction occurs when a user writes a post to the sub-reddit. In the Wikipedia dataset, users and pages are nodes, and an interaction represents a user editing a page. In both aforementioned datasets, the interactions are represented by text features (of a post or page edit, respectively), and labels represent whether a user is banned. Both interactions and labels are time-stamped.

The Twitter dataset is a non-bipartite graph released as part of the 2020 RecSys Challenge (Belli et al., 2020). Nodes are users and interactions are retweets. The features of an interaction are a BERT-based (Wolf et al., 2019) vector representation of the text of the retweet.

Node features are not present in any of these datasets, and we therefore assign the same zero feature vector to all nodes. Moreover, While our framework is general and in section 3.1 we showed how it can process any type of event, these three datasets only contain the edge creation (interaction) event type. Creating and evaluation of datasets with a wider variety of events is left as future work.

The statistics of the three datasets are reported in table 4.

To generate the Twitter dataset we started with the snapshot of the Recsys Challenge training data on 2020/09/06. We filtered the data to include only retweet edges (discarding other types of interactions) where the timestamp was present. This left approximately 10% of the edges in the original dataset. We then filtered the retweet multi-graph (users can be connected by multiple retweets) to only include the largest connected component. Finally, we filtered the graph to only the top 5,000 nodes in-degree and the top 5,000 by out-degree, ending up with 8,861 nodes since some nodes were in both sets.

A.4 Additional Experimental Settings and Results

For the all datasets, we use the Adam optimizer with a learning rate of $0.0001$, a batch size of $200$ for both training, validation and testing, and early stopping with a patience of 5. We sample an equal amount of negatives to the positive interactions, and use average precision as reference metric. Additional hyperparameters used for both future edge prediction and dynamic node classification are reported in table 5. For all the graph embedding modules we use neighbors sampling (Hamilton et al., 2017b) (i.e. only aggregate from $k$ neighbors) since it improves the efficiency of the model without losing in accuracy. In particular, the sampled edges are the $k$ most recent ones, rather than the traditional approach of sampling them uniformly, since we found it to perform much better (see Figure 5). All experiments and timings are conducted on an AWS p3.16xlarge machine and the results are averaged over 10 runs. The code will be made available for all our experiments to be reproduced.

Our results for GAE (Kipf & Welling, 2016), VGAE (Kipf & Welling, 2016), DeepWalk (Perozzi et al., 2014), Node2Vec (Grover & Leskovec, 2016), GAT (Velickovic et al., 2018) and GraphSAGE (Hamilton et al., 2017b), CTDNE (Nguyen et al., 2018) and TGAT (Xu et al., 2020) are taken directly from the TGAT paper (Xu et al., 2020). For Jodie (Kumar et al., 2019) and DyRep (Trivedi et al., 2019), in order to make the comparison as fair as possible, we implement our own version in PyTorch as a specific case of our tgn framework. For Jodie we simply use the time embedding module, while for DyRep we augment the messages with the result of a temporal graph attention performed on the destination’s neighborhood. For both we use a vanilla RNN as the memory updater module.

A.4.1 Neighbor Sampling: Uniform vs Most Recent

When performing neighborhood sampling (Hamilton et al., 2017a) in static graphs, nodes are usually sampled uniformly. While this strategy is also possible for dynamic graphs, it turns out that the most recent edges are often the most informative. In Figure 5 we compare two TGN-attn models (see Table 1) with either uniform or most recent neighbor sampling, which shows that a model which samples the most recent edges obtains higher performances.