FedNSAM:Consistency of Local and Global Flatness for Federated Learning

Il paper propone FedNSAM, un nuovo algoritmo per l'apprendimento federato che armonizza la piattezza locale e globale integrando il momento Nesterov globale negli aggiornamenti locali per migliorare la generalizzazione del modello in scenari con dati eterogenei.

L'essenziale

Il Problema: La "Crisi di Identità" nella Scuola Federata

Immagina un grande progetto scolastico chiamato Federated Learning (Apprendimento Federato).
Invece di mettere tutti gli studenti in una sola classe enorme (che violerebbe la privacy, perché i loro quaderni resterebbero pubblici), ogni studente studia a casa propria con i propri appunti. Ogni tanto, si incontrano con il professore centrale per dire: "Ehi, ho imparato questa cosa, ecco il riassunto". Il professore poi unisce tutti i riassunti per creare un unico "Libro della Conoscenza" globale.

Il problema:
Spesso, gli studenti hanno appunti molto diversi (alcuni studiano solo matematica, altri solo storia). Quando il professore cerca di unire questi riassunti, il risultato finale è confuso e "punto" (in termini tecnici: il modello globale finisce in una posizione "acuta" e instabile, che non generalizza bene). È come se il libro finale avesse pagine strappate e incollate male.

Per risolvere questo, gli studiosi hanno inventato un metodo chiamato SAM (Minimizzazione Consapevole della "Acutezza"). L'idea è: "Non fermarti solo dove pensi di aver finito, ma guarda intorno. Se la strada è liscia e piatta, fermati lì. Se è una montagna ripida, continua a camminare". Questo aiuta a trovare soluzioni più robuste.

La scoperta degli autori:
Gli autori di questo paper (Liu e colleghi) hanno notato un difetto fatale nel metodo SAM quando usato in questa scuola federata:

• Ogni studente, studiando da solo, trova una "strada piatta" perfetta per i suoi appunti.
• Ma quando il professore unisce le strade di tutti gli studenti, queste "strade piatte" non si allineano!
• Immagina che lo studente A trovi un piano perfetto per la matematica, e lo studente B trovi un piano perfetto per la storia. Quando il professore li unisce, il piano si rompe e diventa una montagna ripida. Il modello globale diventa di nuovo "acuto" e poco affidabile.

Hanno chiamato questo divario "Distanza di Piattezza" (Flatness Distance). Più gli studenti sono diversi (dati eterogenei), più grande è questa distanza e peggio funziona il libro finale.

---

La Soluzione: FedNSAM (Il "Corriere Nesterov")

Per risolvere questo, gli autori hanno creato un nuovo algoritmo chiamato FedNSAM.

L'analogia del "Corriere con la Mappa Anticipata":
Immagina che ogni studente debba inviare un messaggio al professore.

• Metodo vecchio (FedSAM): Ogni studente guarda solo il suo tavolo e decide dove andare. Non sa cosa fanno gli altri. Risultato: caos.
• Metodo nuovo (FedNSAM): Il professore invia a ogni studente un "Corriere Speciale" (chiamato Momentum Nesterov).
  • Questo corriere non porta solo un messaggio, ma porta una mappa aggiornata di dove sta andando il gruppo nel suo insieme.
  • Prima che lo studente faccia un passo, il corriere gli dice: "Ehi, guarda dove stiamo andando noi tutti insieme! Allinea il tuo passo con quello del gruppo".
  • In pratica, lo studente non guarda solo il suo tavolo, ma "anticipa" il movimento globale.

Cosa succede magicamente?
Grazie a questo corriere, quando ogni studente cerca la sua "strada piatta", la cerca nella direzione giusta per il gruppo. Le strade piatte di tutti gli studenti iniziano a sovrapporsi e ad allinearsi.
Quando il professore unisce i pezzi, invece di trovare una montagna, trova un grande, liscio piano di atterraggio dove il modello globale può atterrare in sicurezza.

---

Perché è importante? (I Risultati)

Gli autori hanno fatto degli esperimenti (come testare questo metodo su immagini di gatti, cani e auto) e hanno scoperto che:

