Subspace Optimization for Efficient Federated Learning under Heterogeneous Data

1. Enunciato del Problema

Il paper affronta la "triplice sfida" che il Federated Learning (FL) deve affrontare nell'era dei modelli su larga scala (da centinaia di milioni a miliardi di parametri):

1. Computazione: I passaggi di addestramento locale sono computazionalmente pesanti.
2. Memoria: Archiviare stati completi del modello e variabili ausiliarie su dispositivi periferici (edge) con risorse limitate è proibitivo.
3. Comunicazione: Trasmettere aggiornamenti a dimensionalità completa su bande passanti limitate è inefficiente.

Queste sfide sono esacerbate dall'eterogeneità dei dati (Non-IID). In contesti Non-IID, gli obiettivi locali divergono dall'obiettivo globale, causando il "client drift".

• Soluzioni Esistenti e Limiti:
  • Correzione dell'Eterogeneità (es. SCAFFOLD): Utilizza variabili di controllo per correggere il drift ma richiede l'archiviazione e la comunicazione di stati ausiliari a dimensionalità completa, annullando i guadagni di efficienza per i modelli di grandi dimensioni.
  • Compressione (es. Quantizzazione): Riduce i bit trasmessi ma spesso richiede buffer di errore a dimensionalità completa, non riuscendo a ridurre l'impronta di memoria.
  • Metodi a Sottospazio/Basso Rango: Riducono la dimensionalità per l'efficienza ma tipicamente mancano di meccanismi principiali per correggere il client drift. Combinazioni naive spesso scartano le informazioni storiche di correzione quando il sottospazio cambia, portando a instabilità.

Il Divario: Non esiste alcun algoritmo esistente che esegua una correzione nativa dell'eterogeneità all'interno di un sottospazio a bassa dimensionalità preservando al contempo la robustezza dei metodi a dimensionalità completa come SCAFFOLD.

2. Metodologia: Subspace-SCAFFOLD (SSF)

Gli autori propongono SSF, un algoritmo che integra la correzione dell'eterogeneità direttamente in un sottospazio di ottimizzazione a bassa dimensionalità.

Meccanismi Principali

1. Proiezione nel Sottospazio:

   • Ad ogni round $t$, viene generata una matrice di proiezione ortonormale casuale condivisa $P_t \in \mathbb{R}^{r \times d}$ (dove $r \ll d$).
   • L'ottimizzazione (aggiornamenti locali e aggregazione) avviene strettamente nel sottospazio $r$-dimensionale.
   • I gradienti vengono proiettati: $g_{proj} = P_t g$.

