FedXDS: Leveraging Model Attribution Methods to counteract Data Heterogeneity in Federated Learning

2026-06-30

AI Summary

This work addresses the performance degradation in federated learning caused by client data heterogeneity by introducing explainable artificial intelligence (XAI) into the federated framework. It proposes a propagation-based feature attribution mechanism that selectively shares only the attribution information of task-critical features, thereby mitigating heterogeneity while integrating metric privacy techniques to provide formal privacy guarantees. Experimental results demonstrate that the proposed method consistently achieves significant improvements in model accuracy and convergence speed across diverse heterogeneity settings and varying numbers of clients. Furthermore, it exhibits strong robustness against membership inference and feature inversion attacks, highlighting its effectiveness in preserving both utility and privacy in heterogeneous federated environments.

Abstract

Explainable AI (XAI) methods have demonstrated significant success in recent years at identifying relevant features in input data that drive deep learning model decisions, enhancing interpretability for users. However, the potential of XAI beyond providing model transparency has remained largely unexplored in adjacent machine learning domains. In this paper, we show for the first time how XAI can be utilized in the context of federated learning. Specifically, while federated learning enables collaborative model training without raw data sharing, it suffers from performance degradation when client data distributions exhibit statistical heterogeneity. We introduce FedXDS (Federated Learning via XAI-guided Data Sharing), the first approach to utilize feature attribution techniques to identify precisely which data elements should be selectively shared between clients to mitigate heterogeneity. By employing propagation-based attribution, our method identifies task-relevant features through a single backward pass, enabling selective data sharing that aligns client contributions. To protect sensitive information, we incorporate metric privacy techniques that provide formal privacy guarantees while preserving utility. Experimental results demonstrate that our approach consistently achieves higher accuracy and faster convergence compared to existing methods across varying client numbers and heterogeneity settings. We provide theoretical privacy guarantees and empirically demonstrate robustness against both membership inference and feature inversion attacks. Code is available at https://github.com/MaxH1996/FedXDS.

Problem

• Federated Learning
• Data Heterogeneity
• Statistical Heterogeneity
• Data Sharing
• Model Performance Degradation

Innovation

• Explainable AI
• Federated Learning
• Feature Attribution
• Data Heterogeneity
• Metric Privacy