Federated Out-of-Distribution Generalization: A Causal Augmentation View

TL;DR

This work tackles federated learning under out-of-distribution conditions by addressing spurious background correlations that hinder generalization. It introduces FedCAug, a causal augmentation framework with two modules: CRL, which localizes and sharpens causal regions in images, and CA, which creates counterfactual samples by fusing objects with diverse backgrounds while preserving data privacy. The method yields consistent improvements across NICO-Animal, NICO-Vehicle, and ColorMNIST, reducing background reliance and promoting causal feature learning. As a model-agnostic, privacy-preserving approach, FedCAug can be integrated with existing methods to enhance OOD robustness in real-world federated settings.

Abstract

Federated learning aims to collaboratively model by integrating multi-source information to obtain a model that can generalize across all client data. Existing methods often leverage knowledge distillation or data augmentation to mitigate the negative impact of data bias across clients. However, the limited performance of teacher models on out-of-distribution samples and the inherent quality gap between augmented and original data hinder their effectiveness and they typically fail to leverage the advantages of incorporating rich contextual information. To address these limitations, this paper proposes a Federated Causal Augmentation method, termed FedCAug, which employs causality-inspired data augmentation to break the spurious correlation between attributes and categories. Specifically, it designs a causal region localization module to accurately identify and decouple the background and objects in the image, providing rich contextual information for causal data augmentation. Additionally, it designs a causality-inspired data augmentation module that integrates causal features and within-client context to generate counterfactual samples. This significantly enhances data diversity, and the entire process does not require any information sharing between clients, thereby contributing to the protection of data privacy. Extensive experiments conducted on three datasets reveal that FedCAug markedly reduces the model's reliance on background to predict sample labels, achieving superior performance compared to state-of-the-art methods.

Figures (6)

• Figure 1: FedCAug can learn better causal representations by reducing background noise interference through causal decoupling and causal augmentation, thereby providing instructive knowledge for the image classification task.
• Figure 2: Illustration of the framework of FedCAug, it sharpens images and locates class-relevant regions by using the Causal Region Localization (CRL) module; then it fuses the images with common sense background through the Causal Augmentation (CA) module, providing guiding information for the model to learn causal features.
• Figure 3: The prediction confidence obtained from background images after processing with FedAvg and FedCAug methods.
• Figure 4: Visualization of the visual Attention. (a) The FedCAug corrects errors in individual clients. (b) The FedCAug improves the aggregated model by correcting errors in each client, even when both clients make mistakes. (c) The FedCAug increases the model's confidence in the ground-truth.
• Figure 5: A comparative demonstration of causal augmented samples and diffusion model samples.