1. Funziona anche quando gli studenti sono molto diversi: Anche se i dati sono molto disordinati (ad esempio, uno studente ha solo foto di gatti neri e un altro solo di cani bianchi), FedNSAM riesce a trovare un accordo.
2. È più veloce: Il modello impara in meno tempo. È come se il corriere facesse risparmiare agli studenti ore di studio inutile.
3. Funziona con i "Giganti": Hanno provato questo metodo anche con modelli di intelligenza artificiale molto grandi (come quelli che usano per vedere le immagini), e ha funzionato meglio di tutti i metodi precedenti.

In Sintesi

Il paper dice: "Non basta che ognuno di noi sia bravo da solo. Se vogliamo costruire un'intelligenza artificiale globale che funzioni bene, dobbiamo assicurarci che i nostri 'piani di studio' siano allineati."

FedNSAM è il nuovo metodo che usa un "corriere intelligente" (il momentum di Nesterov) per assicurarsi che, anche se ognuno studia da casa propria, tutti stiano camminando nella stessa direzione, trovando così un terreno solido e piatto dove tutti possono stare bene.

Sommario tecnico

1. Il Problema: Incoerenza della Piattezza nell'Apprendimento Federato

L'apprendimento federato (FL) mira ad addestrare modelli globali mantenendo i dati locali sui client per preservare la privacy. Tuttavia, due fattori critici degradano le prestazioni del modello globale:

1. Eterogeneità dei dati (Non-IID): I dati distribuiti tra i client non sono indipendenti e identicamente distribuiti.
2. Aggiornamenti locali multi-step: I client eseguono più iterazioni di aggiornamento locale prima di comunicare con il server.

Questi fattori tendono a far convergere il modello globale verso minimi acuti (sharp minima), che generalizzano male. Per risolvere ciò, algoritmi popolari come FedSAM (Sharpness-Aware Minimization) sono stati adottati per cercare minimi piatti (flat minima) durante l'addestramento locale.

La scoperta critica dell'articolo: Gli autori dimostrano che in scenari ad alta eterogeneità, la ricerca di minimi piatti a livello locale non garantisce minimi piatti a livello globale.

• In contesti a bassa eterogeneità, le regioni piatte dei vari client si sovrappongono, permettendo al modello globale di trovarsi in una zona piatta.
• In contesti ad alta eterogeneità, le regioni piatte dei client si allontanano e diventano disgiunte. Di conseguenza, anche se ogni client trova un minimo piatto localmente, la media aggregata (modello globale) finisce in una regione acuta, degradando la generalizzazione.

2. Metodologia e Proposta: FedNSAM

Per affrontare questo problema, gli autori introducono il concetto di "Flatness Distance" (Distanza di Piattezza) e propongono un nuovo algoritmo chiamato FedNSAM.

A. Concetto Chiave: Flatness Distance ($\Delta_{\mathcal{D}}$)

Gli autori definiscono formalmente la Flatness Distance come la discrepanza tra le regioni piatte dei modelli locali e quella del modello globale.
$$\Delta_{\mathcal{D}} = \frac{1}{N} \sum_{i=1}^{N} \mathbb{E} \| \theta_{i, K}^t - \theta^{t+1} \|^2$$
Dove $\theta_{i, K}^t$ è il modello locale del client $i$ dopo $K$ iterazioni e $\theta^{t+1}$ è il modello globale aggregato.

• Una distanza elevata indica che le regioni piatte dei client sono disallineate, portando a minimi globali acuti.
• L'obiettivo è minimizzare questa distanza per allineare le regioni piatte.

B. L'Algoritmo FedNSAM (Federated Nesterov Sharpness-Aware Minimization)

FedNSAM risolve il problema di allineamento integrando la momentum di Nesterov globale all'interno degli aggiornamenti locali. A differenza di FedSAM che usa solo il gradiente locale per la perturbazione, FedNSAM utilizza una direzione di perturbazione globale stimata.

Meccanismo di funzionamento:

1. Stima del Gradiente Globale: Poiché i gradienti globali non sono accessibili direttamente durante l'addestramento locale, FedNSAM stima la direzione globale utilizzando un momento esponenziale mobile ($m_t$) basato sulle differenze dei modelli inviati dai client.
   $$m_t = \lambda m_{t-1} + \Delta_t$$
   dove $\Delta_t$ rappresenta la media degli aggiornamenti dei client.
2. Allineamento della Perturbazione: Invece di perturbare il modello locale nella direzione del gradiente locale (come in FedSAM), FedNSAM sposta il modello nella direzione del momento globale $m_t$. Questo forza ogni client a cercare un minimo piatto che sia coerente con la direzione di aggiornamento globale.
3. Estrapolazione di Nesterov: L'algoritmo applica un passo di estrapolazione di Nesterov prima del calcolo della perturbazione, accelerando la convergenza e migliorando la stabilità.

