Technische Samenvatting: Federated Learning via Variational Bayesian Inference

1. Probleemstelling

Federated Learning (FL) maakt gedistribueerd machine learning mogelijk terwijl de privacy van cliënten wordt gewaarborgd. De standaard FL-benaderingen, zoals Federated Averaging (FedAvg), lijden echter vaak aan prestatievermindering wanneer ze worden geconfronteerd met twee kritieke uitdagingen in de echte wereld:

1. Dataheterogeniteit (Non-i.i.d.): De datadistributies van cliënten variëren aanzienlijk door verschillen in gebruikersgewoonten, locaties en voorkeuren.
2. Dataschaarste: Cliënten beschikken vaak over beperkte hoeveelheden data vanwege de hoge kosten van verzameling, wat leidt tot overfitting van het model.

Bestaande oplossingen proberen deze problemen aan te pakken via Personalized FL (PFL) en Clustered FL (CFL). PFL-modellen overfitten echter vaak op onvoldoende data, en standaard Bayesian FL (BFL)-modellen hebben moeite met extreme heterogeniteit. Bovendien is de inferentiekosten van Bayesian Neural Networks (BNNs) hoog vanwege de noodzaak van meerdere forward passes om onzekerheid te schatten. Dit papier adresseert de behoefte aan een framework dat gelijktijdig omgaat met heterogeniteit, dataschaarste en efficiëntie bij inferentie.

2. Methodologie

De auteurs stellen een tweelaags optimalisatieframework voor dat BNN's integreert in PFL en CFL via Variational Bayesian Inference (VI). In tegenstelling tot traditionele VI waarbij een vaste prior wordt verondersteld voor parameters, gebruikt dit framework de getrainde globale distributie van de server als de priorverdeling voor elke cliënt. Dit voorkomt een mismatch tussen veronderstelde priors en werkelijke distributies.

Het framework bestaat uit drie specifieke algoritmen:

2.1 pFedBayes (Personalized Bayesian FL)

• Mechanisme: Ontworpen voor scenario's met beperkte data. Het neemt aan dat alle netwerkparameters onafhankelijke en identiek verdeelde (i.i.d.) Gaussische distributies volgen.
• Optimalisatie: Cliënten minimaliseren de som van de reconstructiefout over hun lokale data en de Kullback-Leibler (KL) divergentie tussen hun lokale distributie en de globale distributie. De server aggregeert lokale distributies om de globale distributie bij te werken.
• Doel: Personalisatie bereiken terwijl de globale prior wordt benut om overfitting te voorkomen.

2.2 sFedBayes (Sparse Personalized Bayesian FL)

• Mechanisme: Ontworpen om de efficiëntie van inferentie te verbeteren. Het introduceert sparsity door aan te nemen dat parameters een Bernoulli-Gaussian distributie volgen (een "spike-and-slab" prior).
• Optimalisatie: Vergelijkbaar met pFedBayes maar bevat een leerbare sparsity-structuur. Parameters worden gepruned als de inclusiekans laag is.
• Doel: Opslag- en computationele kosten tijdens inferentie verminderen door redundante parameters te verwijderen, terwijl de modelprestaties behouden blijven.

2.3 cFedBayes (Clustered Bayesian FL)

• Mechanisme: Ontworpen voor extreme statistische diversiteit. In plaats van één enkele globale prior te leren, leert de server meerdere globale distributies (één per cluster).
• Optimalisatie: Cliënten selecteren dynamisch de globale distributie (cluster) die het beste past bij hun private data (minimaliserend van de KL-divergentie) en werken hun lokale distributie bij op basis van die specifie af specifieke prior.
• Doel: Cluster-niveau adaptieve garanties bieden in plaats van een enkele uniforme grens, wat beter aansluit bij cliënten met zeer verschillende datadistributies.

3. Belangrijkste Bijdragen

