Paper: LegoMT2: Non-Blocking Federated Learning for Massive Multilingual Machine Translation (ICLR-2024)

Reviewer comment (Novelty):
The authors should present the key features of the traditional federated learning methods in the related works. The authors claim an efficient approach with a tailored model architecture for massive multilingual neural machine translation. What are the key attributes of the tailored model? In other words, what is the key difference between the federated learning used in this paper compared to the traditional federated method?

Author response:
>**Question 1: What is the key difference between the federated learning used in this paper compared to the traditional federated method?**

The key difference between the traditional federated method and our method is stated in Appendix A.  The main difference includes three aspects:  the data construction, the model architecture, and the communication method.

|  | Traditional Federated Method | LegoMT2 |
| --- | --- | --- |
| Data Construction | Domain [1] | Language-centric grouping |
| Model Architecture | Centralized  | Detachable |
| Communication Method* | Synchronous | Asynchronous |


One key factor in recent advanced FL algorithms is communication compression [2,3,4,5,6,7]. The traditional implementation of FL requires that each client send a full model (or a full model update) back to the server in each round. For large models, this step is likely to be the bottleneck of FL for a few reasons.

**Reason 1:** internet speeds are usually faster for downloading than uploading. For example, Xfinity offers 125Mbps for downloads and only 15Mbps for uploads.

**Reason 2:**  The size of the model is quite large. Therefore, there are numerous strategies available to compress these models or to reduce the bandwidth required for downloading the current model.

**Reason 3:** Different clients may experience a considerable amount of idle time while waiting for the server to collect all models and generate a new global model.

[1] Training mixed domain translation models via federated learning
[2] Federated learning with compression: Unified analysis and sharp guarantees.
[3] Federated learning: Strategies for improving communication efficiency
[4] Communication-efficient adaptive federated learning
[5] AutoFedNLP: An efficient FedNLP framework 
[6] FS-Real: Towards Real-World Cross-Device Federated Learning
[7] FedED: Federated Learning via Ensemble Distillation for Medical Relation Extraction

Reviewer comment (Experiments):
The experimental results are somewhat less convincing. Actually, the model size of the model should be viewed as 10.4B rather than 1.6B. And the final model used in inference is the averaged version of the 8 local models. Therefore, the model should be compared to the same-size finetuned model.

Author response:
>**Question 2-1:  The final model used in inference is the averaged version of the 8 local models.**

There seems to be a misunderstanding here: during inference (the results obtained in Table 1), we utilize only a single flow (mix-flow, 1.6B), not the averaged version of the 8 local models.

>**Question 2-2:  The model size of the model should be viewed as 10.4B rather than 1.6B. Therefore, the model should be compared to the same-size finetuned model.**
Yes,  we agree with your point. A fair comparing base model would be training 10.4B model and distilling into 1.6B model.

**1) another valid comparison**

We’ve given careful consideration to your suggestion and believe that, in addition to the model of the same size already listed in the table, there is indeed another model that could serve as a valid comparison:

+ LegoMT2: In training, the model of the whole system is 10.4B; but in inference, each flow can work independently, and the model size is only 1.6B.

+ Valid-Comparion:  We train a 10.4B model with centralized method, and then distilled this model to 1.6B.

**2) It can be anticipated that this method will lag behind LegoMT2 in terms of both speed and performance.**

+ Time: As evident from Table 3, training a 10.4B model is significantly slower, specifically 16.2 times slower, than LegoMT2. Furthermore, this pipeline also necessitates the distillation of the 10.4B model to a 1.6B model. This process requires the inference of the 10.4B model to obtain the output of each sample, which significantly increases the time requirement.

|  | Lego-MT2 | Valid-Comparison |
| --- | --- | --- |
| Training Speed | 16.2$\times$ | 1.0$\times$ |
| distillation | - | Inference 10.4 to get the output of each sample |

+ Performance:  The 54.5B model’s average performance surpasses LegoMT-2 by one point. Therefore, the upper limit of the 10.4B model has been established, and it’s important to note that the distillation process inevitably leads to performance losses.

Hence, when compared to the solution of training a 10.4B model and then distilling it, LegoMT2 proves to be both faster and more efficient.

Reviewer comment (Experiments):
Why the model is finetuned from the pre-trained model? Why not training from scratch?

Author response:
>**Question 3: Why is the model finetuned from the pre-trained model? Why not training from scratch?**

As the primary objective of LegoMT2 is to construct a system that supports 435 languages, it’s noteworthy that more than 300 of these are low-resource languages.

However, controlling the quality of low resources is not easy, as shown in Table 2 of the  NLLB paper[1]:
|   |  |
|-------------------------------------------|-----|
| # of Languages requiring Re-translation   | 10  |
| Avg # of Re-translations                  | 1   |
| Max # of Re-translations                  | 2   |
| Avg # of Days to Translate 1 language     | 42  |
| Avg # of Days to align                    | 28  |
| Avg # of Days for 1 language              | 119 |
| Shortest Turnaround (days) for 1 language | 70  |
| Longest Turnaround (days) for 1 language  | 287 |
|   |  |