Pseudo-codice semplificato del passo locale:

• Calcola il punto di estrapolazione: $\theta_{i, k+1/4} = \theta_{i, k} + \lambda m_t$.
• Calcola la perturbazione globale: $\delta_{i, k} = \rho \frac{-m_t}{\|m_t\|}$.
• Aggiorna il modello perturbato e calcola il gradiente per l'aggiornamento finale.

3. Contributi Principali

1. Identificazione del Limite di FedSAM: Dimostrazione teorica ed empirica che la minimizzazione della sharpness locale non è sufficiente per garantire la piattezza globale in presenza di alta eterogeneità dei dati.
2. Definizione della Flatness Distance: Introduzione di una metrica quantitativa per misurare la divergenza delle regioni di ottimizzazione tra client e server.
3. Proposta di FedNSAM: Sviluppo di un nuovo algoritmo che allinea le regioni piatte locali tramite momentum di Nesterov globale, senza richiedere comunicazioni aggiuntive di gradienti completi.
4. Analisi Teorica:
   • Dimostrazione di un limite di convergenza più stretto rispetto a FedSAM: $O(\frac{\sqrt{LF}}{\sqrt{TKS}(1-\lambda)})$.
   • Prova che FedNSAM riduce il limite superiore della Flatness Distance rispetto a FedSAM, specialmente quando l'eterogeneità ($\sigma_g^2$) è alta.
5. Risultati Sperimentali: Validazione su modelli CNN (ResNet, VGG) e Transformer (ViT, Swin) su dataset come CIFAR-10, CIFAR-100 e Tiny ImageNet.

4. Risultati Sperimentali

Gli esperimenti confermano la superiorità di FedNSAM rispetto allo stato dell'arte (FedAvg, SCAFFOLD, FedSAM, MoFedSAM, FedGAMMA, FedLESAM):

• Accuratezza Superiore: Su CIFAR-100 con alta eterogeneità (Dirichlet-0.1), FedNSAM raggiunge un'accuratezza del 58.53%, superando significativamente FedSAM (40.18%) e FedLESAM (48.74%).
• Efficienza di Convergenza: FedNSAM converge più velocemente. Ad esempio, su ResNet-18 con CIFAR-100, raggiunge il 55% di accuratezza in soli 316 round, mentre FedSAM ne richiede oltre 1000 e FedSAM impiega circa 3 volte più tempo per raggiungere prestazioni simili.
• Robustezza all'Eterogeneità: La performance di FedNSAM degrada molto meno rispetto agli altri metodi man mano che l'eterogeneità dei dati aumenta (da Dirichlet-0.6 a 0.1).
• Scalabilità su Transformer: L'algoritmo dimostra efficacia anche su modelli grandi come ViT-Base e Swin-Base, ottenendo la massima accuratezza con il minor numero di round di comunicazione.
• Analisi del Landscape: Le visualizzazioni delle superfici di perdita (loss landscapes) mostrano che FedNSAM riesce a trovare minimi globali molto più piatti rispetto a FedSAM, che tende a rimanere in zone acute in scenari Non-IID.

5. Significato e Impatto

Questo lavoro è significativo perché:

• Ridefinisce la strategia SAM nel FL: Sposta il focus dalla semplice ricerca di piattezza locale all'allineamento globale delle regioni di ottimizzazione.
• Risolve un paradosso: Spiega perché tecniche di regolarizzazione locali falliscono in FL e offre una soluzione pratica basata sul momentum globale.
• Efficienza Computazionale: A differenza di altre soluzioni che richiedono calcoli aggiuntivi pesanti o comunicazioni extra, FedNSAM utilizza il momentum già calcolato per la sincronizzazione, rendendolo efficiente in termini di comunicazione e calcolo.
• Applicabilità: La metodologia è valida sia per reti neurali convolutive tradizionali che per architetture Transformer moderne, rendendola rilevante per le applicazioni FL di prossima generazione.

In sintesi, FedNSAM rappresenta un avanzamento fondamentale nell'ottimizzazione federata, garantendo che la ricerca di minimi piatti sia coerente tra il livello locale e globale, migliorando così la generalizzazione del modello in scenari reali e complessi.