Client Contribution Normalization for Enhanced Federated Learning

There exists a vast amount of literature in deep learning to handle statistical heterogeneity using approaches like multitask learning and meta-learning. Recently, these methods of handling non-IID data have been incorporated in FL settings [18, 22, 4]. For instance, MOCHA allows personalization in FL by learning separate but related models at each participating device using multitask learning [26]. Despite these advances, the key challenge of statistical heterogeneity still remains open. Hsu et al. explore the performance of FL approaches on classification tasks and show that FedAvg starts performing poorly with increasing non-IIDness [10]. Other works focus on heterogeneity in computation cost in FL [27, 1]. Additionally, Luo et al. introduced the EdgeFed (Edge FL) approach, which employs an edge server to aggregate weight updates between the centralized server and the clients [17].

Although various methods have targeted data heterogeneity in FL, most of them introduce additional regularizers to handle the non-IID data at the client end. FedProx uses a dynamic regularizer to handle non-IID data at the client end [14]. SCAFFOLD highlights the drawbacks of FedAvg and proposes the use of control variate (variance reduction) to handle client drift in the presence of heterogeneous data, [11]. The recently introduced FedBABU algorithm uniquely updates the model’s body (extractor) during federated training, leaving the head unaltered and randomly initialized [21]. FedNova addresses heterogeneity by normalizing the number of local iterations at each client, assigning weights to clients during aggregation based on these normalized iterations [28]. In FL, Stochastic Gradient Descent with Momentum (SGDM) [30] is employed to optimize model parameters across distributed clients.

In FL, the contribution factor plays a pivotal role in the aggregation process, influencing how individual models contribute to the global model. FedAvg utilizes the number of samples from each local model as a key factor in determining its contribution value during the aggregation process [2]. Various similarity-based approaches have been proposed to tackle issues posed by non-IID data and malicious clients. The adaptive federated averaging method, proposed by Muñoz-González et al., calculates the cosine similarity between local model gradients and the weighted average of all local models in each round and aggregates the global model using contribution values [20]. Similarly, the FLTrust approach introduced by Cao et al. measures cosine similarity between a local model and the server model trained with a small root dataset [3].

While these approaches have made significant progress, several limitations persist. Many methods rely on specific model architectures or introduce substantial computational overhead. Some require modifications to local training procedures or model structures, limiting their applicability. The trade-off between addressing heterogeneity and maintaining model generalization remains a challenge.

Our work addresses these gaps by proposing a flexible, computationally efficient plug-in solution that integrates seamlessly with existing FL algorithms. By utilizing mean latent representations, we offer a versatile approach to normalize client contributions without modifying local training procedures or model architectures. This not only enhances the performance of FL algorithms in non-IID settings but also provides a practical framework for real-world applications where data heterogeneity is prevalent. Through extensive experiments, we demonstrate significant improvements in model accuracy and consistency across diverse datasets and skewed distributions.

At each client, we utilize the activations derived from the final hidden layer of the local model as a representation of semantic features, capturing valuable information about the variation in data across different clients. Each client transmits both its model weights $w_{r}$ and mean latent representation $z_{r}$ to the centralized server. During aggregation, the server calculates similarity among the $z_{r}$ values of all participating clients using the Cosine similarity metric:

The similarity matrix $S$ is defined as follows:

The contribution factor $\Lambda_{r}$ for client $r$ is defined as:

The proposed approach handles variability in data distribution among participating clients using $\boldsymbol{\Lambda}=[\Lambda_{1}\hskip 2.84544pt\Lambda_{2}\hskip 2.84544pt\Lambda_{3}\hskip 2.84544pt\dots\hskip 2.84544pt\Lambda_{R}]^{T}$. It normalizes the contribution of individual clients during aggregation by integrating $\boldsymbol{\Lambda}$ into the mainstream FL framework and modifying the global objective as follows:

where $\nu_{r}$ is the importance function defined by the mainstream FL framework, and $\mathbf{\Lambda}\cdot\boldsymbol{\nu}$ is the dot product of the importance function vector ($\boldsymbol{\nu}$) and the proposed $\boldsymbol{\Lambda}$.

