Neural Collapse-Inspired Multi-Label Federated Learning under Label-Distribution Skew

Dit artikel introduceert FedNCA-ML, een nieuw federatief leerframework dat inspiratie haalt uit Neural Collapse en een attentiemodule gebruikt om effectief multi-label classificatieproblemen op te lossen onder omstandigheden van ernstige label-distributie-scheefheid.

De kern

Stel je voor dat je een groep artsen uit verschillende ziekenhuizen over de hele wereld wilt samenwerken om een slimme computer te trainen die ziektes kan herkennen. Maar er is een groot probleem: ze mogen hun patiëntgegevens niet delen vanwege privacywetten. Dit is wat Federated Learning (Federatief Leren) doet: de computer leert van iedereen, zonder dat de data ooit de ziekenhuizen verlaat.

Maar nu komt de echte uitdaging: Meerdere labels en scheve verdelingen.

Het Probleem: Een Verwarde Puzel

In de echte wereld hebben patiënten vaak meerdere ziektes tegelijk (bijvoorbeeld een longaandoening én een hartprobleem). Bovendien is de data niet eerlijk verdeeld:

• Ziekenhuis A ziet alleen veel patiënten met longziektes.
• Ziekenhuis B ziet alleen patiënten met huidziektes.
• Ziekenhuis C ziet bijna alleen gezonde mensen.

Als je deze artsen nu gewoon laat samenwerken, leert de computer zich te specialiseren in wat elk ziekenhuis ziet. Het resultaat? De computer wordt heel goed in het herkennen van de veelvoorkomende ziektes, maar vergeet de zeldzame ziektes volledig. Het is alsof je een groepje mensen vraagt om een kaart van de wereld te tekenen, maar iedereen tekent alleen de stad waar hij woont. Het eindresultaat is een rommelige, onvolledige kaart.

De Oplossing: FedNCA-ML (De "Neural Collapse" Methode)

De auteurs van dit paper, onderzoekers van de Universiteit van Oxford, hebben een slimme oplossing bedacht genaamd FedNCA-ML. Ze gebruiken een concept uit de wiskunde dat "Neural Collapse" heet. Laten we dit uitleggen met een creatieve analogie.

1. De "Perfecte Zaal" (Neural Collapse)

Stel je voor dat je een grote, ronde zaal hebt met tafels. Elke tafel staat voor een specifieke ziekte (bijv. "Longontsteking", "Huidkanker").

• Het oude probleem: Elke arts (client) plaatst de stoelen rondom hun eigen tafel wat willekeurig neer. De arts die alleen longziektes ziet, duwt de stoelen voor "Longontsteking" heel dicht bij elkaar, maar vergeet de stoelen voor "Huidkanker" helemaal.
• De nieuwe methode: De onderzoekers zeggen: "Laten we een perfecte, vaste indeling maken." Ze stellen een onzichtbare, perfecte geometrische structuur op (een Simplex ETF). Denk hierbij aan een perfecte sterrenhemel waar elke ster (ziekte) op een exacte, evenwijdige afstand van elkaar staat.
• Het effect: Alle artsen krijgen nu dezelfde blauwdruk. Ze moeten hun stoelen (de data van hun patiënten) precies op die perfecte plekken zetten. Hierdoor wordt de "Longontsteking"-arts gedwongen om ook aandacht te besteden aan de "Huidkanker"-tafel, omdat die taak in de blauwdruk vaststaat. Dit voorkomt dat de computer zich alleen richt op wat lokaal bekend is.

2. De "Speciale Brillen" (Attention Module)

In een meervoudige ziekte-situatie is het lastig om één beeld te hebben dat alles dekt.

• De analogie: Stel je voor dat een patiënt een foto is van een drukke markt met veel kraampjes.
• Het probleem: Als je één grote foto maakt van de hele markt, verdwijnen de details van de specifieke kraampjes (de ziektes) in de chaos.
• De oplossing: FedNCA-ML gebruikt een speciale bril (een Attention-module). Deze bril laat de computer niet naar de hele markt kijken, maar focust zich op één kraampje tegelijk.
  • "Kijk nu alleen naar de viskraam (Longontsteking)."
  • "Kijk nu alleen naar de bloemenkraam (Huidkanker)."
    Hierdoor kan de computer voor elke ziekte apart een duidelijk beeld vormen, zelfs als die ziekte zeldzaam is.

3. De "Discipline" (Extra Regels)

Om te voorkomen dat de computer toch weer gaat rommelen, voegen ze twee extra regels toe:

