Q -Learning Scheduler for Multi Task Learning Through the use of Histogram of Task Uncertainty (2022)

Abstract: Simultaneous training of a multi-task learning network on different domains or tasks is not always straightforward. It could lead to inferior performance or generalization compared to the corresponding single-task networks. An effective training scheduling method is deemed necessary to maximize the benefits of multi-task learning. Traditional schedulers follow a heuristic or prefixed strategy, ignoring the relation of the tasks, their sample complexities, and the state of the emergent shared features. We proposed a deep Q-Learning Scheduler (QLS) that monitors the state of the tasks and the shared features using a novel histogram of task uncertainty, and through trial-and-error, learns an optimal policy for task scheduling. Extensive experiments on multi-domain and multi-task settings with various task difficulty profiles have been conducted, the proposed method is benchmarked against other schedulers, its superior performance has been demonstrated, and results are discussed.

## introduction multi-task learning aims to jointly improve the generalization of several classification and regression tasks. it does so by sharing the domain-centric information of each task, reducing trainable parameters, focusing attention on relevant features amid noisy or high dimensional data, regularizing other tasks, and exploiting the relations among tasks @xcite . the tasks can be defined as applying the same model on different data (also known as multi-domain learning) @xcite , or on various problems in a linear @xcite or hierarchical manner (e.g., named entity recognition, entity mention detection, and relation extraction in hierarchical mtl-hmtl @xcite . typical mtl is trained with mini-batches of every task intermittently, while the order of the training is usually uniform or related to the tasks' database size @xcite . this simplistic approach leads to unnecessary computations, due to redundancy in tasks and samples @xcite , or due to the fact that some tasks are prerequisites for learning others @xcite . additionally, task imbalances deteriorate appropriate training because they lead to imbalances between back-propagated gradients @xcite . therefore, researchers have proposed methods to train the tasks and samples in a certain order, a process called "mtl scheduling". effective scheduling increases accuracy, reduces overfitting across multiple tasks, avoids catastrophic forgetting, improves the low-resource task accuracy. meanwhile, it keeps the high resource task accuracy intact, provides extensive control over the training dynamics of mtl, and exploits task relations to learn required features for upsteam tasks. early scheduling approaches include non-adaptive heuristics and fixed strategies to order the training. such scheduling was applied to one target task (single scheduling, figure 1(a )) or was used by scaling per-task learning rates @xcite . as these methods were not adequate to handle more complicated mtl systems (e.g., uber-net @xcite ), more flexible and adaptive scheduling methods started to emerge. they use learning progress signals such as training loss @xcite , validation loss @xcite , and uncertainty @xcite to tailor the schedule accordingly (figure 1(b) ). advanced schedulers are expected to monitor task learning progress, emergent shared feature representation, and sample complexity of each task to be able to provide a suitable strategy for task orders. additionally, when the underlying multi-task models learn to improve the performance of harder tasks, they may hit a plateau; as a result, simpler (or data-poor) tasks can be overtrained @xcite . needless to say that some tasks may be forgotten if the schedule is improper (catastrophic forgetting). using a prefixed model to handle all these factors (model bias problem) without considering instantaneous feedback from the network (lack of temporal monitoring) is the main challenge for many of the current schedulers. to address the challenges, we propose using reinforcement learning (rl) method to learn a potentially complex strategy. this allows for handling different states of the mtl training and avoids catastrophic forgetting. it also enables learning long-term temporal effects of task selection on the network's performance using the intrinsic delayed reward handling mechanism of rl. moreover, it uncovers task relations with no explicit modeling (model-free), merely based on trial-and-error and receiving (delayed) feedback from the mtl training (figure 1(c) ). we use deep q-net (dqn) @xcite to map the state of the mtl to the desired actions (i.e., which task to train on next?), and propose the histogram of task uncertainty to describe the mtl state for the algorithm. our contributions are: in the hard parameter sharing architectures, shared parameters provide a global feature representation, while task-specific layers further process these features or provide a complimentary feature set for a specific task. mtl methods assume that learning easy tasks is the prerequisite for learning complex ones @xcite , hence put tasks in hierarchies @xcite or group similar tasks to form group-specific shared layers @xcite . scheduling mtl: to train a deep neural network on a battery of tasks simultaneously, the tasks are sampled uniformly or in proportion to their dataset size @xcite . since this may offer limited control over the performance trade-offs (e.g., accuracy vs. overfitting), task scheduling methods were proposed to improve the mtl's performance compared to single-task networks, static heuristics, and grid search methods. thus, non-adaptive (fixed strategy) and adaptive methods were used, and var-ious classes of scheduling emerged, such as: sample schedulers try to over-sample tasks with worse results compared to the baseline (kiperwasser and ballesteros, 2018) or down-weight easier samples to focus on harder ones for training @xcite . yet, such strategies fail if a task is highly over-sampled or datasets are imbalanced. task difficulty schedulers are usually built upon the notion of curriculum learning that favors smaller and easier tasks to learn first @xcite , aligned with the training of natural intelligence in human babies. this idea falls into two forms, task hierarchies (e.g., @xcite ) and task prioritization @xcite @xcite . in the latter case, if the data of the tasks are coming from significantly different distributions (e.g., domain adaptation @xcite @xcite ), the assumptions of curriculum learning do not hold @xcite , and prefixed schedulers might not be effective. task weighting is another approach for scheduling problem. mtl is sensitive to the task weights @xcite . these weights scale the loss term of each task in the total loss of the network. static task weights can be selected by hyperparameter tuning @xcite , yet this approach is suboptimal in the presence of less important or redundant tasks/samples @xcite . non-adaptive weighting schemes fetch tasks intermittently early in training and gradually weigh more on a specific task @xcite . adaptive schedulers dynamically prioritize different tasks by monitoring measures such as performance @xcite and homoscedastic uncertainty @xcite . another method is to dynamically tune the gradient magnitudes @xcite . here, we use q-learning scheduler which can dynamically schedule the learning of long-term temporal effects of task selection. as a result of scheduling, qls selects the task to draw samples for the next training episodes. ## multi-task classification we selected text classification, in which a document needs to be assigned to a set of classes. the solutions range from hand-crafting good features to be used in convolutional neural networks (nns) for word-level @xcite and character-level encoding @xcite , recurrent nns @xcite and convolutional recurrent nns @xcite . here, we address the problem of text classification using long-short term memory networks (lstm) with a variant explored in @xcite to facilitate further analysis of the effects of the proposed method in the representation. generally, this method can be applied to any encoder-decoder-based nlp task that uses lstm as the encoder. the text sequence of words @xmath0, w 2 , . . . , w t } is converted to a sequence of word embeddings x i and is given to an lstm layer. each cell of lstm layer at time t, includes input, forget and output gates, a memory and a hidden state h t . the lstm memory chain updates as here, α k is the importance of each task k and l(ŷ, y) is defined in eq(2). a scheduler, when applied to this mtl frameworks, adopts one of the following modifications: change the task weight (α k ), modify its gradient in the back-propagation, re-weight the samples (x t ), or select a task to draw samples for the next training episodes. here, for simplicity, we select the latter case to focus more on the idea. yet, the extension of the proposed idea to all different approaches is straightforward. ## experiment our proposed scheduling method is benchmarked under three different scenarios: (i) a multi-domain setting, in which a similar task is performed on different datasets, (ii) a simple multi-task setting, where the network does three different but slightly related tasks, and (iii) a complex multi-task setting in which task relations are complex (e.g., one task can be a prerequisite for another, improving one may hamper the performance of another, etc.). we have conducted an extensive analysis of the first task to clarify our proposed method's internal mechanism and advantages. ## conclusion we augment the fully-shared mtl framework with a reinforcement-learning-based scheduling scheme that obtains an optimal scheduling policy for tasks through trial and error. the scheduler detects the task states through a novel state definition: histogram of task uncertainty. it adjusts the scheduling policy to improve the training and validation accuracy of the mtl, enhances generalization of the emergent shared features, and handles different relationships among tasks. the proposed method is also capable of leveraging unlabeled data, obtaining highly-nonlinear strategies, and tackling different sources of task uncertainty.