This integration allows our approach to be easily incorporated into existing FL algorithms. For example, in FedAvg, we replace the standard aggregation step with our normalized version:

where $w_{t+1}$ is the global model at round $t+1$, and $w^{r}_{t}$ is the local model of client $r$ at round $t$.

Fig. 1 (a) illustrates the behavior of our proposed approach, controlled by the contribution factor, in a simulated data setting with three clients. Two clients have data from eight CINIC-10 classes, while the third has data from the remaining two. We compare FedAvg and Normalized FedAvg by plotting the similarity between the local and global weight matrices after each round, as shown in Fig. 1(b). The proposed normalization approach prioritizes client 3, bringing its model closer to the global model, unlike FedAvg, which treats all clients equally, resulting in client 3’s model being further from the global model. The inset highlights this effect, demonstrating the effectiveness of our approach in achieving a more balanced weight distribution among clients. Algorithm 1 outlines the methodology employed in the proposed approach.

Our approach introduces three main computational steps: (1) mean latent representation computation, (2) inter-client similarity calculation, and (3) contribution factor derivation. These processes are efficiently integrated into the FL pipeline, occurring once per communication round without requiring significant additional data transfer.

To reduce abrupt changes in contribution values across consecutive rounds, we use temperature scaling. We modify equation (5) as:

where $\Lambda_{r_{T}}$ is the normalized contribution factor for node $r$ with temperature scaling. A temperature value less than 1 results in refinement of the contribution value obtained using the mean latent representation. We perform all our experiments using this temperature scaling approach to ensure consistent and stable evaluation of client contributions.

We utilize a VGG-16 model architecture for all experiments. Each client executes two local epochs per communication round, with a total of 200 communication rounds. The Adam optimizer is employed with a batch size of 64. The mean latent representation is derived from the final hidden layer of VGG-16, projecting the output into a 128-dimensional vector. Our comparative analysis involves six baseline FL techniques: FedAvg [19], FedProx [14], FedBABU [21], SCAFFOLD [11], SGDM [30], and FedNova [28].

We conduct experiments on CIFAR-10 [13], FEMNIST [5], and CINIC-10 [6]. CIFAR-10 is used to evaluate performance under controlled non-IID settings with varying degrees of heterogeneity. FEMNIST is selected to assess the approach’s efficacy on naturally heterogeneous data with imbalanced class distributions. CINIC-10 is employed to test robustness against augmented and transformed data distributions. For CIFAR-10 and CINIC-10, we generate non-IID distributions using a Dirichlet distribution with concentration parameter $\alpha$. Lower $\alpha$ values induce higher data heterogeneity, allowing us to systematically evaluate our method’s performance across a spectrum of non-IID scenarios.

Our experimental results demonstrate the efficacy of the normalization approach across diverse data settings and baseline FL methods. The normalization approach enhances accuracy across all tested datasets, FL methods, and heterogeneity levels, showing particular effectiveness in highly heterogeneous settings (low $\alpha$ values). This improvement is crucial for real-world FL applications where data distributions across clients are often highly skewed. The consistent superior performance of SGDM with normalization suggests a complementary relationship between momentum-based optimization and our contribution-based normalization, potentially opening avenues for further research into hybrid optimization strategies in FL.

The approach’s effectiveness scales well with increasing client numbers, as evidenced by the FEMNIST results, indicating its suitability for large-scale FL deployments. Normalized methods consistently achieve faster convergence, particularly in early training stages, potentially reducing communication rounds in practical FL deployments. This acceleration in convergence is especially valuable in bandwidth-constrained or high-latency network environments, where minimizing communication overhead is critical.

The performance improvements across different methods provide insights into the interaction between our approach and existing FL techniques. For instance, the significant improvement in FedNova’s performance when normalized suggests that our approach effectively complements methods designed to handle heterogeneity. Similarly, the enhanced performance of FedProx with normalization indicates that our method can work synergistically with proximal term regularization to manage client drift. These findings underscore the proposed approach’s potential in enhancing FL performance across diverse real-world scenarios characterized by data heterogeneity.