Entropy-Regularized Probabilistic Gates for Sparse Model Discovery in Scarce-Data Federated Learning

Technische Samenvatting: Entropie-geregulariseerde Probabilistische Poorten voor Ontdekking van Sparse Modellen in Data-arme Federated Learning

Probleemstelling
Federated Learning (FL) staat voor aanzienlijke uitdagingen bij het opereren in scenario's met een hoge dimensionaliteit en beperkte hoeveelheden data, waarbij het aantal parameters ($d$) de totale hoeveelheid samples ($N$) ver overstijgt. In dergelijke instellingen ($d \gg N$) leidt standaard optimalisatie vaak tot parameterconfiguraties die falen in het generaliseren naar ongeziene testdata. Hoewel sparse modellen voordelen bieden op het gebied van communicatie- en computationele efficiëntie voor FL, is het ontdekken van deze sparse structuren moeilijk onder condities van data-heterogeniteit en gedeeltelijke participatie van clients. Bestaande benaderingen, zoals magnitude-gebaseerde pruning of Federated Iterative Hard Thresholding (Fed-IHT), falen vaak in het adequaat verkennen van de parameterruimte, wat leidt tot een te vroege toezegging aan suboptimale sparse supports. Bovendien vertrouwen standaard regularisatietechnieken (L1/L2) op parameter-magnitudes in plaats van het af te dwingen van een specifieke parameterdensiteit, wat het ontdekken van een model met een gewenste sparsiteit $\rho$ uitdagend maakt.

Methodologie
Het artikel stelt E-FLoPS voor (Entropy-Regularized Federated Learning with Probabilistic Gates), een gedistribueerd algoritme dat is ontworpen om sparse modellen te leren onder een $L_0$ densiteitsrestrictie, terwijl de onzekerheidsexploratie behouden blijft.

1. Probabilistische Poorten en Reparametrisatie: De methode hanteert het kader van stochastische poorten $z \in \{0,1\}^d$, waarbij effectieve parameters worden gedefinieerd als $\theta = \tilde{\theta} \odot z$. Om gradiëntgebaseerde optimalisatie mogelijk te maken, worden de discrete poorten benaderd met een Hard Concrete distributie, een continue relaxatie van de Bernoulli-distributie. Hierdoor kan de $L_0$ pseudo-norm (het aantal niet-nul parameters) worden benaderd door het verwachte aantal actieve poorten, $\mathbb{E}_q[z_j]$.
2. Entropie-regularisatie: Een cruciale innovatie is de introductie van een entropie-regularisatieterm in de optimalisatiedoelstelling. De Lagrangian bevat een Kullback-Leibler (KL) divergentieterm, $KL(q(z_j) || p(z_j))$, die de relatieve entropie van de poortverdelingen bestraft. Dit mechanisme voorkomt dat de poortverdelingen te vroeg instorten naar deterministische 0/1-toestanden, waardoor de exploratie van concurrerende sparse configuraties tijdens de training wordt gestimuleerd. De temperatuurparameter $T$ regelt de sterkte van deze straf.
3. Federated Optimalisatie: Het algoritme werkt onder een centrale serverorkestratie. Clients voeren lokale stochastische gradiëntafdaling (SGD) updates uit op de ruwe parameters $\tilde{\theta}$, de poortparameters $\phi = \log \alpha$, en de Lagrange-multiplier $\lambda$ (die een target densiteit $\rho$ afdwingt).
   • De server aggregeert updates synchroon.
   • Om communicatie-efficiëntie te garanderen, gebruikt het algoritme een pruning-mechanisme tijdens de training. Alleen de top-$m$ indices (waarbij $m = \lfloor \rho \cdot |\theta| \rfloor$) van de parameters en hun bijbehorende poortstatistieken worden geserialiseerd en gecommuniceerd.
   • Bij ontvangst reconstrueert de server het globale model en herstelt de poortparameters uit de geaggregeerde poortwaarden, waarbij de ruiscomponent wordt genegeerd.

Belangrijkste Bijdragen

• Formulering: De auteurs formuleren een differentieerbare, entropie-geregulariseerde, $L_0$-geconstraineerde doelstelling die specifiek is ontworpen voor de federated setting. Deze formulering adresseert de niet-differentieerbaarheid van de $L_0$ norm en de instabiliteit van optimalisatie in de $d \gg N$ regio.
• Algoritme (E-FLoPS): Het artikel introduceert een variant van Federated Averaging die probabilistische poorten integreert met entropie-regularisatie. Dit maakt een gebruikersgedefinieerde target sparsiteit mogelijk, terwijl de onzekerheidsexploratie behouden blijft om voortijdige convergentie naar slechte lokale minima te voorkomen.
• Communicatie-efficiëntie: De methode bereikt aanzienlijke communicatiecompressie door alleen de sparse support en noodzakelijke poortstatistieken te verzenden, wat vergelijkbaar is met iterative hard-thresholding methoden, maar zonder de post-training pruning overhead.

Experimentele Resultaten
De auteurs evalueerden E-FLoPS op synthetische en reële benchmarks, waarbij ze het vergeleken met Fed-IHT en dense Federated Averaging (FedAvg) met post-training pruning.

• Synthetische Lineaire Regressie: In hoog-dimensionale settings met variërende sample-naar-parameter ratio's ($N/|\theta|$), behaalde E-FLoPS consistent hogere test $R^2$ scores en een betere nauwkeurigheid van de sparse recovery (gemeten als de fractie van de ware features die zijn teruggevonden) vergeleken met Fed-IHT en FedAvg, met name in de meest data-arme regimes.
• MNIST Beeldclassificatie: Op een CNN-taak met 97,5% target sparsiteit demonstreerde E-FLoPS een test-prestatie en communicatie-efficiëntie vergelijkbaar met Fed-IHT. Het vertoonde echter een grotere reductie in Multiply Accumulate Operations (MACs) tijdens inferentie, wat wijst op superieure computationele efficiëntie.
• Leukemie Classificatie: Gebruikmakend van de Golub gen-expressie dataset (72 samples, 3571 features, $N/|\theta| \approx 0,02$), presteerde E-FLoPS significant beter dan Fed-IHT en FedAvg op test-nauwkeurigheid over 30 verschillende willekeurige seeds. Cruciaal was dat de genen die door E-FLoPS werden geselecteerd (bijv. cystatin C en myeloperoxidase) overeenkwamen met eerder experimenteel onderzoek naar kanker-metastase, terwijl andere methoden minder biologisch relevante features selecteerden.

Betekenis en Claims
Het artikel claimt dat E-FLoPS een robuust mechanisme biedt voor de ontdekking van sparse modellen in FL, waarbij specif de instabiliteit wordt aangepakt die wordt gevonden in kleine-sample, hoog-dimensionale regimes. Door gebruik te maken van entropie-regularisatie voorkomt de methode vroege toezegging aan suboptimale sparse supports, wat leidt tot betere generalisatie en een nauwkeurigere recovery van de werkelijke onderliggende sparse structuur. De auteurs stellen dat hun aanpak deze statistische winsten behaalt terwijl de communicatiekosten vergelijkbaar blijven met bestaande sparse methoden zoals Fed-IHT, wat het een levensvatbare oplossing maakt voor privacy-bewuste, resource-beperkte leeromgevingen. Het werk benadrukt het belang van onzekerheidsexploratie in sparse federated optimalisatie, een factor die vaak over het hoofd wordt gezien door magnitude-gebaseerde pruning technieken.