CA-HFP: Curvature-Aware Heterogeneous Federated Pruning with Model Reconstruction

Dit paper introduceert CA-HFP, een praktisch framework voor federated learning dat door middel van krommingsbewuste, apparaatspecifieke pruning en lichtgewicht reconstructie persoonlijke compressie mogelijk maakt op heterogene randapparaten zonder de aggregatiecompatibiliteit of convergentie te schaden.

De kern

Stel je voor dat een groep vrienden samen een groot, complex recept (een kunstmatig intelligentie-model) wil leren koken. Maar er is een probleem: ze wonen allemaal in verschillende huizen met verschillende keukens. Sommige huizen hebben een dure, professionele oven en een groot aanrecht (krachtige telefoons), terwijl anderen maar een kleine magnetron en een klein aanrecht hebben (zwakke apparaten). Bovendien heeft elke vriend een heel ander voorraadkastje met ingrediënten (verschillende data).

In de wereld van Federated Learning proberen deze vrienden samen te werken zonder hun eigen voorraadkastjes te verplaatsen naar één centrale keuken. Ze sturen alleen hun "kookadviezen" (updates) naar een centrale chef.

Het probleem is dat de vrienden met de kleine keukens de hele taak niet aankanen. Als ze allemaal hetzelfde grote recept moeten leren, gaan ze vastlopen of hun batterij leeglopen.

Hier komt CA-HFP (Curvature-Aware Heterogeneous Federated Pruning) om de hoek kijken. Het is als een slimme, nieuwe manier om samen te koken die rekening houdt met ieders beperkingen. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Iedereen krijgt een aangepast recept (Persoonlijke "Pruning")

In plaats van dat iedereen hetzelfde zware recept moet leren, mag elke vriend een versneden versie van het recept maken die past bij hun keuken.

• De vriend met de grote oven mag bijna het hele recept doen.
• De vriend met de kleine magnetron mag alleen de belangrijkste stappen doen en de rest overslaan.
• De slimme truc: De centrale chef kijkt niet alleen naar hoeveel stappen je overslaat, maar ook naar welke stappen het belangrijkst zijn. Hij gebruikt een soort "krachtmeter" (de kromming of curvature) om te zien welke ingrediënten echt nodig zijn voor de smaak en welke je kunt weglaten zonder dat het gerecht rot wordt.

2. Het probleem van de "ontbrekende stukjes"

Stel je voor dat Vriend A alleen de saus heeft gemaakt en Vriend B alleen het vlees. Als ze hun resultaten naar de centrale chef sturen, kan hij ze niet zomaar samenvoegen. Het zijn twee verschillende gerechten! In de oude methoden was dit een groot probleem: de chef wist niet hoe hij deze verschillende stukken weer tot één compleet gerecht moest maken.

3. De "Reconstructie": De chef vult de gaten in

Dit is de magische stap van CA-HFP. Voordat de chef de resultaten van iedereen samenvoegt, doet hij een reconstructie:

• Hij neemt het versnede recept van Vriend A en vult de ontbrekende stappen in met wat hij al weet van het grote recept.
• Hij doet hetzelfde voor Vriend B.
• Nu hebben ze allemaal weer een compleet, volledig recept (al zijn sommige stappen nog steeds "inactief" of minder belangrijk).
• Pas nu kan de chef ze eerlijk samenvoegen tot één nieuw, beter recept voor de volgende ronde.

Waarom is dit zo goed?

• Snelheid en Energie: De vrienden met kleine keukens hoeven niet meer te worstelen met zware taken. Ze werken sneller en besparen batterij.
• Beter Samenwerken: Omdat de chef de gaten slim invult, verliezen ze geen smaak (accuraatheid). Het eindresultaat is net zo lekker als als iedereen het hele recept had gedaan, maar dan veel sneller.
• Rechtvaardig: Zelfs als de vrienden heel verschillende ingrediënten hebben (ongelijke data), zorgt de slimme "krachtmeter" ervoor dat ze niet de verkeerde stappen weglaten.

Samengevat in één zin:

CA-HFP is als een slimme chef die ervoor zorgt dat elke kok in het team alleen de taken doet die bij zijn keuken passen, en die vervolgens op een slimme manier de ontbrekende stukjes invult voordat hij het gezamenlijke recept verbetert. Zo leren ze allemaal samen, snel en zonder dat de zwakke koks het opgeven.

Technische samenvatting

Probleemstelling

Federated Learning (FL) staat voor grote uitdagingen bij de implementatie op heterogene randapparaten (zoals mobiele telefoons en IoT-sensoren). De paper identificeert twee hoofdproblemen die vaak samen voorkomen:

1. Systeem-heterogeniteit: Apparaten hebben verschillende rekenkracht, geheugen en bandbreedte. Dit leidt tot "stragglers" (trage apparaten) en instabiele deelname, waardoor het onpraktisch is om alle clients volledige modeltraining te laten uitvoeren.
2. Statistische heterogeniteit (Non-IID): Data is niet onafhankelijk en gelijkmatig verdeeld over de clients. Dit veroorzaakt divergerende lokale updates en verslechtert de globale convergentie.

