FedSKD: Aggregation-free Model-heterogeneous Federated Learning via Multi-dimensional Similarity Knowledge Distillation for Medical Image Classification

FedSKD is een nieuw aggregatievrij federatief leerframework voor medische beeldclassificatie dat model-heterogeniteit mogelijk maakt via een ronde-om-circulatie van modellen en multi-dimensionale kennisdistillatie, waardoor schaalbaarheid en personalisatie worden verbeterd zonder centrale server.

De kern

🏥 Het Probleem: De "Geheime Kookpotten"

Stel je voor dat er tientallen ziekenhuizen zijn die allemaal hun eigen unieke recepten hebben voor het diagnosticeren van ziektes (zoals autisme of huidkanker). Ze hebben allemaal geweldige data, maar door privacywetten mogen ze hun patiëntgegevens niet met elkaar delen. Het is alsof elke kok in een gesloten keuken staat en niemand mag proeven wat de ander maakt.

Traditionele methoden proberen dit op te lossen door een centrale chef-kok (een server) aan te stellen. Deze chef vraagt iedereen om hun recepten op te sturen, mixt ze in een grote pot (aggregatie) en stuurt het nieuwe "super-recept" terug.

• Het nadeel: Dit werkt niet goed als elke kok een heel ander type keukenapparatuur heeft (verschillende computermodellen) of als ze heel verschillende ingrediënten gebruiken (verschillende patiëntpopulaties). De grote pot wordt dan een rommelige soep die niemand lekker vindt.

🚀 De Oplossing: FedSKD (De "Ronde Tafel" Methode)

De auteurs van dit paper, Ziqiao Weng en collega's, hebben een nieuwe manier bedacht die geen centrale chef nodig heeft. Ze noemen het FedSKD.

Stel je voor dat de koks niet naar een centrale tafel gaan, maar een ronde tafel vormen waar ze direct met elkaar praten. Maar er is een probleem: als kok A zijn recept aan kok B geeft, en die aan kok C, vergeten ze vaak wat ze van kok A leerden. Dit noemen ze in de vaktaal model drift (je raakt je kennis kwijt) en kennisverdunning (de goede ideeën worden verwaterd).

FedSKD lost dit op met twee slimme trucs:

1. De "Ronde Tafel" (Zonder Server)

In plaats van een centrale server, geven de ziekenhuizen hun modellen (hun "recepten") in een ronde gang door aan elkaar.

• Ziekenhuis A leert van Ziekenhuis B.
• Ziekenhuis B leert van Ziekenhuis C.
• Ziekenhuis C leert weer van A.
  Elk ziekenhuis mag zijn eigen unieke keukenapparatuur gebruiken (hete modellen). Ze hoeven niet allemaal hetzelfde type computer te hebben.

2. De "Meer-dimensionale Smakenproeverij" (Knowledge Distillation)

Dit is het hart van de uitvinding. Als twee koks met elkaar praten, kijken ze niet alleen naar het eindresultaat (de soep), maar analyseren ze de smaken op drie niveaus:

• Niveau 1: De "Bak" (Batch-wise): Kijken ze naar de hele groep patiënten? "Hebben jullie ook gemerkt dat patiënten met blauwe ogen vaker hoofdpijn hebben?" Dit zorgt voor een globaal begrip.
• Niveau 2: De "Pixel" (Pixel/Voxel-wise): Kijken ze naar de details? "Kijk eens naar deze specifieke vlek op de huid." Dit zorgt dat ze de fijne details niet vergeten.
• Niveau 3: De "Regio" (Region-wise): Kijken ze naar samenhangende delen? "In het brein werken deze twee gebieden altijd samen." Dit helpt bij complexe ziektes zoals autisme.

Door op al deze niveaus te "proeven" en te vergelijken, leren ze van elkaar zonder hun eigen unieke stijl te verliezen. Ze houden hun eigen "vork" vast (hun eigen model), maar nemen een hapje van de "schotel" van de ander.

🧪 De Resultaten: Wat leverde het op?

De wetenschappers hebben dit getest op twee echte medische taken:

1. Autisme-diagnose: Gebaseerd op hersenscans (fMRI) van mensen uit 17 verschillende ziekenhuizen in Europa en de VS.
2. Huidkanker: Het herkennen van verschillende soorten moedervlekken.

De uitkomsten waren indrukwekkend:

