FedNNNN: Norm-Normalized Neural Network Aggregation for Fast and Accurate
Federated Learning

Federated Learning (FL) [10] is a distributed learning protocol that enables one to train a model with a massive amount of data obtained by IoT devices or smartphones without expensive training on a centralized server. Instead of collecting the data and training the model on a single machine, FL server only combines models locally trained by edge devices to proceed the learning process, and this combination procedure is known as model aggregation. The properties of FL is extensively studied over the past few years [10, 7, 6, 2, 1, 17, 19, 14, 12, 15, 16, 11], where we see discussions on aspects of FL such as communication efficiency [7], adaptation to heterogeneous systems [12], performance over non-standard data distributions [19, 14, 5, 15, 16, 11], security properties [6, 2, 1], and many more. In this work, we focus on improving the prediction accuracy of FL over non-IID distributions, which is one of the main problems associated with existing FL frameworks.

In this section, we outline the Federated Averaging (FedAvg) framework [10]. The overview is depicted by Figure 2.

In FedAvg, we assume that clients are in possession of the data, and a server aggregates the models from the clients. The learning is proceeded by the following steps.

• •

  Step ➀: Let the total number of clients to be $K$. The server first picks $m=\max(C\cdot K,1)$ clients out of the total $K$ clients for some real number $C\in[0,1]$. In the first round, the server initialize a model $\bm{w}_{0}$ and distribute the model to the selected $m$ clients. Otherwise, the server distributes the aggregated model $\bm{w}_{t}$. We consider the model distribution to be the start of the $(t+1)$-th round of communication.

• •

  Step ➁: Upon receiving the server model, each client locally trains the their own models using $\bm{w}_{t}$ as the initial model for $\mathcal{E}$ epochs. The local model trained on the $k$-th client in the $(t+1)$-th communication round is referred to as $\bm{w}^{k}_{t+1}$.

• •

  Step ➂: Client $k$ returns its trained model $\bm{w}^{k}_{t+1}$ to the server.

• •

  Step ➃: Upon receiving $m$ locally trained models from the clients, the server aggregates the models by simply taking the weighted average as

  $$\bm{w}_{t+1}\leftarrow\sum_{k=1}^{m}\frac{n_{k}}{n}\bm{w}^{k}_{t+1},$$

  (1)

  where $n_{k}$ is the size of the dataset on the $k$-th client, and we have that $n=\sum_{k=1}^{m}n_{k}$.

• •

  Step ➄: The server evaluates a unified model with test data. Additional learning steps can be carried out by repeating the steps from ➀ to ➄.

As mentioned, the averaging aggregation utilized in FedAvg results in significant accuracy degradation during server model evaluation if the clients possess non-IID datasets. A line of works [19, 14, 5, 15, 16, 11] are proposed to address this problem. Here we give a brief review on the existing methods.

The work in [19] is one of the first to point out that non-IID datasets result in drastically different local models, and these models become difficult to aggregate with simple averaging. However, [19] only proposes to share auxiliary datasets to each client so that each client obtains a more IID dataset. [14] tries to adjust the loss function to improve upon [19], but our analysis shows that the improvements are not consistent across datasets. In Section 3.2, we take a deeper look at the exact reason for the degraded accuracy of FL over non-IID datasets.

Other optimization approaches include [5, 15, 16]. In [5], clients are grouped based on the label distribution of their datasets. In [15], the server aggregates the client models whenever a single weight update occurs on some particular client. These approaches clearly incur a large amount of communications between the server and the clients. FedCurv [16], which imposes complex loss function on clients, also induces communication overheads. In summary, existing works generally require additional datasets or complex communication protocols to directly solve the non-IID problem on the data level. In what follows, we introduce quantitative analyses and normalization-based techniques to mitigate the impact of non-IID datasets without relying on auxiliary datasets or complex protocol modification.

We start by defining a per-client ${\Delta\bm{w}}^{k}_{t+1}$ for the $(t+1)$-th communication round

where $\bm{w}_{t}$ is the weight vector (i.e., neural network model) distributed to each client in the start of the $(t+1)$-th communication round, and $\bm{w}^{k}_{t+1}$ is the locally learned weights that will be returned to the server in the $t+1$-th round. Hence, we can simply interpret the weight difference ${\Delta\bm{w}}^{k}_{t+1}$ as the amount of learning proceeded in a single round of communication for client $k$.

Using Equation (2), we can re-formulate Equation (1) the aggregation procedure in FedAvg as

In other words, since all clients share the same $\bm{w}_{t}$, we can express the weight-averaging procedure in FedAvg as the sum of the distributed model $\bm{w}_{t}$ and the averaged sum of the local updates from each of the clients.

