FLASH: Flexible Learning of Adaptive Sampling from History in Temporal Graph Neural Networks

1 Introduction

Dynamic graphs provide a natural framework for modeling real-world systems where entities and their interactions evolve over time. They underpin a wide range of applications, including social and communication networks [Shetty and Adibi, 2004; Kumar et al., 2019], user-item recommendation systems [Kumar et al., 2019], and financial or knowledgeintensive platforms [Ni et al., 2019; Shamsi et al., 2022]. Predicting future interactions in these settings has emerged as a central learning task, leading to the development of Temporal Graph Neural Networks (TGNNs) - models specifically designed to learn from sequences of timestamped events represented by dynamic graphs.

TGNNs process dynamic graphs by encoding temporal interaction patterns into node representations, enabling them to predict future links. A common challenge in these models is how to efficiently aggregate information from a node's history of interactions, which can grow unbounded over time. Processing complete histories quickly becomes computationally prohibitive, especially in high-frequency interaction settings. To address this, existing models such as TGN [Rossi et al., 2020], TGAT [Xu et al., 2020], DyGFormer [Yu et al., 2023], GraphMixer [Cong et al., 2023], FreeDyG [Tian et al., 2023], and others [Zhou et al., 2022; Zou et al., 2024; Gravina et al., 2024; Xu et al., 2024] adopt memory-efficient heuristics. These typically include strategies like uniform sampling, time-decay weighting, or truncating to the k most recent interactions. While effective at reducing computational overhead, these approaches are static and fail to adapt to the local graph structure or task-specific temporal signals. As shown in Figure 1, such heuristics apply uniform or truncated selection schemes that overlook potentially informative neighbors, whereas adaptive strategies can prioritize structurally meaningful interactions.

Static heuristics like uniform sampling [Rossi et al., 2020], truncation [Cong et al., 2023] or hybrid approaches [Luo and Li, 2024] are appealing due to their simplicity, but they treat all interactions as equally informative or rely solely on recency. This ignores the fact that some neighbors may be more relevant than others due to their position in the graph or their interaction patterns. Moreover, the optimal sampled neighborhood may vary across time, nodes, and tasks, making fixed strategies fundamentally limited. These shortcomings are further amplified in heterogeneous or rapidly evolving graphs, where structural context can shift dramatically over time. This motivates the need for a learnable, structure-aware neighborhood selection mechanism that can adaptively prioritize informative past interactions.

To address these limitations, we propose FLASH - Flexible Learning of Adaptive Selection from History for Temporal Graph Neural Networks, a learnable and graphadaptive neighborhood selection mechanism for TGNNs. FLASH replaces static heuristics with a data-driven approach that learns to prioritize historically informative neighbors based on their structural and temporal context. Crucially, because the true importance of neighbors is not known a priori, our method is trained using a self-supervised ranking objective that encourages selecting neighbors most predictive of future interactions. FLASH is lightweight, general-purpose, and integrates seamlessly into a wide range of existing TGNN architectures, including TGNNs with non-differentiable feature extractors [Yu et al., 2023; Tian et al., 2023]. This allows it to improve predictive performance without requiring architectural changes. Our key contributions are as follows:

• We propose FLASH, a novel graph-adaptive, learnable neighborhood selection mechanism that seamlessly integrates with any TGNN.
• We design a self-supervised training objective based on ranking loss, enabling our method to learn informative neighbor selection without access to ground-truth labels.
• We provide a theoretical analysis showing that FLASH is provably more expressive than the common heuristics of recent neighbors selection and uniform sampling.
• We conduct extensive experiments across multiple dynamic graph benchmarks, demonstrating consistent performance gains across diverse TGNN backbones compared to common neighbor sampling baselines.

2 RelatedWork

