FedFMC: Sequential Efficient Federated Learning on Non-iid Data

As a mechanism for devices to update a global model without sharing data, federated learning bridges the tension between the need for data and respect for privacy. However, classic FL methods like Federated Averaging struggle with non-iid data, a prevalent situation in the real world. Previous solutions are sub-optimal as they either employ a small shared global subset of data or greater number of models with increased communication costs. We propose FedFMC (Fork-Merge-Consolidate), a method that dynamically forks devices into updating different global models then merges and consolidates separate models into one. We first show the soundness of FedFMC on simple datasets, then run several experiments comparing against baseline approaches. These experiments show that FedFMC substantially improves upon earlier approaches to non-iid data in the federated learning context without using a globally shared subset of data nor increase communication costs.


Introduction
Over the past decade, improvements in machine learning accuracy has met rising consumer demand for use of learning models. Models were shown to be successful in a plethora of learning tasks, ranging from prediction-based recommendation to a variety of image processing tasks. Standard machine learning applications flourished in optimized processes where data would flow from user devices to centralized data centers for analysis.
However, the rise of mobile and IoT (internet-of-things) devices has come with a changing dynamic resulting in new needs and constraints for machine learning. These include: (1) tighter restrictions [1] on usage of data in relation to user privacy, (2) demand for lower latency in many applications such as remote surgery, and (3) constraints of poor internet connectivity for accessibility in under-developed countries and offsite rigs (for research or oil exploration).
A wide variety of approaches have strived to address these issues. Most early work focused on privacy conscious learning, some of which formed the foundational work of using a centralized server to learn local work without sacrificing privacy [2]. In 2016, Federated Learning (FL) was introcued by McMahan et al. [3] to address these needs. It outlined an approach of keeping data distributed on edge devices, running a global model on each device, and incorporating model updates from each device through a novel method named Federated Averaging. Federated Learning has proven to be groundbreaking in many applications, leading to the birth of many avenues for further research.
Although immediately useful in many contexts, Federated Learning and in particular Federated Averaging depend on an iid global dataset for success in learning. Indeed, McMahan et al. cited non-iid data as a major obstacle in their seminal paper [3]. Historically, non-iid data has not been a problem. In normal machine learning workflows, data centers have access to the entire global dataset and thus distributes data among machines in a carefully controlled iid manner.
Unfortunately, the central server in Federated Learning does not have the same luxury as normal data centers as it lacks direct access to data. Since there are many real-world scenarios which local edge devices will not have data drawn IID arXiv:2006.10937v1 [cs.
LG] 19 Jun 2020 from some shared global distribution, researchers have struggled to address this problem. Current approaches either make sacrifices by requiring a globally shared subset of data or increasing communication costs with more models.
In this paper, we tackle supervised learning tasks drawing examples (x, y) ∼ P i (x, y) from non-identical edge device distributions. More specifically, our method handles bias in label skew, when the marginal distributions of labels P i (y) varies across edge devices while P(y|x) is identical [4].
We restrict our definition of non-iid to define edge devices belonging to certain "archetypes," defined by a set of labels of data that is over-represented by a quantity represented by a bias parameter. Therefore, these archetypes represent label skew.
Our contributions are as follows: 1. Present FedFMC (Fork-Merge-Consolidate) as a novel approach for learning non-iid data in FL.
This builds off of previous distributed learning work by not only supplying local models but also performing training locally on each device. Hence, FedAvg empowers edge devices to collaboratively learn a shared prediction model while keeping all training data local.
Although successful in classic Federated Learning tasks, FedAvg falls short in many real-world scenarios. On non-iid data, the FedAvg algorithm is prone to conflicting updates and performance oscillations as the averaged model is likely to prefer a certain archetype over another. Updates from devices of one archetype tend to undo the updates from another, and as such the averaging function in FedAvg struggles to learn a significant portion of the labels. As a result, FedAvg's slow convergence and low accuracy deem it unusable in most non-iid contents.

Current Approach to non-iid FL
A simple and successful approach to non-iid federated learning is to create a globally-shared dataset of separate origin from the edge devices [4]. This dataset is globally-shared between all devices. As such, during the FedAvg process devices are able to update to the global model with less conflicts, enabling the network to learn non-iid data. In particular, Zhao et al. showed in [5] that globally sharing 5% of data resulted in a 30% increase in accuracy on experiments with the CIFAR-10 dataset. A version of this globally-shared dataset is found in most state-of-the-art solutions to non-iid FL.
Unfortunately, this approach breaches scenarios of strict Federated Learning. It assumes the accessibility and validity of a such subset of data to generalize to the problem at hand. This is inapplicable to many use cases as there are many legal, social, and technical barriers in sharing data representative of all devices. Thus, this approach may in fact increase the likelihood of unfair device bias and raise susceptibility to attacks through data poisoning.