By leveraging the existing model foundation, we can avoid redundancy in the process and devote more time to focusing on the newly added 200+ languages.

[1] No Language Left Behind: Scaling Human-Centered Machine Translation

Reviewer comment (Presentation):
Extensive experimental results and analyses are not fit in 9 pages. There are some description overlaps in Section 1 and 3 so the authors can move the contents from Appendix to the main pages.

Author response:
Thanks for your suggestion! We will move the analysis of the difference between of traditional Federated Learning and comparison with LLM from the Appendix to the main part.

Reviewer comment (Experiments):
Reg Section 3.3; how helpful is the parameter initialization with NLLB-200-1.3B? Have you ever looked into this effect, without having the NLLB initialization?

Author response:
We have a reasonable belief that NLLB initialization is important. However, it is too costly to try full training of LegoMT2 without NLLB initialization. Therefore, we did not conduct alternative experiments. According to NLLB paper[1] (Table 2), it would take over 200 days to replicate the alignment process of NLLB  to extend from NLLB's 200 languages to LegoMT2's 435 languages. 

|   | |
|-------------------------------------------|-----|
| # of Languages requiring Re-translation   | 10  |
| Avg # of Re-translations                  | 1   |
| Max # of Re-translations                  | 2   |
| Avg # of Days to Translate 1 language     | 42  |
| Avg # of Days to align                    | 28  |
| Avg # of Days for 1 language              | 119 |
| Shortest Turnaround (days) for 1 language | 70  |
| Longest Turnaround (days) for 1 language  | 287 |
|   | |


[1] No Language Left Behind: Scaling Human-Centered Machine Translation



Nevertheless, we would like to underscore that NLLB-200-1.3B only supports 200 languages. Our contribution to utilizing NLLB includes effort in extending the vocabulary from 200 languages to 435 languages. The tokens for an additional 235 languages are initiated randomly, resulting 0.3B extra parameters compared to NLLB. In this way, we can best utilize the pre-trained parameters in NLLB while largely expanding the language support in LegoMT2, in a cost-effective way.

Reviewer comment (Experiments):
Have you ever tried with different language grouping?

Author response:
We conducted an additional experiment with a different language grouping and found further increasing the group number can further improve system performance.

The language information:

| **new language group information = {**                                                                                                                                                                         |
|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
|     "family_1": ['fr', 'es', 'en'],                                                                                                                                                                            |
|     "family_2": ['nl', 'tr', 'pl', 'it', 'de', 'pt'],                                                                                                                                                          |
|     "family_3": ['bg', 'ar', 'ru', 'fa', 'el', 'hu', 'ro', 'cs'],                                                                                                                                              |
|     "family_4": ['sk', 'da', 'uk', 'sl', 'he', 'fi', 'id', 'sv', 'vi'],                                                                                                                                        |
|     "family_5": ['ko', 'sq', 'hr', 'mk', 'sr', 'zh', 'no', 'bs', 'hi', 'ja', "zhtrad", 'zhtw', 'zhyue'],                                                                                                       |
|     "family_6": ['et', 'lt', 'lv', 'ca', 'bn', 'ms', 'th', 'gl', 'is', 'ta', 'sw', 'mt', 'eu', 'eo', 'tl', 'ml', 'af', 'si', 'ur', 'be', 'ne', 'mr', 'nb', 'ze', 'ka', 'sh'],                                  |
|     "family_7": ['az', 'ga', 'te', 'hy', 'mg', 'kk', 'cy', 'lb', 'ht', 'jv', 'ast', 'am', 'tg', 'km', 'mn', 'su', 'my', 'so', 'oc', 'xh', 'ha', 'ky', 'ps', 'pa', 'gu', 'ku', 'nds'],                          |
|     "family_8": ['fy', 'la', 'nn', 'kn', 'mi', 'yo', 'br', 'ig', 'sn', 'ny', 'hne', 'ceb', 'zu', 'uz', 'lo', 'cx', 'sd', 'wa', 'as', 'mai', 'se', 'pes', 'dje', 'pck', 'plt', 'crp'],                          |
|     "family_9": ['kek', 'tt', 'rw', 'fil', 'csb', 'yi', 'or', 'wo', 'ug', 'ilo', 'gd', 'ss', 'dv', 'hsb', 'crh', 'gv', 'tn', 'kab', 'an', 'nso', 'ba', 'kbh', 'syr', 'cop', 'chr', 'rom', 'quw'],              |
|     "family_10": ['cjp', 'dop', 'shi', 'ee', 'nhg', 'amu', 'wal', 'jak', 'usp', 'cni', 'li', 'ppk', 'quc',  'mam', 'hbs', 'ch', 'chq', 'cak', 'djk', 'cb', 'gbi', 'ff', 'fur', 'bsn', 'ojb', 'agr', 'dik'],    |
|     "family_11": ['jiv', 'ake', 'acu', 'lg', 'ia', 'tk', 'fo', 'om', 'ber', 'ln', 'tmh', 'pot', 'dz', 'ang',  'mhr', 'io', 'kg', 'bo', 'ns', 'ti', 'sa', 'kw', 'frp', 'arz', 'ckb', 'bar', 'wae'],             |
|     "family_12": ['ks', 'os', 'ace', 'sc', 'jbo', 'ts', 'arq', 'gom', 'szl', 'sco', 'ay', 'shn', 'zht', 'toki',  'zhs', 'mwl', 'kl', 'qu', 'zza', 'vec', 'ltg', 'aa', 'cv', 'ab', 'tlh', 'shs'],               |
|     "family_13": ['ie', 'tmp', 'bem', 'inh', 'pap', 'cmn', 'prs', 'lfn', 'ak', 'ce', 'lmo', 'tet', 'kr', 'wuu',  'st', 'lld', 'gn', 'rm', "nan", 'mus', 'mfe', 'miq', 've', 'scn', 'av', 'dsb', 'rn'],         |
|     "family_14": ['azb', 'pms', 'iu', 'bi', 'ase', 'bm', 'co', 'gr', 'kj', 'pmy', 'ry', 'lad', 'vo', 'efi',  'xal', 'fj', 'gos', 'nap', 'rup', 'ksh', 'tc', 'zsm', 'bh', 'nus', 'fuv', 'din', 'pi'],           |
|     "family_15": ['bal', 'cbk', 'brx', 'hil', 'cu', 'bua', 'dtp', 'pam', 'qd', 'hai', 'mos', 'cnh', 'avk',  'sml', 'lij', 'bug', 'kzj', 'trv', 'bnt', 'swg', 'to', 'tzl', 'zam', 'que', 'orv', 'kha',  'bho'], |
| }                                                                                                                                                                                                              |

