Every Client Is an Environment: Federated De-confounding for Spatio-Temporal Forecasting

In federated spatio-temporal forecasting (STF), models are trained across distributed nodes (clients) without sharing raw data. A persistent challenge is that each client operates in a unique environment (e.g., different weather, traffic patterns, or local events), leading to heterogeneous data distributions. Existing methods often treat this heterogeneity as an optimization hurdle, resulting in models that capture environment-specific biases and fail to generalize when environmental conditions shift. This paper asks: can we leverage this cross-client heterogeneity as a source of causal evidence to improve forecasting robustness?

The authors propose Atlas, a framework that treats each federated client as a distinct causal environment. By formulating the problem through a Structural Causal Model (SCM), they identify that local models suffer from confounding effects because they only observe a biased subset of environmental strata. Atlas addresses this by:

1. Prototype Codebook Learning: Each client learns a local codebook of prototypes representing different environmental regimes.
2. Codebook Alignment: The server uses the Hungarian algorithm to align these local prototypes across clients, ensuring that corresponding slots in the codebooks represent the same environmental regimes.
3. Federated Aggregation: By aggregating these aligned prototypes, the server constructs a global codebook that approximates the global environmental distribution, effectively performing an implicit back-door adjustment to eliminate client-specific confounding.

The study demonstrates that the federated de-confounding error is linearly controlled by the averaged confounding strength across all clients, rather than being limited by the bias of any single client. Empirical results on five real-world datasets show that Atlas consistently outperforms existing federated baselines. Furthermore, the learned environmental representations are shown to be transferable, interpretable, and communication-efficient, providing a more robust foundation for spatio-temporal forecasting under distribution shifts.