• De "Nee-Regel": Als een patiënt geen longontsteking heeft, moet de computer zeker weten dat hij niet per ongeluk denkt dat het er wel is. Dit wordt een "afwijzingsstraf" genoemd.
• De "Groeps-Regel": Alle patiënten met dezelfde ziekte moeten op de foto heel dicht bij elkaar staan (in de digitale wereld), zodat ze makkelijk te herkennen zijn.

Wat levert dit op?

Door deze methode te gebruiken, wordt de samenwerking tussen de ziekenhuizen veel beter. De computer leert niet alleen de "gemakkelijke" en veelvoorkomende ziektes, maar wordt ook veel beter in het herkennen van de zeldzame en complexe gevallen.

In de tests (op datasets met duizenden medische beelden) bleek dat deze methode tot wel 5% beter presteerde dan andere methoden, vooral bij het vinden van de zeldzame ziektes.

Kort samengevat:
FedNCA-ML is als een super-coördinator die ervoor zorgt dat elke arts in het team, ongeacht welke patiënten ze dagelijks zien, zich richt op een perfecte, gezamenlijke kaart van alle mogelijke ziektes. Hierdoor wordt de diagnose voor iedereen beter, en krijgen de zeldzame ziektes de aandacht die ze verdienen.

Technische samenvatting

Titel: Neural Collapse-Geïnspireerde Multi-Label Federated Learning onder Label-Distributie-Verzadiging (Label-Distribution Skew)

Auteurs: Can Peng, Yuyuan Liu, Yingyu Yang, Pramit Saha, Qianye Yang, en J. Alison Noble (Universiteit van Oxford).

---

1. Het Probleem

Federated Learning (FL) staat toe dat modellen worden getraind op gedistribueerde cliënten zonder dat ruwe data wordt gedeeld, wat essentieel is voor privacygevoelige domeinen zoals medische beeldvorming. Hoewel bestaande FL-methoden vaak gericht zijn op enkelvoudige classificatie, zijn veel real-world toepassingen multi-label (een sample kan meerdere labels hebben, bijv. meerdere ziekten in één röntgenfoto).

Deze studie richt zich op een specifiek en onderbelicht scenario: Multi-Label FL onder Label-Distributie-Verzadiging (Label Skew). Dit introduceert drie fundamentele uitdagingen:

1. Ongelijkmatige Data-distributie: Cliënten hebben vaak extreme onevenwichtigheid in labels (bijv. een ziekenhuis ziet veel "longontsteking" maar zelden "fracturen"). Dit leidt tot overfitting op meerderheidslabels en ondertraining van minderheidslabels.
2. Multi-label Co-occurrence Bias: Labels komen vaak samen voor. Frequent voorkomende labels domineren het trainingssignaal en onderdrukken het leren van discriminatieve features voor zeldzame labels.
3. Inconsistentie tussen Cliënten: Cliënten verschillen niet alleen in de frequentie van labels, maar ook in de afhankelijkheidsstructuren tussen labels (welke ziekten komen samen voor). Dit veroorzaakt optimalisatieconflicten en "client drift", waardoor het global model slecht convergeert.

---

2. Methodologie: FedNCA-ML

De auteurs stellen FedNCA-ML (Federated Neural Collapse Alignment for Multi-Label Learning) voor. Dit is een unificerend representatieleraarframework dat gebruikmaakt van de Neural Collapse (NC) theorie.

Kernconcept: Neural Collapse (NC)
NC beschrijft een ideale geometrie in de latente ruimte van een getrainde classifier:

• Features van dezelfde klas collapse naar hun gemiddelde.
• De klassen-gemiddelden vormen een Simplex Equiangular Tight Frame (ETF): een symmetrische, maximaal gescheiden configuratie waarbij alle klassen evenwijdig en evenwijdig gescheiden zijn.

Architectuur van FedNCA-ML:
Het framework bestaat uit drie hoofdcomponenten:

1. Label-Aware Disentanglement Module (LADM):

   • In multi-label settings is één gedeelde beeldrepresentatie vaak onvoldoende omdat het verschillende semantische bewijzen verstrengelt.
   • LADM gebruikt een attention-based mechanism (geïnspireerd op DETR) om klas-specifieke features te extraheren uit de gedeelde beeldfeatures.
   • Er wordt één vaste "query" per klas gebruikt (gedeeld door alle cliënten) om te zorgen voor consistentie. Dit transformeert het multi-label probleem in een reeks per-klas subproblemen die compatibel zijn met NC-uitlijning.