Reviewer comment (Experiments):
The paper did not conduct specific verification experiments on parameter interference to demonstrate that the performance improvement of LegoMT2 over finetuned NLLB-200-1.3B indeed stems from the alleviation of parameter interference phenomena.

Author response:
Parameter interference is a fundamental problem in multilingual machine translation. 

It refers to the competition between different languages for the limited parameters of a model when we hope to use a single model to handle all translation directions. This can result in good translation results for some languages, while the translation results for other languages may be less satisfactory.

As illustrated in Table 1, directly tuning NLLB-200-1.6B yields results that are inferior to those of NLLB-200-1.3B. However, LegoMT2 effectively enhances performance across over 400 languages, which suggests that LegoMT2 successfully mitigates parameter interference.

Reviewer comment (Experiments):
Which of Single-FT or Single-FT + MoE in Table 3 is used for the experiments in Table 1 and Table 2? Have the translation performance of both been evaluated?

Author response:
The Single-FT is used for experiments in Table 1 and Table 2. We evaluated a pretrained 12B model (M2M-100-12B) but did not fine-tune the 12B model, because the training 12B model is too slow (16x slower in Table 3). It would require 240 days on 64 A100 GPUs to complete the training, which is prohibitively costly.

Reviewer comment (Experiments):
Have any other methods for MERGE operation of non-blocking federated learning, apart from simple averaging, been tried and evaluated?

Author response:
No, we do not try other methods for MEREG. 

1)  At present, this MERGE operation is adequate for LegoMT2.  The existing MERGE operation, which involves replacing the global module on the client side with the average of the parameters obtained from the server, is the standard merge operation in federated learning. 

2) Testing different merge operations is a very resource- and time-consuming operation.  While we could explore more variations, it’s important to note that LegoMT2 represents a substantial pre-training effort.  LegoMT2's training takes 15 days on  64 80G A100 GPUs.

3) We sincerely hope that you reconsider the value of our paper. In this paper, we proposed a novel training framework, LegoMT2,  designed for massive massively multilingual machine translation (MNMT) systems. 

 + We introduced an efficient training framework for MNMT that supports 435 languages, which is more than any other system.
 + LegoMT2 partitions the model and data together and employs an efficient non-blocking algorithm to accelerate training.
 + LegoMT2 is small but powerful, which works only with 1.6B parameters and achieves better results than other models at the same or larger size (only behind NLLB-54B which is 30x larger).

Reviewer comment (Experiments):
How about LLMs for Multilingual Machine Translation?

Author response:
The comparison with ChatGPT is in the Appendix. Both in the En→X and X→En direction, ChatGPT falls behind LegoMT2 even with eight-shot. 

| Model              | X$\rightarrow$En | En$\rightarrow$X | AVG. |
|--------------------|------------------|------------------|------|
| ChatGPT zero-shot  | 27.9             | 23.9             | 25.9 |
| ChatGPT eight-shot | 31.9             | 24.7             | 28.3 |
| LegoMT2            | 38.3             | 31.6             | 35.0 |