Learning Task Sampling Policy for Multitask Learning (2021)

Abstract: It has been shown that training multi-task models with auxiliary tasks can improve the target task quality through cross-task transfer. However, the importance of each auxiliary task to the primary task is likely not known a priori. While the importance weights of auxiliary tasks can be manually tuned, it becomes practically infeasible with the number of tasks scaling up. To address this, we propose a search method that automatically assigns importance weights. We formulate it as a reinforcement learning problem and learn a task sampling schedule based on the evaluation accuracy of the multi-task model. Our empirical evaluation on XNLI and GLUE shows that our method outperforms uniform sampling and the corresponding single-task baseline.

## introduction multi-task learning @xcite has been shown to improve performance of multiple related tasks through cross-task knowledge transfer using just one model, which also drastically improves parameter efficiency @xcite . in the case where the objective is improving specific target tasks, the use of auxiliaries has demonstrated substantial benefits, for example, in sequence-to-sequence learning @xcite , and question answering @xcite . in this paper, we follow this line and consider the goal of training a multi-task model to be the maximizing performance of one target task. while there are numerous efforts that attempted to increase the complexity of multi-task models to match the performance through architectural addition, not a lot of emphasis has been placed on harnessing the potential of a model through exploration of task sampling weights, as we view weighted combination of different tasks as a sampling problem. in our work, we formulate this exploration problem as learning a policy to assign task sampling weights deferentially over time. for example, an optimal policy could be putting more weights on auxiliary tasks with rich data at the beginning of training and shifting to the target tasks near the end for finetuning. while manual hyper-parameter fine-tuning should help find a good policy, it becomes challenging as the number of combinations increases exponentially with respect to the number of tasks. to solve this problem, we propose a reinforcement-learning-based search method that automatically finds the optimal policy to maximize the target task accuracy. different from previous dynamic re-weighting works @xcite , our policy is static and agnostic to model training artifacts such as which training examples are selected and network parameters, and thus the policy can also be viewed as a schedule. this property enables reuse of a fixed sampling schedule once search is done. when we need to retrain model due to system change or minor underlying model change, it allows us to directly retrain without performing policy search again. we conduct experiments by targeting at improving one language's performance in the xnli @xcite dataset or one task's performance in the glue dataset. we empirically demonstrate that multi-task models trained with the proposed method outperform multi-task models learned with various rule-based sampling as well as corresponding single-task baselines. we also conduct necessary ablation study to validate our approach. our main contributions are as follows. first, we are the first work to formulate multi-task exploration as a static sampling policy/schedule learning problem to the best of our knowledge. then, we propose a simple and effective policy optimization algorithm for learning task sampling policy. ## related works existing works @xcite show that, by doing proper architecture changes, we can utilize cross-task transfer of multi-task learning and have multi-task models outperform single-task baselines. however, for different multi-task settings, the optimal architectures may be different and difficult to know a priori. it leads us to prefer an automatic searching approach, e.g. rl-based architecture search methods @xcite were proposed to learn how to share under different settings. while there are many successes in searching architecture for multi-task learning, how to select tasks to train together is underexplored. the current approaches are somehow arbitrary. with a goal of optimizing the performance for one target task, some works @xcite @xcite exhaustively explored the relatedness of tasks in a multi-task training setting by manually picking pair of tasks and found that the multi-task performance gains significantly. this becomes practically infeasible as the number of tasks grow making the possible combinations of tasks exponential. up-sampling strategy @xcite , such as uniform sampling that balances high-resource and low-resource tasks, is another dimension that has been explored. finally, the sampling strategy can also change over time. curriculum learning @xcite can be used to choose to learn hard tasks first and easy tasks later. domain adaption by continuous pre-training on the target domain @xcite ) can be seen as a schedule that first pre-train on a general domain and then on a specific domain. instead of manually trying all the task combination options, reward-based learning methods have been proposed by @xcite ; @xcite to solve the problem to maximize task transfer. however, such methods, which implicitly or explicitly depend on artifacts during training such as evaluation accuracy of the current time stamp or example-level weight, suffer from instability of rewards and randomness of states changes. as a result, it can be difficult to re-train to ensure to get good model quality. unlike previous works, we focus on learning a static sampling schedule, which means our schedule, once learned, should be fixed before any future model training. future model training should not involve policy optimization. that makes it easy to re-train the models when underlying system changes or share the sampling schedule with other teams who do not have policy-optimization expertise in an industrial setting. ## methods in sec. 3.1, we introduce our algorithms for learning task sampling policy. in sec. 3.2, we discuss our exploration strategies. ## experiments we use the public bert-base foot_0 @xcite checkpoints as our base model to run multi-task fine-tuning experiments. we consider two scenarios for knowledge transfer: cross-lingual transfer between languages of the same task and cross-task transfer between tasks of the same language. for the cross-lingual transfer study, we use the cross-lingual natural language inference (xnli) corpus @xcite . the data set asks whether a premise sentence entails, contradicts, or is neutral toward a hypothesis sentence. crowdsourced english data are translated to ten other languages by professional translators and used for evaluation, while the english multinli @xcite training data and its translation to other languages, translate-train @xcite , are used for training. for the cross-task transfer study, we cover multiple tasks from the glue benchmark @xcite . for all of our experiments, we compare two of our proposed methods to three baselines: uniformsampling that uses a uniform weight to sample across all tasks or languages, proportional sampling, and the target task baseline that only uses the labeled data from the target task for training. ## conclusions we propose a simple and effective rl-based algorithm to learn a static sampling schedule for multi-task learning, with the goal to improve the performance of one target task. we apply the algorithms to model fine-tuning and show that the proposed approach outperform popular rule-based baselines.