To quantitatively assess the impact of the updating vector, we define a pair of real scalars $N_{t+1},E_{t+1}\in\mathbb{R}$ using the $L_{2}$ norm $\left\lVert.\right\rVert$ as follows

Here, $N_{t+1}$ in Equation (4) is the distance that the server model $\bm{w}_{t}$ moved from the $t$-th round to the $(t+1)$-th round. Whereas, $E_{t+1}$ in Equation (5) is the average of the norms of each local weight updating vector $\lVert{\Delta\bm{w}}^{k}_{t+1}\rVert$ on the $k$-th client. Consequently, we can think of $N$ as the amount of server model updates, and $E$ as the (average) amount of local updates in clients.

The following proposition expresses the relationship between $N$ and $E$.

As Proposition 1 follows trivially from a recursive application of the Pythagorean inequality, we leave a formal proof to the appendix. The main idea behind the proposition is that, while each client locally proceeds the learning process by an average of $E_{t+1}$, after model aggregation, the server only learns by $N_{t+1}$, which is guaranteed to be less than $E_{t+1}$ by Proposition 1.

The previously described learning behavior can be better illustrated through Figure 2. The important observation here is that, since each client learns locally without online communication, their updating direction diverges. As the server averages these local models in FedAvg, by the formulation of Equation (3), the local models tend to cancel each other out, resulting in a extremely small updating vector on the server (i.e., $N\ll E$). We refer to this phenomenon as the weight updating direction divergence (WUDD) problem.

As a demonstration of the NWDA technique, we show the calculated $N$ and $E$ values using MNIST dataset. Figure 4 show how the sizes of $N$ and $E$ change over the communication rounds (i.e., the learning epochs). Initially, the NN tries to learn the dataset through a series of weight updates, and we observe reasonably large updating norms. Since MNIST is a small dataset, on the single-client case, the learning soon converges, and both $N$ and $E$ become extremely small due to the learning convergence (note that $N=E$ for the single-client case). In contrast, while $E$ is large for non-IID FedAvg, due to the WUDD problem, $N$ is relatively small throughout the communication rounds. We also emphasize on the fact that, as the learning proceeds, $N$ becomes convergent (extremely small $N$) while $E$ remains large. This indicates that some learning updates are still available, but cannot be learned (or aggregated) by the server through the simple averaging technique in FedAvg.

To further study the norm properties of single-client, IID and non-IID FedAvg, we plot $\mathcal{N}$, the integral of the amount of norm updates on the server with respect to $t$, i.e., $\mathcal{N}=\int_{\tau=0}^{t}\left\lVert\bm{w}_{\tau+1}-\bm{w}_{\tau}\right\rVert$ in Figure 4. The single-client case shows a clear trend of fast and convergent learning curve, while both IID and non-IID FedAvg do not show clear signs of convergence across the communication rounds.

In Section 3.1, we analyzed the learning process of FL under FedAvg as a series of aggregations of vectors representing weight updates ${\Delta\bm{w}}^{k}_{t+1}=\bm{w}^{k}_{t+1}-\bm{w}_{t}$. Based on our analysis, we can say that the average of the weights across clients (i.e., $N_{t}$) obtains a reasonable large norm only when all clients update in the same direction. If the weight updates diverge (i.e., the WUDD problem), learning cannot be executed effectively and efficiently in a simple FedAvg setting.

While no existing works explicitly derive a norm analysis, some techniques [19, 5, 15, 16] try to implicitly address the WUDD problem. For example, as mentioned Section 2.3, it is shown that if additional datasets (with different class labels) are distributed to the clients along to be learned along with the local datasets on the client, the accuracy improves [19]. With NWDA, we can simply interpret this method as reducing the divergence between the updating weight vectors across independent clients by introducing common datasets. Based on the observation from [19], in [14], FedProx is proposed to reduce the difference between client models. FedProx proposes to add a normalization term $\frac{\mu}{2}\left\lVert\bm{w}^{k}-\bm{w}_{t}\right\rVert^{2}$ in the loss function on each client as

where $Q$ is the original loss, and $\mu$ is a hyperparameter. The basic idea of FedProx is to penalize weight updates that largely increase the distance (i.e., norm) between the locally learned model $\bm{w}^{k}$ and the distributed model $\bm{w}_{t}$. Essentially, by introducing a penalty term in the loss function, FedProx only reduces the size of $E$, and the diverging direction problem is not addressed. However, $N$ is what actually matters to the evaluation accuracy of the aggregated model. As long as $N$ is much less than the required amount of weight updates to fully perceive a dataset, the aggregated model cannot obtain a reasonable level of accuracy.

The main difference between the FedNNNN protocol and the conventional FL protocol is that, we separate the evaluation model and the aggregated model. As shown in Figure 6, after the clients locally produce their models, there are two aggregation steps.