2. Variabili di Controllo Proiettate (Correzione dell'Eterogeneità):

   • SSF imita il meccanismo SCAFFOLD ma opera su quantità proiettate.
   • I client mantengono variabili di controllo locali $c_i$ e il server mantiene una variabile di controllo globale $c$.
   • La direzione di aggiornamento locale viene corretta utilizzando la differenza tra le variabili di controllo locali e globali proiettate: $P_t g - c_{i, proj} + c_{proj}$.

3. Riempimento Conservativo del Residuo (Innovazione Chiave):

   • A differenza dei metodi a sottospazio standard che scartano le informazioni al di fuori del sottospazio attivo, SSF mantiene uno stato di controllo a dimensionalità completa ($c_i$ e $c$).
   • Regola di Aggiornamento: Quando il sottospazio cambia, le variabili di controllo vengono aggiornate come:
     $$c^{t+1} = (I - P_t^\top P_t)c^t + P_t^\top P_t (\text{nuove informazioni sul gradiente})$$
   • Ciò significa che il complemento ortogonale (residuo) della variabile di controllo viene preservato invariato, mentre viene aggiornato solo il componente del sottospazio attivo.
   • Riempimento (Backfill): Dopo l'aggiornamento del sottospazio, il modello a spazio completo viene ricostruito combinando le nuove coordinate proiettate con il residuo preservato: $x^{t+1} = P_t^\top x^{t+1}_{proj} + x^t_{res}$.

Flusso dell'Algoritmo

1. Decomposizione: Il server divide il modello e le variabili di controllo in parti proiettate e residue.
2. Passaggi Locali: I client ricostruiscono il modello completo utilizzando il modello a sottospazio trasmesso e il residuo del server. Calcolano i gradienti nello spazio completo, li proiettano ed eseguono aggiornamenti corretti per l'eterogeneità nel sottospazio.
3. Aggregazione: Il server aggrega gli endpoint proiettati nel sottospazio.
4. Riempimento e Aggiornamento: Il server solleva il modello di nuovo nello spazio completo e aggiorna le variabili di controllo, preservando i componenti residui.

3. Contributi Chiave

1. Progettazione dell'Algoritmo (SSF):

   • Il primo algoritmo FL a eseguire una correzione nativa dell'eterogeneità in un sottospazio a bassa dimensionalità.
   • Introduce un aggiornamento di tipo riempimento (backfill-style) che mantiene i componenti residui delle variabili di controllo, garantendo che le informazioni storiche di correzione del drift non vengano perse quando il sottospazio ruota o cambia.

2. Analisi Teorica:

   • Dimostra che, sotto standard assunzioni di regolarità e varianza limitata, SSF raggiunge un tasso di convergenza non asintotico di $\tilde{O}(1/T + 1/\sqrt{NKT})$.
   • Questo tasso corrisponde all'accelerazione lineare di SCAFFOLD a dimensionalità completa, dimostrando che i vincoli del sottospazio non degradano il comportamento di convergenza asintotico riguardo al client drift.
   • L'analisi vale senza assumere indipendenza tra il proiettore casuale e i gradienti stocastici.

3. Validazione Empirica:

   • Dimostrato sia su un problema giocattolo controllato di regressione matriciale sia su un benchmark di deep learning (CIFAR-100 con ResNet-110).
   • Ha mostrato che SSF riduce significativamente i costi di memoria, comunicazione e computazione mantenendo un'accuratezza comparabile a SCAFFOLD a dimensionalità completa.

4. Risultati Sperimentali

Problema Giocattolo (Regressione Matriciale)

• Robustezza: Sotto livelli variabili di eterogeneità (basso, medio, alto), SSF ha costantemente superato FedSub (un metodo a sottospazio di base senza conservazione del residuo) e FedAvg.
• Stabilità: FedSub ha mostrato instabilità numerica e divergenza (errori NaN) a dimensioni di sottospazio più elevate ($r=50$) a causa della perdita di informazioni ortogonali. SSF è rimasto stabile su tutte le dimensioni.
• Accuratezza: Con un rapporto di sottospazio di $r/d = 0.2$, SSF ha raggiunto errori relativi molto vicini a Full-SCAFFOLD, significativamente migliori di Full-FedAvg.

Deep Learning (CIFAR-100 / ResNet-110)

• Classifica: Full-SCAFFOLD > SSF > FedAvg > FedSub.
• Prestazioni: SSF ha raggiunto un'accuratezza di test finale del 45.43%, superando FedAvg (34.35%) e FedSub (23.17%) con un ampio margine.
• Efficienza: SSF ha ridotto con successo l'overhead di comunicazione e memoria preservando al contempo la robustezza di SCAFFOLD contro l'eterogeneità dei dati.

5. Significato e Impatto

• Colmare il Divario: SSF risolve la tensione fondamentale tra robustezza (gestione di dati Non-IID) ed efficienza (gestione di modelli di grandi dimensioni). Dimostra che non sono necessari stati ausiliari a dimensionalità completa per correggere efficacemente il client drift.
• Scalabilità: Riducendo l'impronta di comunicazione e memoria da $O(d)$ a $O(r)$ (dove $r \ll d$) mantenendo garanzie di convergenza teoriche, SSF rende fattibile il Federated Learning su modelli di grandi dimensioni su dispositivi periferici con risorse limitate.
• Insight Teorico: Il meccanismo "conservativo del residuo" fornisce un nuovo paradigma per l'ottimizzazione a sottospazio, mostrando che mantenere la memoria a dimensionalità completa per le variabili di controllo (anche se accessibile solo parzialmente) è cruciale per la stabilità in ambienti a sottospazio dinamico.

In conclusione, SSF offre una soluzione pratica e teoricamente solida per il dispiegamento di Federated Learning su larga scala ed eterogeneo, raggiungendo il "meglio di due mondi": l'efficienza dei metodi a sottospazio e la robustezza della correzione tramite variabili di controllo.