• Beter dan de rest: FedSKD deed het beter dan alle andere methoden, zelfs beter dan methoden die een centrale server gebruiken.
• Twee vliegen in één klap: Het systeem werd niet alleen beter in het algemeen (voor alle ziekenhuizen), maar elk ziekenhuis werd ook beter in zijn eigen specifieke geval.
• Veiligheid: Zelfs als één ziekenhuis per ongeluk (of kwaadwillend) slechte data gaf, viel het hele systeem niet uit. Het systeem was robuust.

🌟 Waarom is dit belangrijk?

Stel je voor dat je een wereldwijde groep artsen hebt die samenwerken zonder hun patiëntennaamlijsten te delen. FedSKD zorgt ervoor dat:

1. Iedereen mee kan doen: Zelfs als je een oude computer hebt of een heel nieuw systeem, kun je deelnemen.
2. Niets verloren gaat: De kennis die je van een ander leert, wordt niet overschreven door je eigen werk.
3. Privacy gewaarborgd blijft: Er worden geen patiëntdata uitgewisseld, alleen de "slimheid" van het model.

Kortom: FedSKD is als een slimme, wereldwijde kookclub waar elke kok zijn eigen unieke recept behoudt, maar door slim te proeven op verschillende niveaus, leren ze elkaar zo goed kennen dat ze allemaal betere gerechten kunnen maken dan ze alleen ooit hadden kunnen doen.

Technische samenvatting

Probleemstelling

Federated Learning (FL) staat voor privacy-bevorderende samenwerking zonder directe data-uitwisseling, maar bestaande methoden kampen met twee fundamentele beperkingen in medische beeldvorming:

1. Model-heterogeniteit: Cliënten (ziekenhuizen) hebben vaak verschillende rekenkracht en specifieke behoeften, wat leidt tot het gebruik van modellen met verschillende architecturen. Bestaande Model-Heterogene FL (MHFL) methoden zijn afhankelijk van een centrale server voor aggregatie, wat schaalbaarheid beperkt en vaak vereist dat modellen gedeeltelijk identiek zijn.
2. Peer-to-Peer (P2P) uitdagingen: P2P-FL verwijdert de centrale server, maar bestaande P2P-methoden trainen doorgaans één homogeen model sequentieel door cliënten. Dit leidt tot model-drift (de modelparameters oscilleren door veranderende optimalisatiedoelen) en kennisverdunning (kennis van eerdere cliënten wordt overschreven of vergeten naarmate het model zich aanpast aan nieuwe cliënten).

Er is dus behoefte aan een framework dat volledig heterogene modellen ondersteunt, geen centrale server vereist, en de problemen van drift en kennisverlies in P2P-omgevingen oplost.

Methodologie: FedSKD

De auteurs stellen FedSKD voor, een nieuw P2P MHFL-framework dat twee kerninnovaties combineert:

1. Aggregatie-vrije, ronde-wijze modelcirculatie:
In plaats van een centrale server of sequentiële training van één model, circuleren lokaal getrainde heterogene modellen (genaamd Domain-Adaptive Models of DAM) via een ronde-wijze (round-robin) strategie tussen cliënten.

• In elke ronde ontvangt een cliënt een model van een voorgaande cliënt (de Knowledge-Transit Model of KTM).
• De cliënt voert bidirectionele kennisdistillatie uit tussen zijn eigen lokale DAM en de ontvangen KTM.
• Na training wordt de KTM verworpen en wordt de geüpdatete DAM doorgegeven aan de volgende cliënt. Dit creëert een gesloten lus van kennisversterking.

2. Multi-dimensionale Similarity Knowledge Distillation (SKD):
Om kennisoverdracht mogelijk te maken tussen modellen met verschillende architecturen (zonder gedeelde gewichten), introduceert FedSKD een distillatiemethode die puur op feature-representatieniveau werkt. Deze methode omvat drie niveaus:

• Batch-wise SKD (B-SKD): Aligneert semantische relaties tussen samples binnen een mini-batch (inter-sample correlaties).
• Pixel/Voxel-wise SKD (P-SKD): Aligneert ruimtelijke feature-verdelingen op het niveau van individuele pixels of voxels om lokale structurele patronen te behouden.
• Region-wise SKD (R-SKD): Aligneert representaties binnen anatomische regio's (bijv. hersengebieden in fMRI), wat cruciaal is voor het vangen van functionele connectiviteit.

