Feature Norm Regularized Federated Learning: Utilizing Data Disparities for Model Performance Gains

1 Introduction

Federated learning, a cutting-edge machine learning paradigm, has found extensive applications across domains such as healthcare, finance, and smart devices, enabling collaborative model training while preserving data privacy and security [Banabilah et al., 2022]. In federated learning, a crucial distinction from conducting deep learning on a single node is to aggregate locally updated models from individual participants to obtain a global model [Wang et al., 2019]. Most commonly used aggregation algorithms in federated learning is FedAvg

∗These authors are joint corresponding authors.

†These authors are joint corresponding authors.

[McMahan et al., 2017]. FedAvg performs weighted averaging on locally trained models to construct a global model, which is simple but notably effective under i.i.d. (independently and identically distributed) data [Clauset, 2011].

Despite the wide application of FedAvg, non-i.i.d. data presents challenges [Hsieh et al., 2020]. These issues lead to variations in local data characteristics, hindering the convergence and generalization of global models. Recent studies [Li et al., 2019] have shown that FedAvg performs poorly under non-i.i.d. data, as weighted averaging fails to integrate accurate knowledge from participants with skewed distributions.

In this work, we propose Feature Norm Regularized Federated Learning (FNR-FL), a federated learning framework leveraging the class average feature norm to enhance the performance of the global model under the non-i.i.d. data distribution. Compared to prior efforts dedicated to addressing non-i.i.d. issues, FNR-FL allows for significantly better test accuracy and faster convergence and excels in various noni.i.d. scenarios, instead of excelling in only a limited number of settings.

The main achievements, including contributions to the field, can summarised as follows:

• We proposed a regularization algorithm based on class average feature norms, which serves as a plug-and-play module that seamlessly integrates with existing federated learning algorithms. This modular algorithm allows for enhanced model regularization without the need for substantial modifications to the underlying federated learning framework.
• Building upon this regularization algorithm, we have constructed a federated learning algorithm FNR-FL that achieves significantly superior test accuracy and faster convergence under the non-i.i.d. data distribution compared to other Federated Learning (FL) algorithms.
• Under mixed non-i.i.d. scenarios, FNR-FL has proven to outperform other algorithms, marking the first known federated learning algorithm tested under such mixed conditions. This strategic testing highlights the algorithm's robust performance in complex, real-world conditions.
• We have developed two innovative metrics which innovatively capture the trade-off between accuracy, communication, and computational costs, offering a clear benchmark for algorithm comparison.

2 RelatedWork

McMahan et al. [2017] developed FedAvg, a widely used federated learning algorithm that aggregates model parameters through weighted averaging. The essence of FedAvg involves clients uploading their local model parameters to a server, which computes the average of these parameters, weighting them according to the volume of data on each device. This approach has significantly shaped the development of efficient and scalable federated learning systems.

Li et al. [2020] introduced FedProx, a novel optimization framework tailored for federated networks, uniquely addressing both system and statistical heterogeneity. Its innovation lies in allowing flexible local computations while incorporating a proximal term to maintain overall algorithmic stability across diverse devices.

Karimireddy et al. [2020] developed SCAFFOLD, a pioneering stochastic algorithm that creatively tackles gradient dissimilarity in federated settings. Its key innovation is the incorporation of control variates to significantly reduce the gradient variance, leading to enhanced convergence rates and model performance.

Li et al. [2021a] proposed MOON (Model-Contrastive Learning), a groundbreaking approach in federated learning that introduces contrastive learning at the model level. This technique stands out for its simplicity and effectiveness, particularly in enhancing federated deep learning models' performance on non-i.i.d. datasets by encouraging model-level feature alignment.

Wang et al. [2020] introduced FedNova, a comprehensive theoretical framework designed for heterogeneous federated learning environments. Its main contribution is the natural integration of diverse local update steps and optimization techniques (such as gradient descent, stochastic gradient descent, and proximal updates), ensuring fair and efficient convergence across a wide range of network conditions and device capabilities.

