PRES: Toward Scalable Memory-Based Dynamic Graph Neural Networks
Authors: Junwei Su
Difan Zou
Chuan Wu
Venue: ICLR 2024 poster (2024)
Keywords: Dynamic/Temporal Graph Neural Network
Primary area: learning on graphs and other geometries & topologies
TL;DR: This paper studies scalable training for memory-based dynamic graph neural network

Abstract: Memory-based Dynamic Graph Neural Networks (MDGNNs) are a family of dynamic graph neural networks that leverage a memory module to extract, distill, and memorize long-term temporal dependencies, leading to superior performance compared to memory-less counterparts. However, training MDGNNs faces the challenge of handling entangled temporal and structural dependencies, requiring sequential and chronological processing of data sequences to capture accurate temporal patterns. During the batch training, the temporal data points within the same batch will be processed in parallel, while their temporal dependencies are neglected. This issue is referred to as temporal discontinuity and restricts the effective temporal batch size, limiting data parallelism and reducing MDGNNs' flexibility in industrial applications. This paper studies the efficient training of MDGNNs at scale, focusing on the temporal discontinuity in training MDGNNs with large temporal batch sizes. We first conduct a theoretical study on the impact of temporal batch 
size on the convergence of MDGNN training. Based on the analysis, we propose PRES, an iterative prediction-correction scheme combined with a memory coherence learning objective to mitigate the effect of temporal discontinuity, enabling MDGNNs to be trained with significantly larger temporal batches without sacrificing generalization performance. Experimental results demonstrate that our approach enables up to a 4 $\times$ larger temporal batch (3.4$\times$ speed-up) during MDGNN training.

Published as a conference paper at ICLR 2024
PRES: TOWARD SCALABLE MEMORY-BASED DY- NAMIC GRAPH NEURAL NETWORKS
Junwei Su, Difan Zou' & Chuan Wul
Department of Computer Science, University of Hong Kong (jwsu, dzou, cwu hku.hk
(cc @cs.
ABSTRACT
Memory-basee based Dynamic Graph Neural Networks (MDGNNs) are family dy- namic graph neural networks that leverage a memory module to extract, distill, and memorize long-term temporal dependencies, leading to superior performance compared to memory-less counterparts. However, training MDGNNs faces the challenge of handling entangled temporal and structural dependencies, requir- ing sequential and chronological processing of data sequences to capture accu- rate temporal patterns. During the batch training, the temporal data points within the same batch will be processed in parallel, while their temporal dependencies are neglected. This issue is referred to as temporal discontinuity and restricts the effective temporal batch size, limiting data parallelism and reducing MDGNNs' flexibility in industrial applications. This paper studies the efficient training of MDGNNs scale, focusing on the temporal discontinuity in training MDGNNs with large temporal batch sizes. We first conduct a theoretical study on the impact of temporal batch size on the convergence of MDGNN training. Based on the analysis, we propose PRES, an iterative pedictiooncorrectti with a memory coherence learning objective to mitigate the effect of temporal discontinuity, enabling MDGNNs to be trained with significantly larger tempo- ral batches without sacrificing generalization performance. Experimental results demonstrate that our approach enables up to 4 x larger temporal batch (3.4x speed-up) during MDGNN training.
scheme combined
INTRODUCTION
Graph representation learning has become increasingly important due to its ability to leverage both feature vectors and relational information among entities, providing powerful solutions in various domains (Wu et al. 2020 Hamilton 2020). While initial research (Kipf & Welling 2017 Hamilton al. 2018 Veličković et al.] (017) primarily focuses on static graphs, many real world applica- tions involve dynamic graphs (also referred to as temporal graph) characterized by continuously changing relationships, nodes, and attributes. To address this dynamic nature, dynamic graph neural networks (DGNNs) have emerged as promising deep learning models capable of modelling time- varying graph structures (Kazemi al., 22200; Skkrdiing et al. 2021). Unlike their static counter- parts, DGNNs excel at capturing temporal dependencies and learning spatial representations within the context of dynamic graphs. Consequently, they play a critical role in applications such as social media (Zhang et al. 2021), where communication events stream and relationships evolve, and rec- ommender systems Kumarr et al. 2019), where new products, users, and ratings constantly emerge.
Among DGNNs, Memory-based Dynamic Graph Neural Networks (MDGNNs) have demonstrated superior performance compared to memory-less counterparts (Poursafaei al. 2022). An essential feature of MDGNNs is the integration of a memory module within their architectures. This memory module acts as a sequential filtering mechanism, iteratively learning and distilling information from both new and historical graph data (also referred to events). Consequently, MDGNNs excel at capturing long-range dependencies and achieving state-of-the-art performance across a diverse range of wyynami--grap-rehhie tasks (Zhang et al. 2023 Rossi al. 2021 Xu et al. 2020). The
The implementation is available at: battt::///pth/ii wtt.:l/cwwwwwwwww//eea. B5
t: corresponding authors

---

Published as a conference paper at ICLR 2024
training and inference processes of MDGNNs involve three key steps. First, the incoming events are sequentially processed in their temporal order using the MESSAGE module, which extracts and distils relevant information, generating a set of message vectors. These message vectors are then employed, in conjunction with the previous memory state, by the MEMORY module to update the memory states of the vertices involved in the events. Finally, the updated memory vectors, along with structural information, are fed into the EMBEDDING module to generate dynamic embeddings for the vertices. Fig. I visually depicts this process, illustrating the information flow within MDGNNs.
Because of its effectiveness, there is a re- cent uptick in both theoretical exploration (ex- pressive power) (Souza et al. (2022) architec- tural innovation (Zhang et al. 2023) concerning MDGNN. However, the presence of the mem- ory module poses a challenge for MDGNN training. Capturing accurate temporal patterns requires sequential and chronological process- ing of event sequences (Rossi et al., [202]] Zhou jet al.) 2022). To achieve training efficiency, MDGNNs adopt batch processing, where con- secutive events are partitioned into temporal batches and each batch of events is fed to the model concurrently for processing. However, this can be problematic as the temporal depen- dency between the data points within the same batch cannot be maintained. When events in- volving the same node coexist within batch, batch processing only results in one update for MDGNNs, whereas principally, the MESSAGE and MEMORY should be updated twice based on the chronological order of the events (see Sec. [..] for a detailed discussion and Fig. (2b) a visual illustration). This phenomenon is known the temporal discontinuity which can lead to information loss and performance deterioration. Clearly, this issue will become more severe when using a large batch size, greater number of events (which have temporal order) will be included in the same batch and processed simultaneously. Therefore, in practice, one may need to use a small temporal batch size in training MDGNNs (Rossi et al. 2021 Kumar et al. 2019 Zhou et al. 2022). restricting data parallelism and reducing flexibility in practical applications. In contrast, modern deep learning methods have achieved remarkable success by training models on large amounts of data, leveraging large batch sizes to exploit the computational power of parallel processing (Goyal et al.) 2017 You jet al., 2019) 2020; Rasley et al.) 2020). Addressing the batch size bottleneck crucial to enhance the practicality of MDGNNs, enabling more efficient training, optimized utilization of computation resources, and improved performance and applicability across diverse domains.
Mam at t,
Dynamic Graph with
Memmal 4))
Timeslamped Edges eg!!)
Msg.
Emb.al 4
Return Service
Sequence of Events
Offline Data
Figure I: Illustration of the MDGNN process. Ar- rows of the same colour represent simultancous operations. (1) Temporal events are sequentially processed and transformed into messages. (2) Ar- rived messages update the previous memory. (3) The updated memory is used to compute embed- dings. (4) Time-dependent embeddings can be utilized for downstream tasks within the system.
Sec.B.1l
This paper aims to enhance the efficiency of MDGNN training by investigating the impact of tem- poral batch size on the training procedure. Our contributions are summarized as follows:
We provide formal formulation of the MDGNN training procedure and present, to the best of our knowledge, the first theoretical result on the influence of temporal batch size on MDGNN training. Contrary to the prevailing belief that smaller temporal batches always yield better MDGNN performance, we demonstrate that small temporal batches can introduce significant variance in the gradient (Theorem D. Furthermore, leveraging the concept of memory coherence (Def. B present a novel convergence result (Theorem 2 for MDGNN training, offering insights into the factors impacting the convergence rate.
Building upon the aforementioned analysis and drawing inspiration from the similarity between the memory module and the filtering mechanism in control theory, we propose novel training frame- work for MDGNNs, named PRES PRREdicttoo-Smoothe which offers provable advantages (Propo- scheme
sition Dp. PRES consists of two key components: 1) an iterative that can mitigate the variance induced by temporal discontinuity in large temporal batches, and 2) memory smoothing learning objective aiming at improving the convergence rate of MDGNNs.
ceediction-to-correee
To validate our theoretical analysis and the effectiveness of our proposed method, we conduct an extensive experimental study. The experimental results (Sec. 6.1) demonstrate that our approach
2

---

Published as a conference paper at ICLR 2024
enables the utilization of up to 4 larger temporal batch (3.4x speed-up) during MDGNN training, without compromising model generalization performance.
2 RELATED WORK
Due to space limit, here we discuss papers that are most relevant to the problem we study and provide more comprehensive review in Appendix G
Dynamic Graph Representation Learning. Dynamic graph representation learning has garnered substantial attention in recent years, driven by the imperative to model and analyze evolving relation- ships and temporal dependencies within dynamic graphs (Skarding al. 2021] Kazemi et al.) 2020).
Dynamic Graph Neural Networks (DGNNs), as dynamic counterparts to GNNs, have emerged as promising neural models for dynamic graph representation learning (Sankar et al. [2200] Poursafaci
et al. 2022 Xu et al. 2020; Rossi et al. 2021 Wang et al. 2021 Kumar et al. 2019;)rivee et al.)
2019 Zhang et al.. 2023] Pareja al. Trivedi al., 2017). Among DGNNs, MDGNNs have
demonstrated superior inference performance compared to their memory-less counterparts. Con- sequently, there has been recent surge in both theoretical exploration (expressive power) (Souza et al. 2022) and architectural innovation Rossi et al., 2021 Wang et al.) 2021 Kumar et al., 2019 Trivedi et al. 2019 Zhang et al. (023) related MDGNNs. Additionally, there are works dedi- cated to optimizing both the inference and training efficiency of MDGNNs from a system perspec- tive, employing techniques such as computation duplication (Wang & Mendis 2023), CPU-GPU communication optimization (Zhou et al. 2022), staleness (Sheng et al. 20044) and caching (Wang) ict al. 202I). Despite the recognition of the temporal discontinuity problem (which may be referred to differently) (Zhou et al.) 2022, Rossi et al.
202 Kumar et al. 2019), there are still no theoretical insights or founded solutions addressing the temporal discontinuity issue (limited temporal batch size). Our study addresses this gap in MDGNN research by focusing understanding the impact of temporal batch size on training MDGNNs, aiming to enhance data parallelism in MDGNN training. We adopt more theoretical approach, and our proposed method can be used in conjunction with the aforementioned prior works to further improve training scalability
Mini-Batch in Stochastic Gradient Descent (SGD). should be noted that there is another or- thogonal line of research investigating the effect of mini-batch size in SGD training (Goyal al.,
2017 Qian & Klabjan| 2020 Lin et al. 2018 Akiba et al 2017]] Gower et al. [2019). including
studies the context GNNs (Chen et al.) 201 2017 Ying et al. 2018 Huang al., 2018). in SGD and temporal batch in MDGNNs, as they serve distinct pupposes and bring different challenges. The goal of mini-batches in SGD to obtain a good estimation of the full-batch gradient by downsampling the entire dataset into mini- batches. On the other hand, the temporal batch specifically refers partitioning consecutive events data to ensure chronological processing of events. The temporal batch problem studied in this paper aims to increase the temporal batch size to enhance data parallelism in MDGNN training.
is important to differentiate the concepts of
3 PRELIMINARY AND BACKGROUND
Event-based Representation of Dynamic Graphs. this paper, we utilize event-based repre- sentation of dynamic graphs, as described in previous works (Skarding et al.) 2021 Zhang et al.) 2023 A dynamic graph & in this representation consists of a node set V = [I, NJ and an event set E {ess (t)}. where i, V. The event set & represents a stream of events, where each edge eig (t) corresponds to an interaction event between node and node j at timestamp 2 0. Node features and edge features denoted by (t()) and eig (t), respectively. In the case of non-attributed graphs, we assume va = 0 and eig (t) = O, representing zero vectors.
Memory-based Dynamic Graph Neural Network (MDGNN). We adopt an encoder-decoder formulation of MDGNNs, following the setting in (Rossi et al. 2021 Zhang al. 2023). The encoder takes a dynamic graph as input and generates dynamic representations of nodes, while the decoder utilizes these node representations for various downstream tasks. Given event (1)()), the
3