Technische Implementatie:

• Het model wordt opgesplitst in een feature extractor en een prediction header.
• Tijdens de wederzijdse leerfase worden de parameters van de prediction header van de ontvangen KTM bevroren (frozen), terwijl alleen de feature extractor wordt fijngefineerd. Dit voorkomt dat de KTM te sterk aanpast aan de lokale data (catastrophic forgetting) en behoudt de cross-client kennis.
• De totale loss functie combineert cross-entropy loss met de multi-dimensionale SKD loss, gewogen door een hyperparameter $\gamma$.

Belangrijkste Bijdragen

1. Eerste P2P-MHFL Framework: FedSKD is het eerste framework dat peer-to-peer federated learning toepast op volledig heterogene modellen voor medische beeldclassificatie, zonder afhankelijkheid van een centrale server.
2. Oplossing voor Drift en Verdunning: Door middel van de voorgestelde multi-dimensionale SKD en bidirectionele distillatie, worden de fundamentele problemen van model-drift en kennisverdunning in sequentiële P2P-FL effectief aangepakt.
3. Nieuwe Datasets: De auteurs hebben een geografisch gepartitioneerde dataset gebaseerd op ABIDE (FedASD) gecreëerd, die de eerste is in de literatuur die institutionele heterogeniteit behoudt. Daarnaast zijn benchmarks voor huidlaesies (FedSkin) gebruikt.

Resultaten

FedSKD is geëvalueerd op twee medische taken:

1. Autisme Spectrum Stoornis (ASD) Diagnose: Gebaseerd op fMRI-data (ABIDE dataset, 17 centra).
2. Huidlaesie Classificatie: Gebaseerd op dermoscopische beelden (Derm7pt dataset).

Kernbevindingen:

• Superieure Prestaties: FedSKD presteert significant beter dan state-of-the-art baselines (zowel homogene als heterogene), waaronder FedAvg, FedProx, FedBN, pFedAFM, AlFeCo en FedCross.
  • Op de FedASD dataset (heterogeen) behaalde FedSKD een gemiddelde AUC van 76,66% (lokaal) en 67,10% (globaal), wat een verbetering is van ongeveer 5-6% ten opzichte van de beste concurrenten.
  • Op de FedSkin dataset (onder verschillende non-IID scenario's) overtrof FedSKD de concurrenten met marges tot wel 8-9%.
• Personalisatie vs. Generalisatie: Het framework slaagt erin om zowel de prestaties op de lokale cliënt (personalisatie) als de prestaties op andere cliënten (generalisatie) te verbeteren, wat een zeldzame combinatie is in heterogene FL.
• Robuustheid: FedSKD toonde zich robuust tegenover label-flipping data poisoning attacks, waarbij het de prestaties van andere methoden aanzienlijk overtrof.
• Fairness: Het framework verminderde de dispariteit in prestaties tussen geslachtsgroepen (man/vrouw) aanzienlijk, wat wijst op verbeterde eerlijkheid.
• Kosten: Hoewel FedSKD ongeveer het dubbele van de rekencapaciteit (GFLOPs) vereist omdat twee modellen lokaal worden getraind, is dit een aanvaardbare trade-off voor de aanzienlijke prestatiewinst. Het heeft geen extra serverkosten en dezelfde communicatiebandbreedte als standaard FedAvg.

Betekenis en Impact

FedSKD vertegenwoordigt een paradigmaverschuiving in federated learning voor de medische sector. Het lost het dilemma op tussen privacy, model-flexibiliteit en samenwerking:

• Het maakt het mogelijk voor ziekenhuizen met verschillende IT-infrastructuur en privacybeperkingen om samen te werken zonder data te delen of een centrale server te vertrouwen.
• Het bewijst dat kennisdistillatie op feature-niveau (in plaats van gewicht-aggregatie) een krachtige oplossing is voor heterogene systemen.
• De resultaten suggereren dat dit framework schaalbaar en robuust is voor real-world medische toepassingen, waar data-heterogeniteit en privacy cruciaal zijn.

Samenvattend biedt FedSKD een schaalbare, server-vrije oplossing die de beperkingen van bestaande MHFL-methoden overwint en een nieuwe standaard zet voor privacy-bewuste, collaboratieve medische AI.