• •

  ➃ Average Aggregation: The same aggregation in FedAvg is adopted to produce a single model that is capable of performing highly-accurate inference based on the aggregated models.

• •

  ➅ Normalized Aggregation: The proposed normalized aggregation technique is used to resolve the WUDD problem where $N\ll E$. More details on the complete FedNNNN aggregation is provided in Section 4.2.

All other steps remain the same as in the conventional FL protocol. Lastly, we note that the proposed normalized aggregation technique is only good as a foundation of learning. In fact, due to our modification on the norm of the weights (as seen in the next section), prediction accuracy on the aggregated model from step ➅ remains extremely poor.

As discussed in Section 3.2, according to our NWDA, the main reason that existing FL schemes fall short on non-IID datasets is that averaging models with divergent directions significantly reduces the updating norm of the aggregated model. Hence, in FedNNNN, we propose to normalize the updating norm as

Consequently, the norm of the normalized updating weight vector becomes

In order to control the impact of the norm normalization, we introduce a hyperparameter $\beta$. In the experiment, we empirically decide on the value of $\beta$.

A conceptual illustration of our normalized aggregation is depicted in Figure 6. We assume that the weighted average $\sum_{k=1}^{m}\frac{n_{k}}{n}\Delta\bm{w}_{t+1}^{k}$ gives us a reasonably correct updating direction, but with insufficient norm. We amplify the norm of the updating vector in that particular direction with the size of $E_{t+1}$. The reason that $E_{t+1}$ works as a proper normalization factor stems from our observation in Figure 2, where we see that the size of $E$ remains large after $N$ converges in the non-IID case. As a result, the amount of unfinished learning can be expressed through the size of $E$, and if $E$ approaches to zero, we can safely conclude that no significant client update is available anymore. Note that when $N$ is extremely small, e.g., 0 or close to 0, we do not normalize the updating vector and return $\bm{w}_{t}$ as is.

While the proposed normalized aggregation technique described above is successful in advancing the learning process in a given direction, we still find that, due to its democratic nature, FL is susceptible to local extrema. To avoid the weights being stuck in a low-accuracy state, we add a momentum term to the aggregation process that is reminiscent to the momentum stochastic gradient descent (SGD) method [13]

Here, $\gamma$ is the same hyperparameter as in momentum SGD that expresses the amount of past weight changes to be memorized. Equation (10) defines the $(t+1)$-th round momentum $\bm{d}_{t+1}$ to be a scaled version of the average weight updates (note $\bm{d}_{0}=0$), and these weight updates accumulate as a driving force to prevent the learning from being stuck in non-optimized states. Hence, the aggregation formula of our normalized aggregation with momentum adjustment becomes

and we leave the complete description of the entire protocol to the appendix.

We use the MNIST dataset that contains 10 classes of grey-scale handwritten images of numbers, and the CIFAR10 dataset which classifies RGB images into 10 classes. The parameters used in the following experiments are summarized in Table 1, where the notations are commonly used in most FL methods. Total Rounds represents the total number of communications between the server and the clients. $\eta$ and $\lambda$ are learning rate and weight decay parameters, respectively. We used SGD without momentum for optimization. Due to the limited space, the exact valuations of $\mu$, $\beta$ and $\gamma$ which used in each FL method will be listed in the appendix. We split the dataset into training and test data, where the training data are assumed to be inherently held by the clients and test data are evaluated at server side. The hyperparameters $\eta$, $\lambda$, $\mu$, $\beta$, $\gamma$ were chosen to have the highest accuracy on the test data over the training rounds.

Here, we outline the architecture of CNN models used in this section. For MNIST, we used a CNN architecture with 2 convolutional layers followed by 2 fully connected layers. For CIFAR10, a CNN with 3 fully connected layers following 6 convolutional layers with batch normalization layer is used. Convolutional and fully connected layers of both the MNIST and the CIFAR10 architectures have bias terms. In Table 1, BN indicates whether batch normalization is used. As a separate preprocessing, we normalized input images by subtracting the mean and dividing by the standard deviation.

In this work, we take data distribution into account when performing the experiments. There are four types of data distribution: i) IID, ii) non-IID, iii) balanced, and iv) unbalanced. For i) IID, clients possess images of all of the defined classes in the dataset. In contrast, for ii) non-IID, clients have images of only 2 classes (the minimum number of classes that allow meaningful learning). For the iii) balanced condition, all clients have dataset of the same size. On the other hand, iv) unbalanced condition assumes that the dataset sizes of the clients follows a power-law distribution, where a small number of clients hold most of the training images. For privacy reasons, however, we assumed that the server does not know the data size of the client and set $n_{k}/n=1/n$ for all $k$.

