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

Il paper propone FedNCA-ML, un nuovo framework di Federated Learning multi-etichetta che, ispirandosi alla teoria del Neural Collapse e integrando un modulo basato sull'attenzione, affronta efficacemente l'eterogeneità dei dati e lo skew delle distribuzioni delle etichette allineando le rappresentazioni dei client per migliorare le prestazioni di classificazione.

L'essenziale

Immagina di dover formare un super-medico globale capace di diagnosticare molte malattie diverse (come un tumore alla pelle, un'infezione polmonare o una malattia cardiaca) allo stesso tempo.

Il problema è che questo medico non può vedere i dati reali dei pazienti per motivi di privacy. Deve imparare lavorando con 10 ospedali diversi (i "clienti"), ognuno dei quali ha i suoi pazienti, le sue malattie più comuni e le sue regole interne.

Ecco la situazione difficile (il "problema"):

1. Squilibrio: L'Ospedale A vede solo pazienti con la pelle scura e malattie rare, mentre l'Ospedale B vede solo pazienti con la pelle chiara e malattie comuni.
2. Malattie multiple: Un paziente può avere sia l'asma che il diabete. È difficile insegnare al medico a riconoscere entrambe le cose contemporaneamente senza confondersi.
3. Il conflitto: Se ogni ospedale addestra il suo "assistente" solo sui propri pazienti, quando provano a unire le conoscenze, gli assistenti si scontrano. Quello dell'Ospedale A pensa che una certa macchia sia normale, mentre quello dell'Ospedale B pensa sia pericolosa.

La Soluzione: FedNCA-ML (Il "Guru della Geometria")

Gli autori di questo paper hanno creato un nuovo metodo chiamato FedNCA-ML. Per spiegarlo in modo semplice, usiamo un'analogia con un orchestra e una partitura perfetta.

1. La Teoria del "Crollo Neurale" (Neural Collapse)

Immagina che ogni malattia abbia una sua "forma geometrica perfetta" nello spazio mentale del computer. In un mondo ideale, se tutti i pazienti avessero la stessa malattia, le loro immagini si raggrupperebbero tutte in un unico punto preciso, come un gruppo di amici che si stringono in un abbraccio. E ogni gruppo di malattie diverse dovrebbe stare il più lontano possibile dagli altri, come isole in un oceano.
Questo stato ideale si chiama "Crollo Neurale". È come se tutte le malattie avessero una "casa" perfetta e ben definita.

2. Il Problema della Realtà

Nella realtà, gli ospedali sono disordinati. L'Ospedale A ha solo 3 isole (malattie), l'Ospedale B ne ha 10. Se ognuno disegna la sua mappa, le isole non coincidono. Quando si uniscono, le mappe diventano un caos.

3. La Magia del Metodo (FedNCA-ML)

Il metodo proposto fa tre cose intelligenti:

• La Partitura Fissa (ETF): Invece di lasciare che ogni ospedale inventi la sua forma per le malattie, il server centrale fornisce una "partitura geometrica fissa". È come dare a tutti gli ospedali lo stesso modello di come dovrebbero essere disposte le isole (le malattie) nello spazio. Non importa se un ospedale vede poche isole o molte; devono tutte allinearsi a questa mappa perfetta. Questo impedisce che l'Ospedale A "dimentichi" le malattie rare perché non le ha mai viste.
• Gli Occhiali Speciali (LADM - Modulo di Disentanglement): Spesso, quando guardiamo un paziente, vediamo tutto insieme (pelle, respiro, cuore) e ci confondiamo. Il metodo usa un "occhio intelligente" (un meccanismo di attenzione) che, invece di guardare l'immagine intera, estrae un'etichetta specifica per ogni malattia.
  • Analogia: Immagina di avere un'immagine di una stanza piena di oggetti. Invece di dire "questa è una stanza", il sistema usa un filtro speciale per isolare solo "la sedia", poi solo "il tavolo", poi solo "la lampada". Questo aiuta il computer a capire che una persona può avere sia l'asma che il diabete senza che le due diagnosi si cancellino a vicenda.
• Il Controllo di Qualità (Le Regole di Pulizia): Il sistema aggiunge due regole extra per assicurarsi che il lavoro sia pulito:
  1. Non confondere il rumore: Se un ospedale non ha mai visto una certa malattia, il sistema gli dice: "Non cercare di indovinare, ignora quel segnale confuso".
  2. Stai vicino ai tuoi simili: Se un paziente ha il diabete, assicurati che la sua immagine sia molto vicina ad altre immagini di pazienti con il diabete, e molto lontana da quelle con l'asma.

Il Risultato

Grazie a questo approccio, quando gli ospedali uniscono i loro "assistenti", non ottengono un mostro confuso, ma un super-medico equilibrato.

• Non ignora le malattie rare (quelle che solo un ospedale su dieci vede).
• Capisce bene quando un paziente ha più malattie insieme.
• È molto preciso nel distinguere le diverse condizioni.

In sintesi: Il paper dice che invece di far litigare gli ospedali su come classificare le malattie, diamo loro una mappa geometrica perfetta e condivisa e degli occhiali speciali per vedere ogni malattia singolarmente. Il risultato è un'intelligenza artificiale che impara meglio, più velocemente e senza dimenticare nessuno, anche quando i dati sono molto diversi tra loro.

Sommario tecnico

1. Il Problema

Il lavoro affronta una sfida critica nel campo dell'apprendimento federato (Federated Learning - FL): l'addestramento collaborativo di modelli per la classificazione multi-etichetta in scenari caratterizzati da squilibrio nella distribuzione delle etichette (Label-Distribution Skew).

In molte applicazioni reali, come la diagnostica medica, i dati sono:

• Multi-etichetta: Un singolo campione può contenere più condizioni o oggetti contemporaneamente (es. un paziente con più malattie).
• Non-IID (Non-Independent and Identically Distributed): Ogni client (es. un ospedale) possiede una distribuzione di dati unica, spesso con gravi squilibri (alcune classi sono dominanti, altre rare o assenti) e pattern di co-occorrenza delle etichette diversi rispetto ad altri client.

Le sfide principali identificate sono:

1. Squilibrio dei dati: I client ottimizzano i modelli verso le loro distribuzioni locali, portando a un sovradattamento sulle classi maggioritarie e a un sottotraining delle classi minoritarie.
2. Bias di co-occorrenza: Le etichette frequenti tendono a dominare il segnale di addestramento, sopprimendo l'apprendimento di caratteristiche discriminative per le condizioni rare.
3. Incoerenza cross-client: Le differenze nelle frequenze delle etichette e nelle loro dipendenze strutturali tra i client generano conflitti di ottimizzazione, impedendo la convergenza globale e la generalizzazione del modello.

2. Metodologia: FedNCA-ML

Gli autori propongono FedNCA-ML (Federated Neural Collapse Alignment for Multi-Label Learning), un framework che integra la teoria del Neural Collapse (NC) per allineare le rappresentazioni dei client in uno spazio latente coerente.

Il metodo si basa su tre componenti chiave:

A. Modulo di Disentanglement Consapevole delle Etichette (LADM)

Poiché una singola rappresentazione condivisa è spesso insufficiente per catturare evidenze specifiche per classe in contesti multi-etichetta, il framework introduce un modulo basato sull'attenzione (ispirato a DETR).

• Funzionamento: Estrae rappresentazioni specifiche per ogni classe partendo dalle feature dell'immagine condivise.
• Meccanismo: Utilizza un set di "query" fisse (una per classe) che prestano attenzione ai token dell'immagine.
• Vantaggio: Questo trasforma il problema multi-etichetta in un insieme di sottoproblemi per classe, permettendo di applicare l'allineamento geometrico senza forzare l'esclusività mutua nello spazio delle feature di base (backbone), preservando così le relazioni semantiche.

B. Allineamento Geometrico basato su Neural Collapse (NC)

Il cuore del metodo è l'uso di una struttura geometrica condivisa ispirata al Neural Collapse.

• ETF (Equiangular Tight Frame): Viene definito un classificatore globale fisso basato su una matrice ETF, che rappresenta una configurazione geometrica ideale dove i centroidi delle classi formano un simplettico massimamente separato.
• Ancoraggio: Le rappresentazioni specifiche per classe estratte dal LADM vengono allineate a questa struttura ETF globale. Questo agisce come un "ancoraggio" che previene lo drift dei client verso le loro distribuzioni locali distorte, garantendo che tutti i client imparino in uno spazio di feature coerente e bilanciato.

C. Funzioni di Regolarizzazione

Per migliorare la robustezza e la compattezza delle clusterizzazioni, vengono introdotti due termini di perdita aggiuntivi:

1. Negative Feature Rejection Loss: Penalizza l'allineamento spurio delle feature negative (etichette assenti) con prototipi di altre classi, sopprimendo il rumore.
2. Positive Feature Contrastive Loss: Promuove la compattezza intra-classe e la separazione inter-classe, spingendo le feature positive verso il loro prototipo ETF specifico.

L'obiettivo totale combina la Cross-Entropy binaria (BCE) con questi due termini di regolarizzazione.

3. Contributi Chiave

1. Formalizzazione del problema: Definizione rigorosa del FL multi-etichetta sotto squilibrio di etichette, considerando sia le frequenze che i pattern di co-occorrenza.
2. Framework FedNCA-ML: Introduzione di un approccio di pre-allineamento che utilizza la geometria ETF condivisa per mitigare il drift delle rappresentazioni e migliorare l'apprendimento bilanciato.
3. Meccanismo di Attenzione per Classe: Sviluppo di un modulo (LADM) che permette l'applicazione dell'allineamento NC in scenari multi-etichetta, risolvendo il conflitto tra feature condivise e specifiche per classe.
4. Regolarizzazione Complementare: Integrazione di loss di rifiuto e contrasto per migliorare la separabilità delle classi in distribuzioni eterogenee.

4. Risultati Sperimentali

Il metodo è stato valutato su 5 dataset (CIFAR-10, PASCAL VOC, MS COCO, DermaMNIST, ChestX-ray14) in 9 diversi scenari FL con vari gradi di non-IID (parametro $\beta$) e presenza di classi mancanti (parametro $\gamma$).

• Performance Superiori: FedNCA-ML ha ottenuto i migliori risultati nella maggior parte dei casi, superando gli approcci state-of-the-art (come FedAvg, FedProx, FedLGT).
• Miglioramenti Specifici:
  • Su CIFAR-10 multi-etichetta: miglioramento fino al 3.92% nell'AUC per classe e 4.57% nel F1-score per classe.
  • Su DermaMNIST: miglioramento fino al 4.93% nel F1-score per classe.
  • Su ChestX-ray14: miglioramenti significativi nell'AUC per classe, indicando una migliore capacità di riconoscere le malattie rare (classi minoritarie) rispetto ai metodi esistenti che tendono a favorire la classe "nessuna anomalia".
• Analisi di Ablazione: Gli esperimenti dimostrano che la combinazione di LADM, ancoraggio ETF e regolarizzazione è essenziale. L'uso di query fisse (ETF) invece di query apprendibili si è rivelato cruciale per la stabilità in scenari non-IID.
• Visualizzazione: Le visualizzazioni t-SNE e Grad-CAM confermano che il modello apprende cluster semantici più compatti e si focalizza sulle regioni rilevanti dell'immagine per ogni classe specifica.

5. Significato e Impatto

Questo lavoro è significativo perché:

• Colma un gap teorico: Estende la teoria del Neural Collapse, finora limitata alla classificazione multi-classe, al dominio complesso e poco esplorato del FL multi-etichetta.
• Soluzione pratica per la medicina: Offre un approccio robusto per scenari medici reali dove i dati sono privati, distribuiti in modo eterogeneo e caratterizzati da malattie rare e multiple.
• Generalizzazione: Dimostra che imporre una struttura geometrica globale (ETF) può mitigare efficacemente i conflitti di ottimizzazione causati da dati locali distorti, portando a modelli globali più equilibrati e generalizzabili senza richiedere la condivisione di dati grezzi.

In sintesi, FedNCA-ML rappresenta un avanzamento fondamentale nel rendere l'apprendimento federato praticabile ed efficace per compiti di classificazione complessi e sbilanciati, garantendo equità e accuratezza anche per le classi minoritarie.