FedNSAM:Consistency of Local and Global Flatness for Federated Learning

Dit paper introduceert FedNSAM, een nieuw algoritme voor federated learning dat de generalisatie van het globale model verbetert door Nesterov-momentum te gebruiken om de inconsistentie tussen lokale en globale platheid in heterogene data-omgevingen op te lossen.

De kern

De Probleemstelling: Een Verwarde Groepsreis

Stel je voor dat een groep vrienden (de klanten in Federated Learning) samen een grote puzzel probeert op te lossen. Ze willen één gezamenlijk antwoord vinden, maar ze mogen hun eigen stukjes van de puzzel (hun data) niet aan elkaar laten zien om privacyredenen. Ze sturen alleen hun suggesties naar een centrale leider (de server).

Het probleem is dat elke vriend een heel ander stukje van de puzzel heeft.

• Vriend A heeft alleen rode stukjes.
• Vriend B heeft alleen blauwe stukjes.
• Vriend C heeft alleen groene stukjes.

In het verleden probeerden ze dit op te lossen door te zeggen: "Laten we allemaal zoeken naar de platste plek in de kamer waar we staan." Als je op een platte plek staat, ben je stabiel en val je niet snel om (in technische termen: je komt in een vlakke minimum in plaats van een scherpe piek, wat zorgt voor een beter, algemener model).

Maar hier zit de valkuil:
Omdat iedereen een ander stukje puzzel heeft, vinden ze allemaal een andere "platte plek".

• Vriend A vindt een vlakke plek in de hoek met de rode muur.
• Vriend B vindt een vlakke plek bij het blauwe raam.

Als de server nu al hun suggesties middelt, belandt het gezamenlijke antwoord ergens in het midden, precies op een scherpe piek tussen die twee vlakke plekken. Het resultaat is onstabiel en werkt slecht. Dit noemen de auteurs het probleem van de "Flatness Distance" (het afstandje tussen de vlakke plekken van de verschillende vrienden).

De Oplossing: FedNSAM (De Nieuwe Navigatie)

De auteurs van dit paper, Junkang Liu en zijn team, hebben een nieuwe methode bedacht genaamd FedNSAM. Ze gebruiken een slimme truc om ervoor te zorgen dat iedereen naar dezelfde vlakke plek loopt, zelfs als ze verschillende puzzelstukjes hebben.

Ze gebruiken een concept uit de natuurkunde genaamd Nesterov Momentum.

De Metafoor: De Skiër met een Voorspelling

Stel je voor dat je een skiër bent die een berg afdaalt (het zoeken naar de beste oplossing).

1. De oude methode (FedSAM): De skiër kijkt alleen recht vooruit naar de helling onder zijn voeten en stapt dan. Als de helling plotseling verandert (door de andere skiërs), kan hij uitglijden.
2. De nieuwe methode (FedNSAM): De skiër kijkt niet alleen naar waar hij nu staat, maar kijkt ook een stukje vooruit in de richting waar hij naar toe beweegt. Hij gebruikt zijn eigen snelheid en de gemiddelde snelheid van de hele groep om te voorspellen waar de berg over een paar seconden zal zijn.

In de praktijk doet FedNSAM het volgende:

• De server stuurt niet alleen de huidige positie naar de vrienden, maar ook een "voorspelling" (de Nesterov momentum) over waar de groep waarschijnlijk naartoe gaat.
• De vrienden gebruiken deze voorspelling om hun eigen zoektocht naar een "platte plek" te corrigeren. Ze kijken niet alleen naar hun eigen lokale omgeving, maar proberen de globale vlakke plek te vinden die voor iedereen geldt.

Het is alsof de vrienden niet meer blindelings naar hun eigen muur kijken, maar een kompas gebruiken dat aangeeft waar de werkelijke top van de berg ligt, zodat ze allemaal in dezelfde richting lopen en samen in één grote, veilige, vlakke vallei terechtkomen.

Waarom is dit belangrijk?

