Asyn2F: An Asynchronous Federated Learning Framework with Bidirectional Model Aggregation

The training process of a model with Asyn2F is presented in Algorithm 1, which is executed by the main thread at the server. The server initializes a global model, denoted as $W_{0}^{\text{global}}$. The server notifies all the workers joining the training to download the global model and start their training (i.e., execute notify_workers()). As discussed in the framework architecture, let $Q$ denote the queue that contains the local models of the participating workers. While all the workers execute their local training process, the server waits for local models sent from the workers and keeps checking the aggregation condition (i.e., execute check_aggregation_cond($Q$)). If the aggregation condition is satisfied, the server performs model aggregation and produces a new global model, versioned as $W_{i}^{\text{global}}$, $i=1,2,\dots$. The aggregation can be triggered periodically or based on the size of the queue that contains notifications received from the workers. For instance, the server updates the global model whenever there is a new local model submitted. The server can also update the global model when there are $k$ models received from $k$ workers (i.e., the case of FedAvg [7]). If we set $k$ to the number of workers joining the training process, the training becomes entirely synchronous training in which the server waits for all the workers to submit their local model before doing aggregation.

The training will be terminated once the stopping condition is satisfied. Conventionally, the training is terminated when it reaches a pre-defined number of training epochs. In practice, users may consider the trade-off between the model performance and the training cost (e.g., time or computing resources of the workers). Thus, we incorporate two other different stopping conditions including: (i) maximum training duration, and (ii) desired performance of the model.

The training process is illustrated in Fig. 2. The five workers receive the (initialized) global model (illustrated by light green squares) and start the training process at different time instants. Combined with the heterogeneity in computing resources and the size of training datasets, the workers complete their training epoch and send the local model (represented by pink squares) to the server at different time instants. For instance, Worker 2 finishes its training epoch very early compared to Worker 1 and Worker 3 while Worker 4 spends a longer time to finish its training epoch. Without waiting for the new global model, those workers completing their training epoch early will immediately start a new training epoch with their current local model (e.g., Worker 2 and Worker 3). This significantly reduces the idle time of workers and improves resource utilization. Given that the aggregation condition is satisfied (periodically triggered at time $t_{1}$, $t_{2}$, and $t_{3}$ as shown in Fig. 2), the server performs global model aggregation and creates a new version of the global model. The server will then notify all the workers to download the updated global model. In Fig. 2, Worker 1 just finished its training epoch, thus it will take the updated global model and start the next training epoch. Worker 2 and Worker 3 started the subsequent training epoch earlier, thus requiring a local model aggregation to accommodate the updated global model to the local model that is currently being trained. The local model aggregation process is denoted as yellow circles. The local models obtained from the local model aggregation process (denoted as purple circles in Fig. 2) at the workers will also be sent to the server for the next global aggregation. In this way, our framework avoids the problem of obsolete information at the worker that takes a longer time for a training epoch while the global model has been updated many times by the local models obtained from other workers that complete the training faster.

Given that the aggregation condition is satisfied, the server starts the global model aggregation process shown in Algorithm 2. We assume that the server has the metadata of the private datasets of the workers such as the size and quality of the dataset. Worker $k$ has a private dataset $D_{k}$, which has its quality denoted by $Q^{\text{data}}_{k}$, and the number of data samples is denoted as $s_{k}$. It is to be noted that there exist many tools for determining the quality of data and the server can rely on them to request the report of data quality from the workers. Furthermore, metrics for the quality of data are different for different applications. There are many metrics that have been discussed extensively in the literature, such as class overlap, feature relevance, timeliness, completeness, and duplicates explained in detail in [23], which can be used to estimate the quality of data. Technically, the implementation of data evaluation can be diverse [24].

For a local model received from a worker (denoted as $W_{k}^{\text{local}}$ for the local model received from worker $k$), the algorithm computes its contribution to the global model, which depends on the quality of data of the worker, the size of the dataset (i.e., the number of samples), the value of the loss function that is used in the training at the worker and the delay of the local model compared to the global model. The contribution ratio is computed as follows:

where $Q_{k}^{\text{data}}$ and $s_{k}$ are the quality and size of the dataset of worker $k$. $L_{k}$ is the value of the loss function obtained after the training at worker $k$. $i_{k}$ is the version index of the global model that worker $k$ has downloaded and used in its training process, and $i$ is the version index of the resulting global model after aggregation as shown in Algorithm 1. Thus, $i-i_{k}$ is the version delay at worker $k$ in incorporating the newly-released global model in its training. For instance, $i-i_{k}=1$ means that worker $k$ has used the latest version of the global model in its training. Thus, the rationale behind Eq. (1) is that if the value of the loss function and the version delay is large, the contribution of the local model obtained from the worker to the global model should be small. This allows the training to avoid the problem of incorporating obsolete local models received from slow workers into the global model that has been updated by other faster workers. The contribution of each local model will be normalized before being used in updating the global model as shown in the last step of Algorithm 2.

