Overcoming Catastrophic Forgetting beyond Continual Learning: Balanced Training for Neural Machine Translation (2022)

Abstract: Neural networks tend to gradually forget the previously learned knowledge when learning multiple tasks sequentially from dynamic data distributions. This problem is called catastrophic forgetting , which is a fundamental challenge in the continual learning of neural networks. In this work, we observe that catastrophic forgetting not only occurs in continual learning but also affects the traditional static training. Neural networks, especially neural machine translation models, suffer from catastrophic forgetting even if they learn from a static training set. To be specific, the final model pays imbalanced attention to training samples, where recently exposed samples attract more attention than earlier samples. The underlying cause is that training samples do not get balanced training in each model update, so we name this problem imbalanced training . To alleviate this problem, we propose Complementary Online Knowledge Distillation (COKD), which uses dynamically updated teacher models trained on specific data orders to iteratively provide complementary knowledge to the student model. Experimental results on multiple machine translation tasks show that our method successfully alleviates the problem of imbalanced training and achieves substantial improvements over strong baseline systems.

## introduction neural machine translation (nmt) has achieved impressive translation performance on many benchmark datasets in the past few years @xcite @xcite . in the domain adaptation task where we have large-scale outdomain data to improve the in-domain translation performance, continual learning, which is also referred to as fine-tuning, is often employed to transfer the out-domain knowledge to in-domain @xcite . after fine-tuning, the model performs well in in-domain translation, but there is significant performance degradation in out-domain translation because it "forgets" the previously learned knowledge. this phenomenon is called catastrophic forgetting @xcite and has attracted a lot of attention @xcite @xcite . in this work, we observe that catastrophic forgetting not only occurs in continual learning but also affects the traditional static training. to be specific, the final model pays imbalanced attention to training samples. at the end of training, the recently exposed samples attract more attention and tend to have lower losses, while earlier samples are partially forgotten by the model and have higher losses. in short, training samples receive imbalanced attention from the model, which mainly depends on the time when the model last saw the training sample (i.e., the data order of the last training epoch). the underlying cause of this phenomenon is mini-batch gradient descent @xcite , that is, we do not simultaneously use all training samples to train the model but divide them into mini-batches. therefore, training samples do not get balanced training in each update step, so we name this problem imbalanced training. this problem is less severe in some tasks (e.g., image classification and text classification), but it has a significant impact on nmt as machine translation is a challenging task containing numerous translation rules, which are easily forgotten during the training process. besides, we find that the imbalanced training problem is especially severe and non-negligible on low-resource machine translation. to demonstrate that the imbalanced training problem does affect the model accuracy, we first review a widely used technique called checkpoint averaging technique, which has proved to be effec-tive in improving model accuracy but its internal mechanisms are not fully understood. we analyze it from the perspective of catastrophic forgetting and find that their success can be attributed to the alleviation of imbalanced training. we also notice that checkpoint averaging has some limitations, leaving room for further improvements. inspired by the existing solution of checkpoint averaging which leverages the complementarity of checkpoints to improve model accuracy, we propose complementary online knowledge distillation (cokd) to address the problem of imbalanced training. as the model tends to forget knowledge learned from early samples, the main idea of cokd is to construct complementary teachers to re-provide this forgotten knowledge to the student. specifically, we divide the training set into mutually exclusive subsets and reorganize them in a specific orders to train the student and teachers. we perform cokd in an online manner where teachers are on-the-fly updated to fit the need of student. when training the student on a subset, teachers can always provide the student with complementary knowledge on the other subsets, thereby preventing the student from catastrophic forgetting. experimental results on multiple machine translation tasks show that our method successfully alleviates the problem of imbalanced training and achieves substantial improvements over strong baseline systems. especially, on the low-resource translation tasks that are severely affected by imbalanced training, our method is particularly effective and improves baseline models by about 2 bleu points on average. in summary, our contribution is threefold: ## imblanced training before drawing any conclusions, we first conduct experiments on three different tasks, namely, image classification, text classification, and machine translation, to show that the problem of imbalanced training does exist. for image classification, we conduct experiments on cifar-10 and cifar-100 @xcite , both of which contain 50,000/10,000 training/testing images with 32 × 32 pixels drawn from 10/100 classes. for text classification, we conduct experiments on ag-news, which contains 120,000/7,600 training/testing sentences drawn from 4 classes. for machine translation, we conduct experiments on three translation tasks: wmt14 english-german (en-de), iwslt15 english-vietnamese (en-vi), and wmt17 english-turkish (en-tr). we use the resnet-32 network @xcite for image classification, the vdcnn network @xcite for text classification and transformer-base @xcite for machine translation. all the models are trained using cross-entropy loss. we refer readers to appendix a and section 6.1 for the detailed configurations. we train the model until convergence and then take the last checkpoint to calculate losses of training samples in the data order of the last training epoch. if there is a problem of imbalanced training, then training samples at the end of the epoch, which are recently exposed to the model, will tend to have lower losses. in contrast, training samples at the beginning will tend to have higher losses. for quantitative analysis, we use the spearman correlation coefficient between the data order and loss to measure the degree of imbalanced training. specifically, we assign each batch in the training dataset with a batch-id according to the order they appear in the last training epoch, where batch i is the i-th trained batch. we disable regularization techniques such as dropout and label smoothing and calculate the loss for each batch. the correlation coefficient between the batch-id and the loss is used to measure the degree of imbalanced training, and a large negative correlation coefficient indicates that this problem is severe. figure 1 illustrates the relationship between the batch-id and loss. by comparing the loss curves and correlation coefficients on these six datasets, we obtain the following three main observations. the problem of imbalanced training does exist. among the six datasets in our experiments, only cifar-10 has a positive correlation coefficient. two datasets (i.e., ag-news and wmt14 en-de) have small negative correlation coefficients. three datasets (i.e., cifar-100, iwslt15 en-vi, and wmt17 en-tr) have an apparent decline in losses accompanied by large negative correlation coefficients. therefore, we can conclude that the problem of imbalanced training does exist, but the degree of impact varies. imbalanced training is related to task complexity. intuitively, imbalanced training is more likely to occur on complex tasks where previously learned knowledge may be easily forgotten during the learning of numerous new knowledge. comparing the two image classification datasets, cifar-10 and cifar-100 have the same dataset size but a different number of classes. the correlation coefficient on the complex task cifar-100 is -0.29, while the correlation coefficient on cifar-10 is 0.01. the text classification task, which only contains 4 classes, has a small correlation coefficient -0.04. machine translation is generally considered a complex task with exponential search space and numerous translation rules. notably, wmt17 en-tr has the largest correlation coefficient of -0.64. these results are consistent with our intuition that imbalanced training has a greater impact on complex tasks like machine translation. low-resource translation suffers from imbalanced training. comparing the three machine translation datasets, the imbalanced training problem has a much larger impact on low-resource datasets (i.e., iwslt15 en-vi and wmt17 en-tr), where the high-resource dataset wmt14 en-de is less affected. to eliminate the influence of language, we randomly select 100k sentences from the wmt14 en-de dataset for the training to simulate the low-resource scenario. we show the loss curve in appendix b, where the corresponding correlation coefficient is -0.63, which also supports the conclusion. this is counter-intuitive since when there are many training samples, the early samples seem to be more easily forgotten. actually, as figure 1 shows, the loss curves are generally less steep at the beginning, indicating that early samples are nearly "equally forgotten" by the model. for highresource datasets, most samples are nearly "equally forgotten" and only the losses of a few samples at the end are highly correlated with the batch-id, so the overall correlation is low. in comparison, nearly the whole loss curve of low-resource datasets is steep, so the model may simultaneously overfit recent samples and underfit early samples due to imbalanced training. therefore, the problem of imbalanced training is more serious and nonnegligible in low-resource machine translation. loss rises in the end due to the momentum of optimizer. on cifar-100, iwslt15 en-vi, and wmt17 en-tr, though their loss curves are generally downward, they all have a sudden rise at the end. this abnormal phenomenon is actually consistent with our conclusion. because of the momentum factor in the adam optimizer, the impact of a model update is not limited to the current step. the optimizer retains the gradient in the form of momentum, which will affect the gradient updates in the next few steps. therefore, the impact of momentum is not fully released in the last few training steps, so the loss rises in the end. ## checkpoint averaging checkpoint averaging, which directly takes the average of parameters of the last few checkpoints as the final model, is a widely used technique in nmt @xcite . the averaged checkpoint generally performs better than any single checkpoint. however, to the best of our knowledge, its internal mechanism is not fully understood. in this section, we analyze the success of checkpoint averaging from the perspective of imbalanced training. though training samples receive imbalanced attention from each checkpoint, this imbalance is different among checkpoints. if we understand the imbalanced training as the noise on each checkpoint, noises among different checkpoints can be approximately regraded as i.i.d. random variables. by averaging checkpoints, the variance of random noise is reduced and thereby alleviating the problem of imbalanced training. based on the above analysis, we make the following hypothesis and verify it through experiments. ## approach in this section, we propose complementary online knowledge distillation (cokd) to alleviate the problem of imbalanced training. we apply knowledge distillation with dynamically updated complementary teachers to re-provide the forgotten knowledge to the student model. ## experiments 6.1 setup to evaluate the performance of cokd, we conduct experiments on multiple machine translation tasks. for low-resource translation where the problem of imbalanced training is severe, we run experiments on wmt17 english-turkish (en-tr, 207k sentence pairs), iwslt15 english-vietnamese (en-vi, 133k sentence pairs), and ted bilingual dataset. we also evaluate the high-resource performance of cokd on wmt14 english-german (en-de, 4.5m sentence pairs). for wmt17 en-tr and iwslt15 en-vi, we use case-sensitive sacre-bleu @xcite to report reproducible bleu scores. for ted bilingual dataset, following @xcite , we report the tokenized bleu. for wmt14 en-de translation, we report the tokenized bleu @xcite with compound split. for wmt17 en-tr, we use newstest2016 as the validation set and newstest2017 as the test set. we learn a joint bpe model @xcite with 16k operations. for iwslt15 en-vi, we use the pre-processed data used in luong and manning (2015) 2 . for ted bilingual dataset, we use the pre-processed data used in @xcite 3 . for wmt14 en-de, the validation set is newstest2013 and the test set is newstest2014. we learn a joint bpe model with 32k operations. in the main experiments, we set the number of teachers n to 1 and the hyperparameter α to 0.95. we implemented our approach based on the base version of transformer @xcite . @xcite , we increase the dropout rate to 0.2 on wmt17 en-tr and iwslt15 en-vi. for ted bilingual dataset, we further increase the dropout rate of transformer baseline to 0.3. all models are optimized with adam (kingma and ba, 2014) with the optimizer settings in @xcite . the batch size is 32k for all translation tasks. for inference, we average the last 5 checkpoints and use beam search with beam size 5. the checkpoint interval is 1000 for low-resource tasks and 5000 for wmt14 en-de. ## conclusion in this paper, we observe that catastrophic forgetting will cause imbalanced training, which is severe in low-resource machine translation and will affect the translation quality. we rethink the checkpoint averaging technique and explain its success from the perspective of imbalanced training. we further propose complementary online knowledge distillation (cokd), which successfully alleviates the imbalanced training problem and achieves substantial improvements in translation quality. model and draw the loss curve in figure 5 . from both the correlation coefficient and loss curve, we can see that the imbalance is greatly reduced by fine-tuning. regrading the bleu score, the bleu score after fine-tuning is only 0.25 higher than the baseline, which is much smaller than the improvement of checkpoint averaging. we speculate that it is due to some drawbacks of fine-tuning. for example, finetuning with a low learning rate has a risk of overfitting the dataset, which may influence its performance on low-resource datasets.