🤖 machine learning

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

This paper introduces FedXDS, a novel federated learning approach that leverages explainable AI feature attribution to identify and selectively share task-relevant data elements between clients, thereby mitigating data heterogeneity while maintaining formal privacy guarantees and achieving superior accuracy and convergence compared to existing methods.

Original authors: Maximilian Andreas Hoefler, Karsten Mueller, Wojciech Samek

Published 2026-07-01

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Maximilian Andreas Hoefler, Karsten Mueller, Wojciech Samek

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a group of chefs trying to create the world's best pizza recipe. They all have their own secret ingredients and techniques, but they can't share their actual kitchens or raw ingredients with each other because of strict privacy rules. This is Federated Learning: everyone learns together without ever showing their private data.

However, there's a problem. Chef A only has tomatoes from Italy, Chef B only has peppers from Spain, and Chef C only has cheese from France. Because their ingredients are so different (this is called data heterogeneity), when they try to combine their notes, the final recipe gets confused and tastes bad. They keep arguing over which ingredient is the "most important."

Enter FedXDS, a new method proposed in this paper that acts like a smart, privacy-conscious sous-chef to fix this mess. Here is how it works, broken down into simple steps:

1. The "Highlighter" Trick (XAI)

Instead of sending whole photos of their ingredients (which is risky), the chefs use a special Highlighter Pen (called an Attribution Method).

This pen looks at an image of a pizza and figures out exactly which pixels matter most for the taste.
It highlights the pepperoni and the crust but ignores the background noise or the tablecloth.
The Magic: The chefs only send these "highlighted" parts to the group. They throw away the irrelevant junk. This means the group gets a clearer, more focused picture of what makes a pizza good, without the confusion of different backgrounds.

2. The "Blurry Photo" Safety Net (Privacy)

Even if you only send the highlighted parts, someone might still try to guess what the original photo looked like. To stop this, the paper adds a layer of Digital Fog (called Metric Differential Privacy).

Think of it like taking a photo of the highlighted ingredients and then slightly blurring it just enough so no one can recognize the specific brand of pepperoni, but it still looks like pepperoni.
Why this is smarter: Because the chefs already threw away the irrelevant background junk (Step 1), there is less "fog" needed to hide the secrets. If they tried to blur the whole original image, the fog would be so thick the pizza would be unrecognizable. By focusing only on the important parts first, they can keep the image clear enough to learn from, while still keeping it safe.

3. The "Shared Cookbook" (Training)

Now, the central server collects all these "highlighted and slightly blurred" ingredient photos from every chef.

It mixes them into a Global Cookbook (a shared dataset).
Every chef gets a copy of this cookbook. Now, Chef A (who only had Italian tomatoes) can learn from Chef B's Spanish peppers and Chef C's French cheese, even though they never saw the raw ingredients.
This helps everyone agree on a better, more universal pizza recipe much faster.

Why is this a big deal?

The paper tested this idea on computer vision tasks (like recognizing cats, dogs, or numbers) and found:

It's Faster: The chefs reached a great recipe in fewer meetings (communication rounds) than other methods.
It's Smarter: Even when the chefs had very different ingredients, the final recipe was much more accurate than before.
It's Safer: The paper proved that even if a hacker tried to reverse-engineer the "blurred" photos to steal the original ingredients, they failed much more often with this method than with older methods.

In a nutshell: FedXDS is like a group of people solving a puzzle together. Instead of shouting out their whole puzzle pieces (which is dangerous), they use a highlighter to show only the unique shapes that fit the picture, blur them slightly for safety, and then share those. This allows the group to solve the puzzle faster and more accurately, without anyone losing their privacy.

Technical Summary: FedXDS

Problem Statement

Federated Learning (FL) enables collaborative model training without sharing raw data, yet it faces significant performance degradation when client data distributions exhibit statistical heterogeneity (non-IID). While existing approaches attempt to mitigate this through proximal optimization or loss landscape flattening, they often struggle with update divergence. Conversely, sharing raw data or synthetic features to align distributions improves generalization but introduces severe privacy risks. Standard differential privacy (DP) applied to high-dimensional raw data (e.g., images) requires excessive noise to guarantee privacy, leading to substantial utility loss. The core challenge addressed by this paper is how to retain the performance benefits of data sharing in heterogeneous FL environments while providing strong privacy guarantees and maintaining computational efficiency.

Methodology: FedXDS

The authors propose FedXDS (Federated Learning via XAI-guided Data Sharing), a novel framework that leverages Explainable AI (XAI) to selectively share only task-relevant data features. The methodology consists of three primary stages:

1. Attribution-Guided Feature Selection

Instead of sharing raw images, clients utilize propagation-based attribution methods to identify input features most critical for model predictions.

Warmup Phase: Clients first train a local "warmup" model ( $\theta_{warmup}$ ) using standard FedAvg for a set number of rounds.
Attribution Mapping: For each input sample $x$ , an attribution method $\mathcal{A}$ (e.g., Layer-wise Relevance Propagation, Integrated Gradients) computes a pixel-wise relevance score map $h$ .
Masking: A binary mask $m$ is generated by retaining only the top $s$ -percentile of attribution scores (sparsity level). The input is filtered via element-wise multiplication: $f_A(x) = x \odot m$ . This step discards background noise and spurious correlations, retaining only semantically meaningful features.

