FedAgain: A Trust-Based and Robust Federated Learning Strategy for an Automated Kidney Stone Identification in Ureteroscopy

Probleemstelling

De toepassing van kunstmatige intelligentie (AI) in de medische beeldvorming, en specifiek voor de automatische identificatie van nierstenen tijdens ureteroscopie, wordt geconfronteerd met drie kritieke uitdagingen:

1. Data Privacy en Fragmentatie: Medische data is verspreid over verschillende ziekenhuizen en kan niet centraal worden verzameld vanwege privacywetgeving (zoals GDPR en HIPAA). Federated Learning (FL) lost dit op door modelupdates te delen in plaats van ruwe data, maar de huidige implementaties zijn vaak niet robuust genoeg.
2. Data Heterogeniteit (Non-IID): Ziekenhuizen gebruiken verschillende apparatuur, opnameprotocollen en hebben verschillende patiëntpopulaties. Dit leidt tot niet-identiek verdeelde (Non-IID) data, wat de convergentie van standaard FL-algoritmen (zoals FedAvg) verstoort.
3. Beeldcorruptie en Onbetrouwbaarheid: Endoscopische beelden zijn vaak onderhevig aan artefacten zoals bewegingsonscherpte, spiegelende reflecties, variabele verlichting en bloedingen. Bovendien kunnen sommige deelnemende klinieken (clients) onbetrouwbare updates leveren door slechte datakwaliteit of zelfs kwaadaardig gedrag (Byzantijnse clients). Standaard aggregatiemethoden zijn hier erg gevoelig voor, wat leidt tot een daling in diagnostische nauwkeurigheid.

Methodologie: FedAgain

Het artikel introduceert FedAgain, een op vertrouwen gebaseerde Federated Learning-strategie die specifiek is ontworpen om robuustheid en generalisatie te verbeteren in heterogene en corrupte omgevingen.

Kernmechanisme: Dual Trust Mechanism
In tegenstelling tot traditionele methoden die updates gelijk behandelen of alleen op basis van datasetgrootte wegen, berekent FedAgain dynamisch een vertrouwensscore voor elke client per communicatieronde. Deze score is gebaseerd op twee signalen:

1. Benchmark Fout ($E_k$): Voordat een client lokaal traint, wordt de prestatie van het huidige globale model op de lokale validatie-set van die client gemeten. Een hoge fout wijst op een slechte representativiteit van de data of een slecht gekalibreerd model voor die specifieke client.
2. Parameter Divergentie ($D_k$): Na lokale training wordt de afwijking (divergentie) berekend tussen de geüpdatete lokale modelgewichten en de globale server-gewichten. Een extreme divergentie kan wijzen op overfitting op ruis, slechte convergentie of kwaadaardige updates.

Aggregatie Strategie
De server combineert deze twee signalen tot een trust-score ($T_k$) volgens de formule:
$$T_k = \frac{1}{E_k \cdot D_k + \varepsilon}$$
Waarbij $\varepsilon$ een kleine constante is om deling door nul te voorkomen.

• Clients met een lage fout en lage divergentie (betrouwbare updates) krijgen een hoge trust-score en thus een zwaarder gewicht.
• Clients met hoge fouten of extreme divergentie (onbetrouwbare/corrupte updates) krijgen een lage score en worden afgezwakt of genegeerd tijdens de aggregatie.

De globale update wordt vervolgens berekend als een gewogen som van de client-updates, waarbij de gewichten gebaseerd zijn op deze genormaliseerde trust-scores.

Belangrijkste Bijdragen

1. Eerste Robuuste FL voor Nierstenen: Dit is het eerste federated learning-kader dat specifiek is ontworpen voor de in-vivo herkenning van niersteen-types, met expliciete aandacht voor beeldcorruptie en client-onbetrouwbaarheid.
2. Trust-Gewogen Aggregatie: De introductie van een eenvoudige maar effectieve aggregatieregel die client-bijdragen dynamisch weegt op basis van prestatie en consistentie, zonder de noodzaak van complexe uitbijterdetectie of specifieke aanpassingen voor aanvallen.
3. Uitgebreide Evaluatie: Een reproduceerbare PyTorch-codebase en een evaluatieframework dat zeven stress-tests omvat. Dit omvat variaties in corruptiepercentages (10% tot 50%), ernstniveaus, en client-aantallen.
4. Diverse Dataset Validatie: Het model is getest op twee publieke benchmarks (MNIST, CIFAR-10) en drie klinische niersteen-datasets (Michel Daudon, Jonathan El-Beze, en MyStone), waaronder zowel ex-vivo beelden als gesimuleerde endoscopische beelden.

Resultaten

De experimenten tonen aan dat FedAgain consistent presteert beter dan standaard methoden (FedAvg, FedProx) en robuuste methoden (FedMedian, Bulyan) onder diverse omstandigheden:

• Robuustheid tegen Corruptie: Onder zware corruptie-scenario's (tot 40% corrupte clients met hoge ernst) behield FedAgain een nauwkeurigheid van ongeveer 72-73%, terwijl FedAvg en Bulyan aanzienlijk degradeerden (FedAvg daalde naar ~68%, Bulyan naar ~64%). FedAgain toonde minder variabiliteit in prestaties.
• Non-IID Heterogeniteit: In scenario's met label-skew (waarbij ziekenhuizen gespecialiseerd zijn in bepaalde steentypes) en Dirichlet-verdeling, behaalde FedAgain een hogere nauwkeurigheid en lagere variantie dan FedAvg. Het presteerde vergelijkbaar met of beter dan FedMedian, maar zonder de noodzaak om het aantal kwaadaardige clients vooraf te kennen.
• Schaalbaarheid: Naarmate het aantal deelnemende klinieken toenam (van 5 naar 20), degradeerde FedAvg sneller dan FedAgain. FedAgain behield zijn prestaties beter, wat aantoont dat de vertrouwensmechanismen effectief zijn bij het filteren van onbetrouwbare updates in grote federaties.
• Algemene Validatie: Op MNIST en CIFAR-10 bevestigde FedAgain zijn generalisatievermogen, waarbij het op complexere taken (CIFAR-10) een duidelijke meerwaarde bood ten opzichte van FedAvg onder corruptie.

Significantie

FedAgain vertegenwoordigt een praktische doorbraak in de richting van betrouwbare, privacy-bewuste en klinisch inzetbare AI voor medische beeldvorming.

• Klinische Toepasbaarheid: Door de robuustheid tegen beeldcorruptie (zoals die vaak voorkomt in endoscopie) en heterogene data, maakt het de ontwikkeling van AI-assistenten mogelijk die veilig kunnen worden ingezet in diverse ziekenhuizen zonder dat data centraal hoeft te worden opgeslagen.
• Vertrouwen in AI: De methade lost het "black-box" probleem van onbetrouwbare updates op door transparante, prestatie-gedreven gewichting, wat essentieel is voor de reguliere goedkeuring van medische AI.
• Toekomstperspectief: Het kader biedt een solide basis voor toekomstig onderzoek naar privacy-versterkende technieken (zoals differentieel privacy) en adaptieve leerstrategieën in complexe, dynamische medische omgevingen.

Kortom, FedAgain bewijst dat het actief integreren van betrouwbaarheidsmetingen in de federatieve leerlus cruciaal is voor het overwinnen van de realiteit van verspreide, onzuivere en heterogene medische data.