The FedDF framework [2020] leverages a distillation technique to foster a more homogeneous knowledge transfer among decentralized datasets. This method is notable for its distillation from decentralized to centralized, aiding in overcoming the statistical heterogeneity inherent in FL.

FedGen [2021] by Dong et al. integrates a generative component into the federated learning process. This addition aims to synthesize pseudo-samples that represent the global data distribution, thus enhancing the robustness of the model against data heterogeneity.

3 Motivation

3.1 The Fundamental Reason of Non-i.i.d. FL Performance Decline

The primary cause of performance degradation in federated learning due to non-i.i.d. data lies in the varied data distributions across different participants. This variation leads to divergent update directions when training local models [Zhao et al., 2018]. Simply put, the more diverse the data distribution, the greater the directional deviation in local model updates [Li et al., 2021b], which hampers the effective consolidation of knowledge into the global model [Karimireddy et al., 2020;

Wang et al., 2020; Li et al., 2022; Gao et al., 2022]. Consequently, this disparity in data distribution culminates in a decline in the global model's overall test accuracy.

Implication 1. Our objective is to minimize discrepancies in model update directions among participants in federated learning, ultimately enhancing the overall performance of the global model.

3.2 Quantify the Difference Among Local Model Updates

Acknowledging the critical role of minimizing update discrepancies in federated learning (Implication 1), our focus shifts to quantifying these differences effectively. Wei et al. [2023] suggest that by using a consistent feature extractor, it is feasible to approximate data distribution differences through the variance in feature norms.

Implication 2. In a similar vein, we can postulate that, under the assumption of using the same dataset for evaluation, it is possible to gauge the difference among local model updates by leveraging the disparity among feature norms.

3.3 Promoting Alignment of Model Update Directions Among Participants

To counteract the effects of non-i.i.d. data, FedProx integrates a proximal term in each local model's loss function. This term penalizes deviations from the global model, aligning local updates more closely with it.

Implication 3. Informed by FedProx's methodology, we propose enhancing participant training by integrating a regularization term based on feature norm differences [Brownlee, 2018]. This term serves as a corrective measure for participants with diverse data distributions, encouraging more aligned contributions to the global model.

4 Non-i.i.d. Scenarios

In our study, we have investigated three distinct categories of non-i.i.d. settings: feature distribution skew, label distribution skew, and quantity skew [Zhang et al., 2022a]. Additionally, we have explored two mixed non-i.i.d. settings, which combine label distribution with quantity skew and label distribution with feature distribution skew.

4.1 Feature Distribution Skew

We simulated feature distribution skew by adding Gaussian noise to the data held by each participant. Using CIFAR-10 dataset samples, we show in Figure 1 how applying Gaussian noise with varying standard deviations affects the data. Original samples are in Figure 1a, while Figures 1b and 1c display the impact of different noise levels. As the standard deviation increases, the images become progressively obscured, exemplifying the feature distribution's deviation from its original form.

4.2 Label Distribution Skew

The Dirichlet distribution is utilized for allocating data with label distribution skew to different participants. Each class of samples is subject to a Dirichlet distribution sampling process

(c) noise ∼ Gau(0.5)

individually. Taking the allocation of CIFAR-10 dataset to 10 participants as an example, when the parameter α is set to 0.5, the data allocation is depicted in Figure 2a. Darker colors indicate a higher quantity of samples, while white signifies the absence of samples in that class. From the illustration, it is evident that there are substantial variations in the number of samples from each class within the local datasets held by the participants.

4.3 Quantity Skew

Similar to label distribution skew, quantity skew arises when the sizes of local datasets, symbolized as |Di|, differ among parties. We employ the Dirichlet distribution to distribute varying data sample quantities to each party. However, unlike label distribution skew, we maintain a roughly consistent data distribution across parties in this scenario [Archetti et al., 2023]. This consistency allows us to isolate and analyze the specific impact of quantity skew on global model performance. Figure 2b demonstrates this, showing that within each participant's local datasets, the counts of different labels fall within a specific range.