It is to be noted that if a fast worker submits multiple versions of its local model, the latest version will be considered for the global model aggregation. As shown in Fig. 2, before time $t_{2}$ has submitted two versions of its local model. The first one represented by the purple color square is the result of a local aggregation process when the server releases the global model aggregated at time $t_{1}$. The second one is the result of the subsequent training epoch without aggregating with any new global model.

Given a new version of the global model, the workers use Algorithm 3 to incorporate the global model into their training immediately. In the mini-batch training approach, the training process at the workers takes a batch of data from the dataset and trains the model using stochastic gradient descent. After each batch, the workers check if there is a new version of the global model released by the server. If so, the workers will download the global model and aggregate it into the local model before starting a new batch.

Without loss of generality, we assume that worker $k$ is updating the latest global model denoted as $W_{i}^{\text{global}}$ where $i$ is the version index of the model to the local model that is being trained. $W_{i}^{\text{global}}$ has been computed by the server using Algorithm 2 and it is contributed by a number of workers whose average quality of data is $\overline{Q^{\text{data}}}$. The total number of samples in the datasets of those workers is denoted as $S_{i}$ and the average loss of the local models contributing to $W_{i}^{\text{global}}$ is denoted as $\overline{L}$. We also assume that the local model that is being trained has been initiated from the global model versioned $i_{k}$ ($W_{i_{k}}^{\text{local}}$). This model has been trained on the private dataset of worker $k$ for $j$ data batches.

The algorithm computes a weighted parameter to determine the contribution of the latest global model and that of the local model at batch $j$ to the new version of the local model. The formula used to compute this parameter is defined in Eq. (2) where $\beta\in(0,1)$ is a hyper-parameter that defines the importance of the quality of data or the performance of the models represented by the loss.

The above formula takes into account two important factors which are the quality of data and the model performance represented by the loss. Rationally, the higher the quality of data at worker $k$, the higher the contribution of the current local model ($W_{i_{k},j}^{\text{local}}$) to the updated version. If the loss of the current local model is low, it will also have a higher contribution to the updated version of the local model.

For every parameter (weight) of the model, the algorithm calculates the movement distance and direction between two data batches and between the latest global model and the local model. For instance, between two data batches, the movement distance of weight $w$ is calculated as $e^{\text{local}}(w)=W_{i_{k},j}^{\text{local}}(w)-W_{i_{k},j-1}^{\text{local}}(w)$. If $e^{\text{local}}(w)$ is greater than zero, it means that $w$ increases from batch $j-1$ to batch $j$ and vice versa. Similarly, the movement distance and direction between the latest global version and the local version is calculated as $e^{\text{global}}(w)=W_{i}^{\text{global}}(w)-W_{i_{k},j}^{\text{local}}(w)$. Given the values of $e^{\text{local}}(w)$ and $e^{\text{global}}(w)$, weight $w$ of the local model at batch $j$ (i.e., $W_{i_{k},j}^{\text{local}}(w)$) will be updated accordingly. If $e^{\text{local}}(w)$ and $e^{\text{global}}(w)$ have the same sign (i.e., $e^{\text{local}}(w)\times e^{\text{global}}(w)>0$ or both the local model and global model move in the same direction), $W_{i_{k},j}^{\text{local}}(w)$ will be updated directly to $W_{i}^{\text{global}}(w)$. Otherwise, $W_{i_{k},j}^{\text{local}}(w)$ is computed based on the contribution from both the global model and the current version of the local model as follows:

where $\alpha_{k}^{\text{local}}$ is defined in Eq. (2).

The updated local model returned by Algorithm 3 ($W_{i_{k},j}^{\text{local}}$) continues being trained on the next data batches at worker $k$ until completing the training epoch and notifies the server the availability of a new local model.

As the workers can join in and leave the federated learning framework at any time, a monitoring dashboard that supports real-time visualization of the model performance and/or loss value of each worker is necessary. Due to asynchronous communications between the server and workers, as well as the different sizes of datasets, a training epoch may take a longer time for a particular worker than the other, this monitoring dashboard should also report to end-users the current status of the workers and the server about training time and joined in/out timestamp. This helps end-users to stop training when the trained model attains a desired performance or stop a worker if the contribution of the local model is no longer needed. This feature is useful especially when model training incurs a cost, e.g., the cost of computing and storage resources, communication cost, and the cost of training data.