Lifelong Learning and FL
Another approach is to consider learning a classifier on each edge device as a separate learning problem and apply lifelong learning techniques [4]. Lifelong learning is the task of learning separate tasks f 1 , f 2 , ...f n sequentially using a single model, without forgetting the previously learned tasks.
While there have been many recent developments in lifelong learning, a particularly successful one is Elastic Weight Consolidation (EWC). First introduced by Kirkpatrick et. al, EWC is a method to aid the sequential learning of separate tasks using the same model without "forgetting" tasks that were previously learned [6]. EWC identifies the parameters that are most informative for learning a task and penalizes changing these parameters during the training of future tasks [6]. Specifically, for the lifelong learning of two tasks A and B, the objective of training task B would be where I(θ * A ) is the diagonal of the Fisher information matrix [6,7]. The current approaches to applying lifelong learning techniques require that every device be present in every round of training, and full device participation is not necessarily guaranteed in the real world [4,7]. In addition, these multi-task learning approaches typically involve decentralized or peer-to-peer learning schemes, which also assumes complete edge device-edge device connectivity, which is not possible in most privacy-conscious applications [4,8].

The Federated Fork-Merge-Consolidate (FMC) Algorithm
In the line of decentralized peer-to-peer FL, we model the non-iid learning problem as a lifelong learning task. However, instead of considering each edge device as a separate "task" to learn, we design a scheme, FedFMC, to first group devices with similar archetypes together and treat each group as a task to learn.
The FedFMC algorithm has two phases: forking and merging. Here, we denote a (t) n and l (t) n to be the validation accuracy and validation loss, respectively, of device n at time step t. g (t) n is the group that device n belongs to at time step t and θ (t) j is the averaged model associated with the group j. h f and h m are tune-able hyper-parameters for the forking and merging thresholds, respectively. T is the total number of rounds of forking, K is the number of devices participating in each round of a total of N edge devices, and M is the maximum number of rounds per merging of two models. Algorithm 2 details our methodology for dynamically forking a single global model for T rounds. Building off of FederatedAveraging, we first train a subset of K devices every round, update the global model accordingly, and send updates to each edge device. Then, in certain rounds t eligible for forking, each device dynamically sorts itself into the best model through a multi-part process. (1) It first compares its validation loss to other devices reporting to the same model. (2) If its loss is notably higher, it runs all other global models to try to find another model that it has a lower loss. (3) If all the other models perform even worse than the model it is using currently, it forks into a new model. This process is done for each device in the network.
Since devices of similar data skew (archetype and bias) will behave similarly in terms of validation loss, the FORK algorithm encourages dynamic grouping of devices. As such, at the end of T rounds of forking, the global network will have devices that self-forked themselves into groups of similar archetypes. Since FORK does increase communication costs, a certain number of rounds have to pass between rounds eligible for forking. We analyze this added communication costs in more detail in section 4.5. n then n is placed in a new group Algorithm 3 dictates our approach for merging and consolidating several forked models into one. Note that after FORK, there are several models that each have learned a separate data skew / archetype. Then, our consolidation process starts off with one model and iteratively merges models into the globally shared model. As mentioned above, we employ the use of Elastic Weight Consolidation to remember previously learned archetypes as we incorporate more archetypes. At the end, we will have one globally shared model that learns all archetypes.
Note that our work in merging and consolidate models into one is well motivated. Past work have sometimes utilized a process of giving every device a separate model and then boosting device accuracies. However, if the models are not consolidated into one, many advantages of Federated learning are lost. For instance, the ability to receive globally shared updates prove useful in scenarios where the edge device can benefit from other data sources. Indeed, without consolidation, there is little difference between those scenarios and devices which train a model locally offline.
Algorithm 3: MERGE and CONSOLIDATE forked groups into one global model. Input: T, K, N, and M as defined above. Output: θ consolidated , a single consolidated global model that performs well for data from all devices and their archetypes. active_devices ← (n with g (T ) == 0) for f = 1, 2, . . . , number of groups -1 do i ← 0 active_devices ← active_devices + (n s.t. g for all active_devices do w_avg ← AverageWeights(n ∈ active_devices) All devices n ∈ active_devices ← w_avg round_devices ← a random subset of active_devices EWCTrain round_devices for E local epochs each return θ consolidated 4 Results

Experimental Setup
We used subsets of the CIFAR-10 Dataset [9] with a random group of M = N/2 devices participating in each round of training. The only comparative baseline with similar communication costs and underlying assumptions is the FedAvg algorithm. For the EWC of the Merging phases, we used a λ = 1 Number of Groups . This experiment was our proof of concept to show the necessity of different parts of the FedFMC algorithm. We reduced the CIFAR-10 dataset to a very simple example: three archetypes each with only one label (0, 1, or 2). In this case, the forking method consistently forks correctly, so the fork groups are equivalent to groups of devices with the same archetype.   Figures 1 and 2 depict the test accuracies of each fork group using the single consolidated global model on a dataset that has uniform representation of all three labels. It is clear that there is a stabilizing effect of the EWC: without EWC the per-group performance is not stable while learning other groups. As the merging rounds successively incorporated the different groups into a single global model, their performance varied together when the training incorporated EWC (rather than decreasing or oscillating). Figure 3 depicts the results of the FedFMC algorithm on the simple 3-archetype dataset as compared to the Federated Averaging baseline. It's clear that FedAvg has relatively converged compared to the baseline, which has significant oscillations. Between both the forking stages and the merging stages, the performance of the different devices and archetypes are more in-line with each other rather than jumping significantly from round-to-round. In this setup, we simulated devices that had three archetypes, with data drawn iid from the labels [0,1,2,3], [4,5,6], and [7,8,9]. An example of a real-world situation where this type of data may be present is in regional hospitals in different parts of the country, where the composition of patients is relatively uniform within a subgroup of the overall population.

Communication Costs
The forking rounds has similar communication costs to the FedAvg algorithm with the overhead of determining the best model for devices that are performing poorly under their current group's model. The number of updates is as follows: u = FedAvg Updates · Number of Active Devices · Number of Forking rounds = E * K * T We see that in the worst-case scenario, the communication costs are O(T 2 · N ) as compared to the FedAvg communication costs of O(T K + T N ). In practice, even though the opportunity is presented to fork many times, the devices have been sorted into correct or approximately-correect groups such that a new group is not added at every opportunity.

Model Size
One potential concern over solutions for learning non-iid data is an increase in model size. Other solutions have also utilized multiple models on each device, which linearly increases the amount of memory demanded on devices. This is a problem as many edge devices have limited memory storage.
However, we note that FedFMC does not increase the model size. Through the FORK and MERGE algorithms, each device only stores one model at a time (although that model does change).

Conclusion
Federated learning has shown promise as an approach in recent years to address learning problems with privacy constraints. One of the biggest barriers to adoption is the poor experimental performance of traditional algorithms like FedAvg in situations without full device participation and without peer-to-peer edge device connections. Our algorithm, FedFMC, forks edge devices that perform similarly to update different global models, which are then consolidated using EWC through a merging phase.

Discussion
There were many tuneable parameters of FedFMC algorithm, including, but not limited to: • The criteria for forking eligibility: currently the criteria is defined as at least 5 rounds from the start, more than 5 rounds from the end of forking, and at least 4 rounds between forking milestones.
• The threshold for an edge device moved from its original group: currently, if l (t) i ) then the model is moved to the group with the best-performing model on its local dataset (and if that group is the same as the original group, it gets moved into a new group).
• The criteria at which to stop the merging of a certain group: currently, a group stopped merging into the global model if the average accuracy of the global model over the previous 5 rounds was less than 1 away from the maximum accuracy of the global model over the previous five rounds of merging.
The optimal functions and conditions for these parameters were found through experimentation.

