Entropy Regularization ,  Sparsity ,  Federated Learning ,  Uncertainty ,  Parameter Exploration ,  Probabilistic Gates ,  Entropy Maximization ,  Sparse Federated Learning ,  \(L_{0}\) Constraint

1 Introduction↩︎

Federated Learning (FL) algorithms operate in a distributed machine learning (ML) setting in which multiple clients collaborate to train [1]. 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 [2]. 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 [3] and enhance computational and communication efficiency in FL, thereby posing an additional challenge of discovering sparse models [4].

A common approach to inducing sparsity relies on \(L_1\) and \(L_2\) norms¹ for regularization, which depend directly on parameter magnitudes and offer varying levels of shrinkage. In contrast, using a magnitude-independent \(L_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 the \(L_0\) density-constrained optimization problem in FL can be defined as: \[\label{prb95def} \mathfrak{L}(\theta,\lambda) = \sum_{c=1}^{sec: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],\tag{1}\] where, \(\mathcal{L}^{(c)}(\theta)\) denotes the normalized loss at client \(c\), defined as: \[\mathcal{L}^{(c)}(\theta) = \frac{1}{n_c} \sum_{i=1}^{n_c} \ell\big(h(x_i^{(c)}; \theta), y_i^{(c)}\big).\] Here, \(C\) is number of clients in FL, each holding a local dataset \((D^{(c)})_{c=1}^{sec:C} = (X^{(c)}, Y^{(c)})\), where \(X^{(c)} \in \mathbb{R}^{n_c \times in}\), \(Y^{(c)} \in \mathbb{R}^{n_c\times out}\), and \(\sum_{c=1}^{sec:C} n_c = N\). We assume a model \(h(x;\theta): \mathbb{R}^{in} \rightarrow \mathbb{R}^{out}\) and a loss function \(\ell(h(x;\theta), y)\), where \(x \in \mathbb{R}^{in}\), \(y \in \mathbb{R}^{out}\), and \(\theta \in \mathbb{R}^{d}\). The above formulation with Lagrange multiplier \(\lambda\) leads to a min-max optimization problem that can, in principle, be solved using gradient descent-ascent. However, the non-differentiability of the \(L_0\) pseudo-norm complicates optimization within standard gradient-based learning frameworks.

[5] introduce a reparameterization using stochastic gates \(z \in \mathbb{R}^{p}\) with \(\theta = \tilde{\theta} \odot z\)² as the effective parameters of the model. By assuming a hard Concrete distribution over \(z\) (an approximation of the Bernoulli distribution), the \(L_0\) norm is approximated by the expected number of active gates, i.e., \(\sum_{j} \mathbb{E}_{q(z)}[z_j]\), enabling gradient-based optimization of the \(L_0\)-regularized objective. [7] extended this framework by using an \(L_0\) density constraint and solving the resulting min–max problem, thereby enabling the user to define the target sparsity in a centralized setting. [8] 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 \gg N)\) [9]–[12], where the optimization is unstable, resulting in multiple solutions of competing parameter configurations with poor sparsity recovery and generalization behaviour to unseen test data. [13] 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=|\theta|\) similar to a centralized setting, which is extremely challenging under data and client participation heterogeneity in FL.

[5] points to penalizing relative entropy or the Kullback-Leibler divergence \(\mathrm{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 [14], [15]. It is also used in Bayesian inference in centralized and FL settings [16], [17]. In their approach, the hard concrete gate distributions \(q(z)\) are not directly dependent on the data and can be sampled independently of it post-training. In this work, the formulations at [5], [7] 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 [18] and post-training pruning of a dense model using the classic federated averaging algorithm [1].

The remainder of this paper is organized as follows. We first present the entropy-regularized \(L_0\)-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;\theta): \mathbb{R}^{in} \rightarrow \mathbb{R}^{out}\) and a loss function \(\ell(h(x;\theta), y)\), where \(x \in \mathbb{R}^{in}\), \(y \in \mathbb{R}^{out}\), and \(\theta \in \mathbb{R}^{d}\), consider a centralized dataset \(D = (X, Y)\) with \(X \in \mathbb{R}^{N \times in}\) and \(Y \in \mathbb{R}^{N\times out}\). Using [5], [7], the min–max objective can be defined using the expectation of the loss with respect to the distribution of gates, and the \(L_0\) pseudo-norm approximated by the expected number of active gates, i.e., \(\sum_j \mathbb{E}_{q(z \mid \phi)}[z_j]\) as shown in eq. 2 . Since each \(z_j\) is a deterministic transformation of parameter-free noise, the expectation can be optimized using Monte Carlo sampling and reparameterized gradients [19].

