Class-Grouped Normalized Momentum and Faster Hyperparameter Exploration to Tackle Class Imbalance in Federated Learning

Technische Samenvatting: FedCGNM en FedHOO voor klasse-imbalans in Federated Learning

Probleemstelling
Klasse-imbalans blijft een kritieke uitdaging in Federated Learning (FL), met name wanneer de globale labelverdeling een 'long-tailed' karakter heeft. In dergelijke scenario's lijden minderheidsklassen onder een slechte voorspellende prestatie door onderrepresentatie in de trainingsdata. Standaard gecentraliseerde technieken, zoals loss reweighting, geavanceerde sampling of generatieve augmentatie, zijn moeilijk in te zetten in FL. Privacybeperkingen verhinderen gegevensuitwisseling, en synthetische generatoren slagen er vaak niet in om realistische monsters te produceren voor schaarse regimes of domeinspecifieke taken (bijv. defectdetectie in halfgeleiders). Bovereen bestaande optimalisatiegestuurde benaderingen zoals Per-Class Normalization (PCN) kampen met beperkingen in multi-class settings: ze zijn gevoelig voor directionele ruis van minderheidsklassen en introduceren schaalverschillen bij het sommeren van vele eenheidsvectoren, wat de convergentie destabiliseert naarmate het aantal klassen toeneemt.

Methodologie
Het paper stelt een tweeledige aanpak voor om deze problemen aan te pakken: een client-zijde optimizer genaamd FedCGNM (Federated Class-Grouped Normalized Momentum) en een hulp-hyperparameter tuner genaamd FedHOO.

1. FedCGNM (Optimizer):

   • Class Grouping: In plaats van gradiënten voor elke individuele klasse ($C$) te normaliseren, partitioneert FedCGNM klassen in een klein aantal groepen ($H \ll C$, doorgaans meerderheid versus minderheid). De groeperingsregel is afgeleid van het minimaliseren van de variantie binnen de groepen van klasseproporties, een strategie die is aangetoond een betrouwbare proxy te zijn voor een uitputtende threshold-zoektocht.
   • Group-Wise Normalized Momentum: Voor elke groep houdt het algoritme een momentumvector bij. Tijdens lokale training worden stochastic gradients per groep berekend en wordt het momentum bijgewerkt. Cruciaal is dat de momentumvector van elke groep wordt genormaliseerd naar een eenheidslengte (met een stabiliserende constante $\delta$) vóór de sommatie.
   • Update Richting: De uiteindelijke update-richting is de som van deze genormaliseerde groep-momentums. Dit ontwerp egaliseert de gradiëntgrootte tussen meerderheids- en minderheidsgroepen terwijl het de directionele ruis die inherent is aan zeldzame klassen-gradiënten vermindert. Door te opereren op een handvol groepen in plaats van individuele klassen, vermijdt het de schaalmismatch en instabiliteit die geassocieerd wordt met PCN.
   • Client Alignment: Door elke groep-update een eenheidsnorm op te leggen, worden de geaggregeerde update-richtingen over heterogene clients meer uitgelijnd, wat een primaire oorzaak is van prestatieverslechtering in FL.

2. FedHOO (Hyperparameter Tuner):

   • Context: Het selecteren van passende resampling-rates is niet triviaal in FL vanwege de combinatorische zoekruimte over clients en het risico op het destabiliseren van de optimalisatie met fluctuerende schema's.
   • Mechanisme: FedHOO is een X-armed bandit (XAB) gebaseerd algoritme dat ontworpen is voor kleinschalige federaties ($K$ is klein). Het maakt gebruik van de lineaire aard van FL-aggregatie. In elke communicatieronde instrueert de server elke client om te trainen met twee kandidaat-resampling rates.
   • Efficiëntie: Door de geretourneerde client-updates lineair te combineren, kan de server de validatieprestaties van alle $2^K$ mogelijke combinaties van lokale rates afleiden zonder dat hiervoor $2^K$ aparte trainingsrondes nodig zijn. Dit maakt het mogelijk om snel effectieve sampling rates te identificeren vroeg in de training, terwijl de privacy behouden blijft (clients retourneren enkel scalaire validatiemetrieken).