To provide real-time performance monitoring, a testing worker is needed as the server usually does not have a test dataset. This testing worker can belong to the user who requests the model such that the obtained model attains the desired performance of the requester.

The communication cost (either time or monetary cost) is a critical problem of federated learning if models are sent back and forth between the server and the workers frequently. To reduce the impact of this problem, the federated learning framework should allow the workers to train the local model independently using their data until reaching a desired performance or a required number of iterations (i.e., epochs), pre-defined as a training parameter. If the local model of a worker attains this threshold, the worker can start exchanging its local model with the server for global aggregation. This is particularly practical, especially for the models that are trained from scratch and normally have low performance in the first few epochs.

Due to the heterogeneity of computing resources and the size of the training dataset of each worker, the training time for an epoch is different among workers. Since the learning rate is decayed depending on the number of epochs performed, the above issue leads to the fact that the local models are trained at the workers with different learning rates. Those local models when aggregated in the global model affect the stability and convergence speed of the global model. To address this issue, the decay function of each worker must be updated according to the number of training epochs of the global model.

As shown in Fig. 1, there are 4 components: server, worker, storage service, and queue in our federated learning framework. In Fig. 3, we present the sequence diagram showing the interaction among components during the entire process of model training. The server starts the training process by creating a bucket on the storage service for model exchange and registers the queue channel to publish/subscribe the messages from/to the workers.

As discussed earlier, the workers can join the platform and start the training process at any time. When a worker joins, it first registers the queue for message publishing/subscribing, and once the queue registration is successful, it sends an init message to the server via the queue. The format of the init message of the worker is presented in Listing 1. With this message, the worker informs the server about its role in the framework (i.e., trainer or tester based on the requirement R2), its hardware configuration, and metadata of the dataset (i.e., size/number of samples, quality of data). When the server receives the init message from a worker, it creates a WorkerProxy instance and responds to the worker with necessary information packaged in a JSON message presented in Listing 2. This message includes the following information: i) storage information; ii) the current version of the global model and threshold for model exchange (based on the requirement R3); and iii) the aggregation strategy at the server. It is to be noted that the developed framework supports not only the model aggregation approach proposed in our work but also other existing approaches such as FedAvg [7] and M-Step KAFL [16].

During the training process, whenever the server completes the aggregation of the local models received from the workers, it uploads the new global model to the storage and notifies all the workers about the availability of a new global model by sending the message shown in Listing 3. The message contains information on the current version of the global model, the average quality of data, and the average loss of the models that were aggregated in the new version of the global model. This information is needed for the workers to perform local aggregation of the global model to the local model as shown in Algorithm 3. Upon receiving this message, all the workers will need to download the new version of the global model, perform local model aggregation, and continue the training process.

Upon completion of a training epoch and the local model attains the performance threshold for contribution to the global model, the worker uploads the local model to the storage server and uses the message shown in Listing 4 to notify the server. This message contains information about the performance and loss of the local model, the version of the global model that has been used for training the local model. It is to be noted that during a training epoch of the local model, multiple versions of the global model could have been used at different local model aggregations. When notifying the server, the worker takes the latest version of the global model that has been downloaded to include in the notification message.

Abstract Class for Model Definition. To support the training of various deep learning models defined by the users of the framework, we define an abstract class, namely Model, which allows users to define their own neural network architecture and implements 3 required methods: get_weight(), set_weight(), and train() - for one mini-batch. An optional method for a tester is also defined, namely evaluate() allowing users to implement the test function based on their performance metrics.

Messaging System for Communications. We used RabbitQM¹¹1\urlhttps://www.rabbitmq.com/ in our framework. This avoids the use of socket programming and ephemeral ports, which are usually blocked by security systems, e.g., firewalls. All delivered messages are executed by the routing_key concept of RabbitMQ. Each worker has two connections to the queue. The first one asynchronously subscribes to the queue using its own queue name and the same routing_key to listen to the messages from the server. The second one is used to asynchronously publish messages to the server. Similarly, the server also uses two connections to the queue for receiving/publishing messages from/to workers.

Object Storage. Our framework supports both S3²²2\urlhttps://aws.amazon.com/s3/ and Minio³³3\urlhttps://min.io/. In addition to the two main functions that are uploading and downloading models, our framework also supports several additional functions such as the creation of buckets, and deletion of models that are no longer needed. This is useful in the case of S3 as it is chargeable.

