CatFedAvg: Optimising Communication-efficiency and
Classification Accuracy in Federated Learning

Federated learning (McMahan et al. 2016) is a decentralized machine learning method, which uses distributed training on clients without sending the raw data to the central server. We can split this into three categories of vertical federated learning, horizontal federated learning and federated transfer learning.

There have been improvements proposed which, for example, improve privacy protection by the usage of differential privacy techniques (Abadi et al. 2016) or remove the bias of client models using agnostic federated learning (Mohri, Sivek, and Suresh 2019).

In FL, there are several rounds of communication between the clients and the server during training. We have DL models that have millions of parameters. Thus, the larger size of the updates will lead to higher communication costs, which can be a training bottleneck. Further, unreliable network conditions and asymmetric net connection speeds will result in delays from the client.

We can use model compression (Wang et al. 2018) which is used widely in distributed learning where techniques like sparsification, quantization or subsampling (Stich, Cordonnier, and Jaggi 2018) are to reduce the model update size. However, our goal is to maintain the quality of the trained models (Caldas et al. 2018).

Alternatively, we can use selective communication where only relevant or essential updates are transmitted back (Tao and Li 2018). This technique saves communication costs and at times, improves the Global performance. We see the usage of random subsampling of clients in the sketched updating method (Konecný et al. 2016) to improve communication efficiency.

Federated learning has a few statistical challenges. In traditional distributed ML, the central server has access to the complete dataset. Allowing the server to split it into subsets with similar distributions; we transfer these to the clients for distributed training. However, this approach is impractical for FL, as the dataset is local to the client.

In this setting, the clients may have datasets with variable distributions, i.e. the datasets are non-IID. We have cases (Zhao et al. 2018) where FedAvg is unable to achieve the desired accuracy in these conditions. When the dataset is highly skewed and non-IID, we have a solution where a common dataset is sent to the clients by the central server. However, this may not be possible in all cases.

We also see the challenge of global imbalance where the collection of data held across all the clients is class imbalanced, which leads to a decrease in model accuracy. We see the Astrea framework (Duan 2019) trying to solve this problem, where clients initially send their data distribution to the server. Then we have a rebalancing step where each client performs data augmentation on minority classes. There is a mediator that selects the best clients that contribute to the uniform distributions. This approach has shown accuracy improvements.

There have been solutions like FEDPER (Arivazhagan et al. 2019) using concepts from Multi-task learning to learn separate but structurally related models for each participant. We see the training of another set of personalization layers being trained by the client locally on its data. This has been used for building recommender systems, and testing on Flickr-AES dataset (Deng et al. 2017) has shown it to outperform FedAvg. Even though federated multi-task learning has been shown to be an effective paradigm for real-world datasets, it has been applied only to convex models. We compare our algorithm performance with VIRTUAL algorithm (32) with that has been designed for the federated multi-task learning with non-convex models in non-IID conditions.

We have a federated system with $M$ clients with datasets $D_{1},D_{2},D_{3}\dots D_{M}$ with a central server $S$. A fixed number of $K$ clients can be sampled from $M$ clients to participate in updating the global cloud model. The global loss function $F_{s}$ is related with $K$ participating clients with loss function $F_{i}$ for $i_{th}$ client as shown below.

Let $F_{i}c_{j}$ denote loss contribution of $j^{th}$ class id from $i^{th}$ client. The loss contribution from $i^{th}$ client can be expressed in terms of losses of each category for a global problem with C categories, as shown below.

Using eqs. (1, 2) we have,

Ideally, $KF_{s}$ should have loss contribution from all categories, but in real-world skewed Non-IID conditions, some classes may not be covered by the usual random sampling technique of FedAvg technique. Hence, there may exist cases where $F_{i}c_{j}$ is not defined for any value of $i$ $\in[1,K]$

We define a novel method of aggregation which is focused on maximizing category coverage. Let $G_{ji}$ denote loss contribution of $j^{th}$ class from $i^{th}$ client and let $G_{j}$ denote the total loss of $j^{th}$ class id for all clients $i\in[1,K]$.

We select clients in such a way that $G_{j}$ is defined for atleast one client with id $i\in[1,K]$. This is ensured using bitmasking strategy. Hence, global loss function of server $G_{s}$ can be represented as shown below.

Using eqs. (4, 5) we have,

We can state that (3, 5) are equal. The inherent nature of averaging the loss contribution remains unchanged with our formulation.

We have primarily used vision datasets, and all details of the datasets used for our experiments have been summarized in Table 1.

• •

  MNIST: (Deng 2012) The dataset contains a total of 70,000 small square 28x28 pixels grayscale images of handwritten single digits between 0 and 9. The subdivision of dataset is 60,000 training samples and 10,000 testing samples. We federated the dataset amongst the clients to simulate client conditions.

