Causality-inspired Federated Learning for Dynamic Spatio-Temporal Graphs

Abstract: Federated Graph Learning (FGL) has emerged as a powerful paradigm for decentralized training of graph neural networks while preserving data privacy. However, existing FGL methods are predominantly designed for static graphs and rely on parameter averaging or distribution alignment, which implicitly assume that all features are equally transferable across clients, overlooking both the spatial and temporal heterogeneity and the presence of client-specific knowledge in real-world graphs. In this work, we identify that such assumptions create a vicious cycle of spurious representation entanglement, client-specific interference, and negative transfer, degrading generalization performance in Federated Learning over Dynamic Spatio-Temporal Graphs (FSTG). To address this issue, we propose a novel causality-inspired framework named SC-FSGL, which explicitly decouples transferable causal knowledge from client-specific noise through representation-level interventions. Specifically, we introduce a Conditional Separation Module that simulates soft interventions through client conditioned masks, enabling the disentanglement of invariant spatio-temporal causal factors from spurious signals and mitigating representation entanglement caused by client heterogeneity. In addition, we propose a Causal Codebook that clusters causal prototypes and aligns local representations via contrastive learning, promoting cross-client consistency and facilitating knowledge sharing across diverse spatio-temporal patterns. Experiments on five diverse heterogeneity Spatio-Temporal Graph (STG) datasets show that SC-FSGL outperforms state-of-the-art methods.

• The paper introduces SC-FSGL, a causality-driven framework that disentangles shared and client-specific features in dynamic spatio-temporal graphs.
• It leverages a client-conditioned soft mask and a causal codebook with contrastive and IRM losses to simulate interventions and improve invariant feature extraction.
• Empirical results on heterogeneous traffic datasets show stable low MAE and robust generalization compared to traditional federated methods.

Causality-Inspired Federated Learning for Dynamic Spatio-Temporal Graphs: A Technical Analysis

Introduction and Motivation

The paper introduces SC-FSGL, a causality-inspired federated learning framework for dynamic spatio-temporal graphs (FSTGs). The work addresses the limitations of existing federated graph learning (FGL) methods, which are predominantly tailored for static graphs and assume homogeneous knowledge transferability across clients. These traditional assumptions fail in scenarios where data possess marked spatial and temporal heterogeneity, leading to representation entanglement, client-specific interference, and negative transfer, thereby degrading model generalization.

Figure 1: Spatial and temporal heterogeneity across clients and shared causal patterns.

Theoretical Formulation and Causality Perspective

The approach formalizes each client in the federated system as an independent environment, endowed with distinct spatial (graph topology) and temporal (observed time series) data-generating mechanisms. The critical insight is the decomposition of client observations into shared causal variables (stable across environments) and client-specific variables (encapsulating local or ephemeral patterns). Adopting the Structural Causal Model (SCM) formalism, the model seeks to approximate interventional distributions in latent space because direct interventions are infeasible in federated scenarios.

A learnable, client-conditioned soft mask is introduced to simulate interventions in the latent space, attenuating spurious or client-specific effects and promoting extraction of invariant, transferable features. This is further reinforced by an objective rooted in Invariant Risk Minimization (IRM), enforcing feature representations that induce optimal predictors across all client distributions.

Figure 2: A conceptual SCM-inspired illustration on client $k$ showing the relations among observables and latent components.

SC-FSGL Framework: Architectural Innovations

SC-FSGL constructs a client-server federated protocol with the following innovations:

• Client-Side Disentanglement: Each client extracts spatio-temporal features and applies a conditional separation module to segment representations into shared (transferable) and client-specific (non-transferable) causal components.
• Soft Causal Intervention: Learnable soft masks simulate representation-level interventions, aligning extracted features with shared causal structures and suppressing client-specific confounding signals.
• Causal Codebook: A shared, learnable codebook of causal prototypes serves as semantic anchors. Client models align local shared causal variables with the codebook via contrastive learning, facilitating cross-client knowledge consistency.
• Loss Composition: Local training optimizes a combined objective integrating MSE for prediction accuracy, a contrastive loss to tie representations to the codebook, and a differentiable IRM-inspired gradient penalty to promote transferability.

  Figure 3: The overall SC-FSGL architecture, showing client-side causal separation and codebook-based global aggregation.

Experimental Results and Empirical Evaluation

Extensive experiments are conducted across five heterogeneous real-world traffic datasets (METRLA, PEMSD4, PEMSD7(M), PEMSD8, PEMSBAY), each representing a federated client with distinct spatial and temporal characteristics. The results demonstrate:

• Superior Predictive Accuracy: SC-FSGL consistently achieves the lowest MAE and MAPE across all prediction horizons (15, 30, and 60 minutes) when compared to FedAvg, FedProx, Moon, FedRep, FedStar, GMAN, MegaCRN, and FUELS.
• Robustness to Heterogeneity: The model maintains stable, low MAE throughout training and across clients, outperforming baselines which exhibit instability or substantial performance degradation under increased spatial or temporal heterogeneity.
• Causal Feature Disentanglement: t-SNE visualizations of learned representations illustrate clear separation between shared and client-specific components, evidencing successful disentanglement and the benefits of codebook-based alignment.

Figure 4: Performance analysis on the PEMSD4 dataset: (a) prediction results, (b) MAE curves, (c) t-SNE visualization highlighting disentangled causal features.

Figure 5: t-SNE visualization of representation evolution across communication rounds.

Figure 6: RMSE and MAPE curves of SC-FSGL, showing its training stability compared to baselines.

Ablation and Analysis

Ablation studies confirm the indispensability of the causal codebook and the conditional separation (CS) module. Removal of either component leads to substantial increases in MAE; benefit is maximal when both modules operate in tandem. Hyperparameter analysis further reveals optimal performance for balanced contributions from the contrastive and IRM loss components.

Figure 7: Ablation study showing MAE sensitivity to module removal.

Figure 8: Hyperparameter study for contrastive and IRM loss weights.

Convergence and Efficiency

The causal codebook aggregation is theoretically proven to converge under convexity and bounded gradient conditions. Communication efficiency is a prominent feature: only the compact codebook is exchanged between clients and server, mitigating common bandwidth bottlenecks in federated graph learning.

Implications and Future Prospects

SC-FSGL demonstrates that explicit modeling of causal structures in federated, dynamic, heterogeneous graph environments significantly improves generalization and predictive accuracy. Practically, the framework enables effective knowledge transfer, robust out-of-distribution prediction, and efficient communication, making it suitable for real-world decentralized spatio-temporal applications, such as urban mobility forecasting, large-scale sensor networks, and smart infrastructure.

Theoretically, this work reinvigorates the connection between causality and federated learning. Future research can broaden the expressivity of the codebook (e.g., via hierarchical or adaptive prototypes), or extend SC-FSGL to tackle causal drift and more pronounced non-stationarities.

Figure 9: Prediction performance visualization on heterogeneous clients.

Conclusion

SC-FSGL introduces a rigorous, causality-driven solution to federated learning on dynamic spatio-temporal graphs. By disentangling transferable and client-specific factors and enforcing representation-level interventions and alignment, the model reliably outperforms prior approaches. This paradigm shift—from implicit distributional alignment to explicit causal disentanglement—offers a robust foundation for future advances in federated learning on complex, heterogeneous, and dynamic data structures.