Theoretische Analyse
Het paper biedt een convergentieanalyse voor FedCGNM die expliciet rekening houdt met tijdvariërende resampling rates. Onder standaardveronderstellingen van gladheid (smoothness), begrensde gradiënten en onbevooroordeeldheid (unbiasedness), bewijzen de auteurs dat FedCGNM een standaard $O(T^{-1/2})$ stationair punt-rate bereikt. De analyse benadrukt dat voor convergentie te houden, de cumulatieve variatie van de sampling rates over rondes ($V_T$) aan specifieke voorwaarden moet voldoen, wat impliceert dat agressieve maar instabiele resampling de convergentie kan hinderen.

Belangrijkste Bijdragen

1. FedCGNM Optimizer: De eerste client-zijde optimizer die klassen groepeert en een unit-norm momentum per groep toepast. Het balanceert de invloed van meerderheid/minderheid, vermindert ruis en mitigeert schaalmismatch.
2. Variance-Aware Grouping: Een geprincipeerde, lichtgewicht strategie die klasse-partitionering frameert als de minimalisatie van variantie binnen groepen, wat optimale splitsingen oplevert die vergelijkbaar zijn met een uitputtende zoektocht.
3. Convergentie Analyse: Een nieuw theoretisch kader dat FL, groep-normalisatie en dynamische sampling rates combineert, en de impact van rate-schema's op convergentie aantoont.
4. FedHOO: De eerste optimalisatiegestuurde algoritme om combinatorische lokale sampling rates in FL te bepalen met behulp van een X-armed bandit framework dat federatieve parallellisme gebruikt voor efficiënte exploratie.

Experimentele Resultaten
De methoden werden geëvalueerd op vier publieke long-tailed benchmarks (CIFAR-10-LT, CIFAR-100-LT, Adult Income, UNSW-NB15) en een eigen dataset van halfgeleider-chipdefecten (CDD).

• Prestaties: FedCGNM presteerde consistent beter dan de baselines (inclusief FedAvg, FedProx, CLIMB, FedGraB en PCN-varianten) in alle settings. In kleine-client regimes ($K=5$) leverde de combinatie van FedCGNM met FedHOO de beste resultaten op, met verbeteringen tot 29% ten opzichte van FedAvg met resampling.
• Industriële Toepassing: Op de CDD-dataset (1,7% defect rate) bereikte de voorgestelde methode een verbetering van 16% ten opzichte van de beste baseline, wat de effectiviteit in real-world, privacy-gevoelige industriële settings demonstreert.
• Robuustheid: De aanpak toonde veerkracht bij grote-client regimes met gedeeltelijke participatie en onder variërende graden van data-heterogeniteit (Non-IID).
• Ablatie: Experimenten bevestigden dat twee groepen ($H=2$) de beste bias-variance trade-off bieden en dat gematigde momentumfactoren ($\beta \approx 0,5$) optimale prestaties opleveren.

Betekenis en Claims
Het paper claimt dat FedCGNM en FedHOO een praktisch framework vormen voor het verzachten van globale klasse-imbalans in privacy-bewarende federated settings zonder te vertrouwen op gegevensuitwisseling of synthetische generatie. De betekenis ligt in het adresseren van klasse-imbalans op het niveau van de optimalisatie, wat algemener toepasbaar is dan data-niveau trucs die beperkt worden door privacy of dataset-kenmerken. De auteurs benadrukken dat hun aanpak bijzonder effectief is in kleinschalige federaties (gangbaar in industriële implementaties) waar gezamenlijke hyperparameter tuning haalbaar en cruciaal is, terwijl de groeperingsstrategie zorgt voor stabiliteit en uitlijning in multi-class scenario's waar eerdere normalisatietechnieken falen. Het werk draagt bij aan het vakgebied door zowel een robuuste optimizer als een theoretisch gefundeerde, efficiënte methode te bieden voor het tunen van sampling-strategieën in de aanwezigheid van klasse-imbalans.