Technische Samenvatting: Federated Learning met Onzekerheid en Personalisatie via Efficiënte Tweede-orde Optimalisatie

Probleemstelling

Federated Learning (FL) maakt collaboratieve modeltraining mogelijk over gedecentraliseerde cliënten zonder ruwe gegevens te delen. De standaard FL-benaderingen worden echter geconfronteerd met twee primaire uitdagingen:

1. Dataheterogeniteit: Wanneer de data van cliënten niet-IID (niet-onafhankelijk en niet identiek verdeeld) is, faalt een enkel globaal model vaak in het effectief representeren van individuele cliëntdistributies, wat leidt tot suboptimale prestaties.
2. Gebrek aan Onzekerheidskwantificatie: Standaard FL-methoden leren doorgaans punt-schattingen van modelparameters. Dit maakt ze vatbaar voor overfitting in scenario's met weinig data en niet in staat om voorspellende onzekerheid te bieden, wat cruciaal is voor taken zoals Out-of-Distribution (OOD) detectie.

Hoewel Bayesian FL een gefundeerde oplossing biedt door posterieure distributies te leren om onzekerheid te vangen en personalisatie mogelijk te maken (via hiërarchische frameworks), lijden bestaande methoden onder aanzienlijke knelpunten. Huidige Bayesian FL-benaderingen vertrouwen vaak op computationeel dure routines (bijv. MCMC-sampling, Laplace-benadering die expliciete Hessian-berekening vereist, of deep ensembles) en veroorzaken hoge communicatiekosten door volledige posterieure distributies te verzenden. Deze factoren beperken de schaalbaarheid, met name voor hulpbronnenbeperkte cliënten.

Methodologie: FedIvon

De auteurs stellen FedIvon voor, een nieuw Bayesian FL-algoritme dat gebruikmaakt van Improved Variational Online Newton (IVON) om efficiënte tweede-orde variational inference uit te voeren. De methode balanceert de voordelen van Bayesian inference met de computationele en communicatieve efficiëntie die vergelijkbaar is met eerste-orde optimizers zoals Adam.

Kerncomponenten

1. Client-zijde Variational Inference (IVON):

   • In plaats van een punt-schatting te leren, benadert elke cliënt zijn lokale posterieure distributie $q_k(\theta)$ met een Gaussische verdeling met een diagonale covariantie-matrix: $N(\theta | m_k, \text{diag}(\sigma_k^2))$.
   • Tweede-orde Efficiëntie: In tegenstelling tot Adam, dat leersnelheden aanpast op basis van gekwadrateerde gradiënten, gebruikt IVON een herparameterisering om de diagonale Hessian ($h$) impliciet te schatten zonder expliciete berekening. De variantie wordt afgeleid als $\sigma^2 = 1/(h + \delta)$.
   • Optimalisatie: Het algoritme werkt zowel de gemiddelde vector ($m$) als de varianties ($\sigma^2$) bij met behulp van natuurlijke gradiënten. Dit levert op natuurlijke wijze onzekerheidsschattingen op en vangt de geometrie van de verliesfunctie om snellere convergentie te bewerkstelligen, terwijl de computationele kosten vergelijkbaar blijven met die van Adam.

2. Server-zijde Posterieure Aggregatie:

   • De server aggregeert cliënt-posterieuren om een globale posterieure te berekenen. Omdat het product van Gaussische verdelingen ook Gaussisch is, wordt de globale posterieure benaderd als een Gaussische verdeling $N(\theta | \mu, \Lambda^{-1})$.
   • Aggregatiestrategie: De globale gemiddelde ($\mu$) en precisie-matrix ($\Lambda$) worden berekend als gewogen sommen van lokale gemiddelden en precisies, waarbij de gewichten evenredig zijn aan de datasetomvang van de cliënt ($N_k$). Deze strategie is computationeel efficiënt omdat er alleen de gemiddelde waarde en diagonale variantie (of Hessian) verzonden hoeven te worden in plaats van volledige modelgewichten of complexe distributies.

