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

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.

• The paper introduces an attribution-guided, differential privacy protocol that mitigates data heterogeneity in federated learning.
• It leverages Layerwise Relevance Propagation to extract sparse, task-relevant features, reducing privacy risks and computational overhead.
• Empirical results show improved accuracy and faster convergence over FedAvg and competing DP methods across multiple benchmarks.

Attribution-Guided Data Sharing for Federated Learning under Heterogeneity and Privacy Constraints

Motivation and Context

Data heterogeneity among clients in federated learning (FL) presents a persistent obstacle to efficient and generalizable model aggregation. Classical algorithms such as FedAvg degrade when confronted with highly non-IID or distribution-shifted local data. While sharing raw or generative features can help, such approaches introduce severe privacy vulnerabilities and often entail significant computational overhead. Differential privacy (DP) mechanisms, though effective in theory, frequently destroy model utility when applied indiscriminately to high-dimensional client data. The work "FedXDS: Leveraging Model Attribution Methods to counteract Data Heterogeneity in Federated Learning" (2606.31742) introduces a novel solution rooted in explainable AI (XAI) methods to address these challenges systematically.

Methodological Contributions

FedXDS proposes a privacy-preserving data sharing protocol centered around attribution-guided feature extraction. The process is comprised of four steps: (1) an initial warmup phase to establish a reasonable local model, (2) computation of pixel-level relevance maps via a backward XAI technique (with Layerwise Relevance Propagation, LRP, proving optimal empirically), (3) creation of a binary mask that isolates the most task-relevant input features at a prescribed sparsity level, and (4) addition of Gaussian noise to masked features, calibrated according to metric-space DP constraints, before sharing with the federation.

Figure 1: Illustration of DP attribution-guided feature extraction with computation of attributions, masking, and DP noise addition.

This selective strategy achieves two critical effects. Firstly, it minimizes sensitivity to privacy attacks by reducing the input dimensionality to only those regions indispensable for the FL task—thus optimizing the privacy-utility tradeoff. Secondly, it enables efficient data sharing that directly targets statistical heterogeneity without reliance on potentially leaky or computationally intense generative models.

Privacy Analysis and Attack Robustness

The privacy argument is rigorously constructed under $(\varepsilon, \delta)$-metric DP, leveraging the fact that attribution masking is a contractive operator with $\ell_2$-sensitivity at most 1, i.e., no expansion over the input space even in the worst case. The DP noise scale is then matched to this global Lipschitz bound, which empirically results in less required noise versus standard applications of DP to raw features. Sensitivity is driven further down in practice by sparsity in the mask, further reducing reconstruction risk.

A comprehensive empirical evaluation is conducted using both membership inference attacks (MIAs) (Figure 2) and feature inversion attacks (Figures 3–5). FedXDS demonstrates measurable reductions in MI recall and inversion fidelity (as proxied by SSIM) relative to both raw feature sharing and state-of-the-art DP-distilled feature protocols such as FedFed.

Figure 2: MI attack recall decreases significantly for masked, privacy-preserved representations as opposed to raw DP-protected features, particularly at higher sparsity.

Figure 3: Denoising inversion attack on FedXDS-shared representations reveals weak recoverability of the original input.

Figure 4: Baseline inversion attack on DP-protected raw features shows severe privacy leakage.

Figure 5: Inversion attack on FedFed shows partial reconstruction, indicating that generator-based approaches are also susceptible.

Qualitative visualization (Figure 6) confirms that while raw and FedFed features frequently reveal human-identifiable information, FedXDS features, even after aggressive inversion, remain obfuscated.

Figure 6: Visualization of original images for comparison with reconstructed samples from various inversion attacks.

Empirical Performance and Convergence

The framework is benchmarked across standard image datasets (CIFAR-10/100, Tiny-ImageNet) and real federated splits (CelebA, FEMNIST). FedXDS achieves superior accuracy and more rapid convergence than all evaluated baselines; for instance, on CIFAR-10 ($K=100$, $\alpha=0.05$), FedXDS reaches $77.02\%$ accuracy, while FedFed achieves $75.35\%$ and FedAvg just $55.53\%$. Communication efficiency gains are similarly pronounced: FedXDS achieves $70\%$ accuracy in as few as 14 rounds with $K=10$, compared to 49 rounds for FedAvg.

Training progression curves demonstrate accelerated loss minimization and smoother convergence, further substantiating these claims.

Figure 7: Training progression of FedXDS on CIFAR10, $K=10$, $\ell_2$0, showcasing faster and more stable convergence compared to baselines.

Ablations: Attribution Methods, Sparsity, and Privacy-Utility Tradeoffs

A detailed ablation over attribution techniques and sparsity (fraction of retained pixels) underscores the superiority of LRP. While gradient-based methods degrade sharply at high sparsity, LRP maintains >81% accuracy even with just 15% of original pixels (s=85). The coherence of LRP’s explanations, which yield structured contiguous regions, is critical for retaining predictive context post-masking.

Privacy-utility ablations confirm that increasing the DP budget $\ell_2$1 monotonically boosts accuracy, but even stringent settings (e.g., $\ell_2$2) with LRP remain competitive relative to naive masking or generator-based approaches. Varying the global knowledge weight $\ell_2$3 in the composite loss reveals optimal performance at intermediate values, further highlighting the importance of balanced local and shared knowledge integration.

Implications and Theoretical Impact

The central insight of FedXDS is that attribution-guided selection—backed by formal DP analysis—enables feature sharing protocols that are simultaneously private, efficient, and robust to data drift. This approach reframes the role of XAI: rather than mere post hoc interpretability, attribution methods become an operational tool for federated data shaping.

Practically, this work informs the design of federated systems in resource-constrained or privacy-critical domains (e.g., edge healthcare, personalized devices) where central raw-data sharing is infeasible. Theoretically, it advances the intersection of interpretable ML, privacy theory, and distributed optimization by showing that structure-aware sparsification dovetails with privacy noise to deliver Pareto-optimal tradeoffs.

Open directions include instance- or region-adaptive DP noise schedules based on per-sample attribution maps, extension to multi-modal or sequential data, and exploration of more sophisticated structural priors in the masking operator.

Conclusion

FedXDS delivers a systematic procedure for privacy-preserving, communication-efficient federated learning under severe heterogeneity by leveraging attribution-based feature selection and metric differential privacy. Empirical results establish state-of-the-art performance in both accuracy and privacy leakage resistance across multiple datasets and adversarial settings. The methodology paves the way for integrating XAI mechanisms as primary drivers of both utility and privacy in distributed machine learning settings (2606.31742).