# Entropy-Regularized Federated Optimization for Non-IID Data
###### Abstract
Our motivation is to remedy these limitations with an entropy-based regularizer that encourages each client's update to spread its \"mass\" more uniformly across parameters. By maximizing the entropy of the update vector, we promote diverse yet compatible local adjustments that align better with the global objective--without introducing substantial communication or computational overhead. In this work, we propose a novel approach to address these issues by incorporating an entropy-based regularization into the federated optimization process. Our method, called entropy-regularized federated optimization (ERFO), augments the local training objective on each client with a term that encourages high-entropy model updates. Intuitively, instead of allowing a client's update to be overly concentrated on a small subset of model parameters (which can cause the global model to latch onto niche local patterns and potentially forget broader knowledge), we encourage the update to be spread more broadly across the model's parameters. By penalizing very low-entropy (highly peaked) update distributions, we nudge each client to make more distributed changes that better align with the global objective and other clients' updates. The main contributions of this work are as follows:  
1. We introduce entropy-regularized federated optimization (ERFO), augmenting each client's local objective with a Shannon entropy term on per-parameter update magnitudes. This mitigates client drift under highly non-IID data without altering server aggregation or adding communication overhead. 2. We derive a closed-form gradient for the entropy regularizer and integrate it into standard SGD on the client. Unlike FedProx's proximal constraint or SCAFFOLD'scontrol variates, our approach requires only one extra gradient computation per local epoch and no extra vectors to transmit. 3. We conduct extensive experiments on the UNSW-NB15 intrusion detection benchmark with Dirichlet-partitioned non-IID clients, demonstrating that ERFO achieves 81.1% accuracy and 0.791 macro-F1--1.6 pp and 0.008 higher than FedProx--while yielding smoother, more stable convergence. 4. We validate ERFO on the PneumoniaMNIST chest X-ray classification task with balanced client splits, achieving 90.3% accuracy and 0.878 macro-F1--2.8 pp and 0.022 higher than FedAvg--and further show via ablation on the initial entropy weight and a learning-rate vs. entropy-weight sweep that performance is robust across these hyperparameters. The remainder of this paper is organized as follows: Section 2 reviews related work in federated optimization under non-IID settings, as well as entropy-based regularization methods in machine learning. Section 3 presents the ERFO methodology, including the definition of the entropy regularizer, its gradient derivation, and the modified federated training procedure. Section 4 describes the experimental setup, datasets (network intrusion and medical imaging), partitioning scheme, and evaluation metrics. Section 5 reports the empirical results, comparing ERFO with baseline methods and analyzing convergence behavior. Section 6 discusses the insights from the experiments, including the effectiveness of a fixed regularization schedule. Finally, Section 7 concludes the paper with discussions on generalizability, limitations, and future work.