Neighborhood selection in static graphs Existing sampling techniques for large static graphs often rely on substructure sampling (e.g., nodes or edges), as employed by GraphSAGE [Hamilton et al., 2017] and FastGCN [Chen et al., 2018], or utilize random walks, as in PinSage [Ying et al., 2018]. Other methods, such as GraphSAINT [Zeng et al., 2020] and Cluster-GCN [Chiang et al., 2019], are specifically designed to facilitate efficient training on large graphs. However, these approaches typically do not address inference or the temporal nature of dynamic graphs. In many TGNNs, uniform sampling can be viewed as dynamic extension of GraphSAGE-like methods, in which each neighbor in the historical neighborhood is sampled with equal probability. On the other hand, truncation sampling retains only the most recent neighbors, and its stochastic variant of timeweighting, reduces the sampling probability of older neighbors over time, effectively treating recency as a measure of importance. This parallels importance sampling in static graphs (e.g., FastGCN), where recent interactions are explicitly prioritized. However, our experiments show that recency alone is insufficient to capture the complexities of temporal interactions. Instead, incorporating node-specific contextual information at each interaction point, as done by FLASH, proves crucial for achieving robust and accurate performance in dynamic graph settings.

Learning from Large Historical Neighborhoods Learning from large historical neighborhoods in dynamic graphs poses significant challenges in computational cost and capturing long-term dependencies. To address the latter, the patching technique introduced in [Yu et al., 2023], adopted in other recent studies [Tian et al., 2023; Ding et al., 2025]. The method splits the historical neighborhood into chronological patches, each linearly projected into a single representative vector. Since historical neighbors must be encoded (e.g., via co-occurrence encoding [Yu et al., 2023]), neighborhood selection precedes it to reduce computational overhead. Thus, as proposed by [Yu et al., 2023] for DYGFORMER, patching is a complementary strategy for processing large historical neighborhoods, in future link prediction tasks.

Weighing and Selection Mechanisms. Some approaches, such as TGAT [Xu et al., 2020], employ attention mechanisms to weight neighbors, but only after the neighborhood has been sampled. Other works use reinforcement learning [Wang et al., 2022] to learn selection strategies, but they require task-specific reward design and extensive training. Methods based on hard negative mining or curriculum sampling [Younesian et al., 2024] share our motivation of identifying informative samples, but are not tailored for selfsupervised learning and are not readily applicable to TGNNs. In contrast, FLASH learns to sample informative neighbors in a self-supervised, fully differentiable manner.

3 Background

Continuous Time Dynamic Graph (CTDG) is represented as a sequence of time-stamped events G = {(x1, t1), (x2, t2) · · · } where t1 ≤ t2 ≤ · · · . Each event (xi, ti) represents an action on the graph that occurred at time ti. Each such action can be either node addition, node removal, edge addition, or edge removal. We denote Gt = (Vt, Et) the snapshot of CTDG at time t, which is the graph received by applying all the events in G that occurred until time t, where G0 = (∅,∅). We denote FV : V × R+ → RdV and FE : E × R+ → RdE as the functions that map a node or an edge to their features, at a specific point in time. For a given node v and timestamp t, we denote the historic neighborhood Hv,t as the set of all nodes that interacted with v before time t:

Certain TGNNs allow the same node u to appear multiple times in Hv,t, where each occurrence of u is associated with a distinct timestamp that records the time v and u interacted. TGNNs construct node representations by sampling a subset Sv,t(k) ⊆ Hv,t of k historical neighbors. Existing approaches employ various static selection strategies:

where ru is the rank of u with respect to the sequence of sorted neighbors from Hv,t in a decreasing order by their timestamp. Some TGNNs [Rossi et al., 2020] allow sampling beyond the 1-hop historical neighborhood, either by allowing the sampling technique to sample farther nodes in advance or by applying the sampling technique recursively, i.e., sampling from the neighborhood of the sampled neighbors.

Given the sampled neighborhood Sv,t(k), the representation of node v at time t is computed as:

where ψ(·) is a parametric function of the TGNN that maps node to vector based on its sampled neighborhood.

For the task of future link prediction, given a pair of nodes (vi, vj) and their appropriate representations ztvi and ztvj at time t, a TGNN assigns a probability to the existence of a future edge between them using a learnable MERGE function:

To train and evaluate TGNNs for future link prediction, the common approach is to split the entire sequence of interactions into two consecutive non-overlapping segments: a training prefix and an evaluation suffix. All interactions within the training prefix serve as positive examples, indicating node pairs that do form an edge at a specific time. For negative examples, random node pairs are sampled. The TGNN parameters are then updated by minimizing a binary classification loss (e.g., cross-entropy) that distinguishes positive from negative edges:

