Communication-Efficient Federated Learning for Neural Machine Translation

Traditionally, in machine learning problems, a model is trained and validated using data from a particular domain. Such a model usually does not generalize well on domains outside of its training dataset. In order to address this issue, transfer learning was proposed (K. Weiss and Wang, 2016). The idea is to train a neural model in a particular domain, and then use the model in a different domain. Transfer learning has worked well in many fields, and the trained models in one domain can mostly generalize to other domains either straight away or with some fine-tuning. However, when fine tuning a model in a new domain, all neural model parameters are usually updated, which might not be necessary. Houlsby et al. (2019), Pfeiffer et al. (2020), and Rücklé et al. (2020) proposed a new set of architectures, called Adapters, to address this shortcoming. Adapters are low-cost plug-ins that are mounted on pre-trained models inside their internal layers. When fine-tuning the model, only these adapter components are updated, which makes training more parameter-efficient and less expensive. However, adding adapter layers requires changing the internal architecture of the model, and the question of where to exactly insert these adapters to obtain performance improvement is not easily answered (Rücklé et al., 2020).

We argue that using FL framework to train mixed-domain NMT models (the scope of this paper) is similar to transfer learning. Collecting parameters from clients of different domains, aggregating their information, and sending back aggregated knowledge to clients is a form of transfer learning. With this understanding, we ask how can we make this transfer learning more communication and computation efficient? Given that adding adapter layers is not straight forward, we propose a simple yet effective Controller alternative which we cover in detail in Section 3.

To tackle the parameter selection problem, we propose a simple yet effective architectural addition, which we call Controllers. The Controller is another neural layer inserted in between two existing model layers. The goal of these dedicated layers is to adapt and promote information flow from the layers below them, and communicate that information to the central server. During the local update phase, these Controllers will disseminate the aggregate information from the server to the local models. Figure 1 illustrates this concept.

In the model with Controllers, only the parameters of these particular layers are communicated with the server. Not only the communication overhead is reduced, there are additional privacy gains from this approach. Moreover, Controllers help manage the computation cost on the client side. Our Controller mechanism can be used in any neural architecture, but in our experiments, we embedded them into the Transformer (Vaswani et al., 2017) architecture since it is the current state-of-the-art and a well-studied model in the community, and this makes our work easy to replicate by others.

In our models, we treat the Controller layers with two strategies during training. In the first set, we update all model parameters in conjunction with the Controller layers. In the second set, we freeze the original parameters of the model and only update the Controller layers. This is to compare the performance of the model in both scenarios and understand how much we lose in performance, if we reduce the layers involved in the training process. Regardless of which option we choose, the fact that we only communicate the Controller layer parameters to the server, and only need to update them when pushing back the information to clients, helps greatly to reduce the communication overhead. Moreover, if we freeze the original layers and only train the Controller layers, we are able to reduce the training time and computation overhead. Usually, clients in FL settings have limited resources and might not be able to train large models. Our proposed Controller layers provide benefits of an in-domain large model without incurring the computation cost on edge devices.

Rows 1 to 5 in Table 1 report BLEU (Papineni et al., 2002) scores obtained from our stand-alone engines. These models are trained in a non-FL setting with their own training sets, e.g. 6E-6D WMT is a Transformer with 6 encoder and 6 decoder layers trained on the WMT dataset. As expected, the model shows its best performance on in-domain data but it fails to successfully handle other domains’ data. The lowest BLEU score belongs to the PHP test set. Clearly, technical terms and special form of sentences in the PHP dataset is not fully understandable for the WMT model and makes it challenging for the model to generate high-quality translations. We see a similar trend for other models too, namely they perform well on in-domain data but could barely work for out-of-domain samples. However, for datasets from similar domains this performance gap is small(er), e.g. the OS model is able to handle the TED testset relatively well as there are lots of similarities between the OS (film subtitles) and TED (TED talks) sets. PHP is an extreme case in our pool; due to its size and technical genre with unique terms and special grammatical constituents, it cannot consume other domains’ sentences, so the BLEU scores are almost zero for all out-of-domain datasets.

In our setting, a single model trained on in-domain data cannot be a good translation engine to support all domains. Therefore, we train a mixed-domain engine to address this problem. Rows 6 and 7 report related results. In the Fine Tuning approach (Row 6), we combined all datasets to create a huge corpus of training and development sets. We then train a translation engine using such a collection which leads to high-quality results. The average BLEU score of this new model is 32.00 (higher than all others), which is a clear indication of its success. This model is capable of handling different domains quite well. However, this approach relies on a strong assumption, i.e. that all datasets are fully accessible at training time in a centralized repository, which might not be the case in real-world scenarios, and especially in distributed and private settings such as FL. Therefore, we propose another alternative, namely Chained Training (Row 7), where we start from a well-trained model and gradually fine-tune it with other datasets.

For Chained Training, we assume that the WMT model is our base model. We share it with a client and it can fine-tune it with its own data. After each phase of fine-tuning, the updated model is shared with another client, so that it can also add its own data. In this scenario, there is no central training and the model is gradually updated by different clients. As seen in Row 7, results obtained are not as competitive as those in Row 6. This shows how model training can be dramatically impacted in the presence of non-centralized data repositories. It should be noted that, Row 7 does not even address the privacy aspect of the problem and only tries to gradually access data from different clients. Rows 6 or 7 can be considered as a baseline for our FL model, where Row 7 is a more relevant and fair candidate. The setting defined in Row 6 is only suitable for research purposes which discards complications of real-world settings.