• •

  FMNIST Fashion MNIST (Xiao, Rasul, and Vollgraf 2017) abbreviated as FMNIST is a dataset of Zalando’s article images consisting of a training set of 60,000 examples and a test set of 10,000 examples. Each example is a 28x28 grayscale image, associated with a label from 10 classes. It is relatively challenging dataset and simulates the everyday computer vision tasks. The classes consist of fashion objects T-shirt/top, Trouser, Pullover, Dress, Coat, Sandal, Shirt, Sneaker, Bag and Ankle Boot.

• •

  KMNIST-10: Kuzushiji-MNIST-10 (KMNIST-10) (Clanuwat et al. 2018)consists of 70,000 28x28 pixels grayscale images of handwritten Japanese characters of Hiragana. The subdivision of this dataset is 60,000 training samples and 10,000 testing samples. It is a well-balanced dataset across all classes and is relatively difficult to learn as compared to MNIST. We use this to simulate the federated settings.

• •

  FEMNIST: The FEMNIST dataset is the federated version of EMNIST dataset (Cohen et al. 2017). It is a set of handwritten character digits derived from the NIST Special Database 19 and converted to a 28x28 pixel image format and dataset structure that directly matches the MNIST dataset. We used the balanced version with ten single-digit classes between 0 and 9 inclusive, 26 uppercase alphabets and 11 lowercase alphabets. We federated this 47 classes dataset for the simulation of real-world federated settings.

• •

  KMNIST 49: Kuzushiji-MNIST 49 is a much larger, but imbalanced dataset containing 48 Hiragana characters and one Hiragana iteration mark. It is an excellent example to simulate the universal problem of an imbalanced class problem under federated settings. It consists of 49 classes 28x28 grayscale images. It is a comparatively much larger dataset.

We simulate various client data distribution conditions to compare the algorithm. We take a total of 100 clients in the Federated system. We vary the number of clients having a particular category. In all of our experiments, all clients have a proper subset of categories varying from clients having a single category to clients having half the number of total categories. There are five distribution conditions (D₁, D₂, D₃, D₄, D₅) which we use and have been shown in figure 1. We share specific details in the Appendix.

All networks employed are multilayer perceptrons (MLP) with two hidden dense Flipout layers with 100 units in case of MNIST, single hidden layer of 512 units in case of FMNIST and KMNIST-10 and a single hidden layer of 784 units in case of FEMNIST-47 and KMNIST-49. Activation function Rectified Linear Unit (ReLU) is employed between the hidden layers and input layers, while the softmax activation function is employed in the final output layer. The final output layer consists of 10 units in case of MNIST, FMNIST and KMNIST-10. In FEMNIST and KMNIST-49 number of final layer units are 47 and 49 respectively.

We define two modes of experiments per algorithm where we set the value of the maximum number of clients to sample. We define Mode A where N is kept at ten and Mode B, where we keep N as equal to the number of categories in the dataset. This allows us to explore conditions where we may want to sample as few clients as possible. We will denote the algorithms for performance strategy to be A1 and cost strategy to be A2.

Each client has 600 training samples, and batch size is 32 is used in the execution of subroutine of local refinements for all experiments. In the case of a random client selection method, which is the default FedAvg case, a total of 10 clients are selected for participation in global training.

For all cases, the learning rate is used as 0.003, and a total of 50 rounds of global refinements are performed using Stochastic Gradient Descent (SGD) with Federated Averaging (FedAvg) method of updating the global model.

Here we describe the various insights about the experimental results as presented in Tables 2 and 3.

We discuss the number of clients selected, as that reflects the communication efficiency along with the model accuracy numbers under various conditions. We further performed an ablation study which has been shared in the Appendix.

As described in 2 and 3, the performance of Algorithm 1 under mode B is the best across the entire spectrum under all conditions.

If we denote the number of categories in global problem begin solved as C, the relative performance of proposed algorithms is as below.

If N$\geq$C: Mode A-A1 $\approx$ Mode A-A2 $\approx$ Mode B-A1 $\approx$ Mode B-A2

If N$<$C: Mode B-A1 $\geq$ Mode B-A2 $\geq$ Mode A-A2 $\geq$ Mode A-A1

Communication cost is related to the efficiency of a particular algorithm when it comes to deployment, and the cost of refining the global model increases with the number of clients participating in the training round. 4 describes the number of clients selected in various distributions.

Considering all distributions, $D_{1}\dots D_{5}$, we can see that number of clients participating is least in case of Mode A-A2. We limit the clients participating in mode A-A1 to N. In contrast; Mode B-A2 has a comparatively higher number of participating clients when the number of global categories is less than N in mode A-A2.

Let C be the number of global categories. The following is the relative order of algorithms in terms of communication cost expense.

If N$\geq$C: Mode A-A1 $\approx$ Mode B-A1 $\geq$ Mode A-A2 $\approx$ Mode B-A2

If N$<$C: Mode B-A1 $\geq$ Mode B-A2 $\geq$ Mode A-A1 $\geq$ Mode A-A2

We can see that the better performance of mode B-A1 comes at a higher communication cost. If we want the best performance, we need to compromise with communication cost. Alternatively, if we desire a communication efficiency, then Mode A-A2 is desirable. A better compromise is achieved with both performance and communication cost may suggest Mode B-A2 as the best one.