Riepilogo Tecnico: Allineamento della Distribuzione per l'Apprendimento Federato One-Shot tramite Trasporto Ottimale

1. Definizione del Problema

Il documento affronta le sfide inerenti all'Apprendimento Federato One-Shot (OSFL), un regime di comunicazione in cui i client interagiscono con un server centrale esattamente una volta. Mentre l'Apprendimento Federato (FL) standard mitiga l'eterogeneità dei dati attraverso l'ottimizzazione iterativa client-server, l'OSFL manca di questa capacità di raffinamento.

Il problema centrale identificato è l'effetto cumulativo di due tipi specifici di eterogeneità statistica:

1. Spostamento di Dominio (Domain Shift): I client osservano distribuzioni di input $P_k(x)$ differenti a causa di variabili condizioni di acquisizione o ambienti diversi.
2. Spostamento di Etichetta (Label Shift): I client presentano marginali di etichetta $P_k(y)$ differenti, portando a proporzioni di classe variabili.

In assenza di correzione iterativa, questi spostamenti inducono rappresentazioni di caratteristiche disallineate e funzioni decisionali della classe-posteriore inconsistenti tra i client. I metodi OSFL esistenti (che si affidano a distillazione, sintesi di dati o ensembling) spesso assumono rappresentazioni allineate o affrontano questi spostamenti separatamente. Di conseguenza, le discrepanze nella geometria delle caratteristiche e nei bias condizionali delle classi rimangono non corrette, degradando le prestazioni del modello globale aggregato.

2. Metodologia: SLOT-Align

Gli autori propongono SLOT-Align (Single-round, Learning-free Optimal Transport Alignment), un framework consapevole della geometria progettato per armonizzare le distribuzioni delle caratteristiche dei client prima dell'addestramento OSFL a valle. Il metodo opera come uno step di pre-processing e si basa su un encoder condiviso e congelato $f_\theta$ (pre-addestrato su un dataset esterno) per mappare tutti gli input locali in uno spazio latente comune.

Il flusso di lavoro consiste in quattro fasi chiave:

A. Estrazione delle Statistiche Locali delle Caratteristiche

Ogni client $k$ calcola statistiche compatte di primo e secondo ordine delle caratteristiche $z = f_\theta(x)$ estratte dai propri dati locali:

• Media: $\mu_k = \mathbb{E}[f_\theta(x)]$
• Covarianza: $\Sigma_k$ (regolarizzata tramite shrinkage di Ledoit–Wolf per garantire la definitività positiva).
  Queste statistiche fungono da proxy Gaussiano $N(\mu_k, \Sigma_k)$ per la distribuzione delle caratteristiche indotta, catturando le principali discrepanze in posizione e dispersione causate dagli spostamenti di dominio e di etichetta.

B. Costruzione del Riferimento Globale

Il server aggrega queste statistiche locali per costruire una distribuzione di riferimento globale:

• Media Globale: Una media ponderata delle medie dei client.
• Covarianza Globale: Calcolata come il baricentro Bures–Wasserstein delle covarianze dei client. Questo viene risolto tramite un'iterazione a punto fisso per trovare l'unico minimizzatore della somma delle distanze quadrate di Bures, bilanciando efficacemente la variabilità specifica del client con la coerenza globale.

C. Mappatura di Trasporto Geodetico

Ogni client costruisce una mappa di trasporto per allineare la propria approssimazione Gaussiana locale al riferimento globale.

• Mappa Ottimale: Sotto costo quadratico, la mappa di trasporto ottimale tra due Gaussiane è affine e ha una soluzione in forma chiusa che coinvolge le radici quadrate delle matrici di covarianza.
• Interpolazione: Per evitare la sovra-correzione (poiché la vera distribuzione delle caratteristiche non è strettamente Gaussiana), il metodo introduce un parametro di interpolazione $\tau \in [0, 1]$. La trasformazione finale $T^{(\tau)}_k$ è una combinazione convessa della mappa identità e della mappa di trasporto ottimale completa. Ciò crea un'interpolazione di spostamento di Wasserstein, muovendo la distribuzione del client lungo una geodetica a velocità costante verso il riferimento globale.

D. Integrazione

