Asynchronous Convergence in Multi-Task Learning via Knowledge Distillation from Converged Tasks (2022)

Abstract: Multi-task learning (MTL) aims to solve multiple tasks jointly by sharing a base representation among them. This can lead to more efficient learning and better generalization, as compared to learning each task individually. However, one issue that often arises in MTL is the convergence speed between tasks varies due to differences in task difficulty, so it can be a challenge to simultaneously achieve the best performance on all tasks with a single model checkpoint. Various techniques have been proposed to address discrepancies in task convergence rate, including weighting the per-task losses and modifying task gradients. In this work, we propose a novel approach that avoids the problem of requiring all tasks to converge at the same rate, but rather allows for “asynchronous” convergence among the tasks where each task can converge on its own schedule. As our main contribution, we monitor per-task validation metrics and switch to a knowledge distillation loss once a task has converged instead of continuing to train on the true labels. This prevents the model from overfitting on converged tasks while it learns the remaining tasks. We evaluate the proposed method in two 5-task MTL setups consisting of internal e-commerce datasets. The results show that our method consistently outperforms existing loss weighting and gradient balancing approaches, achieving average improvements of 0.9% and 1.5% over the best performing baseline model in the two setups, respectively.

## introduction over the past few years, large pretrained models have achieved great success on a variety of tasks in natural language processing @xcite @xcite @xcite @xcite . most work in this area typically follows the pretraining-finetuning paradigm, in which the model is first pretrained on large text corpora using a self-supervised language modeling objective and then finetuned using supervised data from a target task. in particular, when evaluating on a benchmark that contains multiple tasks such as glue @xcite and superglue @xcite , a common method is to finetune a separate model for each task in order to achieve the best performance. although such an approach produces excellent results, it has several drawbacks, including incurring the costs of repeated finetuning efforts and precluding the possibility of knowledge sharing among related tasks. in addition, when it comes to deploying these models for real-world applications, a separate model is deployed for each individual task which can pose challenges for model hosting and maintenance. an alternative approach is multi-task learning (mtl), which solves multiple tasks together by sharing representations among them. this not only offers benefits in computational and storage efficiency, but also makes it possible to share knowledge among related tasks and encourages the model to learn more robust and generalizable representations @xcite @xcite ). however, one of the biggest challenges in mtl is to balance the convergence schedule across tasks. differences in task difficulty can result in faster convergence on some tasks over others. as a result, the naive approach which simply adds together the losses of all tasks is typically sub-optimal @xcite , since the final model may overfit the tasks that have converged early on during training, while underfitting the others. to tackle this, a large body of work has explored various loss and gradient balancing strategies @xcite @xcite @xcite , in order to enforce the same convergence speed across tasks. despite these efforts, the problem still remains un-solved. as we show in our experiments, existing approaches can still fail to balance learning across different tasks in practice. in this work, we propose a different approach for coordinating the learning across tasks in mtl. instead of artificially forcing all tasks to converge at the same rate, we allow each task to converge according to its own schedule. we call this asynchronous convergence, as opposed to previous methods which seek to achieve a synchronous convergence schedule across tasks. after each task converges, we avoid overfitting by switching to a knowledge distillation loss for that task for the remaining training steps. the intuition is that continuing to train using the true labels can lead to overfitting. instead, the knowledge distillation loss encourages the model to maintain its output distribution and thus its performance on the converged tasks at the best level, while the model learns the remaining tasks. we evaluate the proposed method in two different 5-task mtl setups. our results show that existing loss and gradient balancing approaches fail to produce meaningful improvements over the simple baseline which sums the losses from all tasks or lead to more costly and less efficient training. in contrast, our method achieves consistent improvements. in particular, when comparing the final checkpoint performance, our best performing approach achieves an average improvement of 0.9% over the best baseline model in the first setup and 1.5% in the second setup. ## related work as already mentioned, one key challenge in mtl is to balance the learning speed across tasks. some existing methods address this by applying static weights to the losses of different tasks @xcite @xcite , where the weights are usually determined through extensive hyper-parameter search. however, such an approach tends to be sub-optimal @xcite . one line of research improves this by designing algorithms to automatically determine the weights and dynamically adjust them during training @xcite @xcite . going beyond the loss weighting approaches, there is also work that leverages gradient information and proposes to manipulate the magnitude and/or direction of the gradients from different tasks in order to better coordinate the learning among the tasks @xcite @xcite . recently, @xcite shows improved results by combining loss weighting and gradient manipulation approaches. in all above methods, the goal is to enforce roughly the same convergence speed across tasks, so that the final model fits all tasks well. however, we hypothesize that imposing such an artificial constraint leads to optimization challenges. instead, we propose a simpler method which allows for asynchronous convergence among the tasks, where each task converges on its own schedule. we focus on avoiding overfitting after a task has converged, which is achieved by distilling from the task's best checkpoint for the remaining training steps. on the subject of maintaining the performance of converged tasks, a related research area is continual learning (cl) @xcite . cl studies the problem of learning tasks in a sequential manner with the goal of avoiding catastrophic forgetting @xcite of previous tasks while learning new tasks. replaybased cl approaches are most relevant to our work. the idea is to periodically present the model with examples from past tasks to avoid forgetting @xcite @xcite . in particular, @xcite proposes to avoid forgetting by adding a distillation loss formulated using a small subset of examples from previous tasks. in our case, we experiment with a similar setup where we sequentially add one task at a time while using a distillation loss to preserve performance on converged tasks. however, because our focus is on mtl where we have access to the data of all tasks throughout training, we do not down-sample the data of converged tasks. additionally, our use of the distillation loss is to not only avoid catastrophic forgetting but also overfitting. finally, our work is also related to the field of knowledge distillation (kd) @xcite , where the goal is to transfer the knowledge of one network to another by training the latter network to mimic the predictions of the former network. it is widely used to distill the knowledge of a large teacher model to a small student model @xcite @xcite but has also been applied in mtl @xcite and cl @xcite . in particular, @xcite train a multi-task model by distilling from multiple single-task models and show this outperforms directly training a multi-task model. the main difference between this approach and our work is that we do not require separate teacher models per-task, but rather use intermediate checkpoints as teachers as the model converges on each task. another related work by @xcite also similarly uses intermediate checkpoints for distillation. however, their focus is on training a single machine translation model, whereas we use the technique to train a multi-task model. ## proposal: asynchronous convergence via knowledge distillation as the training set of task t, with @xmath0, . . . , t } and n t being the total number of training examples for task t. let f represent a neural network with parameters θ. in standard supervised training, we would train the network on task t by minimizing a loss l st t (where st stands for supervised training) formulated as specifically, let θ t denote the parameters of the checkpoint when the model converges on task t. we first use the checkpoint to run inference on the task's training set d t to obtain d @xmath1, the model would be trained on task t using a kd loss l kd t formulated as essentially, after the model has converged on task t, we no longer train on the true labels of the task. rather, we use the kd loss to encourage the model to mimic the predictions from the checkpoint where the best performance is achieved. as we show in our experiments, this method effectively maintains the model's performance on converged tasks without overfitting, while the model continues to learn the remaining tasks. to summarize, we propose to minimize the following loss using this method, we experiment with two different training settings. the first is called the joint setting, which is similar to the conventional mtl setup. the model is trained on all tasks together, and we swap in the kd loss as the model converges on different tasks. training stops when all tasks converge. the second setting is called the sequential setting and is inspired by the typical cl setup. here we start training on a single task and then add one new task at a time after the previous task converges. following our proposal, we use the kd loss for all converged tasks, while training the model on the true labels of the new task. the process continues until all tasks converge. one additional hyper-parameter in the sequential setting is the order in which to train the tasks. we experiment with a few different orders, including ordering from (1) the smallest to the largest task by dataset size (sequential (small)), (2) the largest to the smallest dataset size (sequential (large)), and (3) the easiest to the hardest (sequential (easy)). for lack of a better heuristic, sequential (easy) uses the order in which the tasks converge in the ## additional discussions in this section, we provide more discussions on the effects of the design choices in both our joint setting and sequential setting. we illustrate the effects through two additional experiments on the related 5-task setup. in the first experiment, we compare our joint setting against an alternative approach where we take the respective best checkpoints for each task from the uni. weight baseline, and distill them together into a single multi-task model. we call this approach bam (@xmath2. through this experiment, we seek to compare our way of continued training with a mixture of kd and supervision from true labels against pure kd training. we note that our setting is more challenging as the model needs to learn all tasks from scratch, whereas bam (@xmath3. also, we train bam (@xmath4. however, since we need to obtain the best checkpoint for each task from uni. weight, which is also trained for 8200 steps, the total training budget for bam (@xmath5. the results of this experiment are shown in in the second experiment, we compare our sequential setting with the continual mtl method proposed in @xcite . it is similar to our sequential setting in that it adds one new task at a time. the main difference is that they always use real labels for training, whereas we use kd to avoid overfitting on converged tasks. through this experiment, we seek to understand the effects of the kd loss in our sequential setting. specifically, we choose sequential (small) and train continual mtl using the same task ordering. the results are shown under exp. 2 in ## conclusion in this work, we propose a new approach to tackle the challenge of task convergence in mtl. in contrast to conventional loss and gradient balancing methods which attempt to enforce a synchronous convergence schedule, we allow the tasks to converge on asynchronous schedules and use a kd loss to maintain the performance on converged tasks while the model learns the remaining tasks. we show that the proposed method consistently outperforms existing loss and gradient balancing approaches. for future work, we will explore strategies to make model performance invariant of task ordering in the sequential setting, or alternatively, explore strategies to optimally determine task ordering. additionally, we will investigate techniques to improve the efficiency of our method, such as dynamically adjusting the learning rate during training and the batch proportion of converged tasks.