Entropy-Regularized Probabilistic Gates for Sparse Model Discovery in Scarce-Data Federated Learning Krishna Harsha Kovelakuntla Huthasana Alireza Olama Andreas Lundell
arXiv:2607.00275v1 [cs.LG] 30 Jun 2026
Department of Engineering and Information Technology, Åbo Akademi University {kkovelak, alireza.olama, andreas.lundell}@abo.fi July 2, 2026 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 ) 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 an L0 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 , L0 Constraint
1
Introduction
Federated Learning (FL) algorithms operate in a distributed machine learning (ML) setting in which multiple clients collaborate to train [16]. 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]. 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. 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] and enhance computational and communication efficiency in FL, thereby posing an additional challenge of discovering sparse models [25]. A common approach to inducing sparsity relies on L1 and L2 norms[1] for regularization, which depend directly on parameter magnitudes and offer varying levels of shrinkage. In contrast, using a magnitudeindependent L0 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 ρ. The Lagrangian for the L0 density-constrained optimization problem in FL can be defined as: L(θ, λ) =
C X nc c=1
[1] For θ ∈ Rd , L
1 norm is ∥θ∥1 =
N
Pd i=1
L(c) (θ) + λ (∥θ∥0 − ρ|θ|) ,
∥θ∥0 =
|θ| X j=1
|θi | and L2 norm is ∥θ∥2 =
1
P
d θ2 i=1 i
1/2
.
I[θj ̸= 0],
(1)
Server
Client 1
Client 2
Client 3
Client 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.
where, L(c) (θ) denotes the normalized loss at client c, defined as: L(c) (θ) =
nc 1 X (c) (c) ℓ h(xi ; θ), yi . nc i=1
(2)
(c) Here, C is number of clients in FL, each holding a local dataset (D(c) )C , Y (c) ), where X (c) ∈ c=1 = (X P C in Rnc ×in , Y (c) ∈ Rnc ×out , and → Rout and a loss c=1 nc = N . We assume a model h(x; θ) : R in out d function ℓ(h(x; θ), y), where x ∈ R , y ∈ R , and θ ∈ R . The above formulation with Lagrange multiplier λ leads to a min-max optimization problem that can, in principle, be solved using gradient descent-ascent. However, the non-differentiability of the L0 pseudo-norm complicates optimization within standard gradient-based learning frameworks. Louizos et al. [14] introduce a reparameterization using stochastic gates z ∈ Rp with θ = θ̃⊙z [2] as the effective parameters of the model. By assuming a hard Concrete distribution over z (an approximation of the P Bernoulli distribution), the L0 norm is approximated by the expected number of active gates, i.e., j Eq(z) [zj ], enabling gradient-based optimization of the L0 -regularized objective. Gallego-Posada et al. [7] extended this framework by using an L0 density constraint and solving the resulting min–max problem, thereby enabling the user to define the target sparsity in a centralized setting. Huthasana et al. [10] 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 ) [6, 18, 19, 22], where the optimization is unstable, resulting in multiple solutions of competing parameter configurations with poor sparsity recovery and generalization behaviour to unseen test data. Bao et al. [3] 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 size N relative to the number of parameters d = |θ| similar to a centralized setting, which is extremely challenging under data and client participation heterogeneity in FL. Louizos et al. [14] points to penalizing relative entropy or the Kullback-Leibler divergence KL(q(z)∥p(z)) in the optimization where p(z) is the prior and 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, 5]. It is also used in Bayesian inference in centralized and FL settings [1, 26]. In their approach, the hard concrete gate distributions q(z) are not directly [2] ⊙ denotes the elementwise product [9].
2
dependent on the data and can be sampled independently of it post-training. In this work, the formulations at Gallego-Posada et al. [7], Louizos et al. [14] 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] and post-training pruning of a dense model using the classic federated averaging algorithm [16]. The remainder of this paper is organized as follows. We first present the entropy-regularized L0 constrained formulation for federated optimization, then the proposed distributed algorithm, followed by experiments in heterogeneous FL settings, and conclude.
2
Formulation
Assuming a model h(x; θ) : Rin → Rout and a loss function ℓ(h(x; θ), y), where x ∈ Rin , y ∈ Rout , and θ ∈ Rd , consider a centralized dataset D = (X, Y ) with X ∈ RN ×in and Y ∈ RN ×out . Using [7, 14], the min–max objective can be defined using the expectation of the loss with respect to the distribution of P gates, and the L0 pseudo-norm approximated by the expected number of active gates, i.e., j Eq(z|ϕ) [zj ] as shown in eq. 3. Since each zj is a deterministic transformation of parameter-free noise, the expectation can be optimized using Monte Carlo sampling and reparameterized gradients [17, ch. 3.3.3]. " # |θ| R N X 1 X 1 X ℓ h(xi ; θ̃ ⊙ z (r) ), yi + λ Eq(z|ϕ) [zj ] − ρ|θ| . (3) L̂(θ̃, ϕ, λ) = R r=1 N i=1 j=1 (c) In a federated learning (FL) setting with C clients holding datasets (D(c) )C , Y (c) ), we c=1 = (X in out consider a reparameterized linear model h(x; θ̃ ⊙ z) : R → R and a loss function ℓ(h(x; θ̃ ⊙ z), y). PC Here, X (c) ∈ Rnc ×in , Y (c) ∈ Rnc ×out , c=1 nc = N , x ∈ Rin , y ∈ Rout , and θ = θ̃ ⊙ z ∈ Rin . The gate parameters are defined as ϕ = log α ∈ Rin . The Lagrangian corresponding to the entropy regularized L0 density-constrained optimization problem is: |θ| |θ| C X X X nc (c) L (θ̃, ϕ) + λ Eq(z|ϕ) [zj ] − ρ|θ| + T L̂(θ̃, ϕ, λ) = KL(q(zj |ϕ)||p(zj |ϕinit )). (4) N c=1 j=1 j=1
Here, p(z|ϕinit ) is a prior and is also a hard concrete distribution, and L(c) (θ̃, ϕ) denotes the Monte Carlo estimate of the normalized loss at client c, defined as: L(c) (θ̃, ϕ) =
nc R 1 X 1 X (c) (c) ℓ h(xi ; θ̃ ⊙ z (r) ), yi . R r=1 nc i=1
(5)
The stochastic gates z are sampled using the Hard Concrete distribution, applying a hard-sigmoid transformation to a stretched Binary Concrete random variable [14, 15], defined as: u log 1−u + log α s=σ , u ∼ U(0, 1), β′ s̄ = s(ζ − γ) + γ,
z = min(1, max(0, s̄)).
(6)
The expectation of a gate being active [3] is derived at Louizos et al. [14] using the cumulative distribution function Q(s̄). We introduced T ≥ 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 L̂(θ̃, ϕ, λ). θ̃,ϕ
λ≥0
(7)
The parameters θ̃ and ϕ = log α are jointly optimized using gradient descent with reparameterized gradients. λ is updated via gradient ascent and a restart strategy of resetting its value to 0 as and when the sparsity constraint is satisfied [7]. Since the hard concrete distribution is a continuos approximation of Bernoulli distribution the computation of KL(q(zj )||p(zj )) involves an additional term involving s̄, and thus, z ∈ (0, 1) which can be computed using truncated distribution q(s̄|s̄ ∈ (0, 1)) or a Montecarlo estimate of the same. We used the closed-form expressions presented in the appendix A provided at Louizos et al. [14]. γ [3] ′ Eq(z|ϕ) [zj ] = 1 − Q(s̄j ≤ 0 | ϕj ) = σ log αj − β log − ζ
3
3
Algorithm
We use the notations LCon (ϕ) and LKL (ϕ) for L0 density constraint and the KL(q(z)||p(z)) in (4) [10]. The Lagrangian can then be written as: L̂(θ̃, ϕ, λ) =
C X nc
N c=1
L(c) (θ̃, ϕ) + λ LCon (ϕ) + T LKL (ϕ).
(8)
McMahan et al. [16] 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 across C clients each holding a local dataset 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 iteration t, each client performs the following updates: θ̃ct+1 = θ̃ct − ηθ̃ ∇θ̃ L(c) (θ̃t , ϕt ), (c) (c) t (c) t t t t ϕt+1 = ϕ − η ∇ L ( θ̃ , ϕ ) + λ ∇ L (ϕ ) + T ∇ L (ϕ ) , ϕ ϕ ϕ Con ϕ KL c c (c)
λt+1 = λtc + ηλ LCon (ϕt ). c
(9) (10) (11)
At each round, a fraction γc of clients is selected uniformly at random, resulting in K = ⌊γc C⌋ participating clients. In practice, a few mini-batches of uniform size B sampled iteratively at each client to locally update the model 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 keeping T constant or decaying it according to a predetermined schedule. The Lagrange parameter λ is reset to zero when the constraint is satisfied [7]. The aggregation weights wk 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 an L0 constraint using probabilistic gates. We refer to this variant of FedAvg as E-FLoPS, where E stands for entropy regularization. The learning rates ηθ̃ and ηϕ needs to be appropriately tuned. The learning rate ηλ for the Lagrange parameter updates is set in the order of 1/|θ|. The temperature T ∈ (0, ∞) is also initialized at 1/|θ| and increased if needed to encourage uncertainty exploration as the total sample size N decreases. The gate parameters ϕ = log α are initialized from a normal distribution with mean log ρinit − log(1 − ρinit ) and variance 0.01, where ρinit controls the initial density. The target density is denoted by ρtarg = ρ. The Hard Concrete distribution parameters are set to γ = −0.1, ζ = 1.1, and β ′ is recommended to be set at 0.66, following Louizos et al. [14]. The Lagrange multiplier is initialized as λ = 0. Client participation variability is simulated by randomly selecting a fraction of clients at each training round. E-FLoPS achieves test-time sparsity by using deterministic gates ẑ sampled without noise or smoothing [7]. For exact sparsity, the raw parameters θ̃ for the top-m indices of effective parameters θ are pruned, where m = ⌊ρtarg · |θ|⌋. This pruning mechanism can be applied from a pre-defined threshold of epochs, referred to as the prune start epoch, which is set to 0 for E-FLoPS to improve communication efficiency via serialization for message compression [10]. At this stage, the serialized mean gate values, from repeated sampling, corresponding to top-m indices of gates z are retained, and the rest of the gates are replaced avg with their average, denoted by z−m . Only the pruned θ̃ and z, along with their non-zero indices and avg the scalar z−m , are communicated. Upon reception, parameters are reconstructed via θ = θ̃ ⊙ z and z aggregated, and ϕ is recovered from z using ϕ=β ′ log( 1−z ), ignoring the noise component of the gates. The communication cost reduction by this approach in FL is significant for small ρtarg with minimal meta data overhead.
4
Experiments
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 [4] ⊘ represents an element wise division.
4
Algorithm 1 FedAvg variant E-FLoPS. E and B denote the number of epochs and the mini-batch size. 1: Initialization: (θ̃ (0) , ϕ(0) ), compute θ (0) = θ̃ (0) ⊙ z (0) 2: for epoch b = 1 to E do 3: Sample client subset St 4: 5: 6: 7:
for each client k ∈ St do (k, θ̃k , zk , λk ) ← ClientCompute(k, θ̃(t) , z (t) , λ(t) ) end for Server aggregation: X wk θ̃k θ= k∈St
z=
X
wk zk
k∈St
8:
Reconstruct θ̃ = θ ⊘ z [4] and recover ϕ from z using ϕ = β ′ log
z 1−z
Perform server-side updates on (θ̃, ϕ, λ) Either decay or keep T constant if b > prune start then prune θ̃ using top-m θ top-m indices of θ avg prune z with tail statistic z−m Serialize and communicate 15: end if 16: end for 17: ClientCompute(k, θ̃, z, λ): z 18: Recover ϕk from z using ϕ=β ′ log( 1−z ) 19: Sample mini-batches b ∼ D(k) of size B 20: Perform n(B) local SGD steps on (θ̃k , ϕk ) and ascent for λk 21: return Communicate pruned and serialized (θ̃k , ϕk ) and λk 9: 10: 11: 12: 13: 14:
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 by Tong et al. [24], where a hard-thresholding operation is used to retain only top-m parameters 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. For FedAvg the training time performance on test data is evaluated using the top-m parameters 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 variant E-FLoPS enables 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 |θ|), and FedIter-HT : epochs × 4 · (ρtarg |θ|).
Thus, the communication cost of E-FLoPS is more than FedIter-HT, but significantly lower than dense training via FedAvg. For large model sizes |θ| and small target densities ρtarg , the gap between FedAvg and E-FLoPS becomes substantial, while small with that of FedIter-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).
5
N =0.64 over 30 runs of all the algorithms, Figure 2: The figures show (a) mean test R2 and standard deviation for |θ| 2 N (b) test R over varying |θ| by changing total number of samples available across all clients,(c) test R2 at ηϕ = 0.85 for varying ηθ̃ , and (d) test R2 at ηθ̃ = 0.25 for varying ηϕ .
4.1
Experiments on Synthetic Data
We generated synthetic data for sparse linear regression following the procedure from Bertsimas et al. [4] N for a range of |θ| . For a parameter vector dimension d = |θ|, each row xi ∈ Rd of X ∈ RN ×d is sampled from a zero-mean Gaussian distribution with covariance matrix Σ. We use a Toeplitz covariance structure |i−j| defined as Σij = ρcor , i, j = 1, . . . , d. An m-sparse coefficient vector wtrue ∈ Rd is constructed, where m = ⌊ρ · d⌋. A subset of indices Sm ⊆ {1, . . . , d} is selected uniformly at random, and coefficients are assigned as (wtrue )j ∼ Unif{−1, 1} for j ∈ Sm and (wtrue )j = 0 otherwise. The response vector is generated as y = Xwtrue + ε, where ε ∼ N (0, σ 2 IN ). The signal-to-noise ratio (SNR) is defined as ∥Xwtrue ∥2 ∥ true 2 , and the noise level is set to σ= ∥Xw √ √ 2. SNR= For |θ|=1000, we generate a number of total ∥ε∥2 SNR N 2 N training samples N for varying ratios of |θ| with 5000 samples 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, 21]. A fraction of clients (0.6) is randomly sampled in each epoch to introduce participation heterogeneity [20]. We evaluate performance using mean squared error (MSE) and R2 . Experiments are conducted at a correlation level of ρcor = 0.2 and an SNR of 20, with a true sparsity level of 95% or target density of 5%. N As shown in Figure 2, E-FLoPS achieves a higher R2 on unseen test data than FedIter-HT at all |θ| . E-FLoPS also exceeds top-m model of densely training FedAvg during training and the pruned and finetuned model discovered post training. The Figure 2 also shows the stability of across varying learning rates of θ̃, the raw parameters, and ϕ, the parameters of the gates. For stable training of E-FLoPS, we fix ηλ = 0.01 , and tune learning rates within the ranges ηθ̃ ∈ [0.05, 0.5] and ηϕ ∈ [0.1, 0.9]. Table 1 shows the |θ| P accuracy of sparsity recovery A(θ)= (T PT+F , the fraction of true features recovered. E-FLoPS has higher N) N accuracy than both Fed-IHT and FedAvg except at smallest |θ| .
4.2
Experiments on Real Data
4.2.1
Image Classification
We consider publicly available MNIST data [12]: A multi-class classification dataset of handwritten digits with 28 × 28 gray-scale pixel values for each. We use a CNN with two 5 × 5 convolutional layers of 6 and 16 channels, each followed by 2 × 2 max pooling, and three fully connected layers with widths
6
Table 1: Final support recovery across total sample sizes (N ) at d = |θ| showing mean ± standard deviation of A(θ) 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. N d
E-FLoPS
Fed-IHT
FedAvg
Centralized
0.24 0.44 0.64 1.00 1.36 1.56 2.00
0.140 ± 0.039 0.379 ± 0.082 0.626 ± 0.060 0.828 ± 0.034 0.925 ± 0.026 0.920 ± 0.016 0.941 ± 0.021
0.129 ± 0.039 0.292 ± 0.044 0.490 ± 0.055 0.687 ± 0.072 0.831 ± 0.032 0.874 ± 0.032 0.913 ± 0.028
0.179 ± 0.038 0.427 ± 0.036 0.566 ± 0.038 0.741 ± 0.023 0.803 ± 0.030 0.890 ± 0.110 0.917 ± 0.027
0.195 ± 0.027 0.473 ± 0.027 0.736 ± 0.026 0.840 ± 0.016 0.945 ± 0.010 0.980 ± 0.000 0.987 ± 0.010
Figure 3: The figures show: a) Test accuracy over epochs at 97.5% sparsity at Nθ ∼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.
120, 84, and 10 for this classification task resulting in a non-linear model with 44,426 parameters. The targeted density is 2.5% ( sparsity of 97.5% ). The data is downsampled per class at a ratio of ∼0.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 shows that E-FLoPS and FedIter-HT have similar test-time performance and communication efficiency. However, E-FLoPS shows a greater reduction of Multiply Accumulate Operations (MACs) leading to computational efficiency at inference time. 4.2.2
Leukemia 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]. 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-toN dimension ratio, |θ| =0.02, with dense, high-dimensional gene-expression features. It also represents a
7
N ∼0.02 at a fixed data/split seed Figure 4: The figure shows a) Test accuracy over epochs at 99.9% sparsity at |θ| N ∼0.02 at a fixed run seed and and 30 different run seeds and b) Test accuracy over epochs at 99.9% sparsity at |θ| 30 different data/split seeds.
realistic setting in which data privacy may be a central concern, since the samples correspond to patientlevel 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 of 0.1% (99.9% sparsity), with all clients participating. We do not synthetically alter client data using affine shifts. Figure 4 shows that E-FLoPS outperforms the other variants over 30 runs with different seeds for run-time stochastic processes. We confirmed that the difference between E-FLoPS and Fed-IHT is significant using a paired t-test, with a t-statistic of 3.79 and a p-value of 7 × 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-FLoPS again outperforms Fed-IHT and FedAvg. Table 2 presents the most frequently selected genes across different run seeds. The most frequently selected genes by E-FLoPS, cystatin C and myeloperoxidase, are supported by prior experimental research on their roles in promoting metastasis and cancer spread [13]. 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[5] . E-FLoPS
Fed-IHT
FedAvg
Centralized
cystatin C (0.90) myeloperoxidase (0.80) ferritin, light polypeptide (0.80) interleukin 8 (0.53) defensin, alpha 1 (0.43) azurocidin 1 (0.33) histocompatibility complex (0.23)
ferritin, light polypeptide (0.70) cystatin C (0.67) azurocidin 1 (0.57) histocompatibility complex (0.57) hemoglobin, beta (0.57) interleukin 8 (0.47) myeloperoxidase (0.43)
azurocidin 1 (0.33) interleukin 8 (0.33) ferritin, light polypeptide (0.30) myeloperoxidase (0.27) cystatin C (0.23) interleukin 8 (0.17) elastase 2, neutrophil (0.17)
interleukin 8 (0.73) ferritin, light polypeptide (0.70) myeloperoxidase (0.67) cystatin C (0.63) glycoprotein B (0.53) azurocidin 1 (0.47) histocompatibility complex (0.37)
5
Conclusion
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 ) regime under heterogeneous FL conditions remains underexplored. We study entropy regularization of probabilistic gates for sparse FL under an L0 density constraint. The proposed federated averaging algorithm, E-FLoPS, extends sparse federated optimization with a mechanism for uncertaintydriven exploration of sparse parameter configurations. This is particularly relevant in (d ≫ N ) regime, where optimization may otherwise commit early to suboptimal sparse configurations with poor generalization. We formulate a differentiable, entropy-regularized, L0 -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 that E-FLoPS consistently 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 dense FedAvg. These gains are achieved with communication efficiency comparable to Fed-IHT. [5] The gene "major histocompatibility complex, class I, B" is represented as "histocompatibility complex" for legibility in the Table 2.
8
We show that E-FLoPS is 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] Zeinab Abboud, Herve Lombaert, and Samuel Kadoury. Sparse bayesian networks: efficient uncertainty quantification in medical image analysis. In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 675–684. Springer, 2024. [2] Sungtae An, Nataraj Jammalamadaka, and Eunji Chong. Maximum entropy information bottleneck for uncertainty-aware stochastic embedding. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3809–3818, 2023. [3] Yajie Bao, Michael Crawshaw, Shan Luo, and Mingrui Liu. Fast composite optimization and statistical recovery in federated learning. In international conference on machine learning, pages 1508–1536. PMLR, 2022. [4] Dimitris Bertsimas, Jean Pauphilet, and Bart Van Parys. Sparse regression: Scalable algorithms and empirical performance. Statistical Science, 35(4):pp. 555–578, 2020. ISSN 08834237, 21688745. URL https://www.jstor.org/stable/26997931. [5] Christopher P Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in β-vae. arXiv preprint arXiv:1804.03599, 2018. [6] Alyson K Fletcher, Sundeep Rangan, and Vivek K Goyal. Necessary and sufficient conditions for sparsity pattern recovery. IEEE Transactions on Information Theory, 55(12):5758–5772, 2009. [7] Jose Gallego-Posada, Juan Ramirez, Akram Erraqabi, Yoshua Bengio, and Simon Lacoste-Julien. Controlled sparsity via constrained optimization or: How i learned to stop tuning penalties and love constraints. Advances in Neural Information Processing Systems, 35:1253–1266, 2022. [8] Todd R Golub, Donna K Slonim, Pablo Tamayo, Christine Huard, Michelle Gaasenbeek, Jill P Mesirov, Hilary Coller, Mignon L Loh, James R Downing, Mark A Caligiuri, et al. Molecular classification of cancer: class discovery and class prediction by gene expression monitoring. science, 286(5439):531–537, 1999. [9] Roger A Horn. The hadamard product. In Proc. symp. appl. math, volume 40, pages 87–169, 1990. [10] Krishna Harsha Kovelakuntla Huthasana, Alireza Olama, and Andreas Lundell. Federated learning with l0 constraint via probabilistic gates for sparsity, 2025. URL https://arxiv.org/abs/2512. 23071. [11] Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. Foundations and trends® in machine learning, 14(1–2): 1–210, 2021. 9
[12] Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/exdb/mnist/, 1998. [13] Yong Liang, Cheng Liu, Xin-Ze Luan, Kwong-Sak Leung, Tak-Ming Chan, Zong-Ben Xu, and Hai Zhang. Sparse logistic regression with a l1/2 penalty for gene selection in cancer classification. BMC bioinformatics, 14(1):198, 2013. [14] Christos Louizos, Max Welling, and Diederik P Kingma. Learning sparse neural networks through l_0 regularization. arXiv preprint arXiv:1712.01312, 2017. [15] Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv:1611.00712, 2016. [16] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273–1282. PMLR, 2017. [17] Rajesh Ranganath. Black Box variational inference: Scalable, generic Bayesian computation and its applications. PhD thesis, Princeton University, 2017. [18] Galen Reeves and Michael Gastpar. The sampling rate-distortion tradeoff for sparsity pattern recovery in compressed sensing. IEEE Transactions on Information Theory, 58(5):3065–3092, 2012. [19] Galen Reeves and Michael C Gastpar. Approximate sparsity pattern recovery: Information-theoretic lower bounds. IEEE Transactions on Information Theory, 59(6):3451–3465, 2013. [20] Amirhossein Reisizadeh, Farzan Farnia, Ramtin Pedarsani, and Ali Jadbabaie. Robust federated learning: The case of affine distribution shifts. Advances in neural information processing systems, 33:21554–21565, 2020. [21] David Solans, Mikko Heikkila, Andrea Vitaletti, Nicolas Kourtellis, Aris Anagnostopoulos, Ioannis Chatzigiannakis, et al. Non-iid data in federated learning: A survey with taxonomy, metrics, methods, frameworks and future directions. arXiv preprint arXiv:2411.12377, 2024. [22] Thomas James Thomas and J Sheeba Rani. Recovery from compressed measurements using sparsity independent regularized pursuit. Signal Processing, 172:107508, 2020. [23] Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Society Series B: Statistical Methodology, 58(1):267–288, 1996. [24] Qianqian Tong, Guannan Liang, Jiahao Ding, Tan Zhu, Miao Pan, and Jinbo Bi. Federated optimization of l0-norm regularized sparse learning. Algorithms, 15(9):319, 2022. [25] Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H. Brendan McMahan, Blaise Aguera y Arcas, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, Suhas Diggavi, Hubert Eichner, Advait Gadhikar, Zachary Garrett, Antonious M. Girgis, Filip Hanzely, Andrew Hard, Chaoyang He, Samuel Horvath, Zhouyuan Huo, Alex Ingerman, Martin Jaggi, Tara Javidi, Peter Kairouz, Satyen Kale, Sai Praneeth Karimireddy, Jakub Konecny, Sanmi Koyejo, Tian Li, Luyang Liu, Mehryar Mohri, Hang Qi, Sashank J. Reddi, Peter Richtarik, Karan Singhal, Virginia Smith, Mahdi Soltanolkotabi, Weikang Song, Ananda Theertha Suresh, Sebastian U. Stich, Ameet Talwalkar, Hongyi Wang, Blake Woodworth, Shanshan Wu, Felix X. Yu, Honglin Yuan, Manzil Zaheer, Mi Zhang, Tong Zhang, Chunxiang Zheng, Chen Zhu, and Wennan Zhu. A field guide to federated optimization, 2021. URL https://arxiv.org/abs/2107.06917. [26] Xu Zhang, Wenpeng Li, Yunfeng Shao, and Yinchuan Li. Federated learning via variational bayesian inference: Personalization, sparsity and clustering. arXiv preprint arXiv:2303.04345, 2023.
10
A
KL-Divergence for Hard Concrete Gates
The stochastic gates follow the Hard Concrete distribution [14], obtained by stretching a Binary Concrete random variable s ∈ (0, 1): s̄ = s(ζ − γ) + γ,
z = min(1, max(0, s̄)),
(12)
with γ < 0 and ζ > 1. This induces a mixed distribution over z with point masses at 0 and 1, and a continuous component on (0, 1): q(z | ϕ) = Qs̄ (0) δ(z) + (1 − Qs̄ (1)) δ(z − 1) + (Qs̄ (1) − Qs̄ (0)) qs̄ (z | s̄ ∈ (0, 1)), s̄ − γ ϕ and, Qs (s | ϕ) = σ (β ′ (log s − log(1 − s)) − ϕ) . where, Qs̄ (s̄ | ϕ) = Qs ζ −γ
(13) (14)
we define a hard concrete prior p(z) with the same support and then the KL-divergence decomposes as: Qs̄ (0) 1 − Qs̄ (1) + (1 − Qs̄ (1)) log Ps̄ (0) 1 − Ps̄ (1) + (Qs̄ (1) − Qs̄ (0)) Eqs̄ (z|s̄∈(0,1)) log qs̄ (z) − log ps̄ (z) .
KL(q(z)∥p(z)) = Qs̄ (0) log
(15) (16)
The first two terms correspond to masses at z = 0 and z = 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 at Louizos et al. [14].
11