2. NC-Geïnspireerde Feature Uitlijning (ETF Anchoring):

   • In plaats van een leerbare classifier te gebruiken, wordt een globaal gedeelde, vooraf gedefinieerde ETF-matrix gebruikt als de classifier en als de query-matrix voor de LADM.
   • Dit fungeert als een vaste geometrische "anker" voor alle cliënten. Het dwingt de klas-specifieke features om zich te aligneren met deze ideale geometrie, wat client drift vermindert en zorgt voor een consistente decision boundary, ongeacht de lokale data-distributie.

3. Regularisatie Verliezen:
   Om de clustering in de latente ruimte te verbeteren, worden twee extra loss-functies toegevoegd:

   • Negative Feature Rejection Loss: Onderdrukt ruis door te straffen als een feature voor een negatief label (afwezig) per ongeluk hoog lijkt op een ander klassen-prototype.
   • Positive Feature Contrastive Loss: Bevordert compacte intra-klassen clustering en duidelijke inter-klassen scheiding door positieve features dichter bij hun eigen prototype te duwen dan bij andere.

Het totale trainingsdoel is een som van de Binary Cross-Entropy (BCE), de negatieve rejection loss en de positieve contrastive loss.

---

3. Belangrijkste Bijdragen

• Probleemformulering: De eerste formele behandeling van multi-label FL onder zowel label-frequentie skew als label-afhankelijkheids skew.
• FedNCA-ML Framework: Een uniek framework dat NC-theorie toepast op multi-label FL door een gedeelde ETF-geometrie te gebruiken om representatie-drift te mitigeren.
• Klas-specifieke Attention: Een mechanisme (LADM) dat multi-label leren herkwalificeert naar per-klas subproblemen, waardoor NC-uitlijning mogelijk wordt zonder de semantische relaties in de gedeelde backbone te verstoren.
• Robustheid: Introductie van complementaire regularisatietermen die de robuustheid verhogen onder heterogene label-distributies.

---

4. Resultaten

De methode is geëvalueerd op vijf datasets (CIFAR-10, PASCAL VOC, MS COCO, DermaMNIST, ChestX-ray14) onder negen verschillende FL-configuraties met variërende mate van non-IID data en label-skew.

• Prestaties: FedNCA-ML behaalde consistent de beste class-wise (macro) prestaties, wat aangeeft dat het model beter is in het leren van zeldzame klassen in plaats van alleen de meerderheid.
  • Op multi-label CIFAR-10: Verbetering van 3,92% in class-wise AUC en 4,93% in class-wise F1-score ten opzichte van de beste concurrenten.
  • Op DermaMNIST (medisch): Verbetering van 0,77% in AUC en 4,93% in F1-score.
  • Op ChestX-ray14: Verbetering van 1,21% in class-wise AUC, wat cruciaal is voor het detecteren van zeldzame ziekten.
• Ablatie Studies:
  • Het gebruik van alleen een ETF-classifier zonder LADM presteerde slecht (features waren te verstrengeld).
  • Het toevoegen van LADM en de ETF-anker verbeterde de prestaties aanzienlijk.
  • De regularisatietermen zorgden voor verdere verbetering in balans en discriminatie.
• Visualisatie: t-SNE en Grad-CAM visualisaties tonen aan dat FedNCA-ML beter gescheiden clusters leert (op basis van semantiek in plaats van label-aantallen) en zich richt op de juiste semantische regio's in de afbeeldingen.

---

5. Betekenis en Impact

Dit werk is significant omdat het een oplossing biedt voor een van de meest realistische maar moeilijke scenario's in Federated Learning: gedistribueerd leren met ongelijke en complexe multi-label data.

• Medische Toepassingen: Het is direct toepasbaar in de gezondheidszorg, waar ziekenhuizen verschillende patiëntpopulaties hebben en ziekten vaak samen voorkomen. De methode zorgt ervoor dat zeldzame ziekten niet worden genegeerd door het model.
• Theoretische Vooruitgang: Het breidt de toepassing van Neural Collapse (eerdere beperkt tot enkelvoudige classificatie) uit naar complexe multi-label scenario's.
• Privacy en Samenwerking: Het bewijst dat hoogwaardige, gebalanceerde modellen kunnen worden opgeleid zonder dat data wordt gedeeld, zelfs wanneer de lokale data sterk vertekend is.

Kortom, FedNCA-ML biedt een robuust en theoretisch onderbouwd kader om de "client drift" in heterogene multi-label FL-systemen op te lossen, waardoor het een belangrijke stap voorwaarts is voor privacy-bewuste AI in complexe domeinen.