\[\label{lag95den95con951} \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).\tag{2}\]

In a federated learning (FL) setting with \(C\) clients holding datasets \((D^{(c)})_{c=1}^{sec:C} = (X^{(c)}, Y^{(c)})\), we consider a reparameterized linear model \(h(x;\tilde{\theta} \odot z): \mathbb{R}^{in} \rightarrow \mathbb{R}^{out}\) and a loss function \(\ell(h(x;\tilde{\theta} \odot z), y)\). Here, \(X^{(c)} \in \mathbb{R}^{n_c \times in}\), \(Y^{(c)} \in \mathbb{R}^{n_c\times out}\), \(\sum_{c=1}^{sec:C} n_c = N\), \(x \in \mathbb{R}^{in}\), \(y \in \mathbb{R}^{out}\), and \(\theta = \tilde{\theta} \odot z \in \mathbb{R}^{in}\). The gate parameters are defined as \(\phi = \log \alpha \in \mathbb{R}^{in}\).

The Lagrangian corresponding to the entropy regularized \(L_0\) density-constrained optimization problem is: \[\label{prb95def2463} \hat{\mathfrak{L}}(\tilde{\theta}, \phi, \lambda) = \sum_{c=1}^{sec: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})).\tag{3}\] Here, \(p(z|\phi_{init})\) is a prior and is also a hard concrete distribution, and \(\mathcal{L}^{(c)}(\tilde{\theta}, \phi)\) denotes the Monte Carlo estimate of the normalized loss at client \(c\), defined as: \[\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).\] The stochastic gates \(z\) are sampled using the Hard Concrete distribution, applying a hard-sigmoid transformation to a stretched Binary Concrete random variable [5], [20], defined as: \[\begin{align} \label{eq4mainp} s &= \sigma\left( \frac{\log\frac{u}{1-u} + \log \alpha}{\beta'} \right), \quad u \sim \mathcal{U}(0, 1),\notag \\ \bar{s} &= s(\zeta - \gamma) + \gamma, \quad z = \min(1, \max(0, \bar{s})). \end{align}\tag{4}\] The expectation of a gate being active ³ is derived at [5] using the cumulative distribution function \(Q(\bar{s})\). We introduced \(T\geq 0\), which we treat as a constant with or without decay to penalize entropy. The resulting min–max optimization problem is: \[\label{centralFL} \tilde{\theta}^*, \phi^*, \lambda^* = \arg \min_{\tilde{\theta}, \phi} \arg \max_{\lambda \geq 0} \hat{\mathfrak{L}}(\tilde{\theta}, \phi, \lambda).\tag{5}\] The parameters \(\tilde{\theta}\) and \(\phi = \log \alpha\) are jointly optimized using gradient descent with reparameterized gradients. \(\lambda\) 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(z_j)||p(z_j))\) involves an additional term involving \(\bar{s} \text{, and thus, }z\in(0,1)\) which can be computed using truncated distribution \(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 6 provided at [5].

3 Algorithm↩︎

We use the notations \(\mathcal{L}_{\text{Con}}(\phi)\) and \(\mathcal{L}_{\text{KL}}(\phi)\) for \(L_0\) density constraint and the \(KL(q(z)||p(z))\) in 3  [8]. The Lagrangian can then be written as: \[\label{cent95adapt} \hat{\mathfrak{L}}(\tilde{\theta}, \phi, \lambda) = \sum_{c=1}^{sec:C} \frac{n_c}{N} \mathcal{L}^{(c)}(\tilde{\theta}, \phi) + \lambda \, \mathcal{L}_{\text{Con}}(\phi)+T \, \mathcal{L}_{\text{KL}}(\phi).\tag{6}\] [1] 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: \[\begin{align} \tilde{\theta}_{c}^{t+1} &= \tilde{\theta}_{c}^{t} - \eta_{\tilde{\theta}} \nabla_{\tilde{\theta}} \mathcal{L}^{(c)}(\tilde{\theta}^{t}, \phi^{t}), \\ \phi_{c}^{t+1} &= \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), \\ \lambda_{c}^{t+1} &= \lambda_{c}^{t} + \eta_{\lambda} \, \mathcal{L}_{\text{Con}}^{(c)}(\phi^{t}). \end{align}\] At each round, a fraction \(\gamma_c\) of clients is selected uniformly at random, resulting in \(K = \lfloor \gamma_c C \rfloor\) 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 \(\lambda\) is reset to zero when the constraint is satisfied [7]. The aggregation weights \(w_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 an \(L_0\) constraint using probabilistic gates. We refer to this variant of FedAvg as E-FLoPS, where E stands 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 of \(1/|\theta|\). The temperature \(T\in (0,\infty)\) is also initialized at \(1/|\theta|\) and increased if needed to encourage uncertainty exploration as the total sample size N decreases. The gate parameters \(\phi = \log \alpha\) are initialized from a normal distribution with mean \(\log \rho_{\text{init}} - \log(1 - \rho_{\text{init}})\) and variance \(0.01\), where \(\rho_{\text{init}}\) controls the initial density. The target density is denoted by \(\rho_{\text{targ}} = \rho\). The Hard Concrete distribution parameters are set to \(\gamma = -0.1\), \(\zeta = 1.1\), and \(\beta'\) is recommended to be set at \(0.66\), following [5]. The Lagrange multiplier is initialized as \(\lambda = 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 \(\hat{z}\) sampled without noise or smoothing [7]. For exact sparsity, the raw parameters \(\tilde{\theta}\) for the top-\(m\) indices of effective parameters \(\theta\) are pruned, where \(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 to \(0\) for E-FLoPS to improve communication efficiency via serialization for message compression [8]. 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 with their average, denoted by \(z^{\text{avg}}_{-m}\). Only the pruned \(\tilde{\theta}\) and \(z\), along with their non-zero indices and the scalar \(z^{\text{avg}}_{-m}\), are communicated. Upon reception, parameters are reconstructed via \(\theta = \tilde{\theta} \odot z\) and aggregated, and \(\phi\) is recovered from \(z\) using \(\scriptstyle\phi = \beta' \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 \(\rho_{\text{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 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 [18], 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, \[{!}{ \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 of E-FLoPS is more than FedIter-HT, but significantly lower than dense training via FedAvg. For large model sizes \(|\theta|\) and small target densities \(\rho_{\text{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).

4.1 Experiments on Synthetic Data↩︎

We generated synthetic data for sparse linear regression following the procedure from [21] for a range of \(\scriptstyle \frac{N}{|\theta|}\). For a parameter vector dimension \(d=|\theta|\), each row \(x_i \in \mathbb{R}^{d}\) of \(X \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 \(\Sigma_{ij} = \rho_{\text{cor}}^{|i-j|}, \quad i,j = 1, \ldots, d\). An \(m\)-sparse coefficient vector \(w_{\text{true}} \in \mathbb{R}^{d}\) is constructed, where \(m = \lfloor \rho \cdot d \rfloor\). A subset of indices \(S_m \subseteq \{1, \ldots, d\}\) is selected uniformly at random, and coefficients are assigned as \((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 as \(y = X w_{\text{true}} + \varepsilon\), where \(\varepsilon \sim \mathcal{N}(0, \sigma^2 I_N)\). The signal-to-noise ratio (SNR) is defined as \(\scriptstyle\mathrm{SNR} = \frac{\|X w_{\text{true}}\|_2^2}{\|\varepsilon\|_2^2}\), and the noise level is set to \(\scriptstyle\sigma = \frac{\|X w_{\text{true}}\|_2}{\sqrt{\mathrm{SNR}} \, \sqrt{N}}\). For \(\scriptstyle|\theta|=1000\), we generate a number of total training samples N for varying ratios of \(\scriptstyle\frac{N}{|\theta|}\) with \(\scriptstyle5000\) 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 [22], [23]. A fraction of clients (\(0.6\)) is randomly sampled in each epoch to introduce participation heterogeneity [22]. We evaluate performance using mean squared error (MSE) and \(R^2\). Experiments are conducted at a correlation level of \(\rho_{\text{cor}}=0.2\) and an SNR of 20, with a true sparsity level of \(95\%\) or target density of \(5\%\).

As shown in Figure 3, E-FLoPS achieves a higher \(R^2\) on unseen test data than FedIter-HT at all \(\scriptstyle\frac{N}{|\theta|}\). E-FLoPS also exceeds top-\(m\) model of densely training FedAvg during training and the pruned and fine-tuned model discovered post training. The Figure 3 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 of E-FLoPS, we fix \(\scriptstyle\eta_{\lambda} = \frac{0.01}{|\theta|}\), and tune learning rates within the ranges \(\eta_{\tilde{\theta}} \in [0.05, 0.5]\) and \(\eta_{\phi} \in [0.1, 0.9]\). Table  1 shows the accuracy of sparsity recovery \(\scriptstyle A(\theta)=\frac{TP}{(TP+FN)}\), the fraction of true features recovered. E-FLoPS has higher accuracy than both Fed-IHT and FedAvg except at smallest \(\scriptstyle\frac{N}{|\theta|}\).

4.2 Experiments on Real Data↩︎

4.2.1 Image Classification↩︎

We consider publicly available MNIST data [24]: A multi-class classification dataset of handwritten digits with \(28 \times 28\) gray-scale pixel values for each. We use a CNN with two \(5 \times 5\) convolutional layers of 6 and 16 channels, each followed by \(2 \times 2\) max 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 is \(2.5\%\) ( sparsity of \(97.5\%\) ).

The data is downsampled per class at a ratio of \(\sim\)​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 4 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 [25]. 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, \(\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 of \(0.1\%\) (99.9% sparsity), with all clients participating. We do not synthetically alter client data using affine shifts. Figure 5 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 \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-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 [26].

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 \gg N)\) regime under heterogeneous FL conditions remains underexplored. We study entropy regularization of probabilistic gates for sparse FL under an \(L_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 \gg N)\) regime, where optimization may otherwise commit early to suboptimal sparse configurations with poor generalization.

We formulate a differentiable, entropy-regularized, \(L_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 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.

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.

6 KL-Divergence for Hard Concrete Gates↩︎

The stochastic gates follow the Hard Concrete distribution [5], obtained by stretching a Binary Concrete random variable \(s \in (0,1)\): \[\bar{s} = s(\zeta - \gamma) + \gamma, \quad z = \min(1, \max(0,\bar{s})),\] with \(\gamma < 0\) and \(\zeta > 1\). This induces a mixed distribution over \(z\) with point masses at \(0\) and \(1\), and a continuous component on \((0,1)\): \[\begin{align} 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)), \end{align}\] \[\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' (\log s - \log(1-s)) - \phi \right).\] we define a hard concrete prior \(p(z)\) with the same support and then the KL-divergence decomposes as: \[\begin{align} \mathrm{KL}(q(z)\|p(z)) &= 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)} \\ &\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]. \end{align}\]

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 [5].

References↩︎