I client applicano la trasformazione $z' = T^{(\tau)}_k(z)$ alle loro caratteristiche locali. Queste caratteristiche allineate vengono poi alimentate in qualsiasi algoritmo OSFL esistente (ad es., FedAvg, FedCGS, FedPFT) senza modificare le procedure di ottimizzazione a valle.

3. Contributi Chiave

1. Formalizzazione: Il documento formalizza l'OSFL come un problema di allineamento della distribuzione sotto spostamento congiunto di dominio ed etichetta, evidenziando come questi spostamenti limitino l'efficacia degli approcci one-shot.
2. Proposta di Algoritmo: Introduce SLOT-Align, un framework basato sul trasporto ottimale, learning-free (senza apprendimento), che armonizza le distribuzioni delle caratteristiche dei client tramite un singolo scambio di statistiche compatte. Non richiede dati ausiliari, generatori addestrabili o ottimizzazione iterativa.
3. Validazione Empirica: Esperimenti estesi dimostrano che SLOT-Align migliora costantemente l'accuratezza e la robustezza attraverso molteplici benchmark (Office-Home, Digits, DomainNet), vari backbone pre-addestrati (ViT, ResNet) e i migliori algoritmi OSFL allo stato dell'arte.

4. Risultati Sperimentali

Gli autori hanno valutato SLOT-Align sotto diversi gradi di spostamento di etichetta (controllato dal parametro di concentrazione di Dirichlet $\alpha$) e diversità di dominio.

• Guadagni di Performance: Sotto spostamento congiunto di dominio ed etichetta ($\alpha=0.1$), SLOT-Align ha migliorato costantemente l'accuratezza media macro per tutti i baseline testati (O-FedAvg, FedCGS, FedPFT). Ad esempio, sul dataset Office-Home, ha migliorato O-FedAvg del +7.86% e FedPFT del +6.55%.
• Robustezza: Il metodo è rimasto efficace sotto forte skew di etichetta ($\alpha=0.05$ e $\alpha=0.01$), mostrando guadagni anche in regimi di eterogeneità severa.
• Ablation sull'Intensità di Allineamento: Il parametro di interpolazione $\tau$ è risultato critico. Valori intermedi (ad es., $\tau=0.4$) hanno prodotto il miglior compromesso, mentre $\tau=1.0$ (trasporto completo) ha talvolta degradato le prestazioni, suggerendo che un allineamento parziale sia preferibile a una correzione totale.
• Spostamento di Solo Dominio: In scenari con etichette bilanciate (solo spostamento di dominio), SLOT-Align ha fornito guadagni marginali o prestazioni neutre, confermando che il metodo mira specificamente alla complessa interazione tra spostamenti di dominio e di etichetta.
• Efficienza: Il costo di comunicazione è minimo ($O(m + m^2)$ per client per media e covarianza) e l'overhead computazionale è trascurabile rispetto all'estrazione delle caratteristiche e all'addestramento a valle.

5. Significato e Rivendicazioni

Il documento sostiene che SLOT-Align offra una soluzione modulare e computazionalmente efficiente a un collo di bottiglia critico nell'OSFL: l'incapacità di correggere il disallineamento della distribuzione in un singolo round di comunicazione.

• Modularità: Non sostituisce le pipeline OSFL esistenti ma agisce come uno strato di pre-processing compatibile, rendendolo applicabile a una vasta gamma di metodi attuali e futuri che si basano su encoder congelati.
• Consapevolezza della Geometria: Sfruttando la geometria dello spazio 2-Wasserstein e i baricentri Bures–Wasserstein, corregge esplicitamente le discrepanze sia nella posizione (media) che nella forma (covarianza) delle distribuzioni delle caratteristiche.
• Praticità: Il metodo evita i rischi di privacy e l'overhead computazionale associati alla sintesi di dati o alla distillazione iterativa, affidandosi esclusivamente a riepiloghi statistici compatti.

Gli autori concludono che, sebbene l'OSFL sia limitato da budget di comunicazione rigorosi, SLOT-Align mitiga efficacementamente il degrado delle prestazioni causato dalle distribuzioni di dati eterogenee, in particolare quando gli spostamenti di dominio e di etichetta interagiscono per distorcere le statistiche delle caratteristiche.