Experiments are conducted once for each dataset. For our parameter setting, simulation time for total communication rounds took an hour for MNIST and 3.5 hours for CIFAR10, respectively, using GeForce GTX 1080 Ti.

We evaluated five FL aggregation methods on two datasets with the aforementioned four different types of dataset distributions in the experiments. The accuracy comparisons between various experimental conditions are summarized in Table 2. In the table, Equation (8) is used by Norm-Norm method and Equation (10), (11) are used by Momentum. B and UB indicate conditions iii) balanced and iv) unbalanced, respectively.

Regardless of the datasets and distributions, the proposed methods achieve the best prediction accuracy. Although the momentum appears to be relatively strong, it can be seen that the accuracy can be further improved by combining norm normalization. The changes of accuracy as functions of communication rounds are shown in Figure 9 for non-IID - B CIFAR10. We can find FedNNNN to significantly improve the convergence rate, thus speeding up the learning process.

We expect the amount of server update $\left\lVert\bm{w}_{t+1}-\bm{w}_{t}\right\rVert$ in FedNNNN to be larger than that of FedAvg as discussed in Section 4.2. Figure 9 and 9 confirms this expectation. In these figures, $E$ is an average clients update norm $\sum_{k=1}^{m}\frac{n_{k}}{n}\left\lVert{\Delta\bm{w}}^{k}_{t+1}\right\rVert$ and $N$ is the updated norm on the server side $\left\lVert\bm{w}_{t+1}-\bm{w}_{t}\right\rVert$. We can clearly observe that the amount of update on the server in FedNNNN is larger than the $N$ in FedAvg and FedProx. In Figure 9, we see that $E$ of FedNNNN is smaller than that of FedAvg and FedProx, indicating that the weight updates of clients have become smaller because of better initial models. However, in Figure 9, $E$ of FedNNNN appears to be oscillating. This indicates that the weights of the convolutional layers are less likely to converge when compared to fully connected layers. A thorough evaluation of $E$ and $N$ of the fully connected layer and convolutional layers for all settings will appear in the appendix.

Here, we show the complete protocol of FedAvg. The server and clients operate as shown in Algorithm 1 and Algorithm 2 in FedAvg, respectively. The server executes Algorithm 1 as follows.

• •

  Line 1: The server initializes model $\bm{w}_{0}$ and distribute the model to each client.

• •

  Line 2: For each round $t=0,1,2...$, the following steps are repeated.

• •

  Line 3–4: The number of clients to train in this round, $m$, is calculated, and the selected set of clients $S_{t}$ is generated.

• •

  Line 5: For each client $k\in S_{t}$:

• •

  Line 6: The server sends $\bm{w}_{t}$ to the clients.

• •

  Line 7: After the clients train their local models, the server receives trained model $\bm{w}_{t+1}^{k}$ for each client $k$.

• •

  Line 9: The server aggregates models averaging them as $\bm{w}_{t+1}\leftarrow\sum_{k=1}^{m}\frac{n_{k}}{n}\bm{w}^{k}_{t+1}$.

On the client side, the following steps are implemented.

• •

  Input: The client receives the server model $\bm{w}_{t}$ and its label $k$.

• •

  Line 1: the client first obtains some local datasets $\mathcal{P}_{k}$ (Note that this dataset is not from the server). The client splits the entire dataset into mini-batches $\mathcal{B}_{0},\mathcal{B}_{1},\cdots,\mathcal{B}_{B-1}$. We note the set of the mini-batches as $\mathcal{D}=\{\mathcal{B}_{0},\mathcal{B}_{1},\cdots,\mathcal{B}_{B-1}\}$.

• •

  Line 2: For each epoch $i$ from 1 to $\mathcal{E}$:

• •

  Line 3: For each mini-batch $\mathcal{B}\in\mathcal{D}$:

• •

  Line 4: The client trains weight as $\bm{w}\leftarrow\bm{w}-\eta\nabla l(\bm{w};\mathcal{B})$.

• •

  Line 7: The clients send trained weight to the server.

The protocol of FedNNNN is similar to that of the FedAvg protocol. First, we note that the client-side procedure remains unchanged, i.e., exactly same as Algorithm 2. The server-side protocol is also similar but different only for lines 9–12 in Algorithm 3.

• •

  Line 9–10: $N_{t+1}$ and $E_{t+1}$ are calculated.

• •

  Line 11–12: the server adds the update vector $\beta\frac{E_{t+1}}{N_{t+1}}\sum_{k=1}^{m}\frac{n_{k}}{n}(\bm{w}^{k}_{t+1}-\bm{w}_{t})$ to the momentum and updates weight as $\bm{w}_{t+1}\leftarrow\bm{w}_{t}+\bm{d}_{t+1}$.

In this section, we show a formal proof of Proposition 1. Before delving into the proof of Proposition 1, we first outline an important lemma.