5 ProposedMethod

5.1 Class Average Feature Norm

In this section, we describe the computation of the class average feature norm. To facilitate our discussion, we enumerate key notations in Table 1 in the Supplementary Materials. In a federated learning scenario for image classification tasks involving n participants, each participant Pi possesses a local dataset denoted as Di with Ni labeled samples:

Datasets Di contain samples from different classes. Each participant's local model Mi consists of two main components: the feature extractor Mfi and the classifier Mci. For a given sample xij in Di, we define its feature norm as the Euclidean norm (L2 norm) of the extracted features as [Naik et al., 2023]:

where ϑif represents the parameters of the feature extractor Mfi. After calculating the feature norms for each sample, we proceed to compute the average feature norm for each class of participant Pi. Specifically, for the class k in participant Pi, the computation of class average feature normFki is performed as follows:

In this equation, δ[yj i =k] is the indicator function, which is defined as:

This process allows us to obtain class average feature norms for the different classes of participants Pi. The process of computation is summarized in Algorithm 1.

Algorithm 1 Calculate Class Average Feature Norms

1: Input: dataset D 2: Initialize: class feature norms F ← {}, class sample counts C ← {}, class average feature norms avgF ← {} 3: for each batch B in dataset D do 4: for each sample (xi, yi) in batch B do 5: compute the feature normUi for xi using Eq. (2) 6: append Ui to F[yi] 7: C[yi]+ = 1 8: end for 9: end for 10: for each (label k, feature norms Fk) in F do 11: if C[k] > 0 then 12: Calculate the class average feature norm avgf using Eq. (3) 13: avgF [k] ← avgf 14: end if 15: end for 16: Output: class average feature norms avgF

Algorithm 2 Calculate Feature Norms Differences

1: Input: participant Pj, set of all participants P, set of participants to be regularized Pre, feature norm list of all participants Fn 2: ∆j = {} 3: for m ∈ P \ Pre do 4: for label l ∈ Fn[Pj].keys() do 5: ∆temp = (Fn[m][l] - Fn[j][l]) 6: ∆j [l] += ∆temp 7: end for 8: end for 9: return ∆j

5.2 FNR-FL: Feature NormRegularized Federated Learning

Procedure of FNR-FL

In this subsection, we introduce the FNR-FL algorithm, a novel federated learning approach designed to tackle the challenges posed by non-i.i.d. data across different participants. This framework leverages feature norm regularization to align local model updates, thus enhancing global model performance. The detailed steps of this procedure are outlined in Algorithm 3.

In Algorithm 3, The server initializes and distributes the global model to all participant devices. Each participant, P1 through Pn, trains this global model on their own local dataset. After training, they each evaluate their model against a global test dataset. They also calculate the class average feature norms using Algorithm 2. Participants with the relatively lower test accuracy are selected to refine their models further using the FeatureNormRegularization(j, wtj , Dj , ∆j ) in Algorithm 3. The server then collects these local models and performs an aggregation step.

Local training with computation of class average feature norm. In each communication round t ∈ [T], all n participants are active. The server sends the global model wt to each participant. Each participant Pi then updates its local model wt i on local dataset Di for Etrain epochs. Subsequently, each participant is required to compute the class average feature norms F t i (Algorithm 1) and the accuracy Ai of samples in the public dataset Dpublic using its updated local model wt i . The updated local model wt i , the class average feature norms F t i and the accuracy Ai are then sent to the server.

Computation of differences among feature norms. Upon receiving the test accuracies, denoted as A1, ...,An, on the public dataset Dpublic, the server proceeds to select n ∗ p participants with the lowest accuracies to undergo feature norm regularization. Here, p serves as a threshold parameter determining the proportion of participants chosen for regularization.

Algorithm 3 FNR-FL