---

Published as a conference paper at ICLR 2024
eulta)
Negative
hg(t)
P((1.2)44) P022,306
3 loss
dec
ex
samp
emb
Rg(Fa)
ny(ta)
Node Embeddings
Edge probabiiites
Sv(45)
m((s)
ml(z)
Batch
m²(4)
S
s((1)
s(ta)
msg
5
sy(f-)
my(f,)
Positive
sy(+))
m(h)
s(hh)
m-((s)
s((]]
8yl(s)
my((.)
s((5]
Messages
(Updated)
se(t]]
Memory
m(ta)
s(ta)
(a) MDGNNs training
(b) temporal discontinuity
Figure 2: Fig. 2(a) depicts the training flow of MDGNN. The incoming batch serves as training sam- ples for updating the model and memory for the subsequent batch. Fig. 2bb) visualizes the temporal discontinuity that arises from pending events within the same temporal batch and t- indicates the moments before 1. The top section showcases two pending events sharing a common vertex. The middle section demonstrates the transition of memory states when events are sequentially processed according to temporal order. The bottom section illustrates the transition when events are processed in parallel (large batch). The grey colour indicates unobserved or altered memory states and the dotted line indicates missing transition, resulting in temporal discontinuity.
encoder of MDGNN can be defined follows:
mm(((= msg( (s((1 8,((7) (1), At), st) = mem ( (t ),) (t)),
(l = ,
((,()) (1)
sy = mem(s,( (t m,)) (t) = emb(s, (t), hy = emb(s, (tb(s,(t),N,(tt),
where st(t and sy (t are the memory states of nodes and j just before time (i.c., at the time of the previous interaction involving node or j). m,(t) and (y((t) are the messages generated from the event (t). (t) are the temporal neighbours of nodes i and j up to time t, h; h,(t) and (t) are the dynamic embeddings of nodes and] j at time f, and msg(-) (e.g.. MLP), mem(. mem(.)eegg. GRU), and emb((. (e.g, GCN) are learnable functions, respectively, representing the MESSAGE, MEMORY, and EMBEDDING modules discussed earlier.
Training MDGNNs. While MDGNNs possess the theoretical capability of capturing continu- ous dynamics, their training relics on batch processing to account for data discretization and effi- ciency considerations (Rossi et al. 2021 Zhang al. 2023). In the training process of MDGNN on dynamic graph with an event set ε consecutive events are partitioned into temporal batches
B1, ......, Bx of size b [E] temporal relationships. MDGNNNS are commonly trained using the (self-supervised) temporal link prediction task (Rossi al., 2 Zhang et al.) 2023 Xu et al. 2020), which involves binary clas- sification to predict whether an interaction will occur between a given pair of vertices. The event set E provides positive signals (interactions) for the temporal link prediction task. For each batch B;, the negative signals are formed by considering vertex pairs that do not have an event within the time interval of the batch (Xu et al. 2020 Rossi et al. 2021 Zhang et al. 2023). denoted as B,. The com- plete event batch used for training is represented as B, = B, U B. To prevent information leakage of B3 (can predict B, with its own information), MDGNNs adopt a lag-one scheme (Rossi al.) 2021 Zhang et al., 2023) where the temporal batch B,- is used to update the memory state and generate embeddings for predicting B,. The training process in each training epoch with gradient descent can be formulated as the following iterative process:
K These temporal batches are sequentially trained to capture the correct
((((0)) Σ ((00-11)), (((00-1))) = ((B,,0(1-1)) 0(1-1)), 000) 011-1) -V
(2)
where represents the model parameters of the MDGNN, 1j is the learning rate, L(O(V)) is the total loss for the entire epoch with initial parameters ((0), L.(.) is the loss for batch B,, and 2(.)
4

---

Published as a conference paper at ICLR 2024
denotes the loss function (e.g., enosspy (Zhang al. 2023 Xu et al. 2020 Rossi et al, 202I)). Interactions in real-life networks are usually sparse, and it is impractical to process the complete negative event set. Instead, subset of negative events is sampled to facilitate training. We denote the temporal batch with negative sampling B, and its gradient V (0(t-1)) (can be obtained by replacing B; in equation 2 and the training updates become,
g g(i 17 (((i-1))
(3)
We denote VL(O) the gradient of the entire epoch with negative sampling, and make the following assumption on the sampling process of negative events:
Assumption 1. (.(( is unbiased estimate of (L,(0(111 and has bounded variance, i.e. min - E E (000i-1) (C, 11)]) S Fuax
This assumption ensures that the variance of the gradient estimation remains bounded during the sampling process. Fig. 2(a) provides graphical illustration of MDGNN training procedure. Ap- pendix A contains pseudocode for the training procedure, along with more detailed description.
3.1 TEMPORAL DISCONTINUITY AND PENDING EVENTS.
Training MDGNNs using temporal batches introduces challenges accurately capturing temporal patterns. Specifically, events involving the same vertex within a batch are inherently temporal and should be processed chronologically to ensure the memory module accurately captures the temporal patterns, depicted in the central portion of Fig 2(b) Batch processing, however, only triggers one update for the MDGNNs, neglecting the temporal dependency of events, as illustrated in the lower section of Fig 2(b) This phenomenon is termed temporal discontinuity and can lead to information loss and potential noise in the memory state. To formally define this concept, we introduce the terms pending events and pending sets as follows:
Definition 1 (Pending Event). An event e is considered pending on another event e' (denoted as + e) ife and e' share at least one common verlex and 1' < I, where 1 (t') ) is the timestamp of event e (e').
Definition 2 (Pending Set). Given a temporal batch B and an event the pending set P(e, B) {e' + ele'' E B} of an event e is the collection of its pending events within batch B.
In essence, as the temporal batch size increases, it tends to accommodate a greater number of pend- ing events within the batch. As result, training MDGNNs with large temporal batches becomes a formidable task since concurrently processing numerous pending events can result in the loss of critical temporal information within the graph. This limitation imposes constraints on the effective temporal batch size, thereby impeding training parallelization.
THEORETICAL ANALYSIS OF MDGNN TRAINING
In this section, we provide a theoretical analysis of MDGNN training, specifically focusing on how the size of the temporal batch affects the variance and convergent rate of the learning process.
Theorem 1. Let & be given event set and b be the temporal batch size. For given MDGNN parameterized by 9 with training procedure of Eq. 3 have E [[( - VC(0)
18 o2 min
The proof of Theorem I can be found in Appendix B Theorem provides interesting insight into MDGNN training: while is commonly believed that a smaller temporal batch size leads to better information granularity (less subject to pending events), the variance of the gradient for the entire epoch can be significant when the temporal batch size is small. This underscores the unexpected benefit of larger temporal batch sizes, as they exhibit greater robustness sampling noise.
Next, we examine how the size of the temporal batch affects the convergence of the training process. We introduce the concept of memory coherence to measure the loss of temporal dependency in learning dynamics. Specifically,, let s denote the memory state of vertex i after event e (i.e., (e) st(t.)), and we define memory coherence follows.
5

---

Published as a conference paper at ICLR 2024
Definition 3 (Memory Coherence),. The memory coherence of an event Eij By is defined as
(Vilegg (0) ; ,9] (.)) Vi(eij, (ew) ,89 ((e))) Me,, (Bx) := min eEP(eig.BB))
(4)
[Vi(en 8 (ou) (8) (1
where Vileey, Sf sy ( is the gradient incurred by event ey with respect s and of²) (simi- larly for Vi(eiy; (04)) (ew)). The memory coherence Mey captures the minimum coherence of the gradients when fresh memory (s; (o(-u) and past memory from the pending events (s) are used. Intuitively, a positive value of Me., indicates that the directions of these gradients are well aligned. In other words, the convergence properties are less affected by the pending events. is worth noting that the memory coherence Hesy be easily computed empirically during the training process.
Theorem 2. Let & be given event set, b be the temporal batch size, and K [El be the number b
of temporal batches. Let ((())) be the gradient of epoch t when all the events are processed in sequential order (no temporal discontinuity) and V ((t))) be the gradient when events in temporal batches are processed in parallel. Suppose the training algorithm is given by,
8. g(i+11) - 0 (L.+1110 0 = 0 [
(5)
where (L.-) is the gradient of batch B when events are processed parallel and ny the learning rale. Suppose (I) be given in Assumption 7 (2) the loss function is continuously differ- entiable and bounded below, and the grudient L-Lipschitz continuous; (3) )) the memory coherence is lower bounded by some constant ps. Choosing step size T)-
'LVKe' we have
(2/KLIC(O)) - C((0)) + VKo max log 112
1 VT
min
EI [[))))))
(6)
USIST
where 8 is the optimal parameter and 0g is the initial parameter.
The proof of Theorem 2 can be found in Appendix C The assumptions regarding the loss function,
as presented in the theorem, adhere to the standard conventions in convergence analysis (Wan et al, 2022 Chen et al.] 2022; Dai al. 2018). These assumptions can be relaxed with more specific insights into the neural architecture employed in MDGNNS
The insights provided by Theorem 2 offer valuable guidance for the training process of MDGNNs. Firstly, it emphasizes the consideration that the choice of the step size ?/ should account for both the temporal batch size (as captured by K), and the memory coherence μ.. Secondly, the theorem sheds light on the pivotal roles played by memory coherence and variance mmax, as captured by Eq. [6] in the convergence behaviour of MDGNNs. Specifically, it underscores the significance of enhancing memory coherence while simultancously reducing maax. These insights provide strong motivation for the development of our proposed method, as elaborated below.
PREDICT-TO-SMen SMOTHH (PRES) METHOD
We introduce PRES, our proposed method improve MDGNN training with large temporal batch sizes. PRES is based on the theoretical analysis presented in the previous section. It consists of two main components: 1) an iterative prediction-correctin scheme that mitigate the impact of temporal discontinuity in training MDGNNs (optimizing the second term of Eq. 6) and 2) a memory smooth- ing objective based on memory coherence (optimizing the first term of Eq. 6.. An overview of PRES is given in Fig. A.2 The pseudo-code for the complete procedure is provided in Appendix A
5.1 ITERATIVE PREDICTION-CORRECTIOO SCHEMB
To tackle temporal discontinuity, we employ an iterative prediction-correction scheme inspired by filtering mechanisms used in control systems to adjust noisy measurements. The scheme consists of two steps: 1) prediction and 2) update. In the prediction step. a prediction model used to estimate the current state of the system based the previous state information. In the update step,

---

Published as a conference paper at ICLR 2024
the estimation from the prediction model is employed incorporate measurement information and improve the state estimation (reducing the noise to the true value). Here, we treat the memory state with pending events within the temporal batch as noisy measurement.
Our prediction model tasked with modelling the change in the memory state 8; for each vertex i, represented 8xg- We employ Gaussian Mixture Model (GMM), P(6x;) Σj=1 ((.,, 14," x( as the parametric model to estimate the distribution of 8s.- w represents the number of components in GMM, Oj denotes the weights of each components, and ( and x the mean and covariance matrices for the components. We set negative event types in temporal link prediction. Using this GMM, we predict the newest memory state s; (t2) of vertex based on its previous memory state 8; (t2) and the estimated transition:
2 to model the positive and
(t2) st(11) + (12 f1 )5.
To make of the predicted value, we consider the memory state as noisy measurement affected by temporal discontinuity. We then fuse the predicted value and the noisy measurement to obtain better estimate of the memory state. The correction step is given by:
((2) = (1 - 2) (s:((2),
(8)
where y is a learnable variable controlling the fusion the predicted and the current memory state.
Then, we adopt the common approach, Maximum Likelihood Estimation (MLE), for estimating and updating the parameters of GMM. Nevertheless, applying MLE naively would require storing the complete history of each vertex, resulting in substantial memory overhead. To address this, utilize the variance formula Var(X) = E E[x*] E[X]². which allows us to keep track of only a few key values for parameter updates. For each event of type j. the parameter updates are calculated as follows:
((0) = s(12) - (1(12), c = + 56)
()) = + (8(2)) n = ni +1, ( = = 1/s5//nn (1) I - /n - ()
where n -/ (0) are trackers for the number of events, the sum of d( and the sum of (8!())
To summarize, the iterative prediction-correction scheme perceives the memory state impacted by temporal discontinuity as a noisy measurement. It then aims to utilize historical information of the memory state and a GMM prediction model to mitigate this noise/variance, akin the filtering techniques employed in control systems.
5.2 MEMORY COHERENCE SMOOTHING
Based on the preceding analysis, we have identified memory coherence (μ) as a pivotal determinant in the convergence rate of MDGNNs when dealing with large temporal batches. To enhance the statistical efficiency of MDGNN training with larger temporal batches, we propose novel learning objective aimed at promoting larger memory coherence. This learning objective is expressed as:
(()) S(B,) |B-)I ||(().]]
I(Be)
+
(10)
prediction loss
memarr coherence
where (B,) and S(B,) represent the previous and new memory states of vertices within batch B, respectively. The hyperparameter 3 governs the strength of the regularization effect.
Intuitively, if the memory coherence within batch B, is low, Eq. 10 incurs a substantial loss. Conse- quently, this learning objective steers the training process towards parameter values less susceptible to the influence of pending events, effectively enhancing memory coherence. In accordance with Theorem 2 this enhancement is expected to yield superior statistical efficiency.
5.3 THEORETICAL DISCUSSION OF PRES
Variance. Let STANDARD represent the standard temporal batch training with Eq. and PRES denote our proposed framework. We have the following theoretical guarantee for variance reduction.

---

Published as a conference paper at ICLR 2024
9
APPS)
91
APPS)
TGN
91
TGN
JODIE
JODIE
85
APAN
APAN
10
20
40
60
80
100
10
20
40
60
80
100
(a) WIKI
(b) REDDIT
Figure 3: Performance of baselines under different batch sizes. The x-axis represents the batch size,
while the y-axis represents the average precision (AP). The results are averaged over five trials.
86
98
APPN)
APPN)
APP%)
91
TGN TON-PRES
APAN
JODIE
APAN-PRES
87
JODIE-PRES
88
15
20
10
Batch Size
Batch Size (x100)
Batch (x100)
(a) TGN
(b) APAN
(c) JODIE
Figure 4: Performance of baseline methods with and without PRES under different batch sizes on WIKI dataset. The x-axis represents the batch size (multiplied by 100), while the y-axis represents the average precision (AP). The results are averaged over five trials with 8 0.1 for PRES.
Table 1: Performance comparison of existing MDGNNs with and without PRES. The results are averaged over five independent trials with 8 = 0.1 for PRES and 50 epoches. The training efficiency improvement with PRES highlighted in bold. (*s) indicates that it takes * seconds for one epoch.
Datases
REDDIT
WIKI
MOOC
LASTEM
GDELT
Model
Specclup
Speodup
APPS)
Speodup
AP(S)
Speedup
Speedup
TGN
0.14
Ix (321s)
97.8 0.1
1x (87c)
98.8 0.2
1x (212x)
72.1
1x (648s)
96.8 0.1
[x(1325s)
TON-PRES
0.2
34x (91s)
97.5 0.3
3,2x (27k)
97.5.00.11
7122 (71s)
(716)
28x (231s)
96.0
277 (191s)
JODIE
9..
1x (x76)
9.1
1x (62x)
9..0
1x (152x)
1.
1x (152g)
95.1
Ix (1123s)
(876)
0.1
(237))
94.1
(101s)
APAN
98.6 0.1
1x (281s)
99.0 0.1
Tx (71x)
98.5
1x (1736)
69.8 1.6
1x (521s)
0.2
1x (12156)
APAN-PRES
= 0.1
2.2% (127%)
98.5 0.1
2.9x (246)
9.1
2.0x (866)
67.9 1.9
18x (289k)
9..
244x (506s)
Proposition (Informal). If the memory transition with temporal discontimuity can be approximated by linear state-space model with Gaussian noise, PRES achieve smaller noise/variance compared to STANDARD.
The formal version of Proposition T along with its proof provided Appendix D Proposi- tion provides validation PRES. It states PRES can effectively mitigate the variance induced by temporal discontinuity, enabling the training of MDGNN with larger temporal batches. This enhancement, in turn, bolsters the training efficiency without compromising overall performance.
Complexity. The computational complexity of PRES is O( with [B] representing the batch size. As for storage complexity, which arises from the trackers employed in Eq. g it scales at the order of O(V)), where [V] denotes the vertex count. In cases where stringent memory constraints are in place, one can adopt a pragmatic approach by selecting a subset of vertices to serve as an anchor set. This set, in turn, facilitates efficient heuristic for handling the remaining vertices. For a more in-depth exploration of this strategy, please refer to the appendix.
6 EXPERIMENT
We present an experimental study to validate our theoretical results and evaluate the effectiveness of our proposed method, PRES. Due space limit, present subset of results in the main paper and provide a more comprehensive description of the datasets, experimental setup, testbed, and remaining results (such ablation study and GPU memory utilization) in the Appendix
Datasets and Baselines. We use four public dynamic graph benchmark datasets (Kumar et al. 2019), REDDIT, WIKI, MOOC, LASTEM and GDELT. Details these datasets are described in the appendix. We use three state-of-the-art memory-based MDGNNs models: JODIE (Kumar al.
8

---

Published as a conference paper at ICLR 2024
2019). TGN (Rossi et al., 2021) and APAN (Wang al. 2021), and adopt the implementation of these baselines from (Zhou et al. 2022: Rossi et al. 2021), J. We closely follow the settings of (Zhou
et al. 2022) for hyperparameters and the chronological split of datasets, as described in more detail in the appendix. We mainly use average precision as the evaluation metric.
6.1 EXPERIMENTAL RESULTS
Table 2: Performance of MDGNNS w/w.o. PRES
Performance vs. Batch Size. We first ex- amine the relationship between temporal batch size and the performance of baseline methods w./w.o. PRES. Fig. provides insights into this relationship by illustrating the performance of the baselines in the small batch size regime. These results align with Theorem demon- strating that contrary to common belief, smaller temporal batch sizes can lead to larger variance in the training process, and poorer performance (even non-convergence). Fig. further com- pares the baseline methods w./.oo. PRES. The results indicate that 1) the performance of baselines indeed decreases as batch size increases and 2) baselines trained with PRES are less susceptible to performance degradation due increased batch size, validating the effectiveness of PRES.
on node classification task. ROC-AUC (%)
REDDIT
WIKI
MOOC
TGN TGN-PRES
65.6 0.8
86.8
57.8 2.3
65.0 1.2
84.7 + 1.8
56.9 a 3.3
JODIE
6.1
85.2 4: 1.5
66.5 + 2.0
JODIE-PRES
60.9 0.5
85.0 + 1.4
66.0 + 1.7
APAN APAN-PRES
64.3 1.4
90.0 1.7
63.5
6.2
65.5 8.16
98
Overall Training Efficiency Improvement. The next ex- periment demonstrates the potential speed-up and efficiency improvement that baseline methods can achieve when trained with PRES. T Table [lblidd comprehensive com- parisons of the performance of baseline methods with and without PRES. The results clearly illustrate that our approach enables significant increase in data-parallel processing by using larger temporal batch sizes, without sacrificing overall performance in terms of both link prediction and node clas- sification tasks. These findings highlight the effectiveness of method in significantly scaling up MDGNN training while maintaining comparable performance.
APP%)
TGN TON-PRES
100
200
Itertiio
Figure 5: Statistical efficiency of baseline method w./.oo PERS. x- axis is the training iteration and y- axis the average precision. 0.1 is used in PRES.
Effectiveness of Memory Smoothing. Next, show the effectiveness of the memory smoothing objective. Fig. 5 demonstrates that incorporating the memory coherence-basee objective indeed leads better statis- tical efficiency of MDGNN training, validating the effectiveness of our approach.
7 CONCLUSION
This paper studics the effect of temporal batch size on MDGNN training. We present a novel the- oretical analysis of how the size of the temporal batch affects the learning procedure of MDGNNs. Contrary to the common belief, we show a surprising advantage of training MDGNNs in large tem- poral batches (more robust noise). In addition, we present the first convergence result of MDGNN training, illustrating the effecting factors. Based on the analysis and the filtering mechanism, we pro- pose a novel MDGNN training framework, PRES, that can mitigate the temporal discontinuity issue and improve the convergence of MDGNN training with large temporal batch sizes. The scalability enhancements brought about by PRES hold broader implications for both researchers and practi- tioners, as they extend the applicability of MDGNNs to broader spectrum of real-world scenarios featuring large-scale dynamic graphs.
Limitations and Future Work. In this study, our analysis centres on the MDGNN family. How- ever, it is worth noting that the memory module has demonstrated effectiveness in capturing the dynamics of time series data, which exhibits (different) temporal dependencies. An intriguing enue for future research would be to expand our analysis and methods to encompass the realm of time series analysis.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENT
We would like to thank the anonymous reviewers and area chairs for their helpful comments. JS and CW are supported by grants from Hong Kong RGC under the contracts HKU 17207621, 17203522 and C7004-22G (CRF). DZ is supported by NSFC 62306252.
REFERENCES
Takuya Akiba, Shuji Suzuki, and Keisuke Fukuda. Extremely large minibatch sgd: Training resnet- 50 on imagenet in 15 minutes. arXiv preprint arXiv:1711.04325. 2017.
Léon Bottou, Frank Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM review, 60(2):223-3) 331, 2018.
Jianfei Chen, Jun Zhu, and Le Song. Stochastic training of graph convolutional networks with variance reduction. arXiv preprint arXiv:1710. 2017.
Jie Chen, Tengfei Ma, and Cao Xiao. Fastgen: fast learning with graph convolutional networks via importance sampling. arXiv preprint arXiv: 10247, 2018.
Xinshi Chen, Yan Zhu, Haowen Xu, Mengyang Liu, Liang Xiong, Muhan Zhang, and Le Song. Efficient dynamic graph representation learning scale. arXiv preprini arXiv:2112.07701 2021.
Yangrui Chen, Cong Xie, Meng Ma, Juncheng Gu, Yanghua Peng, Haibin Lin, Chuan Wu, and Yibo Zhu. Sapipe: Staleness- aware pipeline for data parallel dnn training. In Advances in Neural Information Processing Systems, 2022.
Wei Dai, Yi Zhou, Nanqing Dong. Hao Zhang, and Eric P. Xing. Toward understanding the impact of staleness in distributed machine learning, 2018.
Swapnil Gandhi and Anand Padmanabha Iyer. P3: Distributed deep graph learning at scale. In OSDI, pp. 551-568, 2021.
Robert Mansel Gower, Nicolas Loizou, Xun Qian, Alibek Sailanbayev, Egor Shulgin, and Peter Richtárik. Sgd: General analysis and improved rates. In International conference on machine learning, pp. 5200-5209. PMLR, 2019.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint orXiv:17066 02677, 2017.
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. Advances in neural information processing systems, 30, 2017.
William L Hamilton. Graph representation learning. Synthesis Lectures Artifical Intelligence and Machine Learning, 14(3):1-159, 2020.
William L. Hamilton, Rex Ying, and Jure Leskovec. Inductive representation learning on large graphs, 2018.
Wenbing Huang, Tong Zhang, Yu Rong, and Junzhou Huang. Adaptive sampling towards fast graph representation learning. arXiv preprint arXiv: SEEiv:180.0055. 05044,,,,,, 2018.
Seyed Mehran Kazemi, Rishab Goel, Kshitij Jain, Ivan Kobyzev, Akshay Sethi, Peter Forsyth, and Pascal Poupart. Representation learning for dynamic graphs: A survey. The Journal of Machine Learning Research, 21(1):2648-272) 2020.
Thomas Kipf and Max Welling. Semi-supervised classification with graph convolutional net- works, 2017.
Srijan Kumar, Xikun Zhang, and Jure Leskovec. Predicting dynamic embedding trajectory in tem- poral interaction networks. In Proceedings of the 25th ACM SIGKDD international conference knowledge discovery & data mining, pp. 1269-1278, 2019.
10

---

Published as a conference paper at ICLR 2024
Tao Lin, Sebastian U Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't use large mini- mini-batches, use local sgd. arXiv preprint 100.::1888..0211111 2018.
Aldo Pareja, Giacomo Domeniconi, Jie Chen, Tengfei Ma, Toyotaro Suzumura, Hiroki Kaneza- shi, Tim Kaler, Tao Schardl, and Charles Leiserson. Evolvegen: Evolving graph convolutional networks for dynamic graphs. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 5363-5370, 2020.
Jingshu Peng, Zhao Chen, Yingxia Shao, Yanyan Shen, Lei Chen, and Jiannong Cao. Sancus: sta le ess-aware с omm u incation-avoiding full-graph decentralized training in large-scale graph neural networks. Proceedings of the VLDB Endowment, 15(99:1937-1950, 2022.
Farimah Poursafaei, Shenyang Huang, Kellin Pelrine, and Reihaneh Rabbany. Towards better evaluation for dynamic link prediction. In Neural Information Processing Systems (NeurIPS) Datasets and Benchmarks, 2022.
Xin Qian and Diego Klabjan. The impact of the mini- batch size on the variance gradients in stochastic gradient descent. arXiv preprint arXiv:2004.13146, 2020.
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System opti-
mizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery Data Mining, pp. -3006, 2020.
Emanuele Rossi, Ben Chamberlain, Fabrizio Frasca, Davide Eynard, Federico Monti, and Michael Bronstein. Temporal Graph Networks for Deep Learning on Dynamic Graphs, In Proceedings of International Conference on Learning Representations, 2021.
Aravind Sankar, Yanhong Wu, Liang Gou, Wei Zhang, and Hao Yang. Dysat: Deep neural rep- resentation learning on dynamic graphs via self-attention networks. In Proceedings of the 13th international conference on web search and data mining, pp. 519-527, 2020.
Mark Schmidt, Nicolas Le Roux, and Francis Bach. Minimizing finite sums with the stochastic average gradient. Mathematical Programming, 162:83-112, 2017.
Guangming Sheng, Junwei Su, Chao Huang, and Chuan Wu. Mspipe: Efficient temporal gnn train- ing via staleness-aware pipeline, 2024.
Joakim Skarding, Bogdan Gabrys, and Katarzyna Musial. Foundations and modeling of dynamic networks using dynamic graph neural networks: A survey. IEEE Access, 9:79143-79168, 2021.
Amauri Souza, Diego Mesquita, Samuel Kaski, and Vikas Garg. Provably expressive temporal graph networks. Advances in Neural Information Processing Systems, 35:32257-32206 2022.
Junwei Su, Difan Zou, Zijun Zhang, and Chuan Wu. Towards robust graph incremental learning on evolving graphs. In International Conference on Machine Learning, pp. 22728-32748,, PMLR, 2023.
Junwei Su, Shan Wu, and Jinhui Li. Mrgglefffective temporal correlation discerning through multi- modal temporal relational graph learning. 2024.
Rakshit Trivedi, Hanjun Dai, Yichen Wang, and Le Song. Know-evolve: Deep temporal reasoning for dynamic knowledge graphs. In international conference on machine learning, pp. 3462-3471. PMLR, 2017.
Rakshit Trivedi, Mehrdad Farajtabar, Prasenjeet Biswal, and Hongyuan Zha. Dyrep: Learning rep- resentations over dynamic graphs. In International conference on learning representations, 2019.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pictro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv:1710.1090) 2017.
Borui Wan, Juntao Zhao, and Chuan Wu. Adaptive message quantization and parallelization for distributed full-graph gnn training. Proceedings of Machine Learning and Systems, 5, 2023.

---

Published as a conference paper at ICLR 2024
Cheng Wan, Youjie Li, Cameron R Wolfe, Anastasios Kyrillidis, Nam Sung Kim, and Yingyan Lin. Pipegen: Efficient full-graph training of graph convolutional networks with pipelined feature communication. arXiv preprint arXivv2203. 2022.
Xuhong Wang, Ding Lyu, Mengjian Li, Yang Xia, Qi Yang, Xinwen Wang, Xinguang Wang, Ping Cui, Yupu Yang, Bowen Sun, et al. Apan: Asynchronous propagation attention network for real-time temporal graph embedding. In Proceedings of the 2021 international conference on management of data, pp. (628-2633) 2021.
Yufeng Wang and Charith Mendis. Tgopt: Redundancy-aware optimizations for temporal graph attention networks. In Pruceedings of the 28th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming, pp. 354-368, 2023.
Blake Woodworth, Kumar Kshitij Patel, Sebastian Stich, Zhen Dai, Brian Bullins, Brendan Mema- han, Ohad Shamir, and Nathan Srebro. Is local sgd better than minibatch sgd? In International Conference on Machine Learning, pp. 10334-10343, PMLR, 2020.
Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Phillip. A comprehensive survey on graph neural networks. IEEE transactions on neural networks and learning systems, 32(1):4-24, 2020.
Da Xu, Chuanwei Ruan, Evren Korpeoglu, Sushant Kumar, and Kannan Achan. Inductive represen- tation learning temporal graphs. arXiv preprint 20X2::0002.07662. 2020.
Rex Ying, Ruining He, Kaifeng Chen, Pong Eksombatchai, William I.. Hamilton, and Jure Leskovec. Graph convolutional neural networks for web-scale recommender systems. Pro- ceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, Jul 2018. doi: 10. 10.1145/3218819332299
URL http: wow...r/g........
321 9819 32 19890
Yang You, Jonathan Hseu, Chris Ying, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large- batch training for Istm and beyond. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-16, 2019.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song. James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes, 2020.
Yao Zhang, Yun Xiong, Dongsheng Li, Caihua Shan, Kan Ren, and Yangyong Zhu. Cope: modeling continuous propagation and evolution on interaction graph. In Proceedings of the 30th ACM International Conference Information & Knowledge Management, pp. 2627-2636, 2021.
Yao Zhang, Yun Xiong, Yongxiang Liao, Yiheng Sun, Yucheng Jin, Xuchao Zheng, and Yangyong Zhu. Tiger: Temporal interaction graph embedding with restarts, 2023.
Hongkuan Zhou, Da Zheng, Israt Nisa, Vasileios loannidis, Xiang Song, and George Karypis. Tgl: A general framework for temporal gnn training on billion-scale graphs. arXiv preprint 2rXhv:2203.14882 2022.
Hongkuan Zhou, Da Zheng, Xiang Song, George Karypis, and Viktor Prasanna. Disttgl: Distributed memory-based temporal graph neural network training. In Proceedings of the International Con- ference for High Performance Computing, Networking, Storage and Analysis, pp. 1-12, 2023.
Difan Zou, Ziniu Hu, Yewen Wang, Song Jiang, Yizhou Sun, and Quanquan Gu. Leyer-dependent importance sampling for training deep and large graph convolutional networks. Advances in neural information processing systems, 32, 2019.
12

---

Published as a conference paper at ICLR 2024
Negative
h144)
hz(4)
2
(((1)))
dec
loss
samp
h2(t))
h²(t))
p((23)]tt)
emb
4
2
Node Embeddings
Edge probabilities
3
4
Batch
(1,
msg
(,,4)
m²(t))
Positive
$2((2)
mz((2)
st((t)
my(tz)
Messages
(Updated) Memory
Figure 6: Illustration of MIGNN Training Procedure. Fig. A.1 depicts the training flow of MDGNN. The incoming batch serves as training samples for updating the model and updating the memory for the subsequent batch.
A ALGORTIHM AND FURTHER DISCUSSION
A.1 TRAINING PROCEDURE MDGNNS
The training procedure of an MDGNN involves capturing temporal dynamics and learning repre- sentations of nodes in dynamic graph. MDGNNs leverage the concept of memory to model the evolving state of nodes over time. The training process typically follows an encoder-decoder frame- work. In the encoder, the MDGNN takes dynamic graph as input and generates dynamic repre- sentations of nodes. This achieved by incorporating memory modules that store and update the past states of nodes based on their interactions and temporal neighbours. The encoder also includes message-passing mechanisms to propagate information through the graph. The decoder utilizes node representations generated by the encoder to perform downstream tasks, such temporal link prediction or node classification. MDGNNs are commonly trained using a self-supervised temporal link prediction task, where the decoder predicts the likelihood of edge between nodes based on their representations.
The training procedure of memory- dynamic graph neural network (MDGNN) involves several steps. First, the dataset is divided into training, validation, and test sets using chrono- logical split. More concretely, suppose given event from time interval [0, 7". chronologi- cal split partition the dataset into [0, Tiram]] (training set), [T [Tirain: T Vaalidationl (validation set), and [T Validatto Tist] (test set). From now on, we focus on the training set and drop the subscript. The training set is then further divided into temporal batches, where each batch consists of consecutive events the dynamic graph. In addition, negative events are sampled from the rest of the graph to provide the negative signal. During training, MDGNNs often adopt lag-one procedure. This means that the model uses the information from the previous batch to update its memory state and generate node embeddings for the current batch. This lag-one scheme helps maintain temporal consistency and ensures that the memory- based model captures the correct temporal patterns. Fig, A. provides graphical illustration of the training flow of a batch. Fig. A.I provides graphical illustration of the training flow between epochs. The podudo-code of the training procedure with cross cross- entropy is summarized in Algorithm
13

---

Published as a conference paper at ICLR 2024
2
t-Batch
t-Batch
t-Batch
C(t-Batch 1,
CfteBa 2,)
C(t Batch K. 0f
Epech
6
2
K
t-Batch
t-Batch
t-Batch K
Clt-aatch 1, O(t-Batch of
C(t-Batch K, of
Epoch
6/"
Figure 7: Illustration of MDGNN Training Procedure. Fig.A.1 A. depicts the training flow of MDGNN for two epoches. The incoming batch serves as training samples for updating the model and updating the memory for the subsequent batch. The model parameter is carried through the second epoch.
Algorithm Standard Training Procedure for Memory- DGNN
Initialization: So +- {Initialize the memory vectors to be zero] for t=1 to T
for B, E B2 Bx do
B, + Sample negative events
B, = B, UB,
B; 1 Temporal batch from last iteration
M, Insg(S;-3) B,-1)
S; mem (S M,)
H; = emb(S,, {where N; is the (Temporal) neighbourhood of vertex
Compute the loss (e.g., binary cross-entropy) and run the training procedure (e.g., backprop- agation)
C(H,, B.)
end for end for
14

---

Published as a conference paper at ICLR 2024
mg(4)
sy(5)
2
milte)
Batch
4 f
Update
8y(6)
-emb
Pg(4))
dec
P((1,2)14 2214)
loss
(((23))ltg
Edge probabilities
8,(4 )
Predict
8z((g)
Figure 8: An illustration of MDGNN training with PRES. A prediction model is inserted into the training procedure to adjust the memory state for the intra-batch dependency, Arrows marked with the same number indicate transitions happening at the same stage.
A.2 PRES
PRES consists two main components: 1) an iterative prediction-correction scheme that incorpo- rates approximation inference into training MDGNNs, and 2) novel memory smoothing objective based on memory coherence. An overview PRES given Fig. A.2 The pseudo-code for the complete procedure is summarized in Algorithm 2
Algorithm PRES
Initialization: So + 0 for t=1 to T do
{. th +- 0
for B, E Bı,..., Bx do
B, +- Sample negative events
B, B, B.
B; +- Temporal batch from last iteration
S Speev (B,-1) <- (,-1(B;_)) {Store the previous memory state of the vertices to be updated)
Mf msg(SS B(-1)
S. = mem 1, M.)
H. emb(S,. A.)
S predict memory state with the prediction model Eq. 7
7 + (1 S
Compute the loss and run the training procedure (e.g., backpropagation)
C(B.) + [1 - Sprev (Bi-1) 5:
pree (B-
||
8g. 5, S
update& &, v. n with Eq.
end for end for
15

---

Published as a conference paper at ICLR 2024
B PROOF OF THEOREM 1
In this appendix, we provide proof of Theorem
Proof. First, recall that for a given event set size M and partition size of k, the training procedure of MDGNN is given by
C(0(0) = =-=, 1,(0--1)),
(1I)
i=1
= - -(B,, 0(1-11)
where K M 0 represents the model parameters of the MDGNN, L()(0)) is the total loss for x. the entire epoch with initial parameters (()), (.)) is the loss for batch B,, and ((.) denotes the loss function. Here, we consider the standard setting where negative event sampling is used to facilitate training. We denote the gradient with sampled negative events as (((6)())) and (L,( is used for update instead.
With the above set-up and by definition of variance, we can write Var[VL((A0)) follows,
K
Var(VL((00)] = B(( (L((i(i-1)
Σ E(V
(12)
By Assumption we assume that the negative sampling is unbiased and have that (L.(-())) an unbiased estimate V (O(0). Then, we can arrange the terms in the above equations and get,
K
K
Var [(O() = B((
V (1.
E(VL, (0
1=1
K
B[(
V (0(i--11))
Σ
(0(--111)
(13)
[
V ( (L((
E ( (L.(0i-11)
Again, by Assumption we have that ((,(0(-11) and (0(i-111)) has bounded difference with- out lower bound of This leads
Var (C(00))
-
(14)
Ke
Substitute in K 21 we have
- Var VC > Ke
M
(15)
>
16

---

Published as a conference paper at ICLR 2024
C APPENDIX.PROOF OF THEOREM 2
In this appendix, we provide proof of Theorem 2 to keep track of the notations and the evolution of the parameters induced by batches and epochs.
The key challenge of the proof for Theorem 2 is
Next, we provide the proof of Theorem 2
Proof. Suppose Bi,.., Bk are the K batch. Let's recall that the gradient of an epoch given as follows.
C(0(0)) Σ 10(0011))
(16)
= 0(1-1)).
Let's start with considering the progression within an epoch, and let 0, 0! denote the model pa- rameter of epoch after being trained on batch B;-
We denote the gradient with sampled negative events as (L,t-(())), and (0((0(()1) is used for update ins the update rule is given by,
0 ((i-1) -1 V(
(17)
equivalently,
8 = ((:(((--))
(18)
By the L-Lipschitz property of the objective function VL, we have the following inequality,
C(0 < ) + (0( of VC(O.)) + 2
$10 - 1000
(19)
Substitute in 0 ofi 1/ V (0(6-11)) have that,
((f(iti)) S ((8(0) + +88 (8( 0 VC(E!)) + L glee 7 6.0012
(20)
= -()))) ne ( (O. v + Lii /V2(0!))
By the premise of the memory coherence in the theorem and Assumption | we have that the inner product of the gradient with and without pending events is bounded by 14, and that the negative sampling is unbiased and has upper bounded variance. Therefore, taking the expectation of each term and substituting the assumption into the expression, we arrive the following inequality.
Lai I (
s B[C(0 - ( (V2(0 VC(O))) +
Lun
s c() M.E 11 ||
+ 2
(L3 (L + (21)
Laf
Luf o²
=( (-TH +
y C (6 |13)
+
2
2
Rearrange the equation above, we get that
(n Lmk (( ((O - L()((tt)) +
Lnf o²
(22)
2
Telescope sum over the entire epoch, namely all the batches, i 1.... K, we have,
Σ
K
(ne - LuZ
]E|
([] < ((, - (( + Σ
K
Laf 0²
(23)
2
2
i=1
17

---

Published as a conference paper at ICLR 2024
By definition, we have that )(0,k)) L(0())).
The above equation equivalent to the following,
K
Lad 2
K Σ
L
Σ
Map
S 2() [(((+1) +
o²,
(24)
2
i=1
Since we are making the learning rate independent of batch number i, we can rearrange the equation above follows,
Lof 2
K
Lmg
EII ((6 - L(0(41) + +K
o²²
(25)
Now telescope sum over the epochs, namely, 1,.., T, we have,
Lug
Lmi
Σ
mpp
E
((0 L(Or) + Σ
K
o²
(26)
2
By set up, have that Ei=1 Bill expression for the above equation as follows.
VC(O.)). Then, we can obtain equivalent
Luf
I 2
T
Σ
[1 =()))))))))))))))))))))))))))))) [(OT) Ko²
(27)
2
Note that the choice of stepsize guarantees that
M Lsn 2
for all Thus, can divide both side by ΣT (Nea
Lu 2
Rearrange the terms and taking the
minimum among & 1,.... T for [ )1(
obtain the following expression,
2() L(Or)
Ko² 2T= n
min EI 11
(28)
Σ - (1:11
Lu Di-1nr
2
Lxd
2
Luf
By the choice of Th, we have that 77/4
2 and this leads to.
2
(C(0))) [((T))
Ko² L ΣT mi
min sst
Σ 11/12
Σ 11//12
(29)
2(( 0 L(0r))
2K 2 LET 1/.
Σ Def
Σ 1//44
Since 8* is the optimal parameter, by definition we have that, C(0") < C(Or).
Substitute this in the equation, we have,
2(((0)) (0) C(OT))
2Ko ΣT TK
min E
Σ11 THM
Σt Thff4
(30)
2(2 (0) -((8))
2Ko L ΣT= 1/2
<
Σ7 11 7//4
Σl= 1: M7/4
Finally, substitute in the stepsize T/t LVKt
get that
2 (VL((C(00)
-((0)
1
min [[
V log T)
(31)
oStST
112
VT
18

---

Published as a conference paper at ICLR 2024
D PROOF OF PRES THEORETICAL GUARANTEE
In this appendix, we provide proofs for the theoretical guarantees of PRES, namely, Proposition
Next, present the formal version of Propostion 1 It should be noted that the notations used here are slightly different from the notations used in the main paper. We denote that s;(() to be the memory state of vertex i time / when processing events sequentially (i.e., there is no intra-batch dependency). We denote (t) to be the memory state of vertex at time t when pending events are processed in parallel. We denote s; similarly for (t-).
be the memory state of vertex from the last update, and
We assume the difference between these two is given by zero-mean Gaussian noise, i.e.,
(t) = (t) + o0)
(32)
where N\O, 01) is the zero-mean Gaussian noise with variance d In other words, we have that
8.. (t) - (1) = (0)
Furthermore, we assume that the state transition for memory state without intra-batch dependency loss is given by the following linear state-space model with Gaussian noise,
(t) = s, st (t -- -t
(33)
Recall that the prediction model we used is,
(t) = st(t (t- (t-
(34)
and the final state estimation is given by.
(t) = (t) + (1 - ())
(35)
Proposition Under the set-up above, we have that
(t) (f) s (())):(1) (t(I]]]-
Proposition that the difference between the memory state estimated with our proposed model and the memory state without intra-batch dependency smaller than the one with the standard batch processing. This means that our proposed method can mitigate the effect brought by batch processing, as it would be closer to the "true" memory state and therefore has a smaller variance as claimed in Proposition| I Now, we show that Proposition 2 is true.
Proof. First, substitute in
(t) 7)((t) + (1-1))))000 (s:((t)....
in
[|=((1) - sx(1l1
we have that
((((()) - s((i)l = (1)8((1) + (1 (1-))(4(0)-
- (1)l
(36)
(1)8:(1) (1-)+++1))(000))
Notice that if 1, then we have
[(: st(1)l = (1)8((1) - s:(() + -1 (6:(00) (t) (1)l
(37)
As y is a learnable parameter between [0, 1). this guarantees that (t) st)ll no worse than
(t) - st) Next, we show that there is indeed a possible improvement. To do so, we substitute
$; (t) = (t) + (1 - (t)
into the equation above, and we get
(1) s.)))] = (1))((() - (1 (1-
(38)
19

---

Published as a conference paper at ICLR 2024
Rearranging the terms, get,
[|x - se (t)l = (l)(() - s; s,(t)) (11 - (t))l
(39)
As y is a learnable parameter between [0, 11, as long show that (š:(1) s.(t)) - 0,
then we have,
(3((1) (t(1) = (l)()(()) s s, (1-7)(8:00) ((())))))))))))))))))))) (())) (l)()((() ((Ti)))
(40)
- (t), st
To show (5:(1) st)) 0, we substitute the modelling assumption and have,
(t) - (t) st (t- (V(, - s;(t)) + - (t- 02) 8..] As the number of samples increase, we have that Blo.,] + (2, and therefore. E(b:(t) s;()) +.
(41)
Therefore, arrive that,
B (t) (t)]ll s (t) st(I]]]].
E EXPERIMENT DETAILS
E. HARDWARE AND SOFTWARE
All the experiments of this paper are conducted on the following machine
CPU: two Intel Xcon Gold 6230 2.1G, 20C/40T, 10.4GT/s, 27.5M Cache, Turbo, HT (125W) DDR4-2933
GPU: four NVIDIA Tesla V100 SXM2 32G GPU Accelerator for NV Link
Memory: 256GB (8 x 32GB) RDIMM, 3200MT/s, Dual Rank
OS: Ubuntu 18.04LTS
E.2 DATASET
E.2.1 DESCRIPTION
We use the following public datasets provided by the authors of JODIE Kumar et al. (2119). (1) Wikipedia dataset contains edits of Wikipedia pages by users. (2) Reddit dataset consists of users' posts on subreddits. with dynamic labels indicating if they get banned after some events. (3) MOOC dataset consists of actions done by students on online courses, and nodes with dynamic labels indicating if students drop out of courses. (4) LastFM dataset consists of events that users listen to songs. MOOC and LastEM datasets are non-attributed. The statistics of the datasets are summarized in Table 3
these two datasets, edges are with 172-d feature vectors, and user nodes are
E.2.2 LISCENCE
All the datasets used in this paper are from publicly available sources (public paper) without a license attached by the authors.
E.3
For the hyper- parameters each experiment, we follow closely the ones used in TGL Zhou et al. 2022). The only parameter, control in the experiment is the temporal batch size used for training.
20

---

Published as a conference paper at ICLR 2024
Table 3: Detailed statistic of the datasets.
Datasets
Wikipedia
Reddit
MOOC
LastFM
GDELT
# vertices
9,227
10,984
7,144
1,980
16,682
# edges
157,474
672,447
411,749
1,293,103
1,912,909
# edge features
172
172
0
0
186
APTR)
91
AP(X)
JODIE
JODIE
85
APAN
APAN
10
40
60
100
10
40
60
100
Batch
Betch
(a) WIKI
(b) REDDIT
70
AP(%)
AP(N)
81
TGN
TGN
JODIE
JODIE
Batch Size
Belch Size
(c) MOOC
(d) LASTEM
Figure 9: Performance of baseline methods with and without PRES under different batch sizes. The x-axis represents the batch size, while the y-axis represents the average precision (AP). The results are averaged over five trials.
F ADDITIONAL RESULTS
In this section, we present additional experimental results.
F.1 BATCHSIZE AND PERFORMANCE
In this subsection, we provide the additional results the remaining dataset for the relation between batch size and the performance of MDGNNs. The results are reported in Fig. 9 this experiment, fix the default hype parameter provided by TGL and change the temporal batch size used in each dataset.
F.2 EFFECTIVENESS OF PRES
In this subsection, we provide the additional results on the remaining dataset for the effectiveness of PRES. The results are reported in Fig. 10 Fig. 11 Fig. and Fig. [3] Similar to the last experiment, we fix the default hype parameter provided by TGL. and change the temporal batch size used each dataset, and compare the performance when training with and without PRES.
F.3 TRAINING EFFICIENCY IMPROVEMENT
In this subsection, we provide the additional results on the remaining dataset for the effectiveness of PRES. The results reported in Fig. 14 Similar the last experiment, fix the default hype parameter provided by TGL as well as the batch size and compare the statistical efficiency when training with and without PRES.
21

---

Published as a conference paper at ICLR 2024
98
98
AP
AP(%)
94
93
TGN
APAN
TGN-PRES
APAN-PRES
90
88
10
15
20
1
10
15
20
Batch Size (x100)
Batch Size (x100)
(a) TGN
(b) APAN
95
AP(%)
91
JODIE JODIE-PRES
87
Batch Size (x100) (c) JODIE
Figure 10: Performance of baseline methods with and without PRES under different batch sizes on WIKI dataset. The x-axis represents batch size (multiplied by 100), while the y-axis represents the average precision (AP). The results are averaged over five trials with 3 = 0.1 for PRES.
F.4 ADDITIONAL EXPERIMENT TO SHOW COMPARISON WITH OTHER EFFICIENT METHODS FROM OTHER DOMAINS
In this section, provide additional experimental results illustrate the relation between relative speedup and the affected performance from the other domains as a comparison.
The studies of efficient method (even at the cost of suboptimal settings for accuracy) is exemplified (but not limited) to the following major lines of research:
1. use of staleness: which sacrifices the "freshness information" in the training process to accelerate computation or communication in training
2. use of quantization in both training and inference: which sacrifice the precision of the weight of the model to accelerate computation
3. simpler model architecture: use simple estimation method to accelerate the efficiency of the model
We have sampled the following methods from each category to create a comparison between the gain efficiency and the effect on accuracy: 1) staleness: PipeGCN Wan al. (2022). SAPipe Chen al. 2022). Sancus Peng et al. (2022), 2) quantization: AdaQP Wan el al. (2023), 3)simpler architecture:
FastGCN Chen et (0118).
The values for each method are obtained in the following ways:
For SAPipe and FastGCN, the values are estimated or taken from the papers. The value
of SAPipe is taken from Chen et al. 2022). The underlying tasks for these values
22

---

Published as a conference paper at ICLR 2024
98
98
AP
AP
94
TGN
APAN
TGN-PRES
APAN-PRES
90
90
10
15
20
1
10
15
20
Batch Size (x100)
Batch Size (x100)
(a) TGN
(b) APAN
98
AP(%)
94
APAN APAN-PRES
90
Batch Size (x100)
(c) JODIE
Figure 11: Performance of baseline methods with and without PRES under different batch sizes on REDDIT dataset. The x-axis represents batch size (multiplied by 100), while the y-axis represents the average precision (AP). The results are averaged over five trials with 3 0.1 for PRES.
image classification tasks and language translation tasks. The value of FastGCN is taken from Chen et al.) (2018) and the underlying task is node classification.
For Sancus, AdaQP, and PipeGCN, the values are obtained from running the open-source code:
1. Saneus node classification with GCN the public OGB-product dataset
A node classification with GCN on the public DGB-product dataset
3. PipeGCN3 node classification with GCN on the public OGB-product dataset
3
PRES(our) are the average values computed from Table
F.5 ADDITIONAL EXPERIMENT ON EXTENDED TRAINING SESSION WITH TGN ON WIKI DATASET
F.6
ADDITIONAL EXPERIMENT ON ABLATION STUDY AND GPU MEMORY UITLIZATION
G
RELATED WORK
this appendix, we provide more comprehensive review discussion of related works.
23

---

Published as a conference paper at ICLR 2024
98
98
AP(%)
AP
94
TGN
APAN
TGN-PRES
APAN-PRES
90
90
10
15
20
1
10
15
20
Batch Size (x100)
Batch Size (x100)
(a) TGN
(b) APAN
98
AP(%)
94
JODIE JODIE-PRES
90
Batch Size (x100) (c) JODIE
Figure 12: Performance of baseline methods with and without PRES under different batch sizes on MOOC dataset. The x-axis represents the batch size (multiplied by 100), while the y-axis represents the average precision (AP). The results are averaged over five trials with 3 = 0.1 for PRES.
G.1 DYNAMIC GRAPH REPRESENTATION LEARNING
Dynamic graph representation learning has received significant attention in recent years, driven by the need to model and analyze evolving relationships and temporal dependencies within dynamic
graphs. Comprehensive surveys Skarding et al. (2021); Kazemi et al. (2020) offer detailed insights
into the existing works. Dynamic Graph Neural Networks (DGNNs), as the dynamic counterparts of
GNNs, have emerged as promising neural models for dynamic graph representation learning (Sankar
ct al.
2020 Poursafaei et al. 2022; Xu et al. 2020; Rossi et al. 2021; Wang et al. 2021: Kuman
al.
2019
Su al. (2024] Trivedi al. 2019
Zhang et al. 2023
Pareja et al. (2020]
Trivedi al.
2017
Among DGNNs, MDGNNS such as
Rossi et al. 2021 Wang et al. 2021 Kumar et al,
2019
Trivedi al. 2019 Zhang al., 2023
have demonstrated superior inference performance
compared to their memory-less counterparts. Most existing works on MDGNNs primarily focus on designing tailored architectures for specific problems settings. To the best of knowledge,
only a few studies have investigated the efficiency of MDGNNs. For example, Wang & Mendis 2023) leverages computation redundancies accelerate the inference performance of the temporal attention mechanism and the time encoder. [Wang et al.) (202I) introduces mailbox that stores neighbour states for each node, accelerating local aggregation operations during online inference. Zhou et al. (2022): (Sheng et al. (2024) provides system framework for optimizing inter-bate batch de- pendency and GPU-CPU communication during MDGNN training. Zhang et al. (2023) proposes restarter module that focuses restarting the model at any time with warm re-initialized memory. EDGE (Chen et al. 2021) focuses on accelerating event embedding computation by intentionally neglecting some updates (i.e., using staleness). Therefore, EDGE's 's focus and objectives differ from ours. DistTGL (Zhou et al. (2023) focuses on the dependency arising from distributed training and addresses efficient scaling concerning the number of GPUs. While these works have addressed the
24

---

Published as a conference paper at ICLR 2024
76
71
AP
72
AP
67
68
63
TGN
APAN
64
TGN-PRES
59
APAN-PRES
60
55
1
10
15
20
1
10
15
20
Batch Size (x (x100)
Batch Size (x100)
(a) TGN
(b) APAN
76
AP(%)
72
68
JODIE JODIE-PRES
64
60
1
10
15
20
Batch Size (x100)
(c) JODIE
Figure 13: Performance of baseline methods with and without PRES under different batch sizes on LASTEM dataset. The x-axis represents the batch size (multiplied by 100), while the y-axis represents the average precision (AP). The results are averaged over five trials with 3 = 0.1 for PRES.
98
98
AP (%)
AP(%)
90
90
APAN APAN-PRES
JODIE JODIE-PRES
82
82
0
100
200
0
100
200
Iteration
Iteration
(a) APAN
(b) JODIE
Figure 14: Statistical efficiency of baseline method w./w.o PERS. x-axis is the training iteration and y-axis is the average precision. B 0.1 used in PRES.
efficiency of MDGNNs in some aspects, none of them provide theoretical insights or directly ad- dress the intra-batch dependency problem. contrast, our study focuses on enlarging the temporal batch size to enable better data parallelism in MDGNN training. We adopt more theoretical ap- proach, and our proposed framework can be used in conjunction with these previous works to further enhance training efficiency. It should be noted that there exists another line of research concerning updating the GNN models on dynamic/cxpanding graphs, referred to as graph continual learning (Su) al. 2023 The objective of this line of research is orthogonal to dynamic representation learning.
25

---

Published as a conference paper at ICLR 2024
0.6
0.8
AdaGP
PipeGCN
FastocN
Accuracy Drop (%)
1.4
1.6
SAPipe
1,8
Sencus
1.6
188
2.0
2.2
2.4
2.6
Relative Achieved (times)
Figure 15: The visualization of the relative speed and affected accuracy of methods from different domains. Note that the value of the y-axis is reversed. The points with higher positions in the y-axis have a smaller accuracy drop. In addition, these methods are from different domains, corresponding different underlying tasks and datasets. This figure is only rough comparison, rigorously they are not directly comparable because of different tasks and datasets. However, this figure demon- strates that the accuracy-speed trade-off of PRES(our) is reasonable
Test Accuracy (%)
Test Accuracy (%)
100
90
94
TGN TGN-PRES
TGN TGN-PRES
0
100
200
300
400
500
100
200
300
400
500
Epoches
Epoches
Wiki
(b) Mooc
Extended training sessions with TGN on WIKI and MOOC Datasets. Figure 16(a) presents the results the TGN model with 500 epochs on the WIKI dataset, and Figure 16(b) presents the results of the TGN model with 500 epochs the MOOC dataset. Fig. [6] (a) illus- trates that with significantly extended training sessions, many of the minor accuracy discrepancies observed in datasets like Wiki between TGN and TGN-PRES can be alleviated attributed to fluctuations arising from distinct fitting processes. Additionally, Fig. [6 (b) demonstrates that the discrepancy gap certain datasets, such MOOC, gradually diminishes training progresses.
Figure 16:
G.2 MINI-BATCH IN STOCHASTIC GRADIENT DESCENT (SGD)
Another line of research investigates the effect of mini-batch size in SGD training (Goyal et al.,
2017 Qian & Klabjan, 2020 Lin et al., 2018 Akiba et al. 2017 Gower et al. 2019 Woodworth al. 2020: Bottou et al. 2018 Schmidt et al.) 2017). Research batch size and SGD has been an active area investigation in the field of deep learning. The choice of mini-bre mini-bbtch size plays crucial role in balancing the computational efficiency, convergence speed, and generalization performance neural networks, Larger mini-batcher batches tend to provide more accurate gradient estimates due to increased sample size, resulting in faster convergence. However,
the relation between mini-
26

---

Published as a conference paper at ICLR 2024
98
AP
90
TGN TGN-PRES-S TGN-PRES-V
TGN-PRES
82
100
200
Iteration
Figure 17: Ablation study the PRES framework. TGN represents the standard implementation de- void of any PRES component. TGN-PRES-S is TGN augmented with memory coherence smoothing only. TGN-PRES-V is TGN implemented solely with the prediction-correttin is TGN combined with both memory coherence smoothing and the prediction-correction scheme. The results validate the distinct operational nature of the two proposed techniques. Memory co- herence smoothing enhances the convergence rate of the learning process, while the prediction- correction scheme attenuates the variance introduced by large temporal batch size. The experiment was conducted on the WIKI dataset with batch size of 1000 and / 0.1 for the PRES scheme.
scheme. TGN-PRES
48
99
Accuracy (%)
Iterations
97
36
95
24
0.01
0.05
0.1
0.5
beta value
Figure 18: Ablation study of the 3 value in PRES framework. Fig. 18 illustrates that increasing leads to faster convergence but can have negative effect on the accuracy. Because of this trade- off, B can not be too large nor small. The plot above motivates our choice of 3 0.1 in the experiments.
27

---

Published as a conference paper at ICLR 2024
3500
Trend TGN
TGN-PRES
3250
8.5% of Total GPU Memory
TGN
3000
TGN-PRES
GPU Memory (MiB)
2750
100
500
1000
1500
2000
Batch Size
Figure 19: This illustration presents the GPU memory utilization of baseline methods, both with and without the PRES implementation, under varying batch sizes on the WIKI dataset. The x-axis denotes the batch size, while the y-axis depicts the GPU memory usage. The findings reveal that the additional GPU memory required by PRES is not influenced by the batch size, indicating its scala- bility with increasing batch sizes. Moreover, the red line demonstrates significant underutilization of the GPU memory in training MDGNNs.
they also require memory and computational resources. On the other hand, smaller mini- batches can introduce more noise into the gradient estimates but may provide better generalization by exploring
diverse set of examples. Researchers have explored the impact of mini-batch size on convergence properties, optimization dynamics, and generalization performance, leading to various recommendations and insights. However, it is important differentiate the concepts of mini-batch in SGD and temporal batch in MDGNNs, as they serve distinct purposes and bring different challenges. The goal of mini-batches in SGD is to obtain good estimation of the full-batch gradient by downsampling the entire dataset into mini- mini-batches. On the other hand, the temporal batch specifically refers partitioning consecutive graph data to ensure the chronological processing of events. The temporal batch problem studied in this paper aims to increase the temporal batch size enhance data parallelism MDGNN training.
G.3 SAMPLING IN GNNS
The full-batch training of typical GCN is employed in Kipf & Welling (2017) which necessities keeping the whole graph data and intermediate nodes' representations in the memory. This is the key bottleneck that hinders the scalability of full-batch GCN training. To overcome this issues, research on vertices and neighbor sampling in graph neural networks (GNNs) has been a topic of significant interest in the field of graph representation learning (Chen al. 2017 Ying et al. 2018 Huang al.
2018 Goyal et al., 2017 Qian & Klabjan 2020 Lin et al. [2L]]2018AAAA
& Iyer 2021 Hamilton al. 2017 Chen et al. 2018 Zou et al. 2019). GNNs operate on graph-structured
et al. 2017,
data and aim to capture the relational information among nodes. One crucial aspect is selecting which nodes to consider during the learning process. Vertices sampling refers to the selection of subset of nodes from the graph for computational efficiency, as processing the entire graph can be computationally expensive for large-scale networks. Different sampling strategies have been explored, including random sampling, stratified sampling based node properties or degrees, and importance sampling based on node importance measures. On the other hand, neighbour sampling focuses determining which neighbours of a node consider during the aggregation step in GNNs.
28

---

Published as a conference paper at ICLR 2024
It addresses the challenge of scalability by only considering subset of node's neighbours in each layer, reducing the computational complexity. Various neighbour sampling techniques have been proposed, such uniform sampling, adaptive sampling based node degrees, sampling proportional to node attention scores. Researchers have investigated the effects of different sampling strategies the expressiveness, efficiency, and generalization capabilities of GNNs, aiming strike a balance between computational efficiency and capturing important graph structures. Essentially, the problem studied in GNN sampling similar to the study of mini-batch in SGD, which to efficiently and effectively estimate the gradient with sampling. Therefore, it is also orthogonal to the problem we studied here.
29


Reviews:
[{"review_id": "e3DXJynxkS", "reviewer": "Reviewer_dnLR", "summary": "This work conducts a theoretical study on the impact of the temporal batch size in MDGNN training. This shows that there can be a significant gradient variance using a small temporal batch, which in turn sheds light on an unexpected benefit of large batch sizes. Next, the authors define memory coherence, which represents the similarity of gradient directions within a temporal batch. Memory coherence is then used to model the upper boundary of gradient. \nWith these theoretical insights, the authors present PRES with two main components: 1) iterative prediction-correction scheme 2) memory coherence smoothing. The former uses a GMM (updated with MLE) to predict newest memory states and fuses it with the calculated memory state to obtain the final ‘corrected’ state. The latter uses a new learning objective to promote larger memory coherence. \nUsing PRES, the authors were able to increase the temporal batch size without compromising overall accuracy.", "questions": "The paper is overall well written. The introduction on MDGNN was easy to follow. Insights from theoretical analyses were well presented. It was also evident how these insights became the main building blocks of PRES. However, my main concern comes from the experiment section.\n\n-\t**Is it really useful to gain speed at the cost of accuracy?**\n\nSo far this is my main concern. At first I thought the authors were trying to achieve SOTA accuracy.\nHowever, what the authors are doing is gaining speedup of around 2x to 3x, at the cost of decreased accuracy (~1.0%).\nI am not so sure about this, considering the effort the community is putting to gain higher accuracy.\nEspecially on the tested datasets, the number of vertices is only around few thousands, which wouldn't take terribly long to train.\nI believe this partially comes from not reporting the training time (only the speedup is reported) and there are not enough baselines to compare. But the bottomline is that a strong justification is needed for this.\n\n-\t**What is the consensus on the ‘optimal’ temporal batch size?**\n\n\tIn Figure 3, the authors show the performance of baselines by increasing the batch size up to 100. In the figure the ‘small batch size’ seems to be ~50. My question is do the majority of MDGNNs use a batch size smaller than 50, or are they already using approximately 500 (which seems to be the optimal size in Figure 4)? If the latter is the case, then personally the insight from theorem 1 (variance of the gradient for the entire epoch can be detrimental when the temporal batch size is small) loses some of its shine. Thus, the authors should try to first do a comprehensive overview on the currently used batch sizes.\n\n-\t**How does PRES differ from other baselines?**\n\tTwo related works came to my mind which are missing in the current paper. “Efficient Dynamic Graph Representation Learning at Scale” (arXiv preprint, 2021, https://arxiv.org/abs/2112.07768) and “DistTGL: Distributed Memory-Based Temporal Graph Neural Network Training” (arXiv preprint, 2023, https://arxiv.org/abs/2307.07649). Both try to increase the temporal batch size without harming the accuracy. The former also uses prediction to utilize data-parallelism, while the latter tries to push the temporal batch size to the extreme for distributed GPU clusters. In my opinion, both (and any other baseline that shares the same goal with this work) should be compared methodologically and speedup-wise (in the current setting). Also, it would be interesting to see if these can also benefit from PRES. \n\n-\t**Why is performance with/without PRES not shown with the LASTFM dataset?**\nLASTFM stands out in that 1) the AP is the lowest 2) the speedup of PRES is the lowest. However, I was unable to find a figure like figure 4 for this dataset. Is there a reason for only leaving this dataset out?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "-\tThe authors provide theoretical results on the influence of temporal batch size on MDGNN training.\n-\tWith memory coherence, the authors effectively define new methods to compensate for the accuracy drop of naively increasing the temporal batch size.", "weaknesses": "- The tradeoff of improved speed at the cost of lower accuracy does not seem to be appealing.\n- Comparison with prior work on increasing temporal batch size is insufficient. \n-\tIn a similar manner, there are only a small number of baselines in the experimental results. \n-\tThe specific results of a dataset (LASTFM) is excluded.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VKcl4CjTIQ", "reviewer": "Reviewer_pqbr", "summary": "This paper aims to offer a scalable training method for Memory-Based Dynamic Graph Neural Networks (MDGNNs) by mitigating the temporal discontinuity issue, thus training MDGNNs with large temporal batch sizes. It consists of two main contributions: 1) conducting a theoretical study on the impact of temporal batch size on the convergence of MDGNN training, and 2) proposing PRES based on the theoretical study, an iterative prediction-correction scheme combined with a memory coherence learning objective to mitigate the effect of temporal discontinuity. The evaluation shows that the proposed approach enables up to 4X larger temporal batch sizes and achieves up to 3.4X speedup during MDGNN training.", "questions": "Overall, this is a solid study with clear innovations. The theoretical study on the impact of temporal batch size on the convergence of MDGNN training is extensive and helpful. My major concerns focus on the evaluation aspects. It would be extremely helpful if the authors could offer more information about these questions:\n\n1. PRES is mainly evaluated on four graph datasets (Reddit, Wiki, Mooc, and LastFM). It seems these graphs are not very large with around 1K to 10K vertices and 400K to 1.3M edges. It would be helpful to justify that these graphs are large enough or PRES’s performance is not affected by the graph size.\n\n2. It would be helpful to report the absolute execution time as well rather than relative speedup only.\n\n3. It would be helpful to discuss if this method can be extended to multi-nodes with more GPUs.\n\n4. It seems Table 1 shows that PRES still sacrifices some precision for performance gains in many cases. Please correct me if I have any misunderstanding here.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ It targets an emerging and important GNN, MDGNNs, and the proposed designs generally make sense. \n+ The problem definition is easy-to-follow.\n+ The introduced concept of memory coherence is interesting.\n+ The code is publicly available.", "weaknesses": "- The trained graphs seem small. It is unclear how PRES performs on large-scale graphs.\n- No absolute execution time is reported.\n- It is tested on four GPUs. Its scalability to multi-nodes (with more GPUs) is somewhat unclear.\n- In many cases, PRES still sacrifices some precision for performance gains.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)