1. Beter Generaliseren: Omdat ze nu in een echte, gezamenlijke vlakke vallei zitten, werkt het eindresultaat veel beter op nieuwe, onbekende situaties (bijvoorbeeld een medisch model dat niet alleen op de data van ziekenhuis A werkt, maar ook op die van ziekenhuis B).
2. Sneller: De skiërs glijden niet meer heen en weer tussen verschillende plekken. Ze komen sneller aan bij het doel.
3. Robuust: Het werkt zelfs als de data heel erg verschillend is (bijvoorbeeld als 90% van de vrienden een heel ander soort puzzelstuk heeft dan de rest).

Samenvatting in één zin

FedNSAM is als het geven van een slimme GPS aan een groep wandelaars met verschillende kaarten; in plaats van dat ze allemaal in hun eigen vallei blijven hangen, gebruikt de GPS hun gezamenlijke beweging om hen allemaal naar één grote, veilige en stabiele vallei te leiden, waar ze samen de beste oplossing vinden.

Technische samenvatting

Titel: Consistentie van Lokale en Globale Vlakheid voor Federated Learning

Auteurs: Junkang Liu, Fanhua Shang, Yuxuan Tian, Hongying Liu, en Yuanyuan Liu.
Conferentie: MM '25 (Dublin, Ierland, 2025).

---

1. Het Probleem

Federated Learning (FL) staat bekend om het trainen van modellen op gedecentraliseerde data zonder dat de ruwe data het apparaat verlaat. Een veelvoorkomend probleem in FL is data-heterogeniteit (non-IID data) en het gebruik van meerdere lokale updates per communicatieronde.

• Scherpe Minima: Deze factoren leiden vaak tot het convergeren van het globale model naar "scherpe" minima (sharp minima) in de verliesfunctie. Modellen die in scherpe minima zitten, generaliseren over het algemeen slechter dan modellen in "vlakke" minima (flat minima).
• Beperkingen van Bestaande Oplossingen (FedSAM): Om dit op te lossen, wordt vaak Sharpness-Aware Minimization (SAM) toegepast in lokale training (bijv. FedSAM). SAM probeert lokale minima te vinden die vlak zijn.
• De Kern van het Probleem: De auteurs stellen vast dat in hoge heterogene settings, het vinden van een vlak lokaal minimum op de client-data niet garandeert dat het geaggregeerde globale model ook in een vlak minimum zit. De "vlakke gebieden" van verschillende clients kunnen zo ver uit elkaar liggen dat hun gemiddelde (het globale model) in een scherp gebied terechtkomt. Bestaande methoden zoals FedSAM falen hierdoor bij hoge heterogeniteit.

2. Methodologie: FedNSAM

Om dit probleem aan te pakken, introduceren de auteurs het concept van Flatness Distance en ontwikkelen ze een nieuw algoritme genaamd FedNSAM (Federated Nesterov Sharpness-Aware Minimization).

A. Het Concept: Flatness Distance

De auteurs definiëren de Flatness Distance ($\Delta_{\mathcal{D}}$) als de kwantitatieve maat voor de discrepantie tussen de vlakke regio's van lokale modellen en het globale model:
$$\Delta_{\mathcal{D}} = \frac{1}{N} \sum_{i=1}^{N} \mathbb{E} \| \theta_{i, K}^t - \theta^{t+1} \|^2$$
Waarbij $\theta_{i, K}^t$ het lokale model is na $K$ iteraties en $\theta^{t+1}$ het geaggregeerde globale model.

• Observatie: Hoe groter de data-heterogeniteit, hoe groter de Flatness Distance. Een grote afstand betekent dat de lokale vlakke regio's niet overlappen, waardoor het globale model niet in een vlak gebied kan vallen.

B. Het FedNSAM Algoritme

FedNSAM lost dit op door Nesterov Momentum te integreren om de consistentie tussen lokale en globale vlakheid te harmoniseren.

1. Globale Momentum Schatting: Omdat de globale gradiënt $\nabla F(\theta)$ niet direct beschikbaar is tijdens lokale training, schat FedNSAM deze met een globale Nesterov momentum ($m_t$). Dit momentum wordt berekend als een exponentiële moving average van de modelupdates ($\Delta_t$) van de clients.
   $$m_t = \lambda m_{t-1} + \Delta_t$$