where yt ij is 1 for observed edges in the training prefix and 0 otherwise. Once trained, the TGNN can be used to predict edges in the evaluation suffix by computing the probability p(vi, vj | t) for new future node pairs.

3.1 Theoretical analysis of Heuristic Neighborhood Samplings

The lack of flexibility in common neighborhood sampling techniques, such as k recent selection (truncation) or uniform sampling, hinder the performance of TGNNs. Specifically, we show that there exist dynamic graphs such that any TGNN relying on these common heuristics cannot learn from them (i.e., it cannot overfit on this graph).

Theorem 1. For any k there exists a dynamic graph on which a TGNN that applies k recent neighbors sampling cannot learn.

Theorem 2. For any k, there exists a dynamic graph on which a TGNN that applies uniform sampling of k historical neighbors cannot learn.

To prove Theorem 1, we find a graph that the k + 1 recent neighbor is required to learn the dynamic behavior of the graph. To prove Theorem 2, we find a dynamic graph that requires consistently selecting the same recent neighbor to learn its dynamic behavior. We provide the full proofs and assumptions for Theorem 1 and Theorem 2 in Appendix D.

Implication. These common sampling heuristics discard potentially crucial historical interactions, limiting the performance of TGNNs and reduce their expressive power. We aim to develop a learnable and adaptive sampling technique that not only generalizes these heuristics but also enables TGNNs to become more expressive.

4 Method

In the previous section, we showed that current neighborhood sampling strategies do not fully account for the graph structure and its features, preventing TGNNs from capturing even simple evolving dynamics. Another limitation of these sampling heuristics stems from the fact that different TGNNs learn temporal dynamics in distinct ways due to their diverse designs. For example, TGN [Rossi et al., 2020] uses memory states, whereas DyGFormer [Yu et al., 2023] employs a Transformer-based architecture [Vaswani et al., 2017] and jointly processes Svi, t and Svj, t. Thus, enforcing the same fixed neighborhood sampling scheme across various TGNNs may undermine their performance. A learnable sampling strategy can adapt not only to the structure of the dynamic graph but also to the way each TGNN architecture exploits temporal signals.

4.1 Desiderata for Adaptive Neighborhood Sampling in TGNNs

We propose an adaptive neighborhood sampling mechanism that satisfies the following desiderata:

1. Adapt to CTDG Dynamics: We seek a mechanism SAMPLE(·) that accounts for node and edge attributes as well as their timestamps, i.e., the learnable parameters θ should be informed by the CTDG's interaction patterns. Formally, for a node v with historical neighborhood Hv,t and a candidate interaction partner v′ ,

where |Sv,t| = k.

1. Generalize Existing Heuristics: SAMPLE(·) should subsume commonly used heuristics (e.g., truncation, uniform sampling) as special cases via appropriate choices of θ.
2. Seamless Integration with TGNNs: Since TGNN backbones encode temporal dynamics differently (e.g., memory modules, attention), the relevance of sampled neighbors may depend on the backbone. Accordingly, the TGNN should inform the sampling procedure and the learning of θ. Moreover, the mechanism should integrate seamlessly with a broad range of TGNNs, including models with non-differentiable feature extractors.
3. Self-Supervised Learning of Neighbor Importance: Because ground-truth neighbor "importance" is unavailable, the mechanism should be learnable in a selfsupervised manner that rewards neighbor subsets conducive to accurate link prediction.

4.2 FLASH

To predict a future link (vi, vj) at time t, we learn how informative each historical neighbor u ∈ Hvi,t is for this prediction. Concretely, SAMPLE(·; θ) assigns a score to each historical neighbor and selects the k neighbors with the highest score.

Learning the scoring function. We first construct an informative representation for each historical neighbor that captures both structural (interaction context) and temporal (time-dependent) information (D1). For a neighbor u ∈ Hvi,t, let tu denote the time at which u and vi interacted, and let ru denote the rank of u when sorting Hvi,t in decreasing order of interaction time. We define:

where ϕ1(·) and ϕ2(·) are Time2Vec representations kazemi2019time2vec, FV maps nodes and timestamps to their corresponding features, FE maps edges and timestamps to their corresponding features (as defined in Section 3), M(·) denotes learnable node features, and ∥ denotes concatenation. Following the standard TGNN pipeline and prior work deng2024taser, we project the neighbor and link representations to a common dimensionality using MLPs, and then combine them via a MERGE operator:

