FedRef: Bayesian Fine-Tuning using a Reference Model to Mitigate Catastrophic Forgetting for Heterogeneous Federated Learning

Riassunto Tecnico: FedRef – Bayesian Fine-Tuning con Modello di Riferimento per il Federated Learning Eterogeneo

1. Il Problema: Catastrophic Forgetting e Unbounded Drift

Nel contesto del Federated Learning (FL), l'obiettivo è addestrare modelli collaborativi preservando la privacy dei dati locali. Tuttavia, due sfide principali emergono quando i dati sono eterogenei (non-IID):

• Catastrophic Forgetting (Dimenticanza Catastrofica): Quando un modello viene addestrato sequenzialmente su dati diversi, tende a perdere le capacità apprese in precedenza.
• Unbounded Drift (Deriva Illimitata): L'eterogeneità dei dati tra i client causa divergenze nelle aggiornamenti dei parametri, portando a una deriva del modello globale che degrada le prestazioni predittive e aumenta l'instabilità.

Le soluzioni esistenti (come FedProx o FedOpt) cercano di mitigare questi problemi introducendo termini di regolarizzazione, ma spesso lo fanno a spese del carico computazionale dei client, che devono eseguire calcoli aggiuntivi (termini proximal) durante l'addestramento locale.

2. Metodologia: L'approccio FedRef

Gli autori propongono FedRef, un metodo di fine-tuning bayesiano che sposta l'onere della regolarizzazione dal client al server.

I pilastri metodologici sono:

• Modello di Riferimento ($\theta_{ref}$): Invece di basarsi solo sul modello globale dell'ultimo round, FedRef costruisce un modello di riferimento che è una combinazione pesata dei modelli globali dei round precedenti. Questo modello cattura una regione di parametri temporalmente stabile, agendo come un "ancoraggio" per evitare derive eccessive.
• Stima MAP (Maximum A Posteriori): FedRef utilizza l'inferenza bayesiana per formulare il problema di ottimizzazione. Il server cerca di massimizzare la probabilità a posteriori, bilanciando la verosimiglianza dei dati dei client con una "prior" derivata dal modello di riferimento.
• Funzione di Costo Regolarizzata: La funzione di costo ottimizzata dal server ($L_{ref}$) integra tre componenti:
  1. La perdita media dei client ($F_k$).
  2. Una regolarizzazione $L_2$ tra il modello globale attuale e quello del round precedente ($\theta_g^r - \theta_g^{r-1}$).
  3. Una regolarizzazione tra il modello globale e il modello di riferimento ($\theta_{ref}$), che mitiga la dimenticanza catastrofica.
• Esecuzione Server-Side: A differenza di FedProx, i client eseguono solo l'addestramento locale standard. Tutti i calcoli complessi legati alla regolarizzazione bayesiana e alla costruzione del modello di riferimento vengono eseguiti esclusivamente sul server.

3. Contributi Chiave

1. Efficienza Computazionale: Riduce il carico sui client (dispositivi edge/IoT) eliminando i termini di regolarizzazione locali, mantenendo lo stesso costo computazionale di FedAvg.
2. Stabilità Temporale: L'uso di un modello di riferimento pesato temporalmente previene l'instabilità causata da modelli in fase di sotto-apprendimento (underfitting) nei primi round.
3. Mitigazione della Deriva: Dimostra matematicamente che l'introduzione della prior bayesiana riduce la probabilità di "unbounded drift" (UDP), mantenendo gli aggiornamenti entro una regione di parametri sicura.

4. Risultati Sperimentali

Il framework è stato testato su tre dataset con diversi livelli di eterogeneità: CINIC-10 e FEMNIST (classificazione immagini) e FeTS2022 (segmentazione di immagini mediche MRI).

• Performance Predittive: FedRef ha superato FedAvg, FedProx e FedOpt in termini di accuratezza, F1-score e Dice Coefficient, specialmente in scenari ad alta eterogeneità (FeTS2022).
• Riduzione della Dimenticanza: La metrica $\zeta$ (che quantifica il bilancio tra miglioramento e degradazione delle prestazioni) è significativamente più bassa in FedRef, indicando una maggiore capacità di preservare la conoscenza acquisita.
• Velocità di Convergenza: FedRef raggiunge le soglie di performance target in un numero di round comparabile o inferiore rispetto ai metodi basati su regolarizzazione locale.
• Efficienza: I test confermano che il costo computazionale lato client è identico a FedAvg, a differenza di FedProx che richiede un overhead del 1-3%.

5. Significato e Conclusioni

FedRef rappresenta un avanzamento significativo per il Federated Learning in scenari del mondo reale (come la medicina o i dispositivi mobili), dove le risorse dei client sono limitate e i dati sono intrinsecamente non uniformi. Spostando l'intelligenza della regolarizzazione sul server tramite tecniche bayesiane, il metodo riesce a stabilizzare l'apprendimento globale senza penalizzare i dispositivi periferici, risolvendo efficacemente il trade-off tra robustezza del modello ed efficienza energetica.