Our Contribution
The goal of this new approach is to develop a new FL scheme that can adequately learn a single model that performs well on all archetypes without significant communication costs.
By terminating the FedFMC algorithm at the end of the Forking stage, one would get personalized models that work well for a particular archetype while still learning from all the data and not explicitly defining which devices fall into which archetype.
By allowing the FedFMC algorithm to continue to completion with the Merging stage, one would get a single model that performs well on all of the individual local archetypes.

Future Work
We also forsee several extensions of our work. Our proposed FedFMC algorithm breaks through current barriers in both Federated Learning and Lifelong Learning by through initial proof of concept experiments. The dynamic nature of FedFMC presents difficulties in using classical analysis but we would like to see further work in formalizing the intuition of forking and consolidation. In particular, developing a rigorous approach to bound the number of rounds needed for FORK and convergence in MERGE-CONSOLIDATE proves to be nontrivial yet useful.
Besides from analysis, there are many promising areas of improvement for the FMC algorithm. On top of finetuning parameters, there is great potential in experimentation with smarter FORK algorithms and optimizing the MERGE-CONSOLIDATE process. Experimentation is needed to find forking algorithms that quicken separation devices of different archetypes -there are a wide variety of forking algorithms to choose form, ranging from heuristic to learned models.
Other avenues of research extending from FedFMC include addressing open FL questions in fairness and robustness. Through the dynamic forking process, FedFMC may hold advantages over existing approaches in ensuring fair device representation. Moreover, FORK presents devices a mechanism to self-distance from malicious devices in the network. This empowers robustness in defense against data poisoning and Byzantine attacks.