3. Personalisatie via Hiërarchisch Bayesian Framework:

   • FedIvon ondersteunt van nature gepersonaliseerde FL. De globale posterieure van de server dient als een prior ($p_k(\theta)$) voor de lokale optimalisatie van de cliënt.
   • De cliënt minimaliseert een gepersonaliseerde verliesfunctie: $L_k(q) = \mathbb{E}_{q_k}[\bar{\ell}_k(\theta)] + \beta D_{KL}(q_k(\theta) \| p_k(\theta))$.
   • De hyperparameter $\beta$ regelt de afweging tussen lokale adaptatie en globale kennis. Wanneer $\beta=0$, is het model volledig gepersonaliseerd; hogere waarden integreren meer globale informatie.

Belangrijkste Bijdragen

• Efficiënte Tweede-orde Bayesian FL: Introductie van FedIvon, dat IVON gebruikt om variational inference uit te voeren met computationele kosten die vergelijkbaar zijn met eerste-orde methoden (zoals Adam), maar die tweede-orde informatie (Hessian-schattingen) impliciet levert.
• Verbeterde Prestaties en Onzekerheid: Demonstratie dat FedIvon superieure voorspellende nauwkeurigheid en betere onzekerheidsschattingen bereikt vergeleken met state-of-the-art (SOTA) Bayesian en niet-Bayesian FL-methoden.
• Natuurlijke Personalisatie: Een framework dat van nature client-niveau modelpersonalisatie mogelijk maakt door de posterieure van de server als prior te gebruiken, waardoor een balans wordt gevonden tussen lokale data-eigenschappen en gedeelde globale kennis.

Experimentele Resultaten

De auteurs evalueerden FedIvon op drie datasets (EMNIST, SVHN, CIFAR-10) onder non-IID instellingen met 200 cliënten, elk met beperkte data (<100 voorbeelden).

Standaard FL Resultaten

• Nauwkeurigheid & Kalibratie: FedIvon presteerde beter dan de baselines (FedAvg, FedLaplace) op het gebied van testnauwkeurigheid, Expected Calibration Error (ECE), Negative Log Likelihood (NLL) en Brier Score over alle datasets.
• OOD Detectie: In Out-of-Distribution detectietaken behaalde FedIvon competitieve of superieure AUROC-scores vergeleken met de baselines, wat de bekwaamheid van de methode om voorspellende onzekerheid effectief te kwantificeren valideert.
• Convergentie: Hoewel FedIvon een iets tragere initiële convergentie vertoonde door de stochastische sampling van gewichten voor gradiëntschatting, overtrof het andere methoden naarmate de training vorderde dankzij een verbeterde Hessian-benadering.

Gepersonaliseerde FL Resultaten

Experimenten op CIFAR-10 (class distribution skew) en CIFAR-100 (class concept drift) vergeleken FedIvon met pFedME, pFedBayes en pFedVEM.

• CIFAR-10: FedIvon bereikte een client-nauwkeurigheid die vergelijkbaar is met pFedVEM, maar demonstreerde significante verbeteringen in de globale server-nauwkeurigheid, wat wijst op een betere aggregatie van heterogene data.
• CIFAR-100: FedIvon liet significante verbeteringen zien ten opzichte van alle baselines in zowel gepersonaliseerde client-nauwkeurigheid als globale server-nauwkeurigheid, wat de robuustheid in het omgaan met concept drift benadrukt.

Betekenis en Claims

Het artikel beweert dat FedIvon de kloof overbrugt tussen de theoretische voordelen van Bayesian FL (onzekerheidskwantificatie, personalisatie) en de praktische beperkingen van federated systemen (computatie, communicatie). Door een efficiënte tweede-orde optimalisatiebenadering te gebruiken, vermijdt de methode de zware computationele overhead van MCMC of expliciete Hessian-berekeningen die bij eerdere werken aanwezig zijn.

De auteurs stellen dat FedIvon de eerste methode is die het volgende biedt:

1. Computationele efficiëntie die vergelijkbaar is met eerste-orde optimizers (zoals Adam).
2. Volledige Bayesian voordelen (onzekerheid en personalisatie) zonder de schaalbaarheidsproblemen van bestaande Bayesian FL-methoden.
3. Superieure prestaties in zowel standaard als gepersonaliseerde FL-settings, met name in scenario's met hoge dataheterogeniteit en beperkte hoeveelheden client-data.

Het werk positioneert FedIvon als een schaalbare, praktische oplossing voor het implementeren van robuuste, onzekerheidsbewuste en gepersonaliseerde machine learning-modellen in real-world federated omgevingen.