1. Nieuwe Algoritmen: Het artikel introduceert drie tweelaagse gepersonaliseerde BFL-algoritmen: pFedBayes, sFedBayes en cFedBayes.
2. Theoretische Garanties:
   • De auteurs bieden generalisatiefoutgrenzen voor alle drie de benaderingen.
   • Ze bewijzen dat de gemiddelde generalisatiefout-rates minimax optimaliteit bereiken tot een logaritmische factor.
   • Specifiek wordt aangetoond dat cFedBayes een cluster-niveau generalisatiefoutgrens bereikt, wat superieure adaptieve garanties biedt ten opzichte van de enkele uniforme grens van pFedBayes.
3. Efficiëntie en Sparsity: sFedBayes wordt voorgesteld om de hoge inferentiekosten van BNN's aan te pakken. Theoretische analyse toont aan dat het minimax optimale convergentierates bereikt in het sparse regime.
4. Stochastische Gradiënt Algoritmen: Bijbehorende stochastische gradiënt algoritmen zijn ontwikkeld voor alle drie de methoden, gebruikmakend van reparameterisatietechnieken (bijv. Gumbel-softmax voor discrete variabelen in sFedBayes) om end-to-end training mogelijk te maken.

4. Experimentele Resultaten

De auteurs evalueerden de voorgestelde methoden op vier benchmark datasets (MNIST, FMNIST, CIFAR-10, SVHN) onder non-i.i.d. en clustering instellingen, waarbij ze werden vergeleken met state-of-the-art baselines (bijv. FedAvg, FedProx, pFedMe, IFCA, FedCFL).

• Prestaties bij Beperkte Data: pFedBayes presteerde aanzienlijk beter dan andere gepersonaliseerde methoden, vooral op kleine datasets. Bijvoorbeeld, op de kleine CIFAR-10 dataset verbeterde pFedBayes de gepersonaliseerde nauwkeurigheid met 12,43% ten opzichte van het beste SOTA-algoritme.
• Voordelen van Sparsity: sFedBayes toonde aan dat matige sparsity (bijv. het behouden van ~35-45% van de parameters) de gepersonaliseerde nauwkeurigheid kan behouden of zelfs licht kan verbeteren vergeleken met het dense pFedBayes-model, wat fungeert als een impliciete regularisatie tegen overfitting.
• Effectiviteit van Clustering: cFedBayes behaalde de hoogste nauwkeurigheid op clustering datasets (bijv. MNIST-MNIST⊥, CIFAR10-CIFAR10⊥), waarbij het clustering-gebaseerde baselines zoals FedCFL en IFCA met ongeveer 1–2% in globale nauwkeurigheid versloeg.
• Onzekerheidsschatting: Het Bayesiaanse karakter van de modellen maakt onzekerheidskwantificatie mogelijk, wat volgens de auteurs gebruikt kan worden om modellen met een lage betrouwbaarheid tijdens aggregatie te filteren.

5. Betekenis en Claims

Het artikel claimt dat de voorgestelde benaderingen effectief de dubbele uitdagingen van heterogene en beperkte data in Federated Learning aanpakken. Door de globale distributie als een prior te gebruiken, vermijden de modellen de valkuilen van het veronderstellen van incompatibele priors, terwijl ze profiteren van de robuustheid van Bayesian inferentie.

De auteurs benadrukken dat hun theoretische analyse minimax optimaliteit vaststelt, wat een rigoureuze fundering biedt voor de geobserveerde empirische winsten. Ze stellen dat cFedBayes een significante vooruitgang biedt door adaptieve, cluster-niveau garanties te bieden in plaats van een "one-size-fits-all" grens. Tot slot positioneren ze sFedBayes als een praktische oplossing voor omgevingen met beperkte middelen door de inferentiekosten te verlagen zonder nauwkeurigheid op te offeren.

Het werk wordt gepresenteerd als een uitbreiding van een eerdere conferentieversie, met belangrijke verbeteringen waaronder de introductie van sparsity (sFedBayes), de ontwikkeling van een geclusterd framework met cluster-niveau grenzen (cFedBayes), en uitgebreide experimentele validatie tegen huidige SOTA-modellen.