Bestaande oplossingen die modelpruning (het verwijderen van parameters) gebruiken om communicatie- en rekentaken te verminderen, kampen vaak met een gebrek aan stabiliteit. Ze struggle om tegelijkertijd communicatie-efficiëntie, aggregatie-robustheid en stabiele convergentie te garanderen onder deze gecombineerde heterogeniteit. Traditionele pruning-methoden negeren vaak de kromming (curvature) van de verliesfunctie, wat leidt tot een suboptimale selectie van parameters om te verwijderen, vooral bij Non-IID data.

Methodologie: CA-HFP

De auteurs stellen CA-HFP (Curvature-Aware Heterogeneous Federated Pruning) voor, een raamwerk dat drie kernstappen combineert:

1. Persoonlijke Gestructureerde Pruning:

   • Elke client past een specifiek "pruning ratio" toe op basis van zijn eigen reken- en communicatiecapaciteiten.
   • In plaats van alleen te kijken naar de grootte van de gewichten (magnitude), wordt een krommingsbewuste significantiescore gebruikt. Deze score is gebaseerd op een tweede-orde Taylor-expansie van de verliesfunctie:
     $$s_{i,t} = \nabla_i F(w_t) w_{i,t} + h_{i,t} w_{i,t}^2$$
     Hierbij vertegenwoordigt de eerste term de gradient en de tweede term de kromming (Hessian-diagonaal). Parameters met een lage score worden als minder belangrijk beschouwd en verwijderd. Dit minimaliseert de verstoring van het verlies door pruning.

2. Lichte Server-side Reconstructie:

   • Omdat elke client een ander submodel heeft (door verschillende pruning-masks), kunnen deze niet direct worden samengevoegd met de standaard FedAvg-methode.
   • De server voert een reconstructie uit voordat hij de modellen aggregeert. Verwijderde filters of neuronen in een client's submodel worden aangevuld met de huidige waarden uit het globale model. Hierdoor hebben alle lokale updates weer dezelfde structuur als het globale model, wat een uniforme aggregatie mogelijk maakt.

3. Theoretische Convergentie-analyse:

   • De auteurs leiden een convergentie-bound af voor federale optimalisatie met lokale SGD-stappen en pruning.
   • De analyse kwantificeert expliciet de impact van lokale berekening, data-heterogeniteit en de ruis die door pruning wordt geïntroduceerd ($e_t$).
   • Hieruit volgt dat een matige hoeveelheid lokale stappen ($E$) en een goed gekozen pruning-strength essentieel zijn om de convergentie stabiel te houden.

Belangrijkste Bijdragen

• CA-HFP Raamwerk: Een nieuw FL-framework dat persoonlijke, gestructureerde pruning toestaat voor resource-beperkte apparaten, terwijl het de aggregatie-compatibiliteit behoudt via reconstructie.
• Krommingsbewuste Pruning-criterium: Een nieuwe methode om parameters te selecteren voor verwijdering die rekening houdt met zowel gradients als de kromming van het verlieslandschap. Dit vermindert de aggregatie-bias die vaak optreedt bij Non-IID data.
• Reconstructie-mechanisme: Een oplossing voor structurele mismatch tussen heterogene submodellen, waardoor synchronisatie in de oorspronkelijke parameter-ruimte mogelijk wordt.
• Theoretische Onderbouwing: Een wiskundige analyse die de convergentie garandeert ondanks de aanwezigheid van pruning-ruis en data-heterogeniteit.

Resultaten

De methode is getest op drie datasets (FMNIST, CIFAR-10, CIFAR-100) met VGG16 en ResNet56 architecturen onder verschillende niveaus van Non-IID data (gecontroleerd via de Dirichlet-parameter $\alpha$).

• Nauwkeurigheid: CA-HFP behoudt een hoge nauwkeurigheid, zelfs onder zware Non-IID condities ($\alpha=0.1$). Het presteert aanzienlijk beter dan bestaande baselines zoals FedMP, PruneFL en DapperFL. Op CIFAR-100 met VGG16 behaalde CA-HFP bijvoorbeeld 73.12% nauwkeurigheid bij $\alpha=1.0$, vergeleken met 76.87% voor FedMP (maar FedMP had een hogere communicatiekost; CA-HFP is efficiënter).
• Efficiëntie: Door gestructureerde pruning reduceert CA-HFP de parameters en FLOPs (rekenkosten) aanzienlijk (tot 90% pruning in extreme gevallen) zonder in te leveren op prestaties. De communicatiekosten (upload/download) dalen evenredig met het pruning-ratio.
• Convergentie: CA-HFP convergeert sneller en bereikt een hogere finale nauwkeurigheid dan concurrenten, vooral bij hoge pruning-ratio's en sterke data-heterogeniteit.
• Ablatie-studie: Het verwijderen van het reconstructie-stap leidt tot een drastische daling in nauwkeurigheid (bijv. van 73.12% naar 62.84% op CIFAR-10 bij $\alpha=0.1$), wat aantoont dat reconstructie cruciaal is voor het oplossen van de aggregatie-bias.

Betekenis en Impact

CA-HFP biedt een praktische oplossing voor de realisatie van Federated Learning in real-world omgevingen waar apparaten en data sterk variëren. Door de combinatie van krommingsbewuste pruning en server-side reconstructie, lost het paper het fundamentele dilemma op tussen efficiëntie (kleine modellen) en effectiviteit (stabiele, nauwkeurige globale modellen). Dit maakt FL toepasbaar op een bredere scala van resource-beperkte randapparaten, wat essentieel is voor de toekomst van privacy-bewuste AI in IoT en mobiele netwerken.