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

2026-06-30

AI Summary

This work addresses the challenge of learning sparse models with strong generalization and accurate structural recovery in federated learning settings characterized by data sparsity, heterogeneity, and partial client participation. The authors propose a novel approach based on a probabilistic gating mechanism, which—by introducing entropy regularization into federated learning for the first time—preserves uncertainty in the sparse structure and prevents premature convergence to suboptimal support sets. Integrating L0 constraints with federated optimization, the method consistently outperforms baseline strategies such as Fed-IHT and post-hoc pruning of FedAvg across both synthetic and real-world datasets, achieving significant improvements in both test performance and accuracy of recovered sparse structures.

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.

Problem

• Federated Learning
• Sparse Model Discovery
• Data Scarcity
• High-Dimensional Regime
• Generalization

Innovation

• entropy regularization
• probabilistic gates
• sparse model discovery
• federated learning
• uncertainty exploration