Sintesi Tecnica: FedProIn – Mitigare il Client Drift per Prototipi Apprendibili nella Medical Imaging Federata

Problematica

L'Apprendimento Federato (FL) nella medical imaging è significativamente ostacolato dall'eterogeneità statistica derivante dalle variazioni nei protocolli degli scanner, nei metodi di acquisizione e nelle diverse popolazioni di pazienti tra i vari centri medici. Questa distribuzione dei dati non-IID (non-Indipendente e Identicamente Distribuita) induce il client drift, dove i modelli locali divergono dal modello globale durante l'ottimizzazione. Nello specifico, istanze identiche vengono mappate in regioni distinte dello spazio delle caratteristiche da diversi client, e le rappresentazioni locali diventano polarizzate verso le caratteristiche specifiche del singolo client.

Gli approcci esistenti presentano limitazioni:

• La regolarizzazione basata su vincoli (es. FedProx, MOON) fatica ad allineare le rappresentazioni semantiche in spazi ad alta dimensionalità.
• I metodi standard basati su prototipi spesso assumono un singolo prototipo globale per classe (tramite vettori di media delle feature), il che è inadeguato per i dati medici eterogenei.
• I metodi a prototipi multipli spesso introducono un'elevata complessità computazionale o si affidano a rappresentazioni unimodali che non riescono a ottimizzare efficacemente i confini decisionali.

Metodologia: FedProIn

Gli autori propongono FedProIn (Federated Learning with Prototype aggregation with Influence), un framework che utilizza Prototipi Apprendibili (LPs) come parametri del modello ottimizzabili anziché come vettori di media statici. Il modello globale consiste in un encoder di feature $f_\theta$ e un insieme di prototipi $P \in \mathbb{R}^{C \times M \times d}$, dove $M$ è il numero di prototipi per classe.

1. Decomposizione del Client Drift

FedProIn affronta esplicitamente il client drift scomponendolo in due componenti e applicando specifiche perdite di regolarizzazione:

• Feature Drift: Affrontato tramite la Feature Divergence Loss (FDL). I client mantengono una copia congelata dell'estrattore di feature globale ($\theta^{t-1}$). L'encoder locale è regolarizzato per minimizzare l'errore quadratico medio tra le embedding delle feature locali e le feature globali di riferimento, garantendo l'allineamento semantico.
• Prototype Drift: Affrontato tramite la Prototype Contrastive Loss (PCL). Questa perdita attira i prototipi locali aggiornati verso i prototipi globali correnti, allontanandoli al contempo dai loro stati locali precedenti. Ciò impedisce ai prototipi di diventare eccessivamente polarizzati verso i cambiamenti dei dati locali.

L'obiettivo locale totale combina la perdita di classificazione primaria ($\ell_{Task}$) con questi termini di regolarizzazione:
$$\ell_{Total} = \ell_{Task} + \lambda_{FDL}\ell_{FDL} + \lambda_{PCL}\ell_{PCL}$$

Per gestire lo sbilanciamento delle classi nei dataset medici a coda lunga (long-tailed), il metodo impiega il campionamento bilanciato per classe e l'augmentation adattiva per garantire che i sottotipi di malattie rare siano adeguatamente rappresentati durante gli aggiornamenti dei prototipi.

2. Normalized Influence Aggregation (NIA)

L'aggregazione standard (es. FedAvg) è insufficiente per i prototipi multipli perché pesa gli aggiornamenti esclusivamente in base alla dimensione del dataset, ignorando l'utilizzo dei prototipi. Ciò può amplificare il rumore proveniente da prototipi sottoutilizzati o sopprimere rappresentazioni specializzate per casi rari.

FedProIn introduce la Normalized Influence Aggregation (NIA):

• Matrice di Influenza ($I_i$): Durante l'addestramento locale, ogni client traccia un punteggio di influenza per ogni prototipo, contando quante volte un particolare prototipo è servito come vicino più prossimo per i campioni locali.
• Aggregazione: I prototipi globali vengono aggregati pesandoli proporzionalmente ai loro punteggi di utilizzo ($I_{k,c}^i$), piuttosto che solo in base alla dimensione del dataset del client. Questo riduce il peso degli aggiornamenti rumorosi da prototipi non utilizzati e preserva i prototipi specializzati rilevanti per specifici subset di client.

Contributi Chiave

1. Formulazione dei Prototipi Apprendibili: Adottare i prototipi come parametri ottimizzabili (congiuntamente con l'estrattore di feature) anziché come vettori di media statici per massimizzare i margini di separazione delle classi.
2. Strategia di Doppia Regolarizzazione: Mitigare esplicitamente il client drift scomponendolo in feature drift e prototype drift, affrontati rispettivamente tramite FDL e PCL.
3. Normalized Influence Aggregation (NIA): Un nuovo meccanismo di aggregazione che pesa adattivamente i prototipi dei client in base alla loro utilità locale, affrontando i problemi di eterogeneità nell'utilizzo e di attivazione sparsa nei contesti a prototipi multipli.
4. Validazione Empirica: Valutazione completa su dataset di medical imaging pubblici che dimostra la robustezza contro l'eterogeneità statistica.

Risultati Sperimentali

Il framework è stato valutato su due dataset di medical imaging a coda lunga: HAM10000 (immagini dermatoscopiche) e Matek-19 (immagini di cellule del sangue periferico). Gli esperimenti sono stati condotti sia in scenari IID che Non-IID (Dirichlet $\alpha=0.5$) con $K=10$ client.

Evidenze delle Prestazioni:

• HAM10000: FedProIn ha raggiunto l'83,5% di accuratezza (IID) e l'81,1% di accuratezza (Non-IID), superando tutti i baseline inclusi FedAvg, FedProx, MOON e altri metodi basati su prototipi (FedProto, FedTGP, FedPLVM).
• Matek-19: FedProIn ha raggiunto il 96,2% di accuratezza (IID) e il 95,8% di accuratezza (Non-IID).
• Robustezza: In condizioni Non-IID, i metodi basati su prototipi esistenti hanno subito cali significativi delle prestazioni (es. una perdita di accuratezza del 6-7% su HAM10000), mentre FedProIn ha limitato le perdite relative a soli il 2,4% e lo 0,4% su HAM10000 e Matek-19, rispettivamente.
• Studio di Ablazione: I risultati hanno confermato che sia FDL che PCL sono necessari per le prestazioni ottimali. Rimuovere FDL ha causato il calo più netto delle prestazioni, indicando che il feature drift è la modalità di fallimento dominante. Il modello completo ha superato significativamente le varianti a componente singola e il caso degenero (equivalente a FedAvg).

Significato e Assertività

Il paper sostiene che FedProIn mitiga efficacemente il client drift nella medical imaging federata sfruttando prototipi apprendibili e un'aggregazione consapevole dell'influenza. Gli autori affermano che il loro approccio:

• Ottiene una generalizzazione superiore rispetto ai baseline di FL esistenti e ai metodi personalizzati, particolarmente in ambienti altamente eterogenei (Non-IID).
• Affronta con successo i limiti delle assunzioni a prototipo singolo e dell'aggregazione non pesata nei contesti a prototipi multipli.
• Fornisce una soluzione robusta per la medical imaging dove l'eterogeneità dei dati e lo sbilanciamento delle classi sono prevalenti.

Gli autori notano che il lavoro futuro esplorerà l'adattamento di FedProIn per la FL personalizzata (rimuovendo la dipendenza dal modello globale) e l'estensione del framework a scenari di FL eterogenei con diverse architetture di modelli. Il codice è reso disponibile per supportare la riproducibilità.