• 1: Input: local datasets Di, number of participants n, number of communication rounds T, number of local training epochs Etrain, number of regularization epochs Ere, public dataset Dpublic, learning rate η, percentage of regularized participants p,
• 2: Initialize: global model w0
• 3: Server execues:
• 4: for t = 1, 2, · · · , T do initialize feature norm list Fn
• 5: for i = 1, 2, · · · , n in parallel do
• 6: send the global model wt to participant Pi

7:

wt

• 8: Fn[i] = F t i

i , F t

• 9: end for
• 10: choose nre = ⌊n ∗ p⌋ participants to refine local model updates
• 11: for j = 1, 2, · · · , nre in parallel do
• 12: compute ∆j using Algorithm 2
• 13: wtj ← FeatureNormRefine(j, wtj , Dj , ∆j )
• 14: end for
• 15: wt+1 ← ∑ni=1 |Di | ∑ni=1 |Di | wt i
• 16: end for
• 17: return wT
• 18: LocalTraining(i,wt ,Di):
• 19: wt i ← wt
• 20: for k = 1, 2, · · · ,Etrain do
• 21: for each batch B in dataset Di do
• 22: L(wk i ;B) = ∑(x,y)∈B ℓ(wk i ; x; y)
• 23: wk i ← wk i - η∇L(wk i ;B)
• 24: end for
• 25: end for
• 26: compute F t i on Dpublic using Algorithm 1
• 27: test wk i on Dpublic and record accuracy Ai
• 28: return wk i , F t i ,Ai
• 29: FeatureNormRegularization(j,wtj ,Dj,∆j):
• 30: for k = 1, 2, · · · ,Ere do
• 31: for each batch B in dataset Dpublic do
• 32: L(wk j ; B) = ∑(x,y)∈B ℓ(wk j ; x; y)

33:

denote unique labels in B as L

• 34: J(j, Fn) = 0.0
• 35: for each label l ∈ L do

36:

ρ = #(y == l)/|B|

• 38: end for

39:

Φ(wk

j ) = L(wk

j ; B) + λ · J(j, Fn)

• 40: wk j ← wk j - η∇Φ(wk j )
• 41: end for
• 42: end for
• 43: return wk j to the server

i , accuracy ← LocalTraining(i, wt , Di)

The set of these selected participants is denoted as Pre and the set of all participants is denoted as P. For each participant Pj in Pre, the server computes the differences of the class average feature norms between Pj and other participants in P \ Pre, denoted as ∆j , using Algorithm 2. For each label l, we use calculate the differences of class average feature norm between participant Pj and other participants in P \ Pre, and record it with its label l as key in ∆j.

Feature norm regularized model update. Consequently, the loss function Φj(wtj ) of participant Pj incorporates its cross-entropy loss L(wk i ,B) and the weighted average sum J(j, Fn) of each element in ∆j with the proportion of the number of label l samples to the batch size |B| as the weight:

The hyperparameter λ controls the trade-off between fitting the local datasets well and encouraging convergence towards a global model with better performance.

Convergence Analysis of FNR-FL

In this section, we delineate the convergence characteristics of the FNR-FL algorithm we proposed. Initially, we establish certain assumptions, commonly adopted in preceding studies [Zhou, 2018; Bao et al., 2022; T Dinh et al., 2020; Haddadpour and Mahdavi, 2019].

Assumption 1 (L-smooth). The functions Φ1, . . . ,ΦN are all L-smooth. This implies that for all x, y, the following inequality holds:

Assumption 2 (µ-strongly convex). The functions are all µstrongly convex, which means for all x, y:

Assumption 3 (Bounded gradients). The stochastic gradients are unbiased and have bounded variance. Specifically, in m-th mini-batch gradient descent step of participant Pi:

and

where bim denotes a random batch in all batches of participant Pi in m-th mini-batch gradient descent step.

Consider the m-th mini-batch gradient descent step: ¯wm+1 = ¯wm-ηmgm. We analyze the distance to the optimal point w∗ in terms of the squared norm:

where gm = ∇Φi(wim, bim), which denotes the gradient of the loss function ∇Φi with respect to the current mini-batch bim; and its average ¯gm = ∇Φi(wim).

Lemma 1. The bound of ∥¯wm -w∗ - ηm¯gm∥2 :

Lemma 1 provides an upper bound of ∥ ¯wm-w∗ -ηm¯gm∥2 characterized by the learning rate ηm, µ and L.

Lemma 2. In order to associate with our assumptions, we consider the expectation of ∥¯gm - gm∥2 :

Lemma 2 bounds the expectation of the gradient difference, which is influenced by the variability in each mini-batch's gradient, as indicated by Assumption 3.

Lemma 3. Based on Lemma 1 and Lemma 2, we get:

Lemma 3 introduces a recursive relationship derived from the previous lemmas.

Letting δm = E∥¯wm - w∗ ∥2 , which denotes the expected squared norm of the deviation of the global model ¯wm from the optimal w∗ at the m-th mini-batch SGD; and thus δm+1 = E∥¯wm+1 - w∗ ∥2 , we can get the following recurrence relation:

This inequality demonstrates the decrement of δm over iterations, thus assuring convergence to the optimal global model w∗ as m increases. The requisite conditions for assuring the convergence of the sequence {δm} are enumerated as follows:

1. The learning rate ηm decays.

Given that all the prescribed conditions are satisfied, we ascertain that δm convergences to 0 as m increases to infinity, signifying the assured convergence of the global model ¯wm to the optimal global model w∗ . While in practice the number of iterations m is finite, the convergence behavior of δm as outlined still assures that we can approximate the desired precision within our computational budget. For the complete proof, see the Supplementary Materials.

6 Evaluations

6.1 Experimental Setup

In our experimental setup, all participants are actively involved in every round of the training process. We use the SGD optimizer [Robbins and Monro, 1951] with a learning rate of 0.1. The batch size of the training data is set to 64 and the test data is set to 32 by default. The number of local training epochs is set to 10 by default and the regularization epochs is set to 5 by default. We conduct each training process for 10 rounds.

6.2 Performance Under Single Types of Non-i.i.d. Scenarios

To investigate the effectiveness of the proposed FNR-FL, we report the performance of FNR-FL, FedAvg, FedProx, SCAFFOLD, MOON, FedNova, FedDyn, FedFTG, FedDF, FedDC, and FedGen when training ResNet-18 [He et al., 2016] or VGG-11 [Krizhevsky, 2009] model on CIFAR-10 [Krizhevsky, 2009]. Experiments are conducted under three single types of non-i.i.d. data distribution.

Test Accuracy

Table 2 shows FNR-FL excelling in non-i.i.d. scenarios on CIFAR-10, notably outperforming other FL methods in feature distribution skew with accuracies of 0.9976 on ResNet-18 and 0.9931 on VGG-11. In label distribution and quantity skew, FNR-FL achieves 0.9970 and 0.9992 for ResNet-18 respectively, demonstrating strong results that underscore its robustness in federated learning.

Further evaluations on benchmark datasets MNIST and FEMNIST are presented in Table 3. FNR-FL achieves accuracies of 0.9953 in feature distribution skew and 0.9944 in quantity skew on MNIST, and 0.9889 and 0.9944 respectively on FEMNIST. These results are comparable to those on CIFAR-10, confirming FNR-FL's consistent effectiveness across diverse data scenarios, further demonstrating its practical utility in federated learning systems.

Convergence

The convergence curves of training ResNet-18 on CIFAR-10 under three single types of non-i.i.d. scenarios are illustrated in Figure 4. We can observe that the FNR-FL algorithm consistently exhibits the same level of convergence speed as other FL algorithms and achieves exceptionally better test accuracies.

Communication and Computational Cost

Metrics for performance evaluation are presented in Table 4, where the Traffic represents the cumulative data exchanged, including both uploads and downloads and Time computes the duration to complete 10 rounds of training. The metrics κ = Accuracy∗104 /Time and ρ = Accuracy∗104 /Traffic were computed to measure test accuracy per unit of time and traffic, respectively, providing a quantifiable tradeoff between effectiveness and efficiency. FNR-FL stands out for its exceptional efficiency in non-i.i.d. scenarios.