where MERGE is an MLP with a single hidden layer, as suggested for combining source-destination information yu2023towards,tian2023freedyg. SCOREu,vi,vj,t measures the utility of neighbor u for predicting the link (vi, vj) at time t, where larger values indicate higher importance. FLASH then selects the k neighbors with the highest scores. Figure 2 provides a detailed schematic of FLASH.

Training FLASH. For each node pair (vi, vj) at time t, let yt ij ∈ {0, 1} indicate whether (vi, vj) is a positive (observed) edge (yt ij = 1) or a negative (unobserved) edge (yt ij = 0). Let a TGNN compute edge probabilities using the subsets selected by FLASH, (Svi,vj,t, Svj,vi,t), and using uniformly sampled subsets (S uni vi ,t , S uni vj,t ). We define the probability difference:

Let svi , svj denote the average FLASH scores over nodes in Svi,vj,t and Svj,vi,t, respectively, and let suni vi , suni vj denote the corresponding averages over S uni vi,t and S uni vj,t . For positive edges, we prefer ∆tij > 0 (our subsets increase the predicted probability), while for negative edges we prefer ∆tij < 0 (our subsets decrease it). We enforce this preference using a pairwise RankNet burges2005learning logistic loss:

Minimizing Ltij encourages the score averages to move in the direction consistent with yt ij without requiring any ground-truth importance supervision (D3). When the condition in Equation (14) holds, the neighborhoods selected by FLASH improve over uniform sampling, and the loss encourages higher scores for selected neighbors over the uniformly sampled neighbors. Otherwise, uniform sampling performs better and the loss correspondingly increases the scores of uniformly sampled neighbors relative to those selected by FLASH. We train FLASH jointly with a given TGNN by optimizing the sum of Equation (14) and Equation (6). Gradients from Equation (6) update only the TGNN parameters, while for Equation (14) we stop gradients through ∆tij so that updates flow only to FLASH parameters (D4). To ensure constant runtime, FLASH first applies a recent-neighbors finder deng2024taser to restrict the candidate set to a fixed size. It then computes scores only within this reduced neighborhood rather than over the full historical neighborhood. In Appendix G, we provide a throughout analysis of FLASH and TASER.

4.3 Theoretical Analysis of FLASH

The components of FLASH adapt to the evolution of the dynamic graph over time and to the interactions being predicted, by incorporating both spatiotemporal features and potentially interacting nodes when constructing each sampled neighborhood. We show that FLASH can replicate both k-recent neighbor selection and uniform sampling. Furthermore, we show that the graphs used in the proofs of Theorem 1 and Theorem 2 can be learned by a TGNN that utilizes FLASH. Therefore:

Theorem 3. FLASH is more expressive than k-recent neighbor selection and uniform sampling.

We provide the full proof of Theorem 3 in Appendix D.

5 Experiments

We evaluate FLASH across multiple dynamic graph benchmarks and compare it with recent state-of-the-art baselines.

Section 5.1 presents our key empirical findings, and Section 5.2 provides additional ablation studies. Additional results are reported in Appendix C. Our experiments aim to answer the following research questions:

• (RQ1) How does our proposed strategy compare to established sampling baselines in predictive accuracy?
• (RQ2) Does our method generalize across varying graph sizes, sparsity levels, and temporal granularities?
• (RQ3) What is the computational overhead of our approach relative to existing methods?

Experimental Setup. We test our method on five common TGNNs: TGAT [Xu et al., 2020], TGN [Rossi et al., 2020], GRAPHMIXER [Cong et al., 2023], DYGFORMER [Yu et al., 2023], and FREEDYG [Tian et al., 2023]. Detailed descriptions of each model are provided in Appendix B. Each model is trained with four neighbor sampling strategies: Truncation [Cong et al., 2023; Yu et al., 2023], Uniform [Rossi et al., 2020], No-Look-Back (NLB) [Luo and Li, 2024], TASER with the recent-neighbors finder and a linear predictor [Deng et al., 2024], and our proposed method. Following standard practice, we use a chronological 70%-15%-15% train-validation-test split. Additional parameters regarding the training scheme and implementation-specific details are provided in Appendix E.

