Federated Causal Representation Learning in State-Space Systems for Decentralized Counterfactual Reasoning

1. Problem Statement

The paper addresses a critical challenge in complex, geographically dispersed industrial systems (e.g., oil refineries, water treatment networks, power grids) where assets (clients) are tightly coupled through physical processes and control inputs.

• The Core Question: How would the output of one client change if another client altered its control input? This requires counterfactual reasoning (e.g., "What would happen to my stability if the upstream pump reduced flow?").
• The Constraints:
  1. Data Privacy & Bandwidth: Client data is high-dimensional, private, and expensive to transmit. Centralizing raw data is infeasible due to privacy regulations and bandwidth limits.
  2. Proprietary Models: Clients possess local, proprietary models (e.g., Kalman filters) trained on their own data that cannot be modified or shared.
  3. Limitations of Existing Methods: Classical tools like Granger Causality are predictive, not interventional, and cannot answer "what-if" questions. Existing causal discovery methods often assume centralized data access or fail to disentangle intrinsic dynamics from control-driven influences.

2. Methodology

The authors propose a Federated Learning framework tailored for Linear Time-Invariant (LTI) State-Space Systems. The goal is to learn the global causal structure (specifically cross-client dependencies) without sharing raw data or modifying local proprietary models.

System Model

The system is modeled as a multi-client LTI state-space system:
$$h_t = A h_{t-1} + B u_{t-1} + w_{t-1}$$
$$y_t = C h_t + v_t$$
Where $h$ is the latent state, $u$ is the control input, and $y$ is the observation. The matrices $A$ (state transition) and $B$ (input) contain both diagonal blocks (local dynamics) and off-diagonal blocks (cross-client causal dependencies).

The Federated Architecture

The framework involves two entities: Clients (holding local data and proprietary models) and a Central Server (coordinating learning).

1. Client-Side (Augmented Model):

   • Clients retain their Proprietary Model (a local Kalman filter using local matrices $A_{mm}, B_{mm}, C_{mm}$) which generates a refined state estimate $\hat{h}_{m,c}$.
   • An Augmented Model is introduced with learnable parameters $\theta_m$ and $\phi_m$.
     • $\theta_m$ adjusts the state estimate based on local measurements.
     • $\phi_m$ acts as a bias term to capture the aggregated influence of other clients' inputs.
   • Clients transmit only low-dimensional latent states and inputs to the server, not raw high-dimensional sensor data.

2. Server-Side (Global Estimation):

   • The server aggregates the low-dimensional states from all clients.
   • It explicitly estimates the off-diagonal blocks of the global matrices ($A_{mn}$ and $B_{mn}$ for $m \neq n$), which represent the causal links between clients.
   • The server computes a loss function that includes a disentanglement penalty. This penalty forces the system to separate the effects of endogenous state propagation ($A_{mn}$) from exogenous input effects ($B_{mn}$).

3. Training Loop:

   • The server sends gradients of the global loss back to the clients.
   • Clients update their local parameters ($\theta_m, \phi_m$) using a combination of their local reconstruction loss and the server's gradients.
   • This iterative process allows clients to implicitly learn cross-client dependencies while the server explicitly learns the structural matrices.

Counterfactual Reasoning

Using the Abduction-Action-Prediction (AAP) framework by Pearl:

• Server: Can predict the counterfactual state change $h_m$ given an intervention $do(u_n = u')$ using the learned $B_{mn}$.
• Client: Can predict the counterfactual output change $y_m$ by leveraging the learned $\phi_m$, which encodes the aggregated causal effects of other clients.

3. Key Contributions

1. First Decentralized Counterfactual Framework: This is the first work to enable decentralized counterfactual reasoning in state-space systems without centralizing raw data or altering proprietary local models.
2. Causal Disentanglement: The authors design an iterative optimization scheme that provably disentangles autoregressive dynamics (endogenous) from control-driven influences (exogenous), a prerequisite for valid causal inference.
3. Convergence to Oracle: The paper proves that the decentralized framework converges to a centralized oracle (a model with access to all raw data) up to a specific bias term, ensuring theoretical soundness.
4. Privacy Guarantees: The framework incorporates Differential Privacy (DP). The authors provide theoretical bounds on the sensitivity of communicated messages and gradients, demonstrating how noise can be added to protect client data while maintaining utility.
5. Transfer Function Identification: The framework implicitly learns the system's transfer function (which encodes causal effects), which is invariant to latent state reparameterization, ensuring the learned causal structure is robust.

4. Results

The authors validated their approach on both synthetic and real-world industrial datasets.

• Synthetic Data:

  • Convergence: The global loss and client residuals decreased over iterations, outperforming the proprietary-only baseline.
  • Disentanglement: The penalty term successfully separated the effects of $A_{mn}$ and $B_{mn}$, as evidenced by the convergence of client parameters to the true aggregated causal effects.
  • Scalability: The method remained stable as the number of clients and measurement dimensions increased, though performance degraded slightly with higher dimensions (expected due to estimation variance).
  • Transfer Function: The estimated poles and zeros closely matched the ground truth, confirming the model captured the underlying stable dynamics.

• Real-World Datasets (HAI and SWaT):

  • HAI (Industrial Control System): The model successfully modeled a 4-process steam turbine system. It showed fast convergence and low residuals, effectively capturing inter-process dependencies.
  • SWaT (Water Treatment Plant): Applied to a 6-stage water treatment plant. The model demonstrated stable convergence and the ability to reason about cross-stage interventions.
  • Robustness: The federated models showed better robustness to non-stationary shifts (mean shifts) compared to the proprietary local models.

5. Significance

This paper makes a significant contribution to the intersection of Causal Inference, Federated Learning, and Industrial Control Systems (ICS):

• Proactive Risk Management: It enables operators to perform "what-if" analysis (e.g., assessing cascading failure risks) without compromising data privacy or intellectual property.
• Scalability: By exchanging only low-dimensional latent states rather than high-dimensional sensor streams, the framework is scalable to large networks of assets.
• Theoretical Rigor: It bridges the gap between practical federated learning and rigorous causal theory, providing convergence proofs and privacy guarantees often missing in applied industrial AI.
• Practical Applicability: The validation on real industrial datasets (HAI, SWaT) demonstrates that the approach is not just theoretical but applicable to critical infrastructure where safety and privacy are paramount.

In summary, the paper presents a robust, privacy-preserving solution for understanding and predicting the behavior of interconnected industrial systems, enabling safer and more efficient decentralized operations.