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

# Entropy-Regularized Probabilistic Gates for Sparse Model Discovery in Scarce-Data Federated Learning
Source: [https://arxiv.org/html/2607.00275](https://arxiv.org/html/2607.00275)
Krishna Harsha Kovelakuntla Huthasana Alireza Olama Andreas Lundell Department of Engineering and Information Technology, Åbo Akademi University \{kkovelak, alireza\.olama, andreas\.lundell\}@abo\.fi

###### Abstract

Federated Learning \(FL\) is a distributed machine learning \(ML\) paradigm with collaboration among multiple clients without sharing data\. FL is challenging under data heterogeneity and partial client participation\. Learning sparse models is useful for communication and computational efficiency in FL, but it is especially difficult in the small\-sample high\-dimensional regime\(d≫N\)\(d\\gg N\)where optimization can yield parameter configurations that fail to generalize to unseen test data\. While magnitude\-based pruning doesn’t account for uncertainty exploration in the parameter space, a formulation with probabilistic gates and anL0L\_\{0\}constraint allows sampling from competing sparse configurations during training\. In this work, we study entropy regularization of gate distributions as a mechanism to maintain uncertainty in sparse federated optimization by preventing early commitment to sparse support\. We examine its impact under data heterogeneity, client participation heterogeneity, and sparsity\. Experiments on synthetic and real\-world benchmarks show consistent improvements over federated iterative hard thresholding \(Fed\-IHT\) and pruning after dense federated averaging \(FedAvg\) training, both in statistical performance on test data and in sparsity recovery accuracy\.

Keywords

Entropy Regularization , Sparsity , Federated Learning , Uncertainty , Parameter Exploration , Probabilistic Gates , Entropy Maximization , Sparse Federated Learning ,L0L\_\{0\}Constraint

## 1Introduction

Federated Learning \(FL\) algorithms operate in a distributed machine learning \(ML\) setting in which multiple clients collaborate to train\[[16](https://arxiv.org/html/2607.00275#bib.bib4)\]\. This framework is characterized by privacy requirements of each client and avoids data sharing\. While not all distributed settings are privacy\-sensitive, FL is still beneficial because it eliminates the need for centralized data\[[11](https://arxiv.org/html/2607.00275#bib.bib5)\]\. FL can be coordinated either by a single server or by clients communicating with one another\. In this work, we study FL with central orchestration by a server to obtain a single global model as shown in Figure[1](https://arxiv.org/html/2607.00275#S1.F1)\.

ServerClient 1Client 2Client 3Client 4

Figure 1:Client–server federated learning architecture with central orchestration\. Solid arrows indicate the aggregated global model distributed to clients, while dotted arrows indicate local model updates sent from clients to the server\.A global model is typically learned by iteratively averaging parameters or gradients from clients and redistributing the global model to clients for further learning\. However, statistical heterogeneity across clients and partial client participation during training pose challenges to the learning process in FL\. Furthermore, sparse training and inference are desirable to improve generalizability\[[23](https://arxiv.org/html/2607.00275#bib.bib6)\]and enhance computational and communication efficiency in FL, thereby posing an additional challenge of discovering sparse models\[[25](https://arxiv.org/html/2607.00275#bib.bib42)\]\.

A common approach to inducing sparsity relies onL1L\_\{1\}andL2L\_\{2\}norms\[1\]\[1\]\[1\]Forθ∈Rd\\theta\\in R^\{d\},L1L\_\{1\}norm is‖θ‖1=∑i=1d\|θi\|\\\|\\theta\\\|\_\{1\}=\\sum\_\{i=1\}^\{d\}\|\\theta\_\{i\}\|andL2L\_\{2\}norm is‖θ‖2=\(∑i=1dθi2\)1/2\\\|\\theta\\\|\_\{2\}=\\left\(\\sum\_\{i=1\}^\{d\}\\theta\_\{i\}^\{2\}\\right\)^\{1/2\}\.for regularization, which depend directly on parameter magnitudes and offer varying levels of shrinkage\. In contrast, using a magnitude\-independentL0L\_\{0\}pseudo\-norm is advantageous because it imposes a constant penalty on nonzero parameters and is useful for learning a model with a desired parameter densityρ\\rho\. The Lagrangian for theL0L\_\{0\}density\-constrained optimization problem in FL can be defined as:

𝔏​\(θ,λ\)=∑c=1CncN​ℒ\(c\)​\(θ\)\+λ​\(‖θ‖0−ρ​\|θ\|\),‖θ‖0=∑j=1\|θ\|𝕀​\[θj≠0\],\\mathfrak\{L\}\(\\theta,\\lambda\)=\\sum\_\{c=1\}^\{C\}\\frac\{n\_\{c\}\}\{N\}\\mathcal\{L\}^\{\(c\)\}\(\\theta\)\+\\lambda\\left\(\\\|\\theta\\\|\_\{0\}\-\\rho\|\\theta\|\\right\),\\quad\\\|\\theta\\\|\_\{0\}=\\sum\_\{j=1\}^\{\|\\theta\|\}\\mathbb\{I\}\[\\theta\_\{j\}\\neq 0\],\(1\)where,ℒ\(c\)​\(θ\)\\mathcal\{L\}^\{\(c\)\}\(\\theta\)denotes the normalized loss at clientcc, defined as:

ℒ\(c\)​\(θ\)=1nc​∑i=1ncℓ​\(h​\(xi\(c\);θ\),yi\(c\)\)\.\\mathcal\{L\}^\{\(c\)\}\(\\theta\)=\\frac\{1\}\{n\_\{c\}\}\\sum\_\{i=1\}^\{n\_\{c\}\}\\ell\\big\(h\(x\_\{i\}^\{\(c\)\};\\theta\),y\_\{i\}^\{\(c\)\}\\big\)\.\(2\)Here,CCis number of clients in FL, each holding a local dataset\(D\(c\)\)c=1C=\(X\(c\),Y\(c\)\)\(D^\{\(c\)\}\)\_\{c=1\}^\{C\}=\(X^\{\(c\)\},Y^\{\(c\)\}\), whereX\(c\)∈ℝnc×i​nX^\{\(c\)\}\\in\\mathbb\{R\}^\{n\_\{c\}\\times in\},Y\(c\)∈ℝnc×o​u​tY^\{\(c\)\}\\in\\mathbb\{R\}^\{n\_\{c\}\\times out\}, and∑c=1Cnc=N\\sum\_\{c=1\}^\{C\}n\_\{c\}=N\. We assume a modelh​\(x;θ\):ℝi​n→ℝo​u​th\(x;\\theta\):\\mathbb\{R\}^\{in\}\\rightarrow\\mathbb\{R\}^\{out\}and a loss functionℓ​\(h​\(x;θ\),y\)\\ell\(h\(x;\\theta\),y\), wherex∈ℝi​nx\\in\\mathbb\{R\}^\{in\},y∈ℝo​u​ty\\in\\mathbb\{R\}^\{out\}, andθ∈ℝd\\theta\\in\\mathbb\{R\}^\{d\}\. The above formulation with Lagrange multiplierλ\\lambdaleads to a min\-max optimization problem that can, in principle, be solved using gradient descent\-ascent\. However, the non\-differentiability of theL0L\_\{0\}pseudo\-norm complicates optimization within standard gradient\-based learning frameworks\.

Louizoset al\.\[[14](https://arxiv.org/html/2607.00275#bib.bib12)\]introduce a reparameterization using stochastic gatesz∈ℝpz\\in\\mathbb\{R\}^\{p\}withθ=θ~⊙z\\theta=\\tilde\{\\theta\}\\odot z\[2\]\[2\]\[2\]⊙\\odotdenotes the elementwise product\[[9](https://arxiv.org/html/2607.00275#bib.bib37)\]\.as the effective parameters of the model\. By assuming a hard Concrete distribution overzz\(an approximation of the Bernoulli distribution\), theL0L\_\{0\}norm is approximated by the expected number of active gates, i\.e\.,∑j𝔼q​\(z\)​\[zj\]\\sum\_\{j\}\\mathbb\{E\}\_\{q\(z\)\}\[z\_\{j\}\], enabling gradient\-based optimization of theL0L\_\{0\}\-regularized objective\.Gallego\-Posadaet al\.\[[7](https://arxiv.org/html/2607.00275#bib.bib13)\]extended this framework by using anL0L\_\{0\}density constraint and solving the resulting min–max problem, thereby enabling the user to define the target sparsity in a centralized setting\.Huthasanaet al\.\[[10](https://arxiv.org/html/2607.00275#bib.bib56)\]further adapted this approach to sparsity learning in the FL context, accounting for heterogeneity in data and client participation\.

However, learning sparse models is challenging, particularly in the small\-sample high\-dimensional regime\(d≫N\)\(d\\gg N\)\[[18](https://arxiv.org/html/2607.00275#bib.bib48),[6](https://arxiv.org/html/2607.00275#bib.bib49),[22](https://arxiv.org/html/2607.00275#bib.bib50),[19](https://arxiv.org/html/2607.00275#bib.bib51)\], where the optimization is unstable, resulting in multiple solutions of competing parameter configurations with poor sparsity recovery and generalization behaviour to unseen test data\.Baoet al\.\[[3](https://arxiv.org/html/2607.00275#bib.bib59)\]work on statistical recovery in a low sample–high dimensional regime, but a low sample size for the individual client relative to the number of parameters is considered, while the total sample size N far exceeds the number of parameters\. We aim to study a low total sample sizeNNrelative to the number of parametersd=\|θ\|d=\|\\theta\|similar to a centralized setting, which is extremely challenging under data and client participation heterogeneity in FL\.

Louizoset al\.\[[14](https://arxiv.org/html/2607.00275#bib.bib12)\]points to penalizing relative entropy or the Kullback\-Leibler divergenceKL​\(q​\(z\)∥p​\(z\)\)\\mathrm\{KL\}\(q\(z\)\\\|p\(z\)\)in the optimization wherep​\(z\)p\(z\)is the prior andq​\(z\)q\(z\)is the approximate posterior for uncertainty exploration\. The concept of entropy regularization is not new, and it is commonly used to encourage diversity and exploration at the decision level, distributions of class or action, in reinforcement learning \(RL\), and adapted to uncertainty exploration in input space, distributions of latent variables dependent on data, in variational auto–encoders\[[2](https://arxiv.org/html/2607.00275#bib.bib52),[5](https://arxiv.org/html/2607.00275#bib.bib53)\]\. It is also used in Bayesian inference in centralized and FL settings\[[1](https://arxiv.org/html/2607.00275#bib.bib54),[26](https://arxiv.org/html/2607.00275#bib.bib55)\]\. In their approach, the hard concrete gate distributionsq​\(z\)q\(z\)are not directly dependent on the data and can be sampled independently of it post\-training\. In this work, the formulations atLouizoset al\.\[[14](https://arxiv.org/html/2607.00275#bib.bib12)\], Gallego\-Posadaet al\.\[[7](https://arxiv.org/html/2607.00275#bib.bib13)\]are adapted with relative\-entropy penalization in a communication\-efficient FL setting to explore uncertainty in non\-zero parameter configurations and show that it consistently outperforms iterative hard\-thresholding\-based pruning during training\[[24](https://arxiv.org/html/2607.00275#bib.bib26)\]and post\-training pruning of a dense model using the classic federated averaging algorithm\[[16](https://arxiv.org/html/2607.00275#bib.bib4)\]\.

The remainder of this paper is organized as follows\. We first present the entropy\-regularizedL0L\_\{0\}\-constrained formulation for federated optimization, then the proposed distributed algorithm, followed by experiments in heterogeneous FL settings, and conclude\.

## 2Formulation

Assuming a modelh​\(x;θ\):ℝi​n→ℝo​u​th\(x;\\theta\):\\mathbb\{R\}^\{in\}\\rightarrow\\mathbb\{R\}^\{out\}and a loss functionℓ​\(h​\(x;θ\),y\)\\ell\(h\(x;\\theta\),y\), wherex∈ℝi​nx\\in\\mathbb\{R\}^\{in\},y∈ℝo​u​ty\\in\\mathbb\{R\}^\{out\}, andθ∈ℝd\\theta\\in\\mathbb\{R\}^\{d\}, consider a centralized datasetD=\(X,Y\)D=\(X,Y\)withX∈ℝN×i​nX\\in\\mathbb\{R\}^\{N\\times in\}andY∈ℝN×o​u​tY\\in\\mathbb\{R\}^\{N\\times out\}\. Using\[[14](https://arxiv.org/html/2607.00275#bib.bib12),[7](https://arxiv.org/html/2607.00275#bib.bib13)\], the min–max objective can be defined using the expectation of the loss with respect to the distribution of gates, and theL0L\_\{0\}pseudo\-norm approximated by the expected number of active gates, i\.e\.,∑j𝔼q​\(z∣ϕ\)​\[zj\]\\sum\_\{j\}\\mathbb\{E\}\_\{q\(z\\mid\\phi\)\}\[z\_\{j\}\]as shown in eq\.[3](https://arxiv.org/html/2607.00275#S2.E3)\. Since eachzjz\_\{j\}is a deterministic transformation of parameter\-free noise, the expectation can be optimized using Monte Carlo sampling and reparameterized gradients\[[17](https://arxiv.org/html/2607.00275#bib.bib36), ch\. 3\.3\.3\]\.

𝔏^​\(θ~,ϕ,λ\)=1R​∑r=1R\[1N​∑i=1Nℓ​\(h​\(xi;θ~⊙z\(r\)\),yi\)\]\+λ​\(∑j=1\|θ\|𝔼q​\(z∣ϕ\)​\[zj\]−ρ​\|θ\|\)\.\\hat\{\\mathfrak\{L\}\}\(\\tilde\{\\theta\},\\phi,\\lambda\)=\\frac\{1\}\{R\}\\sum\_\{r=1\}^\{R\}\\left\[\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}\\ell\\big\(h\(x\_\{i\};\\tilde\{\\theta\}\\odot z^\{\(r\)\}\),y\_\{i\}\\big\)\\right\]\+\\lambda\\left\(\\sum\_\{j=1\}^\{\|\\theta\|\}\\mathbb\{E\}\_\{q\(z\\mid\\phi\)\}\[z\_\{j\}\]\-\\rho\|\\theta\|\\right\)\.\(3\)
In a federated learning \(FL\) setting withCCclients holding datasets\(D\(c\)\)c=1C=\(X\(c\),Y\(c\)\)\(D^\{\(c\)\}\)\_\{c=1\}^\{C\}=\(X^\{\(c\)\},Y^\{\(c\)\}\), we consider a reparameterized linear modelh​\(x;θ~⊙z\):ℝi​n→ℝo​u​th\(x;\\tilde\{\\theta\}\\odot z\):\\mathbb\{R\}^\{in\}\\rightarrow\\mathbb\{R\}^\{out\}and a loss functionℓ​\(h​\(x;θ~⊙z\),y\)\\ell\(h\(x;\\tilde\{\\theta\}\\odot z\),y\)\. Here,X\(c\)∈ℝnc×i​nX^\{\(c\)\}\\in\\mathbb\{R\}^\{n\_\{c\}\\times in\},Y\(c\)∈ℝnc×o​u​tY^\{\(c\)\}\\in\\mathbb\{R\}^\{n\_\{c\}\\times out\},∑c=1Cnc=N\\sum\_\{c=1\}^\{C\}n\_\{c\}=N,x∈ℝi​nx\\in\\mathbb\{R\}^\{in\},y∈ℝo​u​ty\\in\\mathbb\{R\}^\{out\}, andθ=θ~⊙z∈ℝi​n\\theta=\\tilde\{\\theta\}\\odot z\\in\\mathbb\{R\}^\{in\}\. The gate parameters are defined asϕ=log⁡α∈ℝi​n\\phi=\\log\\alpha\\in\\mathbb\{R\}^\{in\}\.

The Lagrangian corresponding to the entropy regularizedL0L\_\{0\}density\-constrained optimization problem is:

𝔏^\(θ~,ϕ,λ\)=∑c=1CncNℒ\(c\)\(θ~,ϕ\)\+λ\(∑j=1\|θ\|𝔼q​\(z∣ϕ\)\[zj\]−ρ\|θ\|\)\+T∑j=1\|θ\|KL\(q\(zj\|ϕ\)\|\|p\(zj\|ϕi​n​i​t\)\)\.\\hat\{\\mathfrak\{L\}\}\(\\tilde\{\\theta\},\\phi,\\lambda\)=\\sum\_\{c=1\}^\{C\}\\frac\{n\_\{c\}\}\{N\}\\mathcal\{L\}^\{\(c\)\}\(\\tilde\{\\theta\},\\phi\)\+\\lambda\\left\(\\sum\_\{j=1\}^\{\|\\theta\|\}\\mathbb\{E\}\_\{q\(z\\mid\\phi\)\}\[z\_\{j\}\]\-\\rho\|\\theta\|\\right\)\+T\\sum\_\{j=1\}^\{\|\\theta\|\}KL\(q\(z\_\{j\}\|\\phi\)\|\|p\(z\_\{j\}\|\\phi\_\{init\}\)\)\.\(4\)Here,p​\(z\|ϕi​n​i​t\)p\(z\|\\phi\_\{init\}\)is a prior and is also a hard concrete distribution, andℒ\(c\)​\(θ~,ϕ\)\\mathcal\{L\}^\{\(c\)\}\(\\tilde\{\\theta\},\\phi\)denotes the Monte Carlo estimate of the normalized loss at clientcc, defined as:

ℒ\(c\)​\(θ~,ϕ\)=1R​∑r=1R1nc​∑i=1ncℓ​\(h​\(xi\(c\);θ~⊙z\(r\)\),yi\(c\)\)\.\\mathcal\{L\}^\{\(c\)\}\(\\tilde\{\\theta\},\\phi\)=\\frac\{1\}\{R\}\\sum\_\{r=1\}^\{R\}\\frac\{1\}\{n\_\{c\}\}\\sum\_\{i=1\}^\{n\_\{c\}\}\\ell\\big\(h\(x\_\{i\}^\{\(c\)\};\\tilde\{\\theta\}\\odot z^\{\(r\)\}\),y\_\{i\}^\{\(c\)\}\\big\)\.\(5\)The stochastic gateszzare sampled using the Hard Concrete distribution, applying a hard\-sigmoid transformation to a stretched Binary Concrete random variable\[[14](https://arxiv.org/html/2607.00275#bib.bib12),[15](https://arxiv.org/html/2607.00275#bib.bib15)\], defined as:

s\\displaystyle s=σ​\(log⁡u1−u\+log⁡αβ′\),u∼𝒰​\(0,1\),\\displaystyle=\\sigma\\left\(\\frac\{\\log\\frac\{u\}\{1\-u\}\+\\log\\alpha\}\{\\beta^\{\\prime\}\}\\right\),\\quad u\\sim\\mathcal\{U\}\(0,1\),s¯\\displaystyle\\bar\{s\}=s​\(ζ−γ\)\+γ,z=min⁡\(1,max⁡\(0,s¯\)\)\.\\displaystyle=s\(\\zeta\-\\gamma\)\+\\gamma,\\quad z=\\min\(1,\\max\(0,\\bar\{s\}\)\)\.\(6\)The expectation of a gate being active\[3\]\[3\]\[3\]𝔼q​\(z∣ϕ\)​\[zj\]=1−Q​\(s¯j≤0∣ϕj\)=σ​\(log⁡αj−β′​log⁡\(−γζ\)\)\\mathbb\{E\}\_\{q\(z\\mid\\phi\)\}\[z\_\{j\}\]=1\-Q\(\\bar\{s\}\_\{j\}\\leq 0\\mid\\phi\_\{j\}\)=\\sigma\\left\(\\log\\alpha\_\{j\}\-\\beta^\{\\prime\}\\log\\left\(\-\\frac\{\\gamma\}\{\\zeta\}\\right\)\\right\)is derived atLouizoset al\.\[[14](https://arxiv.org/html/2607.00275#bib.bib12)\]using the cumulative distribution functionQ​\(s¯\)Q\(\\bar\{s\}\)\. We introducedT≥0T\\geq 0, which we treat as a constant with or without decay to penalize entropy\. The resulting min–max optimization problem is:

θ~∗,ϕ∗,λ∗=arg⁡minθ~,ϕ⁡arg⁡maxλ≥0⁡𝔏^​\(θ~,ϕ,λ\)\.\\tilde\{\\theta\}^\{\*\},\\phi^\{\*\},\\lambda^\{\*\}=\\arg\\min\_\{\\tilde\{\\theta\},\\phi\}\\arg\\max\_\{\\lambda\\geq 0\}\\hat\{\\mathfrak\{L\}\}\(\\tilde\{\\theta\},\\phi,\\lambda\)\.\(7\)The parametersθ~\\tilde\{\\theta\}andϕ=log⁡α\\phi=\\log\\alphaare jointly optimized using gradient descent with reparameterized gradients\.λ\\lambdais updated via gradient ascent and a restart strategy of resetting its value to 0 as and when the sparsity constraint is satisfied\[[7](https://arxiv.org/html/2607.00275#bib.bib13)\]\. Since the hard concrete distribution is a continuos approximation of Bernoulli distribution the computation ofKL\(q\(zj\)\|\|p\(zj\)\)KL\(q\(z\_\{j\}\)\|\|p\(z\_\{j\}\)\)involves an additional term involvings¯​, and thus,​z∈\(0,1\)\\bar\{s\}\\text\{, and thus, \}z\\in\(0,1\)which can be computed using truncated distributionq​\(s¯\|s¯∈\(0,1\)\)q\(\\bar\{s\}\|\\bar\{s\}\\in\(0,1\)\)or a Monte\-carlo estimate of the same\. We used the closed\-form expressions presented in the appendix[A](https://arxiv.org/html/2607.00275#A1)provided atLouizoset al\.\[[14](https://arxiv.org/html/2607.00275#bib.bib12)\]\.

## 3Algorithm

We use the notationsℒCon​\(ϕ\)\\mathcal\{L\}\_\{\\text\{Con\}\}\(\\phi\)andℒKL​\(ϕ\)\\mathcal\{L\}\_\{\\text\{KL\}\}\(\\phi\)forL0L\_\{0\}density constraint and theKL\(q\(z\)\|\|p\(z\)\)KL\(q\(z\)\|\|p\(z\)\)in \([4](https://arxiv.org/html/2607.00275#S2.E4)\)\[[10](https://arxiv.org/html/2607.00275#bib.bib56)\]\. The Lagrangian can then be written as:

𝔏^​\(θ~,ϕ,λ\)=∑c=1CncN​ℒ\(c\)​\(θ~,ϕ\)\+λ​ℒCon​\(ϕ\)\+T​ℒKL​\(ϕ\)\.\\hat\{\\mathfrak\{L\}\}\(\\tilde\{\\theta\},\\phi,\\lambda\)=\\sum\_\{c=1\}^\{C\}\\frac\{n\_\{c\}\}\{N\}\\mathcal\{L\}^\{\(c\)\}\(\\tilde\{\\theta\},\\phi\)\+\\lambda\\,\\mathcal\{L\}\_\{\\text\{Con\}\}\(\\phi\)\+T\\,\\mathcal\{L\}\_\{\\text\{KL\}\}\(\\phi\)\.\(8\)McMahanet al\.\[[16](https://arxiv.org/html/2607.00275#bib.bib4)\]propose federated averaging \(FedAvg\), a distributed algorithm for learning a global model via synchronous updates from clients using gradient averaging\. In this setting, a central server coordinates training acrossCCclients each holding a local datasetD\(c\)D^\{\(c\)\}\. Each client performs stochastic gradient descent \(SGD\) updates locally for few iterations before communicating the parameters to the server, to reduce the communication between server and clients by increasing computation at clients\. The server aggregates the parameters by averaging to obtain the global model\.

At iterationtt, each client performs the following updates:

θ~ct\+1\\displaystyle\\tilde\{\\theta\}\_\{c\}^\{t\+1\}=θ~ct−ηθ~​∇θ~ℒ\(c\)​\(θ~t,ϕt\),\\displaystyle=\\tilde\{\\theta\}\_\{c\}^\{t\}\-\\eta\_\{\\tilde\{\\theta\}\}\\nabla\_\{\\tilde\{\\theta\}\}\\mathcal\{L\}^\{\(c\)\}\(\\tilde\{\\theta\}^\{t\},\\phi^\{t\}\),\(9\)ϕct\+1\\displaystyle\\phi\_\{c\}^\{t\+1\}=ϕct−ηϕ​\(∇ϕℒ\(c\)​\(θ~t,ϕt\)\+λ​​∇ϕℒCon\(c\)​\(ϕt\)\+T​​∇ϕℒKL\(c\)​\(ϕt\)\),\\displaystyle=\\phi\_\{c\}^\{t\}\-\\eta\_\{\\phi\}\\left\(\\nabla\_\{\\phi\}\\mathcal\{L\}^\{\(c\)\}\(\\tilde\{\\theta\}^\{t\},\\phi^\{t\}\)\+\\lambda\\text\{ \}\\nabla\_\{\\phi\}\\mathcal\{L\}\_\{\\text\{Con\}\}^\{\(c\)\}\(\\phi^\{t\}\)\+T\\text\{ \}\\nabla\_\{\\phi\}\\mathcal\{L\}\_\{\\text\{KL\}\}^\{\(c\)\}\(\\phi^\{t\}\)\\right\),\(10\)λct\+1\\displaystyle\\lambda\_\{c\}^\{t\+1\}=λct\+ηλ​ℒCon\(c\)​\(ϕt\)\.\\displaystyle=\\lambda\_\{c\}^\{t\}\+\\eta\_\{\\lambda\}\\,\\mathcal\{L\}\_\{\\text\{Con\}\}^\{\(c\)\}\(\\phi^\{t\}\)\.\(11\)At each round, a fractionγc\\gamma\_\{c\}of clients is selected uniformly at random, resulting inK=⌊γc​C⌋K=\\lfloor\\gamma\_\{c\}C\\rfloorparticipating clients\. In practice, a few mini\-batches of uniform sizeBBsampled iteratively at each client to locally update the modeln​\(B\)n\(B\)times per communication round or a full pass over client data amount to a local epoch\. The server gathers all updates from clients and performs a synchronous update of global model , using the averages of the updates from clients, after each communication round or global epoch keepingTTconstant or decaying it according to a predetermined schedule\. The Lagrange parameterλ\\lambdais reset to zero when the constraint is satisfied\[[7](https://arxiv.org/html/2607.00275#bib.bib13)\]\. The aggregation weightswkw\_\{k\}for clients sampled in an epoch or round can be uniform or proportional to the number of samples each client holds\.

This approach enables learning a global sparse model in FL with entropy regularization and anL0L\_\{0\}constraint using probabilistic gates\. We refer to this variant of FedAvg asE\-FLoPS, whereEstands for entropy regularization\. The learning ratesηθ~\\eta\_\{\\tilde\{\\theta\}\}andηϕ\\eta\_\{\\phi\}needs to be appropriately tuned\. The learning rateηλ\\eta\_\{\\lambda\}for the Lagrange parameter updates is set in the order of1/\|θ\|1/\|\\theta\|\. The temperatureT∈\(0,∞\)T\\in\(0,\\infty\)is also initialized at1/\|θ\|1/\|\\theta\|and increased if needed to encourage uncertainty exploration as the total sample size N decreases\. The gate parametersϕ=log⁡α\\phi=\\log\\alphaare initialized from a normal distribution with meanlog⁡ρinit−log⁡\(1−ρinit\)\\log\\rho\_\{\\text\{init\}\}\-\\log\(1\-\\rho\_\{\\text\{init\}\}\)and variance0\.010\.01, whereρinit\\rho\_\{\\text\{init\}\}controls the initial density\. The target density is denoted byρtarg=ρ\\rho\_\{\\text\{targ\}\}=\\rho\. The Hard Concrete distribution parameters are set toγ=−0\.1\\gamma=\-0\.1,ζ=1\.1\\zeta=1\.1, andβ′\\beta^\{\\prime\}is recommended to be set at0\.660\.66, followingLouizoset al\.\[[14](https://arxiv.org/html/2607.00275#bib.bib12)\]\. The Lagrange multiplier is initialized asλ=0\\lambda=0\. Client participation variability is simulated by randomly selecting a fraction of clients at each training round\.

E\-FLoPSachieves test\-time sparsity by using deterministic gatesz^\\hat\{z\}sampled without noise or smoothing\[[7](https://arxiv.org/html/2607.00275#bib.bib13)\]\. For exact sparsity, the raw parametersθ~\\tilde\{\\theta\}for the top\-mmindices of effective parametersθ\\thetaare pruned, wherem=⌊ρtarg⋅\|θ\|⌋m=\\lfloor\\rho\_\{\\text\{targ\}\}\\cdot\|\\theta\|\\rfloor\. This pruning mechanism can be applied from a pre\-defined threshold of epochs, referred to as the*prune start*epoch, which is set to0forE\-FLoPSto improve communication efficiency via serialization for message compression\[[10](https://arxiv.org/html/2607.00275#bib.bib56)\]\. At this stage, the serialized mean gate values, from repeated sampling, corresponding to top\-mmindices of gateszzare retained, and the rest of the gates are replaced with their average, denoted byz−mavgz^\{\\text\{avg\}\}\_\{\-m\}\. Only the prunedθ~\\tilde\{\\theta\}andzz, along with their non\-zero indices and the scalarz−mavgz^\{\\text\{avg\}\}\_\{\-m\}, are communicated\. Upon reception, parameters are reconstructed viaθ=θ~⊙z\\theta=\\tilde\{\\theta\}\\odot zand aggregated, andϕ\\phiis recovered fromzzusingϕ=β′​log⁡\(z1−z\)\\scriptstyle\\phi=\\beta^\{\\prime\}\\log\\left\(\\frac\{z\}\{1\-z\}\\right\), ignoring the noise component of the gates\. The communication cost reduction by this approach in FL is significant for smallρtarg\\rho\_\{\\text\{targ\}\}with minimal meta data overhead\.

Algorithm 1FedAvg variantE\-FLoPS\.EEandBBdenote the number of epochs and the mini\-batch size\.1:Initialization:

\(θ~\(0\),ϕ\(0\)\)\(\\tilde\{\\theta\}^\{\(0\)\},\\phi^\{\(0\)\}\), compute

θ\(0\)=θ~\(0\)⊙z\(0\)\\theta^\{\(0\)\}=\\tilde\{\\theta\}^\{\(0\)\}\\odot z^\{\(0\)\}
2:forepoch

b=1b=1to

EEdo

3:Sample client subset

StS\_\{t\}
4:foreachclient

k∈Stk\\in S\_\{t\}do

5:

\(k,θ~k,zk,λk\)←\(k,\\tilde\{\\theta\}\_\{k\},z\_\{k\},\\lambda\_\{k\}\)\\leftarrowClientCompute\(k,θ~\(t\),z\(t\),λ\(t\)\)\(k,\\tilde\{\\theta\}^\{\(t\)\},z^\{\(t\)\},\\lambda^\{\(t\)\}\)

6:endfor

7:Server aggregation:

θ\\displaystyle\\theta=∑k∈Stwk​θ~k\\displaystyle=\\sum\_\{k\\in S\_\{t\}\}w\_\{k\}\\tilde\{\\theta\}\_\{k\}z\\displaystyle z=∑k∈Stwk​zk\\displaystyle=\\sum\_\{k\\in S\_\{t\}\}w\_\{k\}z\_\{k\}
8:Reconstruct

θ~=θ⊘z\\tilde\{\\theta\}=\\theta\\oslash z\[4\]\[4\]\[4\]⊘\\oslashrepresents an element wise division\.and recover

ϕ\\phifrom

zzusing

ϕ=β′​log⁡\(z1−z\)\\phi=\\beta^\{\\prime\}\\log\\left\(\\frac\{z\}\{1\-z\}\\right\)
9:Perform server\-side updates on

\(θ~,ϕ,λ\)\(\\tilde\{\\theta\},\\phi,\\lambda\)
10:Either decay or keep

TTconstant

11:if

b\>b\>prune startthen

12:prune

θ~\\tilde\{\\theta\}using top\-

mmθ\\thetatop\-

mmindices of

θ\\theta
13:prune

zzwith tail statistic

z−mavg\\scriptstyle z^\{\\text\{avg\}\}\_\{\-m\}
14:Serialize and communicate

15:endif

16:endfor

17:ClientCompute\(k,θ~,z,λ\)\(k,\\tilde\{\\theta\},z,\\lambda\):

18:Recover

ϕk\\phi\_\{k\}from

zzusing

ϕ=β′​log⁡\(z1−z\)\\scriptstyle\\phi=\\beta^\{\\prime\}\\log\\left\(\\frac\{z\}\{1\-z\}\\right\)
19:Sample mini\-batches

b∼D\(k\)b\\sim D^\{\(k\)\}of size

BB
20:Perform

n​\(B\)n\(B\)local SGD steps on

\(θ~k,ϕk\)\(\\tilde\{\\theta\}\_\{k\},\\phi\_\{k\}\)and ascent for

λk\\lambda\_\{k\}
21:returnCommunicate pruned and serialized

\(θ~k,ϕk\)\(\\tilde\{\\theta\}\_\{k\},\\phi\_\{k\}\)and

λk\\lambda\_\{k\}

## 4Experiments

We include experiments on synthetic and real\-world datasets\. We evaluate true sparsity recovery in linear regression \(LR\) on synthetically generated data, convolutional neural network \(CNN\) on the MNIST digit classification data, and softmax multi\-class classifier on the Golub leukemia cancer classification data\. We compare our method with the federated iterative hard thresholding algorithm \(Fed\-IHT\) proposed byTonget al\.\[[24](https://arxiv.org/html/2607.00275#bib.bib26)\], where a hard\-thresholding operation is used to retain only top\-mmparameters in absolute magnitude at each iteration or epoch along with federated averaging to enforce sparsity, and with classic federated averaging \(Fed\-Avg\) with dense training and pruning after the last training epoch\. ForFedAvgthe training time performance on test data is evaluated using the top\-mmparameters after each epoch though the training is dense without imposing sparsity\. An approximate upper bound is established by centralized training by pooling data from clients with distributional shifts\. A tuning phase, marked by a vertical dotted line in all figures in the experiments section, is conducted to further improve the statistical performance of the models with the fixed sparse support discovered by the end of the training phase\.

In FL, comunication is a bottle neck and needs to be minimized\. The federated averaging variantE\-FLoPSenables sparse communication throughout training without compromising statistical performance\. The theoritical uplink and downlink communication costs can be estimated as multiples of the message size and the number of communication rounds, assuming 4 bytes per parameter and index each\. The total two way communication cost per client in each server round in FL is,

FedAvg:epochs×4​\|θ\|,E\-FLoPS:epochs×4⋅\(2​ρtarg​\|θ\|\),andFedIter\-HT:epochs×4⋅\(ρtarg​\|θ\|\)\.\\resizebox\{310\.4979pt\}\{\}\{$\\texttt\{FedAvg\}:\\text\{epochs\}\\times 4\|\\theta\|,\\text\{ \}\\texttt\{E\-FLoPS\}:\\text\{epochs\}\\times 4\\cdot\(2\\rho\_\{\\text\{targ\}\}\|\\theta\|\),\\text\{ and \}\\texttt\{FedIter\-HT\}:\\text\{epochs\}\\times 4\\cdot\(\\rho\_\{\\text\{targ\}\}\|\\theta\|\)$\}\.Thus, the communication cost ofE\-FLoPSis more thanFedIter\-HT, but significantly lower than dense training viaFedAvg\. For large model sizes\|θ\|\|\\theta\|and small target densitiesρtarg\\rho\_\{\\text\{targ\}\}, the gap betweenFedAvgandE\-FLoPSbecomes substantial, while small with that ofFedIter\-HT\.

The experiments are conducted on an Apple MacBook with an M4 Pro chip \(12\-core CPU\) and 24 GB unified memory, running macOS 15\.5\. The implementation uses Flower \(v1\.29\.0\) framework for simulation of FL and PyTorch \(v2\.7\.0\) framework for training, with Python \(v3\.12\.7\)\.

### 4\.1Experiments on Synthetic Data

We generated synthetic data for sparse linear regression following the procedure fromBertsimaset al\.\[[4](https://arxiv.org/html/2607.00275#bib.bib23)\]for a range ofN\|θ\|\\scriptstyle\\frac\{N\}\{\|\\theta\|\}\. For a parameter vector dimensiond=\|θ\|d=\|\\theta\|, each rowxi∈ℝdx\_\{i\}\\in\\mathbb\{R\}^\{d\}ofX∈ℝN×dX\\in\\mathbb\{R\}^\{N\\times d\}is sampled from a zero\-mean Gaussian distribution with covariance matrixΣ\\Sigma\. We use a Toeplitz covariance structure defined asΣi​j=ρcor\|i−j\|,i,j=1,…,d\\Sigma\_\{ij\}=\\rho\_\{\\text\{cor\}\}^\{\|i\-j\|\},\\quad i,j=1,\\ldots,d\. Anmm\-sparse coefficient vectorwtrue∈ℝdw\_\{\\text\{true\}\}\\in\\mathbb\{R\}^\{d\}is constructed, wherem=⌊ρ⋅d⌋m=\\lfloor\\rho\\cdot d\\rfloor\. A subset of indicesSm⊆\{1,…,d\}S\_\{m\}\\subseteq\\\{1,\\ldots,d\\\}is selected uniformly at random, and coefficients are assigned as\(wtrue\)j∼Unif​\{−1,1\}​for​j∈Sm​and​\(wtrue\)j=0​otherwise\(w\_\{\\text\{true\}\}\)\_\{j\}\\sim\\mathrm\{Unif\}\\\{\-1,1\\\}\\text\{ for \}j\\in S\_\{m\}\\text\{ and \}\(w\_\{\\text\{true\}\}\)\_\{j\}=0\\text\{ otherwise\}\. The response vector is generated asy=X​wtrue\+εy=Xw\_\{\\text\{true\}\}\+\\varepsilon, whereε∼𝒩​\(0,σ2​IN\)\\varepsilon\\sim\\mathcal\{N\}\(0,\\sigma^\{2\}I\_\{N\}\)\. The signal\-to\-noise ratio \(SNR\) is defined asSNR=‖X​wtrue‖22‖ε‖22\\scriptstyle\\mathrm\{SNR\}=\\frac\{\\\|Xw\_\{\\text\{true\}\}\\\|\_\{2\}^\{2\}\}\{\\\|\\varepsilon\\\|\_\{2\}^\{2\}\}, and the noise level is set toσ=‖X​wtrue‖2SNR​N\\scriptstyle\\sigma=\\frac\{\\\|Xw\_\{\\text\{true\}\}\\\|\_\{2\}\}\{\\sqrt\{\\mathrm\{SNR\}\}\\,\\sqrt\{N\}\}\. For\|θ\|=1000\\scriptstyle\|\\theta\|=1000, we generate a number of total training samples N for varying ratios ofN\|θ\|\\scriptstyle\\frac\{N\}\{\|\\theta\|\}with5000\\scriptstyle 5000samples of test data in all cases\. The generated training data is then distributed among clients\.

We employed affine shifting for attribute distribution skew, Dirichlet partitioning protocol \(DPP\) for skew in the number of samples available at a client\[[20](https://arxiv.org/html/2607.00275#bib.bib29),[21](https://arxiv.org/html/2607.00275#bib.bib28)\]\. A fraction of clients \(0\.60\.6\) is randomly sampled in each epoch to introduce participation heterogeneity\[[20](https://arxiv.org/html/2607.00275#bib.bib29)\]\. We evaluate performance using mean squared error \(MSE\) andR2R^\{2\}\. Experiments are conducted at a correlation level ofρcor=0\.2\\rho\_\{\\text\{cor\}\}=0\.2and an SNR of 20, with a true sparsity level of95%95\\%or target density of5%5\\%\.

As shown in Figure[2](https://arxiv.org/html/2607.00275#S4.F2),E\-FLoPSachieves a higherR2R^\{2\}on unseen test data thanFedIter\-HTat allN\|θ\|\\scriptstyle\\frac\{N\}\{\|\\theta\|\}\.E\-FLoPSalso exceeds top\-mmmodel of densely trainingFedAvgduring training and the pruned and fine\-tuned model discovered post training\. The Figure[2](https://arxiv.org/html/2607.00275#S4.F2)also shows the stability of across varying learning rates ofθ~\\tilde\{\\theta\}, the raw parameters, andϕ\\phi, the parameters of the gates\. For stable training ofE\-FLoPS, we fixηλ=0\.01\|θ\|\\scriptstyle\\eta\_\{\\lambda\}=\\frac\{0\.01\}\{\|\\theta\|\}, and tune learning rates within the rangesηθ~∈\[0\.05,0\.5\]\\eta\_\{\\tilde\{\\theta\}\}\\in\[0\.05,0\.5\]andηϕ∈\[0\.1,0\.9\]\\eta\_\{\\phi\}\\in\[0\.1,0\.9\]\. Table[1](https://arxiv.org/html/2607.00275#S4.T1)shows the accuracy of sparsity recoveryA​\(θ\)=T​P\(T​P\+F​N\)\\scriptstyle A\(\\theta\)=\\frac\{TP\}\{\(TP\+FN\)\}, the fraction of true features recovered\.E\-FLoPShas higher accuracy than bothFed\-IHTandFedAvgexcept at smallestN\|θ\|\\scriptstyle\\frac\{N\}\{\|\\theta\|\}\.

![Refer to caption](https://arxiv.org/html/2607.00275v1/LRUD_federated_with_centralized_final_test_score.png)

![Refer to caption](https://arxiv.org/html/2607.00275v1/test_score_vs_n_over_theta.png)

![Refer to caption](https://arxiv.org/html/2607.00275v1/test_score_by_lr.png)

![Refer to caption](https://arxiv.org/html/2607.00275v1/test_score_by_gate_lr.png)

Figure 2:The figures show \(a\) mean testR2R^\{2\}and standard deviation forN\|θ\|=0\.64\\scriptstyle\\frac\{N\}\{\|\\theta\|\}=0\.64over 30 runs of all the algorithms, \(b\) testR2R^\{2\}over varyingN\|θ\|\\scriptstyle\\frac\{N\}\{\|\\theta\|\}by changing total number of samples available across all clients,\(c\) testR2R^\{2\}atηϕ=0\.85\\eta\_\{\\phi\}=0\.85for varyingηθ~\\eta\_\{\\tilde\{\\theta\}\}, and \(d\) testR2R^\{2\}atηθ~=0\.25\\eta\_\{\\tilde\{\\theta\}\}=0\.25for varyingηϕ\\eta\_\{\\phi\}\.Table 1:Final support recovery across total sample sizes \(NN\) atd=\|θ\|d=\|\\theta\|showing mean±\\pmstandard deviation ofA​\(θ\)A\(\\theta\)over 30 different run seeds\. The federated method with the best support recovery is indicated in bold in each row, and the centralized method is shown as a reference upper bound\.
### 4\.2Experiments on Real Data

#### 4\.2\.1Image Classification

We consider publicly available*MNIST*data\[[12](https://arxiv.org/html/2607.00275#bib.bib33)\]: A multi\-class classification dataset of handwritten digits with28×2828\\times 28gray\-scale pixel values for each\. We use a CNN with two5×55\\times 5convolutional layers of 6 and 16 channels, each followed by2×22\\times 2max pooling, and three fully connected layers with widths 120, 84, and 10 for this classification task resulting in a non\-linear model with 44,426 parameters\. The targeted density is2\.5%2\.5\\%\( sparsity of97\.5%97\.5\\%\)\.

The data is downsampled per class at a ratio of∼\\sim0\.02, resulting in 1184 training samples and 10000 test samples\. The data is distributed to clients using DPP for label skew and only a fraction of clients\(0\.6\) are randomly sampled to train in each round\. The affine shifts for further distributional heterogeneity are not performed in this case\. For evaluation, we use cross\-entropy \(CE\) loss with classification accuracy\. Figure[3](https://arxiv.org/html/2607.00275#S4.F3)shows thatE\-FLoPSandFedIter\-HThave similar test\-time performance and communication efficiency\. However,E\-FLoPSshows a greater reduction of Multiply Accumulate Operations \(MACs\) leading to computational efficiency at inference time\.

![Refer to caption](https://arxiv.org/html/2607.00275v1/cnn_federated_with_centralized_final_test_score.png)

![Refer to caption](https://arxiv.org/html/2607.00275v1/communication_compression_bars.png)

![Refer to caption](https://arxiv.org/html/2607.00275v1/total_communication_bytes_bars.png)

![Refer to caption](https://arxiv.org/html/2607.00275v1/final_structured_compact_macs_bars.png)

Figure 3:The figures show: a\) Test accuracy over epochs at97\.5%97\.5\\%sparsity atNθ∼0\.02\\scriptstyle\\frac\{N\}\{\\theta\}\\sim 0\.02, b\) Communication compression compared to dense training, c\) Total communication bytes between server and clients in training and fine\-tuning, and d\) MACs reduction\.
#### 4\.2\.2Leukemia Classification

We use the publicly available Golub leukemia gene\-expression dataset from the study on molecular classification of acute leukemia types, acute lymphoblastic leukemia \(ALL\) and acute myeloid leukemia \(AML\), using gene expression\[[8](https://arxiv.org/html/2607.00275#bib.bib57)\]\. The data consist of 72 patient samples with 3571 features after initial preprocessing\. Each feature represents the measured expression level of a gene/probe in a patient sample; larger values indicate higher abundance of the corresponding gene transcript\. These gene\-expression profiles are used to classify samples as ALL or AML\. This dataset combines a very small sample\-to\-dimension ratio,N\|θ\|=0\.02\\scriptstyle\\frac\{N\}\{\|\\theta\|\}=0\.02, with dense, high\-dimensional gene\-expression features\. It also represents a realistic setting in which data privacy may be a central concern, since the samples correspond to patient\-level biomedical measurements\.

We randomly select 15 samples for the test set and distribute the remaining samples across 5 clients using DPP at a fixed split seed\. We train a softmax classifier at a target density of0\.1%0\.1\\%\(99\.9% sparsity\), with all clients participating\. We do not synthetically alter client data using affine shifts\. Figure[4](https://arxiv.org/html/2607.00275#S4.F4)shows thatE\-FLoPSoutperforms the other variants over 30 runs with different seeds for run\-time stochastic processes\. We confirmed that the difference betweenE\-FLoPSandFed\-IHTis significant using a pairedtt\-test, with att\-statistic of3\.793\.79and app\-value of7×10−47\\times 10^\{\-4\}\. We also performed experiments across 30 train/test splits generated with different split seeds to ensure that the results were not driven by a particularly favorable random partition of the data\. Across these splits,E\-FLoPSagain outperformsFed\-IHTandFedAvg\. Table[\[6\]](https://arxiv.org/html/2607.00275#footnote6)presents the most frequently selected genes across different run seeds\. The most frequently selected genes byE\-FLoPS, cystatin C and myeloperoxidase, are supported by prior experimental research on their roles in promoting metastasis and cancer spread\[[13](https://arxiv.org/html/2607.00275#bib.bib58)\]\.

![Refer to caption](https://arxiv.org/html/2607.00275v1/luk_federated_with_centralized_final_test_score.png)

![Refer to caption](https://arxiv.org/html/2607.00275v1/bysplit.png)

Figure 4:The figure shows a\) Test accuracy over epochs at99\.9%99\.9\\%sparsity atN\|θ\|∼0\.02\\scriptstyle\\frac\{N\}\{\|\\theta\|\}\\sim 0\.02at a fixed data/split seed and 30 different run seeds and b\) Test accuracy over epochs at99\.9%99\.9\\%sparsity atN\|θ\|∼0\.02\\scriptstyle\\frac\{N\}\{\|\\theta\|\}\\sim 0\.02at a fixed run seed and 30 different data/split seeds\.Table 2:The table shows the top genes selected by each method\. Values in parentheses indicate the proportion of runs in which a particular gene is selected\[6\]\[6\]\[6\]The gene "major histocompatibility complex, class I, B" is represented as "histocompatibility complex" for legibility in the Table[\[6\]](https://arxiv.org/html/2607.00275#footnote6)\.\.

## 5Conclusion

Prior work has shown that sparse models can reduce communication overhead, memory footprint, and possibly inference cost, while also improving generalization\. However, the challenge of sparse recovery in\(d≫N\)\(d\\gg N\)regime under heterogeneous FL conditions remains underexplored\. We study entropy regularization of probabilistic gates for sparse FL under anL0L\_\{0\}density constraint\. The proposed federated averaging algorithm,E\-FLoPS, extends sparse federated optimization with a mechanism for uncertainty\-driven exploration of sparse parameter configurations\. This is particularly relevant in\(d≫N\)\(d\\gg N\)regime, where optimization may otherwise commit early to suboptimal sparse configurations with poor generalization\.

We formulate a differentiable, entropy\-regularized,L0L\_\{0\}\-constrained objective using Hard Concrete gates, and show how to optimize it in a federated setting using reparameterized gradients at clients and aggregation of sparse updates\. The resulting method helps achieve a user\-defined target sparsity while improving exploration of the parameter space during learning\.

Experiments on synthetic linear regression, CNN for image classification and softmax classifier for leukemia classification show thatE\-FLoPSconsistently improves sparsity recovery and statistical performance at the target sparsity compared with iterative hard\-thresholding\-based sparse federated training,Fed\-IHT, and post\-training pruning of denseFedAvg\. These gains are achieved with communication efficiency comparable toFed\-IHT\.

We show thatE\-FLoPSis effective for both linear and non\-linear models at high sparsity levels\. Understanding how entropy regularization affects sparse recovery differently in data\-scarce versus overparameterized regimes is an interesting direction for future work, along with extending the method to structured sparsity\.

## Acknowledgments

This work has been supported by FAST, the Finnish Software Engineering Doctoral Research Network, funded by the Ministry of Education and Culture in Finland\.

## Declaration on Generative AI

During the preparation of this work, the author\(s\) used Grammarly in order to: Grammar and spelling check in Overleaf\. After using these tool\(s\)/service\(s\), the author\(s\) reviewed and edited the content as needed and take\(s\) full responsibility for the publication’s content\.

## References

- \[1\]\(2024\)Sparse bayesian networks: efficient uncertainty quantification in medical image analysis\.InInternational Conference on Medical Image Computing and Computer\-Assisted Intervention,pp\. 675–684\.Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p6.4)\.
- \[2\]S\. An, N\. Jammalamadaka, and E\. Chong\(2023\)Maximum entropy information bottleneck for uncertainty\-aware stochastic embedding\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,pp\. 3809–3818\.Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p6.4)\.
- \[3\]Y\. Bao, M\. Crawshaw, S\. Luo, and M\. Liu\(2022\)Fast composite optimization and statistical recovery in federated learning\.Ininternational conference on machine learning,pp\. 1508–1536\.Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p5.3)\.
- \[4\]D\. Bertsimas, J\. Pauphilet, and B\. V\. Parys\(2020\)Sparse regression: scalable algorithms and empirical performance\.Statistical Science35\(4\),pp\. pp\. 555–578\.External Links:ISSN 08834237, 21688745,[Link](https://www.jstor.org/stable/26997931)Cited by:[§4\.1](https://arxiv.org/html/2607.00275#S4.SS1.p1.18)\.
- \[5\]C\. P\. Burgess, I\. Higgins, A\. Pal, L\. Matthey, N\. Watters, G\. Desjardins, and A\. Lerchner\(2018\)Understanding disentangling inβ\\beta\-vae\.arXiv preprint arXiv:1804\.03599\.Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p6.4)\.
- \[6\]A\. K\. Fletcher, S\. Rangan, and V\. K\. Goyal\(2009\)Necessary and sufficient conditions for sparsity pattern recovery\.IEEE Transactions on Information Theory55\(12\),pp\. 5758–5772\.Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p5.3)\.
- \[7\]J\. Gallego\-Posada, J\. Ramirez, A\. Erraqabi, Y\. Bengio, and S\. Lacoste\-Julien\(2022\)Controlled sparsity via constrained optimization or: how i learned to stop tuning penalties and love constraints\.Advances in Neural Information Processing Systems35,pp\. 1253–1266\.Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p4.7),[§1](https://arxiv.org/html/2607.00275#S1.p6.4),[§2](https://arxiv.org/html/2607.00275#S2.p1.11),[§2](https://arxiv.org/html/2607.00275#S2.p4.13),[§3](https://arxiv.org/html/2607.00275#S3.p2.8),[§3](https://arxiv.org/html/2607.00275#S3.p4.17)\.
- \[8\]T\. R\. Golub, D\. K\. Slonim, P\. Tamayo, C\. Huard, M\. Gaasenbeek, J\. P\. Mesirov, H\. Coller, M\. L\. Loh, J\. R\. Downing, M\. A\. Caligiuri,et al\.\(1999\)Molecular classification of cancer: class discovery and class prediction by gene expression monitoring\.science286\(5439\),pp\. 531–537\.Cited by:[§4\.2\.2](https://arxiv.org/html/2607.00275#S4.SS2.SSS2.p1.1)\.
- \[9\]R\. A\. Horn\(1990\)The hadamard product\.InProc\. symp\. appl\. math,Vol\.40,pp\. 87–169\.Cited by:[footnote \[2\]](https://arxiv.org/html/2607.00275#footnote2)\.
- \[10\]K\. H\. K\. Huthasana, A\. Olama, and A\. Lundell\(2025\)Federated learning with l0 constraint via probabilistic gates for sparsity\.External Links:2512\.23071,[Link](https://arxiv.org/abs/2512.23071)Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p4.7),[§3](https://arxiv.org/html/2607.00275#S3.p1.4),[§3](https://arxiv.org/html/2607.00275#S3.p4.17)\.
- \[11\]P\. Kairouz, H\. B\. McMahan, B\. Avent, A\. Bellet, M\. Bennis, A\. N\. Bhagoji, K\. Bonawitz, Z\. Charles, G\. Cormode, R\. Cummings,et al\.\(2021\)Advances and open problems in federated learning\.Foundations and trends® in machine learning14\(1–2\),pp\. 1–210\.Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p1.1)\.
- \[12\]Y\. LeCun\(1998\)The mnist database of handwritten digits\.http://yann\. lecun\. com/exdb/mnist/\.Cited by:[§4\.2\.1](https://arxiv.org/html/2607.00275#S4.SS2.SSS1.p1.5)\.
- \[13\]Y\. Liang, C\. Liu, X\. Luan, K\. Leung, T\. Chan, Z\. Xu, and H\. Zhang\(2013\)Sparse logistic regression with a l1/2 penalty for gene selection in cancer classification\.BMC bioinformatics14\(1\),pp\. 198\.Cited by:[§4\.2\.2](https://arxiv.org/html/2607.00275#S4.SS2.SSS2.p2.6)\.
- \[14\]C\. Louizos, M\. Welling, and D\. P\. Kingma\(2017\)Learning sparse neural networks throughL​\_​0L\\\_0regularization\.arXiv preprint arXiv:1712\.01312\.Cited by:[Appendix A](https://arxiv.org/html/2607.00275#A1.p1.1),[Appendix A](https://arxiv.org/html/2607.00275#A1.p2.2),[§1](https://arxiv.org/html/2607.00275#S1.p4.7),[§1](https://arxiv.org/html/2607.00275#S1.p6.4),[§2](https://arxiv.org/html/2607.00275#S2.p1.11),[§2](https://arxiv.org/html/2607.00275#S2.p4.13),[§2](https://arxiv.org/html/2607.00275#S2.p4.5),[§2](https://arxiv.org/html/2607.00275#S2.p4.7),[§3](https://arxiv.org/html/2607.00275#S3.p3.17)\.
- \[15\]C\. J\. Maddison, A\. Mnih, and Y\. W\. Teh\(2016\)The concrete distribution: a continuous relaxation of discrete random variables\.arXiv preprint arXiv:1611\.00712\.Cited by:[§2](https://arxiv.org/html/2607.00275#S2.p4.5)\.
- \[16\]B\. McMahan, E\. Moore, D\. Ramage, S\. Hampson, and B\. A\. y Arcas\(2017\)Communication\-efficient learning of deep networks from decentralized data\.InArtificial intelligence and statistics,pp\. 1273–1282\.Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p1.1),[§1](https://arxiv.org/html/2607.00275#S1.p6.4),[§3](https://arxiv.org/html/2607.00275#S3.p1.6)\.
- \[17\]R\. Ranganath\(2017\)Black box variational inference: scalable, generic bayesian computation and its applications\.Ph\.D\. Thesis,Princeton University\.Cited by:[§2](https://arxiv.org/html/2607.00275#S2.p1.11)\.
- \[18\]G\. Reeves and M\. C\. Gastpar\(2013\)Approximate sparsity pattern recovery: information\-theoretic lower bounds\.IEEE Transactions on Information Theory59\(6\),pp\. 3451–3465\.Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p5.3)\.
- \[19\]G\. Reeves and M\. Gastpar\(2012\)The sampling rate\-distortion tradeoff for sparsity pattern recovery in compressed sensing\.IEEE Transactions on Information Theory58\(5\),pp\. 3065–3092\.Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p5.3)\.
- \[20\]A\. Reisizadeh, F\. Farnia, R\. Pedarsani, and A\. Jadbabaie\(2020\)Robust federated learning: the case of affine distribution shifts\.Advances in neural information processing systems33,pp\. 21554–21565\.Cited by:[§4\.1](https://arxiv.org/html/2607.00275#S4.SS1.p2.5)\.
- \[21\]D\. Solans, M\. Heikkila, A\. Vitaletti, N\. Kourtellis, A\. Anagnostopoulos, I\. Chatzigiannakis,et al\.\(2024\)Non\-iid data in federated learning: a survey with taxonomy, metrics, methods, frameworks and future directions\.arXiv preprint arXiv:2411\.12377\.Cited by:[§4\.1](https://arxiv.org/html/2607.00275#S4.SS1.p2.5)\.
- \[22\]T\. J\. Thomas and J\. S\. Rani\(2020\)Recovery from compressed measurements using sparsity independent regularized pursuit\.Signal Processing172,pp\. 107508\.Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p5.3)\.
- \[23\]R\. Tibshirani\(1996\)Regression shrinkage and selection via the lasso\.Journal of the Royal Statistical Society Series B: Statistical Methodology58\(1\),pp\. 267–288\.Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p2.1)\.
- \[24\]Q\. Tong, G\. Liang, J\. Ding, T\. Zhu, M\. Pan, and J\. Bi\(2022\)Federated optimization of l0\-norm regularized sparse learning\.Algorithms15\(9\),pp\. 319\.Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p6.4),[§4](https://arxiv.org/html/2607.00275#S4.p1.2)\.
- \[25\]J\. Wang, Z\. Charles, Z\. Xu, G\. Joshi, H\. B\. McMahan, B\. A\. y Arcas, M\. Al\-Shedivat, G\. Andrew, S\. Avestimehr, K\. Daly, D\. Data, S\. Diggavi, H\. Eichner, A\. Gadhikar, Z\. Garrett, A\. M\. Girgis, F\. Hanzely, A\. Hard, C\. He, S\. Horvath, Z\. Huo, A\. Ingerman, M\. Jaggi, T\. Javidi, P\. Kairouz, S\. Kale, S\. P\. Karimireddy, J\. Konecny, S\. Koyejo, T\. Li, L\. Liu, M\. Mohri, H\. Qi, S\. J\. Reddi, P\. Richtarik, K\. Singhal, V\. Smith, M\. Soltanolkotabi, W\. Song, A\. T\. Suresh, S\. U\. Stich, A\. Talwalkar, H\. Wang, B\. Woodworth, S\. Wu, F\. X\. Yu, H\. Yuan, M\. Zaheer, M\. Zhang, T\. Zhang, C\. Zheng, C\. Zhu, and W\. Zhu\(2021\)A field guide to federated optimization\.External Links:2107\.06917,[Link](https://arxiv.org/abs/2107.06917)Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p2.1)\.
- \[26\]X\. Zhang, W\. Li, Y\. Shao, and Y\. Li\(2023\)Federated learning via variational bayesian inference: personalization, sparsity and clustering\.arXiv preprint arXiv:2303\.04345\.Cited by:[§1](https://arxiv.org/html/2607.00275#S1.p6.4)\.

## Appendix AKL\-Divergence for Hard Concrete Gates

The stochastic gates follow the Hard Concrete distribution\[[14](https://arxiv.org/html/2607.00275#bib.bib12)\], obtained by stretching a Binary Concrete random variables∈\(0,1\)s\\in\(0,1\):

s¯=s​\(ζ−γ\)\+γ,z=min⁡\(1,max⁡\(0,s¯\)\),\\bar\{s\}=s\(\\zeta\-\\gamma\)\+\\gamma,\\quad z=\\min\(1,\\max\(0,\\bar\{s\}\)\),\(12\)withγ<0\\gamma<0andζ\>1\\zeta\>1\. This induces a mixed distribution overzzwith point masses at0and11, and a continuous component on\(0,1\)\(0,1\):

q​\(z∣ϕ\)=Qs¯​\(0\)​δ​\(z\)\+\(1−Qs¯​\(1\)\)​δ​\(z−1\)\+\(Qs¯​\(1\)−Qs¯​\(0\)\)​qs¯​\(z∣s¯∈\(0,1\)\),\\displaystyle q\(z\\mid\\phi\)=Q\_\{\\bar\{s\}\}\(0\)\\,\\delta\(z\)\+\\left\(1\-Q\_\{\\bar\{s\}\}\(1\)\\right\)\\,\\delta\(z\-1\)\+\\left\(Q\_\{\\bar\{s\}\}\(1\)\-Q\_\{\\bar\{s\}\}\(0\)\\right\)q\_\{\\bar\{s\}\}\(z\\mid\\bar\{s\}\\in\(0,1\)\),\(13\)where,Qs¯\(s¯∣ϕ\)=Qs\(s¯−γζ−γ\|ϕ\)and,Qs\(s∣ϕ\)=σ\(β′\(logs−log\(1−s\)\)−ϕ\)\.\\text\{where\},\\quad Q\_\{\\bar\{s\}\}\(\\bar\{s\}\\mid\\phi\)=Q\_\{s\}\\left\(\\frac\{\\bar\{s\}\-\\gamma\}\{\\zeta\-\\gamma\}\\,\\middle\|\\,\\phi\\right\)\\quad\\text\{and,\}\\quad Q\_\{s\}\(s\\mid\\phi\)=\\sigma\\left\(\\beta^\{\\prime\}\(\\log s\-\\log\(1\-s\)\)\-\\phi\\right\)\.\(14\)we define a hard concrete priorp​\(z\)p\(z\)with the same support and then the KL\-divergence decomposes as:

KL​\(q​\(z\)∥p​\(z\)\)\\displaystyle\\mathrm\{KL\}\(q\(z\)\\\|p\(z\)\)=Qs¯​\(0\)​log⁡Qs¯​\(0\)Ps¯​\(0\)\+\(1−Qs¯​\(1\)\)​log⁡1−Qs¯​\(1\)1−Ps¯​\(1\)\\displaystyle=Q\_\{\\bar\{s\}\}\(0\)\\log\\frac\{Q\_\{\\bar\{s\}\}\(0\)\}\{P\_\{\\bar\{s\}\}\(0\)\}\+\\left\(1\-Q\_\{\\bar\{s\}\}\(1\)\\right\)\\log\\frac\{1\-Q\_\{\\bar\{s\}\}\(1\)\}\{1\-P\_\{\\bar\{s\}\}\(1\)\}\(15\)\+\(Qs¯​\(1\)−Qs¯​\(0\)\)​𝔼qs¯​\(z\|s¯∈\(0,1\)\)​\[log⁡qs¯​\(z\)−log⁡ps¯​\(z\)\]\.\\displaystyle\\quad\+\\left\(Q\_\{\\bar\{s\}\}\(1\)\-Q\_\{\\bar\{s\}\}\(0\)\\right\)\\mathbb\{E\}\_\{q\_\{\\bar\{s\}\}\(z\|\\bar\{s\}\\in\(0,1\)\)\}\\big\[\\log q\_\{\\bar\{s\}\}\(z\)\-\\log p\_\{\\bar\{s\}\}\(z\)\\big\]\.\(16\)
The first two terms correspond to masses atz=0z=0andz=1z=1, and the last term accounts for the continuous component\. In practice, this divergence is evaluated either using the closed\-form expressions if available or Monte Carlo estimation\. We used the closed form expressions provided atLouizoset al\.\[[14](https://arxiv.org/html/2607.00275#bib.bib12)\]\.