6.3 Performance Under Mixed Types of Non-i.i.d. Scenarios

We evaluated the performance of federated learning algorithms on the CIFAR-10 dataset under mixed scenarios of label and feature distribution skew as well as label and quantity distribution skew. In Table 5 (Supplementary Materials), FNR-FL outshines all FL algorithms in label+feature distribution skew at noise levels of 0.1 and 0.5, with top accuracies of 0.9755 and 0.9111, respectively. Other algorithms show significant performance drops, especially at higher noise. For label+quantity skew, FNR-FL leads with 0.8826 accuracy, while others hover around 0.3000.

As noted by recent studies [Li et al., 2022], mixed types of skew present more complex challenges than single types. FNR-FL's ability to maintain high accuracy levels addresses the critical need for effective algorithms capable of operating under mixed types of non-i.i.d. scenarios.

6.4 Orthogonality of FNR-FL with Existing FL Algorithms

The experimental results in Table 6 in the Supplementary Materials demonstrate the effectiveness of integrating the FNR-FL algorithm with other FL algorithms. Across all scenarios, combinations involving FNR-FL consistently outperform their standalone FL algorithms, indicating a substantial enhancement in handling non-i.i.d. data distributions. These results suggest that FNR-FL's regularization mechanism effectively augments existing FL algorithms [Acar et al., 2021].

6.5 Effect of Feature NormRegularization

The ablation study assesses the impact of feature norm regularization (FNR) within the FNR-FL algorithm on a participant's local model. Notably, classes with lower local sample sizes, such as 'frog' (471) and 'horse' (485), showmore significant accuracy improvements. This observation is aligned with the hypothesis that FNR enables participants to benefit from the shared knowledge across the federation, particularly when their local data is insufficient to capture the complexity of the class. In contrast, classes with more substantial local representations, like 'automobile' (512) and 'ship' (528), exhibit smaller gains. This can be attributed to the participant's local model already achieving a higher accuracy due to the ample class-specific data, hence the marginal utility of the shared knowledge from FNR is less impactful [Du et al., 2022; Shao et al., 2023; Qi et al., 2023].

6.6 Effect of Noise

The experiment evaluates the performance and robustness of FNR-FL against other FL algorithms across different noise levels within feature and label distribution skews. The results are presented in Table 7. FNR-FL demonstrates remarkable noise immunity, preserving high accuracy despite increasing noise. Its robustness to noise underscores its suitability for real-world federated learning applications.

6.7 Effect of Training Hyperparameters

Table 6 shows the impact of various hyperparameters on global model accuracy. Our findings reveal that changes in batch size, learning rate, and the number of training epochs do not significantly affect the test accuracy of the global model. This suggests that our proposed FNR-FL maintains stable performance across a variety of hyperparameter settings.

&RPSDULVRQ RI 0RGHO $FFXUDF\ %HIRUH DQG $IWHU )HDWXUH 1RUP 5HJXODUL]DWLRQ

7 Conclusions and FutureWork

In this paper, we introduce the FNR-FL algorithm as a new approach to enhance performance in non-i.i.d. data scenarios. The core philosophy of FNR-FL revolves around leveraging feature norms as a metric to quantify and mitigate the divergence in local model updates. Our evaluations demonstrate that FNR-FL achieves superior test accuracy without excessive overhead and maintains convergence. An orthogonal experiment highlights its modularity, seamlessly enhancing FL algorithms. The necessity of feature norm regularization is confirmed through an ablation experiment.

For future work, we aim to further optimize FNR-FL for communication cost reduction by integrating model compression techniques. We also plan to explore applying FNRFL to privacy-preserving federated learning and to develop a decision-making algorithm that analyzes data distribution characteristics in federated settings to recommend the most appropriate FL algorithm.