FedX: Explanation-Guided Pruning for Communication-Efficient Federated Learning in Remote Sensing

This paper introduces FedX, a novel explanation-guided pruning strategy for federated learning in remote sensing that significantly reduces communication overhead by identifying and removing task-irrelevant model components at the central server while enhancing global model generalization.

The Gist

Imagine you are the captain of a massive fleet of ships (the clients), and you all need to build a single, super-smart map of the world (the global model) together. But there's a catch: you can't share your actual maps or the raw data you collected because of privacy laws and security risks. Instead, you can only send each other small notes about what you learned ("model updates").

The problem? These notes are huge. Sending them back and forth across the ocean is slow, expensive, and clogs up the communication channels, especially if your ships are small or the internet connection is spotty (like satellites or drones in remote areas).

This paper introduces FedX, a clever new strategy to shrink these notes without losing any important information. Think of it as a "smart editor" that helps you send only the most vital parts of your story.

The Problem: The "Heavy Backpack"

In traditional Federated Learning, every time the ships talk to the captain, they have to send their entire backpack of learning. Even if 90% of the stuff in the backpack isn't actually helping them find the treasure, they still have to carry and ship it. This is the communication overhead.

The Solution: FedX (The "Smart Editor")

FedX is like a super-intelligent editor who doesn't just randomly throw things out of the backpack. Instead, it uses Explanation-Guided Pruning.

Here is how it works, using a simple analogy:

1. The "Why" Matters (Explanation)

Imagine you are teaching a student how to identify a "forest" in a photo.

• Old Way (Random Pruning): You might say, "Let's just delete 50% of the words in your textbook randomly." The student might accidentally delete the word "trees" and keep "blue sky," which doesn't help much.
• FedX Way (Explanation-Guided): FedX asks the model, "Which parts of your brain actually helped you identify the forest?" It uses a technique called Backpropagation (think of it as a "trace-back" flashlight) to light up the specific neurons that were most important for the decision.

2. The "Server" as the Editor

In this system, the central server (the Captain) acts as the editor.

• The ships send their updates to the Captain.
• The Captain combines them into one big model.
• The Magic Step: The Captain takes a small, public set of test images (like a practice quiz) and runs them through the model. Using the "trace-back flashlight," the Captain sees exactly which parts of the model are doing the heavy lifting and which parts are just dead weight.
• The Captain then prunes (cuts out) the useless parts. It's like taking a giant, messy sketch and erasing all the faint, unnecessary pencil lines, leaving only the bold, clear outline.

3. The "Layer-by-Layer" Rule

The paper found a crucial trick: You can't just cut randomly from the whole book. You have to cut a little bit from every chapter (every layer of the network).

• The Analogy: Imagine a book where the first chapter is short and the last chapter is huge. If you try to cut 50% of the total words from the whole book, you might accidentally delete the entire first chapter and leave the last one mostly intact. That ruins the story.
• FedX's Fix: It ensures that every chapter gets trimmed fairly. This keeps the story balanced and the model working, even when you cut out a massive amount of text.

The Results: Faster, Smaller, and Smarter

The researchers tested this on real-world satellite images (like looking at cities or forests from space).

• Less Traffic: They managed to shrink the data sent between ships and the captain by up to 44%. That's like sending a postcard instead of a heavy crate.
• No Loss in Smarts: Surprisingly, the model didn't get dumber. In fact, by removing the "noise" (the useless parts), the model sometimes got smarter and more accurate, similar to how a student learns better when they focus only on the key concepts rather than memorizing every single word in a textbook.
• Works on Any Ship: It worked well whether the ships were small (simple models) or massive (complex AI models like Transformers).

Why This Matters

In the real world, this means we can train powerful AI systems on sensitive data (like satellite images of private properties or borders) without violating privacy laws. We can do this even if the satellites or drones have slow internet connections, because FedX ensures we only send the "essence" of the learning, not the whole library.

In short: FedX is a smart pruning tool that acts like a ruthless but fair editor, cutting out the fluff from AI models so they can communicate faster and cheaper, without ever forgetting what they learned.

Technical Summary

1. Problem Statement

Context: Remote Sensing (RS) image classification increasingly relies on Deep Learning (DL). However, RS data is often distributed across decentralized archives (e.g., satellites, UAVs, ground stations) due to privacy concerns, legal restrictions, and bandwidth limitations. Federated Learning (FL) is a natural solution for training models on this decentralized data without sharing raw images.

Challenge: The primary bottleneck in applying FL to RS is communication overhead. FL requires frequent transmission of large model updates (gradients or weights) between clients and a central server. In RS environments, where clients may be resource-constrained edge devices (e.g., satellites) with limited bandwidth, this overhead slows down training and introduces latency.

Limitations of Existing Solutions:

• Model Compression/Quantization: Can degrade performance by approximating important parameters and adds encoding/decoding computational costs.
• Knowledge Distillation: Requires student models that may underperform and risks leaking information through shared logits.
• Standard Pruning: Often relies on simple heuristics like weight magnitude. In FL, where data is non-IID (heterogeneous), magnitude-based pruning can be unreliable because important parameters for specific clients may have small magnitudes in the aggregated global model, leading to the removal of task-relevant information.

2. Methodology: FedX

The authors propose FedX, a novel strategy that integrates explanation-guided pruning into the FL pipeline to reduce communication overhead while preserving model performance.

Core Concept