2. Lokale Updates met Extrapolatie: In plaats van alleen op de lokale gradiënt te vertrouwen, gebruiken clients de globale momentum $m_t$ op twee manieren:
   • Richting van Extrapolatie: De lokale update start met een Nesterov-extrapolatie stap in de richting van $m_t$ ($\theta_{i, k+1/4} = \theta_{i, k} + \lambda m_t$).
   • Richting van Perturbatie: De perturbatie voor de SAM-stap (het zoeken naar de scherpheid) wordt niet bepaald door de lokale gradiënt, maar door de genormaliseerde globale momentum: $\delta_{i, k} = \rho \frac{-m_t}{\|m_t\|}$.
3. Doel: Door de lokale perturbaties en updates te aligneren met de globale dynamiek, worden de vlakke regio's van de clients dichter bij elkaar getrokken, zodat het globale model binnen deze gezamenlijke vlakke regio kan convergeren.

3. Belangrijkste Bijdragen

1. Concept van Flatness Distance: De auteurs introduceren een nieuwe theoretische maatstaf om de inconsistentie van lokale minima onder data-heterogeniteit te kwantificeren en tonen aan dat dit direct correleert met een verslechtering van de generalisatie.
2. FedNSAM Algoritme: Een nieuw FL-algoritme dat Nesterov-extrapolatie gebruikt om lokale vlakke regio's te aligneren. Het gebruikt globale momentum om de perturbatierichting in SAM te sturen.
3. Theoretische Bewijzen:
   • Ze bewijzen een strakkere convergentiebound voor FedNSAM dan voor FedSAM: $\mathcal{O}(\frac{\sqrt{LF}}{\sqrt{TKS}(1-\lambda)})$.
   • Ze tonen theoretisch aan dat de bovengrens van de Flatness Distance voor FedNSAM lager is dan die van FedSAM, wat betekent dat FedNSAM effectiever is in het handhaven van globale vlakheid.
4. Uitgebreide Experimenten: Validatie op CNN-modellen (LeNet, VGG, ResNet) en Transformer-modellen (ViT, Swin) op datasets zoals CIFAR-10, CIFAR-100 en Tiny ImageNet.

4. Resultaten

De experimentele resultaten tonen aan dat FedNSAM superieur presteert ten opzichte van bestaande methoden (FedAvg, FedSAM, MoFedSAM, FedGAMMA, FedLESAM, etc.):

• Generalisatie: FedNSAM bereikt aanzienlijk hogere testnauwkeurigheid, vooral bij hoge data-heterogeniteit (bijv. Dirichlet-0.1).
  • Voorbeeld: Op CIFAR-100 met ResNet-18 en Dirichlet-0.1, behaalde FedNSAM 58.53% nauwkeurigheid, terwijl FedSAM slechts 40.18% haalde.
• Convergentiesnelheid: FedNSAM convergeert sneller. Het bereikt dezelfde nauwkeurigheid als FedSAM in veel minder communicatierondes (bijv. 316 rondes vs. 1000+ rondes voor ResNet-18 op CIFAR-100).
• Robuustheid: Het algoritme presteert stabiel bij lage participatiepercentages (tot 2%) en bij verschillende modelarchitecturen, inclusief grote Vision Transformers.
• Loss Landscapes: Visualisaties van de loss-oppervlakken tonen aan dat FedNSAM een veel vlakker globaal verlieslandschap bereikt dan FedSAM, wat de superioriteit in generalisatie verklaart.

5. Betekenis en Impact

Dit paper is significant omdat het een fundamenteel inzicht biedt in waarom bestaande SAM-varianten in Federated Learning falen bij heterogene data: het gebrek aan consistentie tussen lokale en globale vlakheid.

• Theoretisch: Het biedt een nieuwe perspectief (Flatness Distance) om het generalisatievermogen in FL te analyseren en bewijst dat het aligneren van lokale dynamiek via momentum essentieel is.
• Praktisch: FedNSAM biedt een efficiënte oplossing die geen extra communicatiekosten vereist (behalve het delen van momentum, wat al standaard is in veel FL-methoden), maar wel aanzienlijk betere prestaties levert. Dit maakt het zeer relevant voor real-world toepassingen zoals gezondheidszorg en persoonlijke diensten waar data sterk gedecentraliseerd en heterogeen is.

Kortom, FedNSAM lost het "vlakheids-probleem" in Federated Learning op door lokale updates te sturen met globale momentum, waardoor het model sneller convergeert naar een robuust, goed generaliserend optimum.