Importance of Small Neighborhoods As discussed in Section 1, processing large neighborhoods is computationally prohibitive in both time and space. Figure 4 reports the throughput of state-of-the-art TGNNs for varying neighborhood sizes. The results show that increasing the neighborhood processed in parallel substantially degrades model throughput. Consequently, TGNNs should employ small neighborhoods (2-4) to achieve peak runtime performance. Accordingly, we evaluate neighborhood selection strategies under a small neighborhood budget (2-4) to determine which strategy is optimal for future link prediction while sustaining peak throughput in TGNNs.

5.1 Results

Evaluation with DyGLib. We use the WIKIPEDIA, REDDIT, MOOC, LASTFM, ENRON, SOCIAL EVO., and UCI datasets from the DyGLib benchmark [Yu et al., 2023]. DyGLib includes datasets spanning a wide range of graph sizes, with some containing over a million edges (RQ2). Full dataset statistics are provided in Appendix A. We evaluate future edge prediction (Equation (6)) under two settings: (1) transductive, where all nodes appear during training, and (2)

inductive, where test nodes are unseen during training. Tables 1 and 2 report results across models and methods. Overall, FLASH consistently outperforms previously suggested neighborhood selection strategies, both in the transductive setting and the inductive setting (RQ1). For example, on SOCIAL EVO., FLASH achieves a ∼8% relative improvement over truncation for TGAT, and on LASTFM, a ∼15% relative improvement for FREEDYG, demonstrating strong effectiveness even for computationally expensive models.

Evaluation with TGB. We evaluate FLASH on the TGB benchmark [Huang et al., 2023] using TGBL-WIKI, TGBLREVIEW, and TGBL-COIN. In contrast to DyGLib, TGB samples multiple random negative edges for each positive edge and is evaluated using Mean Reciprocal Rank (MRR). From Figure 3, FLASH yields performance that is consistently on par with or better than baselines across TGNNs and neighborhood sizes. Notably, on TGBL-REVIEW, our method improves performance by ∼2.67× for GRAPHMIXER when k = 4, highlighting robustness across diverse dynamic graph scenarios (RQ2).

5.2 Ablation Study

FLASH Design Choices. We ablate key components of FLASH, namely the historical position embedding ϕ2, temporal embedding ϕ1, learnable node embedding M, and the link-awareness module MLPlink in Equation (12), to assess the contribution of temporal, structural, and interaction context for future edge prediction (RQ2). Table 3 shows that the positional embedding ϕ2 is particularly important on ENRON and UCI, indicating that relying solely on temporal information is suboptimal. The interaction context module, MLPlink, is consistently beneficial across all datasets, providing reliable performance improvements.

Computational Overhead. Table 4 reports a throughput comparison for FLASH, measured as the number of processed edges per second, against the baselines (RQ3). For a fair comparison, both TASER and FLASH employ the same neighbor finder, which constrains their neighborhoods to an identical size, this benefits both methods relative to uniform sampling. Because all experiments use a unified batch size, model latency is inversely proportional to throughput. Appendix F further provides a time-complexity analysis of FLASH and the baselines.

6 Conclusion

In this work, we introduce FLASH, a flexible, end-to-end learnable neighborhood selection framework for dynamic graph learning. FLASH factorizes the neighbor scoring function into spatial and temporal components, ensuring that neighbor importance reflects both topological structure and temporal context. Our training objective couples the learned scores with downstream predictive performance, enabling effective optimization of the selection policy even when ground-truth importance scores are unavailable and TGNN operations are non-differentiable. This design allows FLASH to integrate seamlessly with a broad range of TGNN architectures, providing an efficient mechanism for selecting the most relevant historical neighbors. We further evaluate FLASH across multiple dynamic graph benchmarks and show that it matches or outperforms the considered baselines.

Contribution Statement

O.F. contributed to the design of the method, implementation, experiments, figure/plot design, and manuscript writing. K.S.I.M. contributed to the design of the method, figure/plot design, and manuscript writing. C.B.S. contributed to manuscript writing. C.B. and M.E. contributed to manuscript writing and equally supervised the work.