{'fraction_non_alphanumeric': 0.02433662449861154, 'fraction_numerical': 0.007713668620796051, 'mean_word_length': 4.919073165163794, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Aergia: Leveraging Heterogeneity in Federated Learning Systems

Federated Learning (FL) is a popular approach for distributed deep learning that prevents the pooling of large amounts of data in a central server. FL relies on clients to update a global model using their local datasets. Classical FL algorithms use a central federator that, for each training round, waits for all clients to send their model updates before aggregating them. In practical deployments, clients might have different computing powers and network capabilities, which might lead slow clients to become performance bottlenecks. Previous works have suggested to use a deadline for each learning round so that the federator ignores the late updates of slow clients, or so that clients send partially trained models before the deadline. To speed up the training process, we instead propose Aergia, a novel approach where slow clients (i) freeze the part of their model that is the most computationally intensive to train; (ii) train the unfrozen part of their model; and (iii) offload the training of the frozen part of their model to a faster client that trains it using its own dataset. The offloading decisions are orchestrated by the federator based on the training speed that clients report and on the similarities between their datasets, which are privately evaluated thanks to a trusted execution environment. We show through extensive experiments that Aergia maintains high accuracy and significantly reduces the training time under heterogeneous settings by up to 27% and 53% compared to FedAvg and TiFL, respectively.


Introduction
Federated Learning (FL) is a decentralized and inherently privacy-preserving learning paradigm where clients collectively train a machine learning model [3,22]. During a learning round, a federator selects a subset of the clients that return an update of the global model computed using their local dataset. Upon receiving client updates, the federator aggregates them into a global model update, which is then shared with all clients. Most of existing aggregation algorithms, including FedAvg [22] and FedProx [21], are synchronous, and require the federator to collect all updates from the selected clients before moving to the next training round.
In a practical FL system, clients might have heterogeneous computational resources and possess data that differ both in quantities and class distribution. It has been shown that both resource and data heterogeneity negatively impact the performance of a FL system [12,14,21,40]. First, relying on a mix of weak and strong clients instead of homogeneous clients to train a model can significantly prolong the training time [6]. Second, a classification model trained with federated learning is less accurate when the client datasets are non independently and identically distributed (non-IID) [7].
To mitigate the impact of weak clients, also called stragglers, the state-of-the-art methods attempt to equalize the learning speed amongst the clients by (i) partitioning them based on offline profiling [6], or by (ii) dropping the updates of stragglers during the training rounds [20,24]. The former approach may fall short in capturing transient heterogeneity caused by applications possibly collocated on the clients, whereas the latter might incur a severe accuracy degradation. Moreover, the impact of stragglers is further aggravated when encountering non-IID data among clients. Indeed, stragglers might possess a unique dataset that is critical to the overall model accuracy. In addition, due to the privacy preserving nature of FL, it is not really possible for the federator to infer the data distribution based only on the clients model updates [21,33]. To limit the risk of model divergence, prior studies aggregate the non-IID client data by adding a regularization term, like in FedProx [21], or by estimating their contributions, like in FedNova [33]. However, these works implicitly assume that the client nodes are homogeneous.
In this paper, we aim to accelerate the FL training of convolutional neural networks (CNN) in presence of stragglers and non-IID data. A CNN is composed of convolutional layers and fully connected layers [18], which respectively learn the representation of local data and map the extracted representation into classes. The local training of CNN entails forward and backward passes on both types of layers.
To retain the representation of the unique datasets of stragglers, we advocate to freeze their convolutional layers, and offload the computing and updating of the convolutional layers to strong clients. We propose Aergia 1 , a federated learning algorithm that monitors the local training of selected clients and offloads part of the computing task of stragglers to strong clients that have spare and idle capacities. Aergia relies on a client matching algorithm that associates a straggler to a strong node based on an estimated performance gain and on the similarity between their datasets, since blindly offloading local models to nodes that have drastically different data distribution leads to weight divergence [7]. To ensure privacy, data similarities are securely evaluated using the clients' local data distributions (i.e., the number of labels per class) in an Intel SGX enclave [8], which is hosted by the federator.
We implement Aergia in PyTorch as a middleware running on top of Kubernetes. We evaluate Aergia on three datasets, namely MNIST, FMNIST, and Cifar-10, on different network architectures against four previous heterogeneity or non-IID aware aggregation solutions [6,21,22,33]. Our FL systems consist of a mix of 24 weak, medium and strong nodes that use a different number of CPU cores. Our evaluation results show that Aergia achieves the highest accuracy within the lowest training time.
In a nutshell, this paper makes the following contributions: • We explain how a straggler can offload the training of its model to a strong client. • We present an algorithm that matches the performance profile and data similarity of clients. • We design Aergia 2 , a federated learning middleware for highly heterogeneous clients and non-IID data that leverages model training offloading and online client matching. Aergia relies on a trusted execution environment (an Intel SGX enclave) so that the federator can evaluate the similarity of client datasets without getting access to their private class distribution.
• We evaluate Aergia on a FL cluster built on top of Kubernetes. Our evaluation results on three datasets and several networks show that Aergia effectively leverages the spare computational capacity of strong clients to achieve high accuracy in low training time.
The remainder of this paper is organized as follows. §2 provides some background on Federated Learning, data and resource heterogeneity, as well as on their impact on training time and accuracy. §3 provides an overview of Aergia, while §4 describes its algorithms and implementations details. §5 presents our performance evaluation. §6 reviews the related work. Finally, §7 concludes this paper.

Background and Motivation
In this section, we first recall necessary background on deep learning models, which are core components of the federated learning paradigm, the practical heterogeneity challenges that federated learning faces and their impact on training time and accuracy.

Premier on Convolutional Neural Networks
The state-of-the-art image classifier follows the structure of convolutional neural networks (CNN) [18], which consist of convolutional and fully connected layers. The former maps the image features into a compact representation, hence they are also referred to as feature layers. The latter are dense fully connected layers that classify the representation into one of the classes. The other difference between these two types of layers is their resource demands [9]: convolutional layers are computationally intensive while fully connected layers are memory intensive. The training time of a client's local model can be divided into two parts: the forward pass that computes the classification outcome of images, and the backward pass that back-propagates the model weights. Consequently, the training time of a typical CNN classifier can thus further be categorized into four parts: (i) ff: forward pass on feature layers, (ii) fc: forward pass on fully connected layers, (iii) bc: backward pass on fully connected layers, and (iv) bf: backward pass on feature layers.

Federated Learning
Federated learning (FL) [12,14,21,40] is an emerging decentralized learning paradigm where clients and a federator jointly train a machine learning model in consecutive rounds while local data stays on premise. In this paper, we specifically consider an image classification model that maps images, to one of labels, denoted by , through a function ( ), parameterized by weights . Prior to the training, the federator initializes the model architecture, the objective function, the training algorithm, the training hyperparameters, and the aggregation protocol for the client's Variance of CPUs Test Accuracy (c) Accuracy in a non-IID scenario without and with deadlines. Figure 1. Heterogeneous computational powers among the clients increase the duration of the FL training process (Figure 1(a)). One could use deadlines so that the federator discards late updates in a round before starting the next one, which effectively reduces the training time (Figure 1(b)). However, using deadlines badly degrades the model accuracy, in particular in non-IID settings (Figure 1(c)).
local update 3 . We consider convolutional neural networks (CNN) [18] as the classifier model. The clients train the classifier model based on their own real data, which never leaves their premises, whereas the central server iteratively aggregates and distributes models submitted from clients until reaching the global model convergence.
Local Training. In each global round , the clients receive the latest aggregated global model ( ( − 1)) from the federator and use their local data to perform local updates for epochs, e.g., stochastic gradient decent (SGD) updates [23]. The cross entropy loss function [29,30,39] is widely adopted for classification problems. Specifically, a client aims to find ( ) that minimizes the loss function: min ( ) ( ( ); , ), using the local data points ( , ), where is an input data, e.g., images, and is the class label. Upon finishing the local training, clients send their local model parameters, i.e., ( ), to the federator.
Model aggregation. After receiving all model updates from clients, the federator aggregates the clients' model parameters into the latest global model that is returned to the clients in the beginning of the next round. Specifically, at each round , a subset of clients is selected to do local training and send back their latest model weights, ( ). The aggregation algorithms differ in the frequency and weights in aggregating the local models. FedSGD [4,23] treats all local models equally, and trains the entire local data in one epoch. The gradients are sent to the federator for aggregation every epoch. To minimize the communication and avoid the divergence of local models, FedAvg [22] lets local models train for multiple epochs and then aggregate the models. Specifically, FedAvg calculates the global model of round 3 We note that a variant of FL [26] does not rely on a federator and aggregates the model in a peer-to-peer manner. We do not consider this case here. as the weighted average of all local model weights:

Sources of heterogeneity
Data heterogeneity. Clients possess different and unique privacy-sensitive datasets. A common assumption in the prior art is that client data are identically and independently distributed, which is the so called independent and identically distributed (IID) case. Taking the image data benchmark Cifar-10 [16] as an example, which contains 60,000 images from 10 classes, in the IID case each client would own an equal amount of images that would be equally distributed across classes. Recent studies point out that in practice distributed datasets are highly non-IID, and differ both in size and in distribution [1,33]. For instance, it is easier to identify clients that own horse images than deer images (both are classes in Cifar-10). Consequently, unique images like deer are owned by a small client subset, whereas common images like horse have a higher probability to be equally distributed across all clients. Such non-IID data distribution, i.e., clients owning data in different quantities and distributions, have been shown to be challenging for FL and detrimental to the accuracy of the global models [12,14,40]. The heterogeneity of a non-IID dataset can be captured by its Earth Mover Distance (EMD) [28]. The higher the EMD of a dataset, the higher the heterogeneity of the client data distribution. To mitigate the accuracy degradation in FL due to non-IID data, related studies [21,33] added regularization terms in the objective function, altering the aggregation algorithms, or augmenting the dataset. Clients resource heterogeneity. Edge devices are highly heterogeneous in their computing and network resources [34]. Their hardware and software stacks evolve after each generation, i.e., every 5 to 6 years. It is challenging to find an optimal (deep) learning model for a diversified set of devices. Due to the differences in the type and number of CPU cores and memory provisioning, the computation time of deep models on edge devices vary a lot. Moreover, the network connectivity of edge devices may often be unstable and expensive [7]. Instead of computation, communication is more a bottleneck for edge devices that have poor connectivity, questioning the effectiveness of aggregation algorithms relying on frequent communication. The heterogeneity of clients' resources essentially leads to (highly) unbalanced local training time and also long global training time, because the federator can only aggregate after receiving all the local models, including the slowest one, during the synchronous training. To alleviate either the communication or computation bottleneck, the prior art proposes to have asynchronous communication across global round [2], i.e., the federator aggregates the local updates as soon as a new update is received. The downside of asynchronous training is the risk of slow convergence and low accuracy of the global model. Thus, the state-of-the-art mainly addresses resource heterogeneity by equalizing the communication and computation time across clients so as to maintain the synchronous training and high model accuracy.

Motivation: Impact of heterogeneity on training
time and accuracy Resource heterogeneity in FL algorithms increases training time. Aergia addresses resource heterogeneity thanks to its model freezing and offloading approach. Several works also documented the negative effect of data heterogeneity among clients on FL accuracy [15,21,33], which further increases the effect of resource heterogeneity. Aergia mitigates this issue by taking into account the similarities between client datasets in its offloading scheduling algorithms.
We evaluate the influence of different degrees of CPU core heterogeneity among clients on the FL training time. We then evaluate the learning degradation that would be incurred when naively equalizing the training time of all rounds based on deadlines. For this experiment, we use the MNIST dataset, and refer the reader to Section 5.1 for additional experimental details. Figure 1(a) shows how the overall training time of various sizes of FL clusters increases with the variance between the clients CPUs. We consider 1 to 13 clients. We set the average computational capacity per client to be 0.5 CPU, and assign cores to clients with a variation shown in the figures. A high variance implies a high difference among slow and fast clients. We compute the overall training time as follows: we add the time required for any pre-training requirements (such as offline profiling if the algorithm demands it) to the time required for all training rounds. The duration of a training round is measured by the Federator, using its local clock, from the moment the request is sent to the clients until the last participating client responds with its results. One can see that a small perturbation in clients CPU core, i.e., standard deviation, can significantly extend the training time because the federator has to wait for the stragglers because of the synchronous training process. The delay that is added to the total training time grows with the variance of capacity as well as with the cluster size.
The state-of-the-art minimizes the impact of resource heterogeneity of clients by equalizing their computation time through estimating clients speed and assigning different amounts of training loads. In Figures 1(b) and 1(c), we evaluate a naive solution -terminating the training per round according to the deadlines. Clients who cannot send back their model weights on time are not included in the aggregation, which effectively bounds the training time but severely degrades accuracy. In contrast, the sophisticated equalization algorithms carefully factor in the trade-off between the accuracy and training time, biasing toward achieving a high accuracy. Motivated by this result, we aim to derive an algorithm that equalizes the training time within and across training rounds.

System Model
We consider a traditional federated learning architecture, which Figure 2 illustrates, where multiple clients are connected to a central server, and are also able to directly communicate with each other. We consider heterogeneous settings in the sense that clients might have different computational powers and be interconnected by network links that differ in terms of bandwidth and latency. We consider that the computational load of each client might evolve with time, e.g., because they are running other applications in addition to the training process that we aim at optimizing. We however assume that the network properties remain stable during the training process. We also assume all parties to be honest. We assume that communications are asynchronous but reliable,  i.e., that there is no bound on the time it takes for messages to reach their destination but all messages eventually arrive. Clients are equipped with local clocks so that they can measure the time they require to train their models. We do not assume these clocks to be synchronized, but they do need to have similar frequencies. We do not require the federator to make use of a clock. We assume the federator to be correct. It is however equipped with an Intel SGX enclave that the clients can authenticate using remote attestation [8], and to which they send their encrypted and privacy-sensitive dataset class distribution. The enclave provides code integrity and confidentiality of the data it manipulates. The federator relies on the datasets similarity matrix that is computed by the enclave to refine its offloading decisions and maintain high accuracy.
Aergia allows the training process to become aware of the shortcomings and strengths of specific clients and organize their collaboration so that they rely on each other to reduce the overall training time. The central server plays a predominant role in the sense that it coordinates the overall training process and periodically identifies the clients that should offload part of their training tasks to more powerful clients through a model freezing method.

Leveraging the Heterogeneity of the Learning
Phases A CNN can be split into two sections. The first set of layers are the feature layers and the second part are the classifier layers. During training, for each minibatch the network performs a forward pass and a backward pass, respectively in the forward and in the backward propagation phases, for both sections of the layer. Figure 3 illustrates the four phases of a training cycle: forward pass feature layers (ff), forward pass classifier layers (fc), backwards pass classifier layers (bc), and backwards pass feature layers (bf).
We have profiled the time required by each of the four phases during the network training process with several datasets. To accurately measure the computing time, our profiling time presented here is under the single client scenario. Our results indicate that the time required to perform a cycle is not uniform across phases. Figure 4 shows that the majority of the time is spent in the backwards pass of the feature layers (from 52% to 75%). Aergia therefore focuses on offloading the backwards pass feature layers phase of the weak clients to stronger ones.
Aergia leverages parameter freezing, which is commonly used for transfer learning [25], personalisation tasks [1,38] and to accelerate training [5,37]. Freezing parameters in layers eliminates the needs to perform back propagation for the frozen parts of the network.

Round training with Model Freezing and Offloading
The client selection and aggregation in Aergia is done in the same way as in typical federated learning. For each local update, the server selects a subset of the devices and sends them the central model. Aergia differs from the standard FL starting from the moment where clients start training locally. In the following, we describe the various tasks executed during a round, starting from the reception of the global model by the clients. Section 4 provides a precise description and the implementation details of these tasks.
Early training and profiling. At the beginning of a training round, clients start by executing complete batches, i.e., they use the 4 training phases, and monitor the speed at which they execute each phase. Clients then report their performance measurements, obtained from the online profiler (which directly uses the clients' clock to measure processing times), to the central server, and continue fully training their model while waiting for the server's instructions.
Centralized scheduling. When it has received the performance measurements of all clients, the server identifies the clients that would slow down the entire training process by sending their updated model later than others, and computes a schedule where slow clients offload part of their training process to more powerful and data compatible clients. This computation takes the clients dataset similarities into account by leveraging an Intel SGX enclave for privacy. The server indicates to the slow clients that they should offload their model to a stronger client, whose IP address is specified and who will train part of their model for them. The server also informs the stronger clients that should train the model of weaker clients. Messages from the server detail the global training round number so that clients can ignore late messages.
Model freezing, offloading and training. Clients are informed of the offloading they might have to execute according to the schedule the server computed, and the time at which all clients are expected to send their updated model to the server. Weak clients that need to offload their models freeze the first layers of their model, send their model to a stronger client and continue updating their model with a lighter procedure. Strong clients that are executing offloaded tasks return their own fully trained model and the (possibly partially trained) model they might have been updating for a weaker node at the time of the deadline.
Model aggregation. Upon receiving all model updates, the server recombines the models of clients that offloaded their training process: the first layers are received from the stronger client to which their training was offloaded, and the remaining layers are received directly from the weaker client. Based on the reconstructed models and based on the models that were fully trained by clients, the server uses the classical FL averaging method to compute the next global model, and start the next global update.

System Details of Aergia
In this section, we detail the key components of Aergia. We first describe Aergia's model freezing and offloading algorithm. We then detail how the federator relies on profiling results communicated by clients to schedule the models freezing and offloading that should happen during a global round. We finally explain how the federator leverages an SGX enclave to compute the clients dataset similarities to refine its schedule. Figure 5 illustrates the potential impact of task offloading on the duration of a learning round with heterogeneous clients, and allows us to explain how model freezing and offloading intervene in Aergia. At the beginning of a round, the server sends the current global model to the clients, after which they compute their local model using their own datasets. In this example, clients 3 and 4 are more powerful than clients 1 and 2, which impacts the time it takes them to execute a training phase (numbered from 1 to 4). In this example, the clients profile their performance over a number of local batch updates (Aergia uses 100 batches), and report their performance metrics to the server afterwards. After having collected all performance measurements, the server determines which clients should freeze the second part of their model and offload its training to other clients. The server then informs the clients to freeze their model and to offload the training of their last layers to a selected strong client. Symmetrically, the federator also communicates with the strong clients about the offloading decision so that they are aware of the identity of the weak client whose model they should train, and know the number of local training batches they should use to train it. While they are waiting for the server's scheduling decision, clients keep training their models using all four phases. In Figure 5, clients 1 and 2 are informed that they should respectively freeze and offload their models to clients 3 and 4. When clients 3 and 4 finish executing their local updates (here after 3 batches), they train the offloaded models they have received on their local datasets. All clients eventually return the local models they train (offloaded or not) to the server, which finally aggregates them. In this example, clients 1 and 2 execute their rounds faster than they would have without offloading because they save the execution of two phases. Note that scheduling decisions are cryptographically signed by the federator for authenticity, and that they contain a monotonically increasing sequence number so that they cannot be replayed and so that messages sent by the federator that arrive late (i.e., in the next round) are ignored.

Online Profiling
Minimizing the maximum training time of all clients during a local update is equivalent to minimizing the variance between the client training times. To reach this goal, Aergia gathers information about the training speed of each client over the first local updates of a global training round. This information is initially not available to the federator, and might also evolve over time as clients might dedicate a various proportion of their computing power to the learning task. During the profiling phase, clients keep track of the execution time of the forward-and backward propagation of each layer. These execution times are then communicated to the server and allows it to obtain a clear picture of the current computational power of the clients, and identify the clients that are slowing down the whole learning process. The profiling phase should include sufficiently enough batches to report accurate numbers, but not too many so that there is still headroom to optimize the remaining part of the learning process (whose entire duration depends on client dataset sizes). Aergia profiles the training process during 100 local batch updates out of 1600 local batch updates 4 .

Centralized Scheduling
The federator in Aergia executes a scheduling algorithm that considers the training process of each client as two consecutive tasks that can potentially be executed on different clients. The first task of each client is always executed locally while the second task can be offloaded to another client. Based on the performance metrics gathered by the profiler, the offloading scheduling algorithm decides whether the second task of each client should be offloaded to another client. In Figure 5, clients 1 and 2 respectively offload the second task of their training process to clients 3 and 4. Clients 3 and 4 execute these offloaded tasks as soon as they receive them and are done computing their own model updates. Overall, thanks to the offloaded tasks, the training time of the overall round can be reduced.
When it has received the performance measurements of all clients, the server identifies the clients that would slow down the entire training process by sending their updated model later than others, and computes a schedule where slow clients offload part of their training process to more powerful and data compatible clients. The server directly informs the slow clients that should offload their training process to another client.
The main objective of the server is to minimize computational variance by correctly matching a weak client with a strong client. The central server is during the local training of the clients informed about the training performance of the clients. The clients gather these performance estimates with the online profiler while training locally on the data. Note that the profile has a very low overhead as it only represents 0.58% of the total training time. Using the performance indicators, the central server matches underperforming clients with strong clients following Algorithm 1. We estimate a mean compute time for each round based on the performance data of each client that is active in the current round. The central server uses the mean compute time ( ) as a target time for the offloading schedule. Using and the performance indicators of each client, the weak and strong clients are identified (Line 13 and 14). Pairs of weak and strong 4 The number of local upates used for the profiler depends on the size of the dataset and the batch size. For our experiments we use 100 local batch updates to limit the profiling overhead and still achieve a good profiling accuracy. : remaining local updates of client .

4:
, : Similarity values between clients and , describing pair wise data similarity. : array of times at which a client must send the partially or completely trained models to the server 9: : an array that details the stronger client to which a client should offload its model to (possibly None). : Estimated duration between two nodes. = _ 19: return , ones because the global training time in a round is determined by the weakest client. For a possible offloading option, the algorithm evaluates the optimal offloading time (Line 23). The optimal offloading time is calculated using the performance indicators of two nodes as can be seen in Algorithm 2. A strong client can only be used one time per round for offloading. Once a strong client is assigned to a weak client, the strong client is removed from the list of possible receiving nodes (Line 29). One can modify Algorithms 1 and 2 to support heterogeneous network transmission latencies and bandwidths, which is straightforward and has been omitted for space reasons.
Globally minimizing the time required to train all client models during a FL round can be seen as a | | job scheduling problem, which can be solved using dynamic programming [17]. Our algorithm is a variant of the greedy longest-processing-time-first (LPT) algorithm that distributes the training tasks over the clients, which has been demonstrated to produce schedules that are close to optimal [11]. In addition, LPT scales linearly with the number of clients, and therefore does not significantly slow down the federator.
Training procedure. The scheduling algorithm is used by the federator to orchestrate the training process. The actions of the Federator during a training round consist of the following steps. First, the client selection, in the same manner as with FedAvg, is performed. The selected clients are by the Federator instructed to start training for rounds with the online profiler active. Each client after number of local batch updates informs the federator about the performance metrics that are gathered with the profiler. Simultaneously, the clients stop the online profiler and continue training until further instruction. With all the performance metrics that the federator receives, an offloading schedule is created using Algorithm 1. The clients are informed by the federator about the offloading schedule. The model aggregation is performed when the federator has received a training result from all selected clients. This process is repeated for rounds.

Refining Schedules with Data Heterogeneity
So far, the scheduling algorithm we have presented does not take into account the fact that clients might possess non-IID local datasets. In practice, offloading the model of a client to a stronger client with a vastly different dataset would be detrimental to the training accuracy. The central server therefore also takes into account the similarity of two datasets to compute the model freezing and offloading that should take place during a round. The similarity between datasets is computed by the federator using the data distribution over the classes. We implemented this process using a trusted execution environment, namely Intel SGX, but it could also be implemented using cryptographic methods such as homomorphic encryption. We demonstrate that using dataset similarities allows Aergia to improve its model accuracy and leave such privacy-preserving extensions to future work. The local data distribution of a client is privacy sensitive, therefore, we calculate the similarity matrix inside a trusted execution environment (i.e., an SGX enclave) of the federator node. Clients send an encrypted vector that contains the number of labels per class to the enclave. The trusted execution environment computes the pair-wise similarity between pairs of clients based on the EMD metric. The result is a similarity matrix , a × triangular matrix where is the number of clients. The similarity matrix enables Aergia to improve the offloading algorithm without publicly sharing sensitive information.
The cost function on line 24 of Algorithm 1 takes the data similarity of two clients into account. The similarity matrix is calculated before the training starts, and it contains the pair wise similarity of the local data using the EMD metric. Variable on line 24 of Algorithm 1 is used to control the effect of the inter client similarity. With = 0, the scheduling only takes the performance metrics into account and ignores the local data similarities, as discussed in section 4.3. As soon as > 0, the inter-client data similarity is taken into consideration while creating the schedule. A larger increases the influence of data similarity when determining the offloading decisions. Using this simple objective function, the scheduling algorithm associates each weak client to a suitable stronger client that will partially train its model on a dataset that is sufficiently similar to its own.

Evaluation
In this section, we summarize the evaluation results of Aergia on a testbed, where heterogeneous clients own diversified data sets and learn the CNN classifier in a federated manner. We compare the overall training time and accuracy against four state of the art solutions on three datasets. We also conduct a sensitivity analysis to demonstrate the effectiveness of Aergia under different degrees of non-IIDness.

Evaluation Setup
The central federator performs the default FL tasks (client selection and model aggregation) and also computes the client matching for the offloading of model training. During the local training, the client sends the federator performance metrics of the local training. With this information, the federator can spot stragglers that slow down the training process. Using these performance metrics, the federator matches stragglers with powerful clients in order to offload computational tasks. The matched clients communicate directly and transfer the tasks to each other. Baselines. We compare Aergia to four baselines: FedAvg [22], FedNova [33], FedProx [21], and TiFL [6]. FedAvg, FedNova, and FedProx implicitly assume that client nodes are homogeneous and that communication is synchronous. FedProx minimizes the amount of drift a client can obtain during local training to improve the convergence when learning over non-IID data. FedNova normalizes the client updates at the aggregation stage. TiFL handles stragglers on a global level by grouping the clients in tiers based on their computation speed.
Testbed. The evaluations are performed on a testbed aimed to run FL systems in a distributed environment. The testbed is implemented in Python and built on top of Py-Torch. Each node is fully isolated from each other and can only communicate through messages. All communication is asynchronous and peer to peer and based on RPC. This means that nodes can message each other directly without relying on the Federator as a relay (Most FL systems assume a star network). In our evaluation we assume a fully connected network. The testbed allows all nodes in the system to run independent and isolated from each other. We use an Aurora R13 desktop with a 5.20 GHz i9 CPU with 24 cores and 64 GB RAM. We use up to one client per CPU core. We implemented the dataset similarity computation that runs in an Intel SGX enclave in C++ and used Graphene [31].
Datasets. We evaluate Aergia using three datasets: MNIST, FMNIST and Cifar-10. The MNIST and FMNIST datasets contain 60,000 training images and 10,000 test images with a dimension of 28x28 pixels. The Cifar-10 dataset has a total of 60,000 images of 32x32 pixels split over 50,000 training images and 10,000 test images.
Networks There are three networks used to evaluate the datasets. For the MNIST dataset we use a three layer CNN with two convolutional layers and a single fully connected layer. For FMNIST we use the same model as used to evaluate MNIST. Lastly, an eight layer CNN is used to evaluate Cifar-10 consisting of six convolutional layers and two fully connected layers.
Heterogeneous Resource Setup. In the real world, machines are rarely equal in computing power. Therefore, using large discrete steps, such as 0.25, 0.5, 0.75 and 1.0 CPU, to define resource heterogeneity between nodes is not realistic. We use a total of 24 cores to evaluate all the algorithms we consider. To reproduce a real-world scenario, the CPU speed of each client is selected uniformly at random as a fraction that ranges between 0.1 and 1.0 of the original CPU speed. We use Docker containers to isolate the nodes on different cores and control their CPU speed. This allows for an realistic approximation of the aforementioned scenario.
Heterogeneous Data Distribution. The clients that participate in the Federated Learning process may have nonuniform data distributions. To evaluate performance under non-IID heterogeneity, clients sample 3 classes out of the 10 available.Clients are then biased towards their local dataset during training. Local client datasets are disjoint, i.e., they do not share images. This approach is used for all three datasets when evaluating the non-IID scenario.

Accuracy and training times
We divide our evaluation into two scenarios: IID and non-IID on three datasets. We report the accuracy and training time in Figures 6-7.
IID: When the local training data has a IID distribution, the effects of the resource heterogeneity are negligible. In figure 6(b)-a we compare the accuracy of FedAvg, FedNova, FedProx, TiFL, and Aergia on the FMNIST dataset. The accuracy after 100 rounds is comparable among the aforementioned algorithms. In contrary to the accuracy values, the results of the real-time duration show that there is a noticeable difference between Aergia and the rest. On average, Aergia is able to do the same amount of training as FedAvg and TiFL in 27% and 45 % less time respectively.
Non-IID: The impact of non-IID data is intensified by the system heterogeneity and leads to a significant longer training times. This effect can be seen in Figure 7(f). Aergia is able to reduce the training time per round and thereby reduce the overall global training time by 27% compared to FedAvg and 53% compared to TiFL. The same trend is visible in Figures 7(d) and 7(e). The accuracy reached in non-IID scenarios is comparable between the Aergia and the other algorithms with the exception of FedNova. Figures 7(e) and 7(f) show that TiFL is unable to prevent the slowdown of the FL system, in contrary to what is reported in [6]. We think this is caused by the higher variation of CPU power in this scenario. TiFL is in this scenario not able to reduce the variance in CPU using their tier solution due to a higher intra-tier CPU variance. Overall we observe that Aergia is able to reduce the training time up to 53% while achieving the comparable accuracy as the state of the art non-IID aware aggregation algorithms.
To better understand the performance advantages of Aergia, we zoom into the detailed performance of FMNIST, i.e., the time distribution per training round. Figure 8 shows the density of the rounds duration during the training process for Aergia and the baselines we consider. Aergia's distribution is shifted to the left compared to all baselines, which indicates its ability to minimize the duration of rounds during training.

Rounds duration and impact of the similarity factor
Aergia uses the similarity factor in Line 24 in Algorithm 1 to control the impact of the inter-client data similarity. A value of 0 means that the inter-client similarity has no effect on the scheduling while > 0 increases the impact of similarity in Algorithm 1. A higher value of restricts the number of favourable strong offloading options. Figure 9(b) shows that the average round time decreases when similarity is not taken into account. A low values of has impact on the global model accuracy as can be seen in Figure 9(a). A positive value of increases the global model accuracy in Aergia.
We evaluate learning FMNIST on 24 clients, 3 of which are selected in each round. Figures 9(a) and 9(b) report the impact of the similarity factor parameter on accuracy and on the training time, respectively. With a similarity factor of 0, the scheduling algorithm ignores the inter-client similarity and uses performance indicators only. With a positive similarity factor, the inter-client similarity has more impact on the scheduling decisions. Increasing the similarity factor can restrict the available offloading options and thereby increases the average round time.

Impact of the Degree of Non-IIDness
The degree of data non-IIDness has an impact on the convergence speed during training. We control the number of classes a client can own as a measure to evaluate and create non-IID data. The lower the number of sampled classes the  Figure 9. Impact of similarity factor on accuracy and training time.
higher the degree of non-IID in the clients local data. For instance, non-IID(2) indicates each client only has 2 classes of data points out of 10 available classes. In Figure 10 we show the effect of the degree of non-IIDness on the Aergia's algorithm. All variations train for the same amount of rounds. The difference in completion time is low but the difference in accuracy is apparent. A high degree of non-IID data results in a decrease of model accuracy. This is comparable to previously reported results [6]. Profiler The profiling time of the online profiler impacts the error in the performance indicators used for scheduling. A longer profiling time offers better performance indicators, but a longer profiling time reduces the training speed and increases the training time. The profiler has a negligible overhead of 0.22% ± 0.09 on average for all models. Test Accuracy IID non-IID (10) non-IID (5) non-IID(2) Figure 10. Test accuracy depending on the level of dataset non-IIDness measured during the training process. The level of non-IID limits the number of sampled classes in a clients' local training data. The levels of non-IIDness follow the setup used in [6].

Related Work
Devices may differ in computational and communication capabilities due to hardware (CPU, memory) and network connectivity. Challenges such as straggler mitigation, client drift, and the effect of non-IID data are exacerbated by the aforementioned system-level characteristics. limits how far the local model of the client is allowed to drift from the global model. The objective function is altered with a parameter to penalize drift from the global model. Their theory shows that choosing > 0 should improve convergence when learning over non-IID data, but practice shows that this is not the case for non-IID label distributions. The solution of Li et al. [21] does not account for stragglers. FedNova [33] limits the effect of heterogeneous clients in the system by altering the global aggregation rule. Clients that perform more steps return a larger update and thus more significantly impact the global model, even when the datasets of all clients have the same size. Although Wang et al. [33] mitigate the effect that heterogeneous nodes have on the model accuracy, it is not able to remove the under representation of stragglers in federated learning systems. Table 1 summarizes our analysis of the existing works that consider heterogeneous client resources or datasets, and shows that Aergia is the only algorithm that minimizes the actual training time, thanks to its model freezing and offloading mechanisms, and maintains high accuracy even in presence of non-IID datasets, thanks to its similarity computation that is executed in an SGX enclave.

Task Offloading
Dong et al. [10] use an edge server equipped with a Trusted Execution Environment (TEE) to offload tasks of edge devices. Because of the limited capabilities of the TEE, this method is only tested on the MNIST dataset on a very small network. Besides that, the data is also randomly (but equally) distributed over the clients. The main goal of this work is to mitigate the straggler problem by utilizing the more powerful edge servers for task offloading. EAFL [13] offloads federated learning tasks from edge clients to nearby edge servers. To lighten to computational burden of weak clients, EAFL partially offloads data to an edge server, while Aergia offloads models, which is more privacy preserving. FedAdapt [35] leverages split learning to allow an IoT device to train its model with the help of the central server by exchanging partially trained models during each round. In comparison, we leverage model freezing and allow clients to offload training tasks to each other, which is more scalable, and consider the data heterogeneity issue, without which accuracy would be decreased.

Federated Learning with deadlines
One solution to mitigate stragglers is for the federated learning system to impose a deadline for the clients to submit their locally trained models. While this reduces the overall idle in the whole system, it often diminishes the contribution of the straggler by dropping late submission. Li et al. [20] use the imposed deadline to tune the local settings such as CPU frequency on each client to meet the deadline with the least amount of energy consumption. This approach lowers the amount of dropped clients, however the effect of non-IID data is not considered. Nishio and Yonetani [24] use resource aware client selection (FedCS) to limit the amount of participating stragglers. They show that this solution works mainly works in an IID data settings in large networks. When dealing with non-IID data, models trained with FedCS incur a significant drop in accuracy.

Stragglers
Different techniques can be used to minimize the impact of stragglers. TiFL [6] groups clients in tiers based on their performance characteristics. Each round a different tier is chosen for client selection. This reduces the variance between training times within each round.
Sageflow [27] uses periodic aggregation rounds to limit the decrease of training time per round caused by stragglers. Late contributions send in by stragglers are incorporated into the global model corrected with a staleness factor. The downside of Sageflow is that the method uses public data distributed across the clients.
Lee et al. [19] use adaptive deadlines to accelerate the convergence to the global model. By calculating a deadline per round based on the participating devices, the mean idle time per device is minimized. This solution still can cause stragglers to be dropped by the system, since the deadline is calculated based on the fastest participating client plus the tolerated waiting time. The authors note that their solution has no guarantees when applied on non-IID data.

Parameter freezing
Parameter freezing is a well known technique in transfer learning that is used to fine tune the model. An alternative use for Parameter Freezing is to speed up the training process. Chen et al. [7] limit the communication between the federated and edge device by freezing parameters in the early stage of the training process. However, challenges are that the local parameters excluded from global synchronization may diverge on different clients, and meanwhile some parameters may only temporally stabilize. Adaptive Parameter Freezing (APF) proposes to fix (freeze) the non-synchronized stable parameters in intermittent periods.
In order to reduce the communication cost, parameter freezing can be used. Brock et al. [5] use parameter freezing to speed up the computations by slowly freezing out the first few layers. This avoids the cost of computing the gradients and speeds up convergence. This approach is coarse and degrades the accuracy of model in spite of a speedup in training. Chen et al. [7] use parameter freezing to lower the communication cost when sending the model of weights. By freezing stale parameters the number model weights that are shared with the server are reduced. While this significantly reduces data transfers, it does not alleviate the computational burden of stragglers.
Veit and Belongie [32] use partial execution of networks to speed up the model execution during inference. By excluding some layers in the network from execution based on the input image, the inference execution time is reduced while retaining accuracy. Similarly, Wu et al. [36] use layer dropping in residual networks to create dynamic execution path to reduce the execution time during inference.

Conclusion
In practical settings, Federated Learning (FL) is largely impacted by the heterogeneity of clients, as each training round needs to wait for slow clients, which are also called stragglers. In this work, we have proposed Aergia a novel FL algorithm that leverages model freezing and offloading. In this protocol, slow clients freeze the first layers of their model and offload it to a faster client that trains these layers using its own dataset. Aergia uses a simple, yet fast and scalable, scheduling algorithm to regularly organize the offloading of models between clients. The central server uses its own dataset to decide which clients have compatible datasets to maintain high accuracy when organizing model freezes and offloadings. Our experiments on three datasets demonstrate that Aergia reduces the overall training time by up to 53% and maintains high accuracy.