Row 8 in Table 1 is our first FL system where we have five clients. Each client’s model is initialized with the WMT engine parameters and is updated by the local data in each round of FL. The client model’s parameters are sent to the server at the end of each round and the server aggregates clients’ information via FedAVG. The average BLEU score of this FL system is 33.83, which was impressive and at the same time unexpected given the fact that the server has no access to clients’ data, and the entire process is carried out in a decentralized and private environment. The communication cost (C-Cost) for each client is “6E-6D” which means it exchanges 6 encoder and 6 decode layers with the server. The training cost (T-Cost) is “6E-6D + W” since every client updates all the network parameters including those of encoder and decoder, and word embeddings.

In Row 9, we introduce a new FL configuration where instead of 6 layers we use 8 (for each encoder and decoder). Controllers define dedicated and extra layers responsible for communication and model updates. This increases the number of layers from 6 to 8, so Row 9 would be a fair baseline for Controller-based models. The expectation for Row 9 was a higher BLEU score than that of Row 8, since we increased the number of layers, and in general the quality of deep learning models should increase proportionally with the increase of their size. The same assumption also holds for our model, but the larger model in Row 9 requires more steps of training to generate better results. For all these models, we ensure that at the end of the training process, either centralized or distributed, each model is only trained for a total of 150K steps.²²2The core model, WMT, is trained for 100K steps then we run extra 50K steps of fine-tuning during FL rounds. In order to have a fair setting, we also trained the 8E-8D/A-A model for 150K steps.

In Row 10, we apply the Controller idea for the first time. We freeze all layers during fine-tuning on the client side and only update Controller layers. The T-Cost for this setup is “2E-2D”, which is significantly lower than all previously reported results. We still communicate all the encoder/decoder layers in-between clients and the server, so C-Cost is still “8E-8D”. In this configuration, the first and second Controller layers are placed after the second and sixth encoder layers, respectively. We have an identical setting for the decoder Controllers. In Row 11, we have another configuration with a different positioning strategy. The average BLEU score of latter configuration is higher than the previous one, and it shows how placing Controllers in the right place can impact model quality. In Row 12, we have the same setting as in Row 10, but this time not only do we freeze Controllers during fine-tuning but we also limit the FL communication to these layers. Both C-Cost and T-Cost in this case is “2E-2D” which is considerably more affordable. Each encoder layer has 3,416,320 and each decoder layer has 4,204,032 parameters. The embedding table is a matrix with 33,116,512 parameters. Considering these numbers, the bandwidth used (for both communication and training) in Rows 9 and 12 are worth 94,079,328 and 15,240,704 parameters, respectively, which means that our Controller-based models are about 6 times less expensive. Another interesting observation is that we are able to achieve better performance than all other 8E-8D models in Row 12, even though we update/communicate fewer parameters.

In Rows 14 and 15, we challenged our proposed solution and instead of assigning additional layers we pick some layers of the 6E-6D model to play the role of Controllers. Our findings show that these 6E-6D models can perform as accurately as their 8E-8D peers, if we pick the right layers as the Controllers. However, 6E-6D models are much more sensitive to the position of Controllers.

Adaptors (Section 2.1) follow a similar logic as our Controllers, but they are solutions for centralized settings. In Adapters, a small unit is implanted inside each layer of a neural model whose job is (i) to connect two consecutive layers to each other in a way that they do not deviate from their main distribution, and (ii) take care of the transition between domains. Such a similar architecture persuaded us to try Adapters in our FL setting but our experiments were not successful. Our investigation showed that when an Adapter unit is sent out to the server side to be combined with other Adaptors, it becomes so distant from its original form that when it is placed back in the body of a client model, it destroys the internal information flow and stops the client from converging. We ran preliminary experiments with Adapters instead of Controllers and obtained gravely lower BLEU scores.

Aside from this qualitative aspect of the problem, we also do not find Adapters suitable for FL settings, because of a design overhead they introduce, e.g. some references recommend that the best position for Adapters to be placed is before Layer Norm (Rücklé et al., 2020), others placed them after Feed Forward sub-layers (Bapna et al., 2019), and the placement options are not limited to only these alternatives. Moreover, each model proposes a different internal neural architecture for Adapters. These design decisions are made through exhaustive empirical investigations. This form of architecture search is not affordable in distributed learning settings like FL. Any exchange of information and re-run of training could be quite costly.

This design dilemma is not unique to Adapters, and our Controllers could also suffer from similar issues, but in the Controller case the situation is far less severe. As Table 1 shows, we have tried different positioning strategies and our findings show that positioning Controllers layers affect the final quality. This also becomes even more important when working with smaller models (6E-6D vs 8E-8E). However, this does not stop us from training clients in the FL setting, because almost any random positioning strategy leads to acceptable results. There is only one exception and that is when we placed the first Controller in between the embedding table and the first layer of the encoder. If we define a configuration such as 8E-8D/A-C (-1-6) where the first layer of an encoder/decoder is a Controller, a client model would never converge. We could not find any mathematical justification for this observation, but we assume that because the first layers plays a critical role in building input representations, any noisy information can affect the entire model’s sustainability. It seems information added from other domains to the first layer after the embedding table, when it is used as a Controller, has a similar effect and stops clients from converging.