2. Privacy-Preserving Mechanism

To protect the remaining features, the framework employs Metric Differential Privacy.

Sensitivity Reduction: By sparsifying the input through attribution masking, the sensitivity ( $\Delta f$ ) of the feature selection function is strictly bounded (proven to be $\le 1$ in the $\ell_2$ norm).
Noise Injection: Gaussian noise is added to the masked features. Because the sensitivity is reduced by discarding irrelevant dimensions, significantly less noise is required to achieve a specific privacy budget ( $\epsilon, \delta$ ) compared to applying noise to the full raw image. This preserves higher utility.
Formal Guarantee: The mechanism satisfies $(\epsilon, \delta)$ -metric privacy, where the privacy loss is bounded relative to the distance between inputs in the metric space.

3. Federated Training with Shared Data

Global Aggregation: Clients upload their private, masked, and noised features to the server, which aggregates them into a global dataset $D_g$ .
Local Optimization: Clients download $D_g$ and optimize a composite objective function that balances local task performance with knowledge from the global shared data:
$\min_{\theta} \left[ \mathbb{E}_{(x,y) \sim D_k}[\ell(f_\theta(x), y)] + \lambda \mathbb{E}_{(x,y) \sim D_g}[\ell(f_\theta(x), y)] \right]$
Here, $\lambda$ controls the trade-off between local specialization and global knowledge integration.

Key Contributions

Novel Algorithm: The introduction of FedXDS, the first approach to utilize XAI feature attribution techniques to identify and selectively share specific data elements between clients to mitigate heterogeneity.
Privacy-Preserving Dimensionality Reduction: A mechanism that uses attribution-guided masking to reduce the dimensionality of raw features before applying noise. This allows for metric differential privacy with lower noise levels, thereby improving utility compared to standard DP on raw data.
Theoretical and Empirical Privacy Validation: The paper provides theoretical bounds on sensitivity and empirically demonstrates robustness against Membership Inference Attacks (MIA) and Feature Inversion attacks, showing that attribution-masked features offer stronger protection than unmasked features under equivalent privacy budgets.
Efficiency: Unlike generator-based approaches (e.g., GANs or VAEs) that require costly retraining, FedXDS computes attribution masks in a single backward pass during a warmup phase, incurring minimal computational overhead.

Experimental Results

The authors evaluated FedXDS on standard benchmarks (CIFAR-10, CIFAR-100, Tiny-ImageNet) and real-world heterogeneous datasets (CelebA, FEMNIST) using the LEAF framework.

Accuracy: FedXDS, particularly when using Layer-wise Relevance Propagation (FedXLRP), consistently outperformed state-of-the-art baselines (including FedAvg, FedProx, FedDyn, and data-sharing methods like FedFed and FedFTG).
- On CIFAR-10 with 100 clients and high heterogeneity ( $\alpha=0.05$ ), FedXLRP achieved 83.46% accuracy, significantly outperforming FedAvg (60.94%) and FedFed (82.58%).
- On real-world datasets (CelebA, FEMNIST), FedXDS achieved the highest accuracy across all tested methods.
Convergence Speed: FedXDS demonstrated superior communication efficiency. For example, on CIFAR-10 with 10 clients, it reached 70% accuracy in 14 rounds, whereas FedAvg required 49 rounds.
Ablation Studies:
- Attribution Method: LRP proved superior to gradient-based methods (Gradient $\times$ Input, Integrated Gradients, SmoothGrad) because it produces structurally coherent relevance maps, preserving meaningful visual patterns even under high sparsity.
- Hyperparameters: An intermediate knowledge weight ( $\lambda=0.5$ ) yielded optimal performance, and a warmup period of 10 rounds was found to be the optimal trade-off for convergence.
Privacy Analysis: Empirical tests showed that FedXDS features were significantly harder to reconstruct (lower SSIM scores in feature inversion attacks) and provided better resistance to membership inference compared to unmasked features or features from other data-sharing methods.

Significance and Claims

The paper claims that FedXDS represents a significant advancement in addressing the "trilemma" of Federated Learning: heterogeneity, privacy, and efficiency.

Bridging XAI and FL: The work demonstrates for the first time that XAI methods, traditionally used for interpretability, can be repurposed as a functional tool to solve the statistical heterogeneity problem in FL.
Utility-Privacy Trade-off: By leveraging the sparsity induced by attribution maps, the method achieves a more favorable privacy-utility trade-off. It argues that discarding task-irrelevant information is not just a filtering step but a privacy-enhancing mechanism that reduces the sensitivity of the data, allowing for less noise injection.
Practicality: The approach avoids the computational burden of generative models, making it suitable for resource-constrained environments while providing formal privacy guarantees.

The authors conclude that attribution methods are a promising tool for improving federated learning, enabling effective knowledge transfer while preserving privacy and computational efficiency.