FedDis: A Causal Disentanglement Framework for Federated Traffic Prediction

TL;DR

This work tackles privacy-preserving traffic forecasting in federated settings where client data are non-IID. It introduces FedDis, a dual-branch causal disentanglement framework that separates globally shared spatial-temporal patterns from client-specific local dynamics using a Global Pattern Bank and a Personalized Pattern Bank, with a mutual information objective enforced via the CLUB bound to encourage orthogonality. Spatial-temporal features are extracted with adaptive graph convolutional recurrent encoders, and federated optimization leverages collaborative pattern sharing and graph attention fusion to exchange knowledge without exposing raw data. Empirical results on four real-world benchmarks show FedDis achieves state-of-the-art performance among federated methods, closely approaching centralized performance while offering robustness to increasing client numbers and clear ablations validating the importance of each component. The approach offers a scalable, interpretable, and privacy-preserving solution for cross-region traffic prediction with strong cross-client knowledge transfer and local adaptability.

Abstract

Federated learning offers a promising paradigm for privacy-preserving traffic prediction, yet its performance is often challenged by the non-identically and independently distributed (non-IID) nature of decentralized traffic data. Existing federated methods frequently struggle with this data heterogeneity, typically entangling globally shared patterns with client-specific local dynamics within a single representation. In this work, we postulate that this heterogeneity stems from the entanglement of two distinct generative sources: client-specific localized dynamics and cross-client global spatial-temporal patterns. Motivated by this perspective, we introduce FedDis, a novel framework that, to the best of our knowledge, is the first to leverage causal disentanglement for federated spatial-temporal prediction. Architecturally, FedDis comprises a dual-branch design wherein a Personalized Bank learns to capture client-specific factors, while a Global Pattern Bank distills common knowledge. This separation enables robust cross-client knowledge transfer while preserving high adaptability to unique local environments. Crucially, a mutual information minimization objective is employed to enforce informational orthogonality between the two branches, thereby ensuring effective disentanglement. Comprehensive experiments conducted on four real-world benchmark datasets demonstrate that FedDis consistently achieves state-of-the-art performance, promising efficiency, and superior expandability.

Figures (7)

• Figure 1: Illustration on PEMS-BAY. (a) The illustration of universal spatial-temporal patterns across clients. (b) Performance comparison of centralized, federated, and our FedDis.
• Figure 2: FedDis framework. Each client adopts a dual-branch architecture with a global branch and a personalized branch, which enables causal decoupling. The Global Pattern Extractor maintains a Global Pattern Bank ($\mathbf{W}$), which captures globally shared patterns, while the Personalized Pattern Extractor maintains a Personalized Pattern Bank ($\mathbf{L}$), capturing local dynamic traffic variations. Client aggregates node embeddings to yield its Graph Prototypes (GP), which indicate the client's overall structural and properties. On the server side, GP guides the aggregation of sharable parameters $\mathbf{\theta_u^{(i,s)}}$, and the $\mathbf{W}$ facilitates global knowledge sharing across clients.
• Figure 3: Personalized pattern extractor.
• Figure 4: Global pattern extractor.
• Figure 5: Ablation study.