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

In dit artikel wordt FedX voorgesteld, een nieuwe strategie voor federatief leren in de aardobservatie die communicatie-inefficiëntie aanpakt door het gebruik van uitleggerichte pruning om de grootte van de uitgewisselde modellen te verminderen zonder in te leveren op de prestaties.

De kern

Stel je voor dat je een gigantische puzzel moet oplossen, maar de stukjes zijn verspreid over honderden verschillende huizen in een dorp. Iedereen heeft een deel van de puzzel, maar niemand mag zijn eigen stukjes naar een centraal kantoor sturen vanwege privacyregels. In plaats daarvan sturen ze alleen een foto van hoe hun stukje eruitziet, zodat een centrale meesterpuzzelaar een groot plaatje kan maken.

Dit is precies hoe Federated Learning werkt, een slimme manier om kunstmatige intelligentie te leren zonder dat data de locatie verlaat. Maar er is een probleem: die "foto's" van de puzzelstukjes zijn vaak gigantisch groot. Het versturen van deze enorme bestanden kost veel tijd en internetbandbreedte, net als het proberen om een hele bibliotheek per post te versturen.

In dit paper introduceren de auteurs FedX, een slimme oplossing die dit probleem oplost. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Probleem: De "Lekke" Internetverbinding

In de wereld van satellietbeelden (zoals foto's van de aarde gemaakt door drones of satellieten) zijn de data vaak verspreid over verschillende landen of organisaties. Ze willen samenwerken om betere kaarten te maken, maar kunnen de foto's niet zomaar delen.
Het probleem is dat de "instructies" (het model) die ze uitwisselen enorm groot zijn. Het is alsof je elke dag een zware koffer vol met onbruikbare spullen moet meenemen, terwijl je alleen een paar belangrijke sleutels nodig hebt.

2. De Oplossing: FedX als een Slimme Editor

FedX is als een zeer slimme redacteur die in het centrale kantoor werkt. In plaats van de hele zware koffer te versturen, doet FedX het volgende:

• De "Waarom"-Vraag: Normaal gesproken kijken computers alleen naar hoe groot een getal is (een gewicht) om te beslissen of het belangrijk is. FedX doet iets anders: het vraagt zich af "Waarom is dit stukje belangrijk voor de oplossing?".
• De "Verklaring" (Explanation): FedX gebruikt een techniek die lijkt op het terugsporen van een spoor. Het kijkt naar een paar voorbeeldpuzzels (die openbaar zijn) en vraagt zich af: "Welke puzzelstukken hebben echt bijgedragen aan het juiste antwoord?"
• Het Weggooien van Ruis: Als FedX ziet dat bepaalde stukjes van de puzzel (de neurale netwerken) eigenlijk niets bijdragen aan het begrijpen van de beelden, dan worden die stukjes weggehaald. Het is alsof je de koffer leegt van alle oude kranten en alleen de sleutels en de foto's van de puzzel overhoudt.

3. De Creatieve Analogie: De Tuin en de Tuinman

Stel je het AI-model voor als een enorme, overwoekerde tuin.

• De Traditionele Manier: De tuinman (het algoritme) kijkt alleen naar hoe groot de struiken zijn. Als een struik klein is, wordt hij weggegooid. Het probleem is dat soms een heel kleine bloem (een klein stukje van het model) de belangrijkste is voor de kleur van de tuin, maar omdat hij klein is, wordt hij per ongeluk verwijderd.
• De FedX Manier: FedX is een tuinman die een verklaring gebruikt. Hij kijkt niet alleen naar de grootte, maar vraagt: "Welke planten zorgen ervoor dat de tuin er mooi uitziet?" Hij gebruikt een speciale bril (de explanation method) om te zien welke planten echt bloeien en welke alleen maar ruimte innemen.
  • Hij knipt alleen de dode takken en onkruid weg.
  • De rest van de tuin blijft intact, maar is nu veel lichter en makkelijker te verzorgen.

4. Waarom is dit zo slim? (De Laag-voor-Laag Benadering)

Een van de belangrijkste ontdekkingen in dit paper is dat FedX niet zomaar over de hele tuin knipt.

• Fout: Als je over de hele tuin één grote schaar gebruikt, verwijder je per ongeluk de diepere lagen van de tuin (diep in de grond) omdat die daar "minder belangrijk" lijken op dat moment, terwijl ze juist cruciaal zijn.
• FedX: FedX kijkt naar elke laag apart. Hij zegt: "In deze laag van de tuin mag je 20% weg, maar in die andere laag mag je 10% weg." Hierdoor blijft de structuur van de tuin perfect in balans, zelfs als je heel veel weghaalt.

5. Het Resultaat: Sneller, Lichter en Slimmer

Door deze methode te gebruiken:

• Minder Verkeer: De bestanden die tussen de satellieten en het kantoor worden verstuurd, worden tot wel 44% kleiner. Het is alsof je van een vrachtwagen vol met zand overstapt op een fiets met een paar sleutels.
• Beter Resultaat: Verwacht je dat het model slechter wordt omdat er minder is? Integendeel! Vaak werkt het model zelfs beter na het weghalen van de ruis. Het is alsof je een student helpt door alle afleidingen uit zijn kamer te halen; hij kan zich beter concentreren op wat echt belangrijk is.
• Privacy: Niemand hoeft zijn eigen foto's of data te delen. Alleen de "schoonmaakregels" worden gedeeld.

Conclusie

FedX is als een slimme, zorgzame editor die voor je zorgt dat je alleen de essentiële informatie deelt. Het maakt het mogelijk om samen te werken aan grote projecten (zoals het monitoren van bossen of steden via satelliet) zonder dat je je internetverbinding overbelast of je privacy opgeeft. Het is een stap in de richting van een efficiëntere, veiligere en snellere toekomst voor kunstmatige intelligentie.

Technische samenvatting

Probleemstelling

Federated Learning (FL) is een veelbelovend paradigma voor het trainen van diepe neurale netwerken op gedecentraliseerde data, wat essentieel is voor Remote Sensing (RS) toepassingen vanwege privacywetgeving en de onmogelijkheid om grote datasets centraal op te slaan. Een kritieke beperking van FL in RS-contexten is echter de communicatie-overhead. De frequente uitwisseling van grote modelupdates tussen clients (zoals satellieten of drones) en de centrale server vereist veel bandbreedte, wat vaak beperkt is in deze omgevingen.

Bestaande oplossingen zoals modelcompressie (kwantisatie) en kennisdistillatie hebben nadelen:

• Compressie kan leiden tot prestatieverlies door het verwijderen van belangrijke parameters en introduceert extra rekentijd voor codering/decodering.
• Kennisdistillatie vereist vaak studentmodellen die minder goed presteren dan docentmodellen en kan gevoelige informatie lekken via gedeelde logits.
• Traditionele pruning (bijv. op basis van gewichtsgrootte) is vaak onbetrouwbaar in FL, omdat updates van heterogene clients elkaar kunnen opheffen, waardoor belangrijke parameters in het globale model een lage magnitude hebben en per ongeluk worden verwijderd.

Methodologie: FedX

De auteurs stellen FedX voor, een nieuwe strategie die uitleg-geleide pruning (explanation-guided pruning) combineert met FL om de communicatiekosten te verlagen zonder de prestaties te schaden.

Kernprincipes:

1. Server-side Pruning: In tegenstelling tot veel andere methoden die pruning op de client uitvoeren, voert FedX de pruning uit op de centrale server na het aggregeren van lokale updates.
2. Explanation-Guided Relevance: In plaats van te vertrouwen op de magnitude van gewichten, gebruikt FedX backpropagation-gebaseerde uitlegmethode (zoals Layer-wise Relevance Propagation - LRP en Integrated Gradients - IG). Deze methoden berekenen "relevantiescores" voor modelcomponenten (bijv. convolutiefilters) op basis van een referentiedataset ($D_{ref}$) die op de server beschikbaar is.
3. Lagen-voor-lagen Pruning (Layer-wise): Een cruciale innovatie is het gebruik van een layer-wise pruning configuratie. Omdat relevantie over de lagen heen behouden blijft (relevance conservation), zouden diepere lagen met meer parameters een lagere gemiddelde relevantie per parameter hebben. Een globale drempel zou deze lagen onevenredig hard treffen. FedX past een drempel toe binnen elke laag afzonderlijk, wat zorgt voor een evenwichtigere en robuustere pruning.
4. Warm-up Fase: Pruning begint pas na een warm-up fase (bijv. 10 communicatierondes) om te zorgen dat het globale model stabiel is voordat de relevantiescores betrouwbaar zijn.
5. Consistent Mask: Eenmaal een pruning-mask (binair masker) is gegenereerd, wordt dit mask behouden en hergebruikt in alle volgende rondes. Dit voorkomt structurele divergentie tussen clients.

Werkingscyclus:

1. Clients trainen lokaal en sturen updates naar de server.
2. De server aggregeert de updates tot een globaal model.
3. De server berekent relevantiescores voor componenten op basis van $D_{ref}$.
4. Op basis van deze scores wordt een masker gegenereerd dat de minst relevante componenten per laag verwijdert.
5. Het gesparseerde (gepruneerde) globale model wordt teruggestuurd naar de clients voor de volgende ronde.

Belangrijkste Bijdragen

• FedX Framework: De eerste FL-strategie die uitleg-geleide pruning gebruikt om communicatie-overhead te reduceren in RS-beeldclassificatie.
• Layer-wise Configuratie: Het introduceren van een lagen-specifieke pruning-drempel om de bias tegen diepere lagen te elimineren die optreedt bij globale pruning.
• Superioriteit boven Magnitude-based Pruning: Het aantonen dat het gebruik van relevantiescores (via LRP/IG) superieur is aan traditionele magnitude-based pruning, vooral bij hoge sparsiteitsniveaus.
• Open Source: De code is publiek beschikbaar gesteld.

Experimentele Resultaten

De auteurs evalueerden FedX op twee datasets: BigEarthNet-S2 (multi-label classificatie, non-IID data) en EuroSAT (single-label classificatie, IID data), met verschillende architecturen (ResNet, ViT, ConvNeXt).

• Prestaties bij Hoge Sparsiteit:

  • FedX (met LRP en layer-wise configuratie) behoudt aanzienlijk betere prestaties dan globale pruning of magnitude-based pruning. Bij 90% pruning behield FedX (LRP, Layerwise) nog steeds ~55% mAP op BigEarthNet-S2, terwijl globale strategieën instortten tot ~33%.
  • FedX presteerde consistent beter dan state-of-the-art methoden zoals FedDST, FedDIP, PruneFL en LotteryFL, vooral bij extreme sparsiteit (80-90%).
  • In sommige gevallen (bijv. ViT-B16) presteerden de gepruneerde modellen zelfs beter dan het ongepruneerde model, wat suggereert dat pruning als een effectieve regularisator werkt.

• Communicatie-efficiëntie:

  • FedX reduceerde de totale communicatiekosten met tot 44% (bij 90% pruning) zonder significante prestatieverlies.
  • Voor grote modellen zoals ViT-B16 (106 GB totaal) resulteerde dit in een besparing van ongeveer 46 GB aan data-uitwisseling.

• Rekenkosten:

  • De berekening van relevantiescores introduceert enige rekenkosten (FLOPs), maar dit gebeurt slechts één keer per pruning-rondje op de server. De kosten voor het toepassen van het masker zijn verwaarloosbaar. De besparing in communicatie weegt ruimschoots op tegen deze extra server-rekenlast.

Significantie en Conclusie

FedX biedt een robuuste oplossing voor de bandbreedte-beperkingen in Federated Learning voor Remote Sensing. Door uitleg-geleide relevantiescores te combineren met een lagen-specifieke aanpak, slaagt het erin om de communicatie-overhead drastisch te verlagen terwijl de generalisatiecapaciteit van het model behouden blijft of zelfs verbetert.

De methode is architectuur-agnostisch en vereist geen wijzigingen in de client-side training, wat het zeer geschikt maakt voor schaalbare implementaties in omgevingen met strikte privacy- en wettelijke beperkingen. De resultaten tonen aan dat FedX een nieuwe standaard kan worden voor communicatie-efficiëntie in gedecentraliseerde RS-toepassingen.