Federated Learning With L0 Constraint Via Probabilistic Gates For Sparsity

TL;DR

This paper introduces a novel federated learning method that enforces sparsity through an L0 constraint using probabilistic gates, improving communication efficiency and statistical performance under data heterogeneity.

Contribution

It proposes a new L0 constraint approach with probabilistic gates for federated learning, enabling controlled sparsity and better performance compared to existing pruning methods.

Findings

Achieves target sparsity levels with minimal performance loss.

Outperforms magnitude pruning in communication efficiency.

Effective on various models and datasets, including CNNs and multi-class tasks.

Abstract

Federated Learning (FL) is a distributed machine learning setting that requires multiple clients to collaborate on training a model while maintaining data privacy. The unaddressed inherent sparsity in data and models often results in overly dense models and poor generalizability under data and client participation heterogeneity. We propose FL with an L0 constraint on the density of non-zero parameters, achieved through a reparameterization using probabilistic gates and their continuous relaxation: originally proposed for sparsity in centralized machine learning. We show that the objective for L0 constrained stochastic minimization naturally arises from an entropy maximization problem of the stochastic gates and propose an algorithm based on federated stochastic gradient descent for distributed learning. We demonstrate that the target density (rho) of parameters can be achieved in FL, under data and client participation heterogeneity, with minimal loss in statistical performance for linear and non-linear models: Linear regression (LR), Logistic regression (LG), Softmax multi-class classification (MC), Multi-label classification with logistic units (MLC), Convolution Neural Network (CNN) for multi-class classification (MC). We compare the results with a magnitude pruning-based thresholding algorithm for sparsity in FL. Experiments on synthetic data with target density down to rho = 0.05 and publicly available RCV1, MNIST, and EMNIST datasets with target density down to rho = 0.005 demonstrate that our approach is communication-efficient and consistently better in statistical performance.