Instead of pruning based on weight magnitude, FedX uses backpropagation-based explanation methods to estimate the task-specific relevance of model components. The pruning decision is made at the central server, ensuring that only the most informative components are retained in the global model sent to clients.

Key Technical Components

1. Server-Side Relevance Estimation:

   • After aggregating local updates from clients, the server uses a small, public reference dataset ($D_{ref}$) to compute relevance scores for model components (e.g., filters, neurons).
   • It employs explanation methods such as Layer-wise Relevance Propagation (LRP) and Integrated Gradients (IG).
   • Relevance is propagated backward from the output layer to the input layer, ensuring relevance conservation across the network.

2. Layer-Wise Pruning Configuration:

   • The paper identifies a critical flaw in "global" pruning: due to relevance conservation, deeper layers (which have more parameters) often show lower average relevance per parameter compared to shallower layers. Global thresholds disproportionately prune these deeper, critical layers.
   • Solution: FedX applies layer-wise pruning. A separate pruning threshold ($\tau^{(l)}_q$) is calculated for each layer $l$ based on the $q$-quantile of relevance scores within that specific layer. This ensures a balanced reduction across the network architecture.

3. Pruning Process:

   • Warm-up Phase: Pruning is delayed for $\upsilon$ communication rounds (e.g., 10 rounds) to allow the global model to stabilize.
   • Mask Generation: A binary mask is generated where components with relevance scores below the layer-specific threshold are marked for removal (0), and others are retained (1).
   • Application: The mask is applied to the global model weights ($w = w \odot M$). The resulting sparse model is distributed to clients for the next round of local training.
   • Consistency: Once a mask is established, it is retained for subsequent rounds to ensure all clients operate on structurally identical models, preventing structural divergence.

3. Key Contributions

• Novel Strategy: Introduction of FedX, the first FL strategy specifically designed for RS that utilizes explanation-guided pruning to reduce communication overhead.
• Architectural Innovation: The proposal of a layer-wise pruning configuration to overcome the structural bias inherent in global pruning when using relevance-based methods.
• Robustness to Heterogeneity: Demonstrating that explanation-guided pruning is more effective than magnitude-based pruning in heterogeneous (non-IID) FL settings, as it identifies task-relevant features rather than just large weights.
• Comprehensive Evaluation: Extensive experiments on two major RS datasets (BigEarthNet-S2 for multi-label and EuroSAT for single-label classification) across six diverse DL architectures (ResNet variants, ConvNeXt, ViT).

4. Experimental Results

The authors evaluated FedX against unpruned baselines, magnitude-based pruning, and state-of-the-art (SOTA) FL pruning algorithms (FedDST, FedDIP, LotteryFL, PruneFL).

• Pruning Configuration:

  • Layer-wise vs. Global: Layer-wise pruning significantly outperformed global pruning, especially at high sparsity (e.g., 90% pruning). Global pruning caused a drastic performance drop (e.g., mAP falling to 33%), while layer-wise maintained high accuracy (55-69%).
  • Explanation Method: LRP generally outperformed Integrated Gradients (IG) in robustness, particularly at extreme sparsity levels.
  • Magnitude vs. Explanation: Magnitude-based pruning failed at sparsity levels >30%, whereas explanation-guided pruning maintained high performance.

• Performance Across Architectures:

  • FedX maintained competitive performance across all tested models (ResNet6 to ViT-B16).
  • Generalization Boost: In several cases (e.g., ViT-B16, ConvNeXt Atto), the pruned models achieved higher accuracy than the unpruned baseline, suggesting that FedX acts as an effective structural regularizer by removing redundant parameters.
  • ViT-B16: Maintained >68% mAP on BigEarthNet-S2 even with extreme compression.

• Communication Efficiency:

  • FedX achieved up to 44% reduction in total communication volume (uplink + downlink) with minimal accuracy loss.
  • For large models like ViT-B16, this saved approximately 46.3 GB of data transmission.
  • Smaller models (ResNet6) also saw significant relative savings (~44%).

• Computational Cost:

  • The relevance estimation step adds computational overhead (e.g., $5.0 \times 10^{13}$ FLOPs for ViT-B16), but this occurs only once per pruning cycle on the server.
  • The pruning mask application is negligible.
  • The trade-off is highly favorable: moderate server-side computation yields massive communication savings.

• Comparison with SOTA:

  • FedX consistently outperformed SOTA algorithms (FedDST, FedDIP, etc.) and random pruning across all sparsity levels.
  • At 90% pruning, FedX maintained competitive performance while other methods degraded significantly.

5. Significance

• Enabling Real-World RS FL: FedX addresses the critical bandwidth bottleneck in Remote Sensing, making FL feasible for satellites and UAVs where communication is expensive and limited.
• Privacy and Efficiency: It achieves communication efficiency without compromising data privacy (no raw data leaves the client) or model performance.
• Methodological Insight: The paper highlights that how pruning is applied (layer-wise vs. global) is as important as the pruning metric itself. It establishes that explanation-guided methods are superior to magnitude-based heuristics in the context of Federated Learning with heterogeneous data.
• Scalability: The architecture-agnostic nature of FedX allows it to be seamlessly integrated into existing FL workflows for various modern deep learning architectures.

In conclusion, FedX represents a significant advancement in making Federated Learning practical for large-scale, privacy-sensitive, and bandwidth-constrained Remote Sensing applications by intelligently pruning models based on task relevance rather than simple weight magnitude.