Monitoring and Report. To meet requirement R1, the monitoring service is implemented by subscribing to messages from workers, the tester, and the server. From these messages, the service extracts necessary information such as worker_id, timestamp, loss value, and performance based on a pre-defined metric, and pushes them to InfluxDB⁴⁴4\urlhttps://www.influxdata.com/ and uses Grafana⁵⁵5\urlhttps://grafana.com/ to visualize the information in real-time during the training process.

Synchronizing Learning Rate. To meet the requirement R4, the decay function is defined by using the version of the global model. Both information (i.e., version of the global model and learning rate) are included in the message that the server uses to notify the workers (see Listing 3).

Supporting Multiple Model Aggregation Strategies. In our framework, WorkerProxy keeps all the necessary information of a training worker. It also provides an abstract class namely Strategy with an aggregate() method and a train() method. With this abstract class, apart from the aggregation algorithms presented in previous sections, our framework additionally implemented the existing algorithms such as M-Step KAFL [16], AsynFedED [17] or FedAvg [7]. Respectively, the server also implements the aggregate() method for the above techniques. This makes our framework effective for not only asynchronous training techniques but also synchronous ones.

We set up the experimental infrastructure that is composed of 1 server, 1 testing worker, and 10 training workers. The details of location and computing resources are given in Table I. All the training techniques (Asyn2F, FedAvg, and M-Step KAFL) use the same infrastructure.

Data Preparation. We use both iid and non-iid random methods to split the CIFAR10 dataset into 10 sub-datasets. When using the iid random method, the sub-datasets are created with two 2 scenarios:

• •

  In the first scenario, each sub-dataset is randomly selected from the original dataset by a random ratio from $20\%$ to $40\%$. This leads to the fact that there is an overlap among 10 sub-datasets, and the total data size of 10 sub-datasets is greater than the original size (i.e., $50,000$ samples).

• •

  In the second scenario, there is no overlap among the 10 sub-datasets. However, we make sure that every dataset should have samples of all the labels.

When using non-iid random method, all the sub-datasets are non-overlapping. The data distribution of one of the experiments is presented in Fig. 4 where the colors represent the labels of the datasets.

Model architecture and hyper-parameter setting. We use ResNet18 model [27], implemented in TensorFlow. We use a stochastic gradient descent optimizer with the momentum parameter set to 0.9. We run the experiments with 3 settings of learning rate (LR).

• •

  Fixed learning rate: All the workers use the same learning rate for the entire training set to 0.01.

• •

  Synchronously-decayed learning rate: The learning rate is decayed along the training process by the server using the cosine decay function with an initial value set to 0.1. The updated learning rate is synchronized with all the workers.

• •

  Asynchronously-decayed learning rate: Each worker uses the cosine decay function by its local training steps with an initial value set to 0.1. The updated learning rate is used locally without synchronizing with other workers.

For each training approach, we run the training process for 3 hours. Each worker trains 5 iteration rounds of local data as an epoch. For Asyn2F, the global model aggregation process at the server is triggered when there are three local models received from three different workers (always get the latest version). Whereas, M-Step KAFL waits for three models without differentiating workers. The client selection strategy of FedAvg is all clients, i.e., 10 workers.

Analysis of Results. In Table II, we present the accuracy of the ResNet18 models trained by our training technique and the other two existing techniques. The experimental results show that Asyn2F achieves the best performance compared to the other two existing techniques. The improvement is more significant (up to 2.90%) in the case of non-overlapping and non-iid sub-datasets, which is usually a practical situation where data providers have different (diverse) datasets. Comparing three methods of learning rate setting, we observe that synchronizing the learning rate among workers achieves the best performance. This reflects the requirement (R4) of the training framework discussed in Section 4.1. It is worth mentioning that the existing works (FedAvg and M-Step KAFL) do not address this issue in their work. When integrating their training techniques into our framework, we additionally implement their techniques with the three learning rate setting strategies, thus providing insightful performance comparison.

Looking into convergence speed, we present in Fig. 5 and Fig. 6 the performance of the global model obtained by the testing worker over training time in case of overlapping and non-overlapping sub-datasets, respectively. The experimental results show that the global model trained by Asyn2F converges faster compared to the models trained by the other two techniques regardless of the learning rate setting methods. This is explained by the fact that the local model aggregation in Asyn2F enables a faster convergence of local models trained by the workers and avoids the information obsolete issue at the workers. In the long run, all the models will achieve maximum performance (which could not be further improved by any chance). But with a desired performance (e.g., $90\%$ of accuracy), a fast convergence of the model benefits a lot. For instance, the users can stop the training when the global model achieves the desired performance. From the training cost point of view, this could lead to a lower training cost for both time and monetary cost of computing resources, especially in the context of federated learning marketplaces.

