CoMemNet: Contrastive Sampling with Memory Replay Network for Continual Traffic Prediction

Abstract: In recent years, the integration of non-topological space modeling with temporal learning methods has emerged as an effective approach for capturing spatio-temporal information in non-Euclidean graphs. However, most existing methods rely on static underlying graph structures, which are inadequate for capturing the continuously expanding and evolving patterns in streaming traffic networks. To address this challenge, we propose a simple yet efficient dual-branch continual learning framework for traffic prediction, named CoMemNet. The fast-converging Online branch undertakes the primary prediction tasks, while the momentum-updated Target branch extracts historical information using Wasserstein Distance features to create a Dynamic Contrastive Sampler (DC Sampler). This sampler selects a node set with significant dynamic network feature changes for training, effectively mitigating the issue of catastrophic forgetting. Additionally, the backbone incorporates a lightweight Node-Adaptive Temporal Memory Buffer (TMRB-N) to consolidate old knowledge through memory replay and address the risk of memory explosion. Finally, we provide two newly curated open-source datasets. Experimental results demonstrate that CoMemNet achieves state-of-the-art (SOTA) performance across all three large-scale real-world datasets. The code is available at: https://github.com/meiwu5/CoMemNet.

• The paper proposes a dual-branch continual learning framework that efficiently adapts to new traffic distributions while retaining historical features.
• It employs a Dynamic Contrastive Sampler and a Node-Adaptive Temporal Memory Replay Buffer to select hard examples and prevent memory explosion.
• Experimental results on CalTrans PeMS datasets show state-of-the-art accuracy with significant reductions in node training and computational overhead.

CoMemNet: Contrastive Sampling with Memory Replay for Continual Traffic Prediction

Problem Formulation and Motivation

Traffic networks evolve continuously with incremental changes in topology, sensor deployments, and flow dynamics. Conventional spatio-temporal models, including STGNNs and attention baselines, assume a static underlying graph, neglecting real-world expansion and distribution drift. This leads to catastrophic forgetting and computational inefficiency in streaming scenarios where retraining or full-history replay is impractical. The paper introduces CoMemNet, which directly addresses the three principal challenges in continual traffic prediction: flexible feature extraction under dynamic graphs, mitigation of catastrophic forgetting, and prevention of memory explosion as node scale grows.

Figure 1: (A) Expansion of the traffic network structure, (B) traffic flow statistics for Station 40009 over multiple years, (C) degradation in prediction performance on old tasks due to continual training on new tasks.

Model Architecture and Methodological Innovations

CoMemNet employs a dual-branch continual learning framework. The Online branch rapidly adapts to new traffic distributions, while the Target branch ensures historical feature retention via exponential moving average (EMA) parameter updates ($\beta$-momentum). The backbone eschews explicit adjacency matrix inputs, focusing on embeddings and leveraging temporal feature replay buffers.

Figure 2: Structural model for traffic signals $X_\tau$ in time period $\tau$.

Central to CoMemNet is the Dynamic Contrastive (DC) Sampler, which performs difficulty-aware online sample selection. DC Sampler computes Wasserstein distances between discretizations of node embeddings across periods, identifying the Top-M nodes with maximal distributional shift. This hard-example curriculum prioritizes nodes exhibiting significant spatio-temporal drift and new arrivals ($V_\tau \setminus V_{\tau-1}$), concentrating learning resources.

Figure 3: Process for screening Top-M nodes with largest Wasserstein distances across time via DC Sampler.

To address memory explosion, CoMemNet introduces a Node-Adaptive Temporal Memory Replay Buffer (TMRB-N). TMRB-N stores low-dimensional representations of critical nodes, selected via feature-difference criteria rather than random or global replay. The buffer incorporates a gating mechanism for update, maintaining synchronization between current and historical embeddings and selecting Top-K nodes by difference magnitude. This yields efficient memory usage, robust consolidation, and minimizes redundant replay.

Figure 4: End-to-end architecture of TMRB-N showing update and storage of temporal features $H_\tau$ in memory replay buffer.

Experimental Evaluation

The paper rigorously benchmarks CoMemNet on three large-scale real-world datasets from CalTrans PeMS (PEMSD3(S), PEMSD4(L), PEMSD8(M)), including two newly curated open-source datasets. CoMemNet achieves SOTA results on all metrics—MAE, RMSE, and MAPE—across 15, 30, and 60 min granularity. Critically, CoMemNet outperforms retrained and existing continual learning approaches, while requiring orders-of-magnitude fewer nodes for training in later years due to DC Sampler efficiency.

In PEMSD3(S), for example, annual node training drops from 655 to just 192 for 2017 (15–30% of network), yet MAE remains superior ($<14.9$ versus baselines above 17.0). On PEMSD4(L), memory and computational cost is significantly reduced as replay is focused only on informative nodes. The method also demonstrates stable performance under ablation, highlighting the necessity of both DC Sampler and TMRB-N selection/update.

Ablation and Sensitivity Analyses

The ablation studies reveal that omitting DC Sampler or TMRB-N impairs accuracy and increases computational overhead, with random or global replay leading to inefficient knowledge consolidation and increased forgetting. Parameter sensitivity tests confirm optimal Top-K selection (K=12), and DC Sampler ratio $\rho$ provides a robust trade-off between accuracy and resource usage.

Figure 5: Sensitivity analysis of hyperparameter K for annual average 12-step MAE.

Discrete probability visualization of node feature distributions over years further supports the curriculum selection logic. Nodes exhibit temporal drift and heterogeneity, with some highly stable and others dramatically shifting, confirming that uniform replay is suboptimal.

Figure 6: Example discrete probability distributions for Node 74 on PEMSD3 (2012), illustrating feature drift targeted by DC Sampler.

Practical and Theoretical Implications

CoMemNet presents definitive practical advances for continual learning in evolving traffic networks: it dynamically adapts to changing topologies without explicit graph input, minimizes unnecessary memory overhead, and operates efficiently in high-dimensional streaming environments. The momentum contrastive dual-branch architecture is generalizable to other spatio-temporal streaming data domains.

Theoretically, the work advances knowledge consolidation schemes via feature-difference-based sampling and buffer updates, suggesting future directions in adaptive replay for continual learning under graph drift. The reliance on Wasserstein distance for curriculum-like node selection holds potential for broader applications in dynamic graph forecasting and lifelong learning frameworks.

Conclusion

The CoMemNet framework provides a robust, efficient solution to continual traffic prediction on expanding networks, combining dual-branch contrastive learning, adaptive memory replay, and difficulty-aware dynamic sampling. Empirical results confirm strong accuracy and scalability. The methodology sets a precedent for embedding-based continual learning without explicit graph structures, with potential for wider adoption in dynamic spatio-temporal forecasting and streaming graph analytics.