A fast convergence of the model also implies a lower communication cost. In Fig. 7, we present the number of training epochs carried out at the workers and the server until the global model achieves $90\%$ of accuracy. The experimental results show that while Asyn2F achieves a faster convergence, all the workers and the server need fewer training epochs compared to M-Step KAFL to converge to the desired performance. This implies that Asyn2F incurs a lower communication cost as the workers and the server need to exchange the obtained models after each training epoch (at the workers) and after each aggregation at the server. The ResNet18 model has a size of $42.7$ MB, which incurs a significant communication cost, especially when we use Amazon S3 as the storage service, which charges monetary cost based on the amount of data uploaded/downloaded from the cloud. It is worth mentioning that the higher the number of models exchanged (stored), the higher the storage cost as well. To reduce the storage cost, those intermediate (local and global) models (obtained prior to the final global model) could be deleted.

It is to be noted that, in Fig. 7, we do not present the number of training epochs of FedAvg, which requires the same number of training epochs for all the workers and the server (10 epochs to attain $90\%$ of accuracy). However, due to the synchronization among workers, the time duration for each training epoch is very long due to the heterogeneity of computing resources of the workers and their dataset. This explains the reason why FedAvg spends a longer time to converge to the stabilized performance of the model. This demonstrates the practicality of our technique and the implemented framework, which can work very well in a heterogeneous infrastructure including both network connectivity and computing resources.

In this section, we perform the experiments with a larger training infrastructure including $20$ heterogeneous workers. The details of the infrastructure are presented in Table III. It is worth mentioning that several workers are running CPU resources, which are much slower than GPU resources when training large deep learning models on large datasets.

Data Preparation: We split the CIFAR10 dataset into 20 sub-datasets with non-iid splitting method. The data size and label distribution of one of the runs are shown in Fig. 8.

Analysis of Results: As shown in Table IV, we obtain the same performance behavior of the proposed technique and the existing ones. Asyn2F outperforms both M-Step KAFL and FedAvg. Obviously, compared to the case with fewer workers, there is a performance drop due to the scarce distribution of the same number of training data samples to a larger number of workers. In Fig. 9 (captured from the monitoring dashboard of the framework), we present the performance evolution of the local models trained by the workers despite small training datasets. All the local models converge to a stabilized performance thanks to the knowledge obtained from the global model, which is in turn aggregated from other local models. The curve represented by the tester is the performance of the global model that stabilizes after a few hours of training. We also observe an interesting behavior of the framework in this experiment. Due to a large number of training workers and small local training datasets, the workers complete their training epoch very fast, especially the workers with GPU resources. These workers quickly submit their local models to the server for aggregation and continue the next training epoch with the new version of the global model aggregated by the server. The issue of obsolete information does not incur with these workers. Nevertheless, the larger the training framework, the more heterogeneous the workers. Those workers with less powerful computing resources will suffer the consequence of the obsolete information issue, thus significantly benefiting when joining our framework.

It is to be noted that M-Step KAFL needs 10 local models for every global aggregation to achieve an accuracy of $92.10\%$. With a fair comparison with 3 local models for every global aggregation, M-Step KAFL achieves an accuracy of only $83\%$. Further experiments with 7 local models for every global aggregation, M-Step KAFL achieves an accuracy of $89\%$. This makes M-Step KAFL less practical in a heterogeneous federated learning environment.

EMBER [26] is a malware dataset including 1 million samples, each being represented by a vector of 2381 feature values. There are 300000 benign samples (labeled as 0), 300000 malicious samples (labeled as 1) and the remaining 400000 samples are undefined. We remove the 400000 undefined samples and use 600000 labeled samples for our experiments. We split the data into 7 sub-datasets using the non-overlapping and non-iid splitting methods. We adopt MalConv model [26], whose detailed parameters are presented in Table V and train it on the training framework with 7 training workers and a tester as shown in Table VI. We note that with only 2381 features representing a sample, the length of the feature vector of the EMBER dataset is much smaller compared to CIFAR10. Even with a deep learning model described in Table V (which is also much smaller than ResNet18 used for CIFAR10), the workers with GPU resources complete a training epoch very fast. This explains why we split the original dataset into only 7 sub-datasets so that the training time of one epoch on a worker will be long enough to demonstrate the heterogeneity of the framework.

The performance reported in Table VII demonstrates the effectiveness of the training framework as well as the aggregation algorithms. We obtain similar performance behavior as with CIFAR10 such that Asyn2F has a superior performance compared to the existing ones. The results also confirm the fact that synchronizing the decayed learning rate among training workers results in the best performance.