Uncertainty Regularized Multi-Task Learning (2022)

Abstract: By sharing parameters and providing task-independent shared features, multi-task deep neural networks are considered one of the most interesting ways for parallel learning from different tasks and domains. However, fine-tuning on one task may compromise the performance of other tasks or restrict the generalization of the shared learned features. To address this issue, we propose to use task uncertainty to gauge the effect of the shared feature changes on other tasks and prevent the model from overfitting or over-generalizing. We conducted an experiment on 16 text classification tasks, and findings showed that the proposed method consistently improves the performance of the baseline, facilitates the knowledge transfer of learned features to unseen data, and provides explicit control over the generalization of the shared model.

## introduction multi-task learning (mtl) is a branch of supervised learning that strives to improve the generalization of the regression or classification task by leveraging the domain-specific information contained in the training signals of related tasks @xcite . mtl has been investigated in various applications of machine learning, from natural language processing @xcite and speech recognition @xcite to computer vision @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 (e.g., named entity recognition, entity mention detection and relation extraction in hmtl @xcite ). when training a multi-task learner, training each task normally increases its accuracy (fine-tuning) and, at the same time, provides more information for the shared representation that affects the accuracy of the rest of the tasks (generalization). balancing the finetuning-generalization trade-off has been the subject of several studies. @xcite adjusts tasks' relative weights in the loss function in proportion to the task uncertainty. @xcite divides the feature space into taskspecific and shared spaces and later employs adversarial learning to encourage shared feature space to contain more common information and no taskspecific information @xcite . @xcite proposed orthogonality constraints to punish redundancy between shared and taskspecific layers. in line with this direction, learning through hints @xcite ) directly trains a network to predict the most important features. yet, none of those methods explicitly balances the fine-tuning of the under-training task with its effect on the other tasks. here, we propose a method that considers the generalization of all tasks along with the taskspecific loss function. as a good indicator of how other tasks are affected by the change of the shared feature space, we proposed measuring the learner's uncertainty on each task. task uncertainty captures the relative confidence between tasks and reflects the inherent uncertainty in each task @xcite . therefore, the objective of the proposed mtl would be to maximally fine-tune each task on its corresponding training data while keeping the uncertainty of other tasks to the minimum. in other words, the mtl is expected to maintain a low level of the overall uncertainty and disentangle the training on the shared layers and task-specific layers. our main contributions are when used with deep learning, mtl models tend to share learned parameters across different tasks through (i) hard parameter sharing @xcite in which the hidden layers are shared between all tasks, while several taskspecific output layers are fine-tuned for each task, (ii) soft parameter sharing, in which each task has its model, and the distance between the parameters of the models for different tasks are regularized to encourage the parameters to be similar using, e.g., ℓ 1 norm @xcite or trace norm @xcite , or (iii) partial parameter sharing, to avoid task interference and leverage task commonalities among a subset of the tasks @xcite . in the hard parameter sharing architectures, shared parameters provide a global feature representation, while task-specific layers further process these features or provide a complementary set of features suitable for a specific task. some mtl approaches are based on the intuition that learning easy tasks is the prerequisite for learning more complex ones @xcite , hence put tasks in hierarchies @xcite @xcite or try to automatically group similar tasks to dynamically form shared layers @xcite . task uncertainty: in an mtl setting, multiple tasks are intermittently trained and modify the shared parameters to minimize their loss (the losses can be back-propagated at once as well). this change affects how other tasks behave in various ways, one of which is the amount of uncertainty that each task bears. uncertainty signals the information that the model lacks, or the sort of information that cannot be inferred from data @xcite . there are various ways to measure uncertainty. @xcite measures uncertainty via drop-out sampling. later, @xcite proposed homoscedastic uncertainty to measure the uncertainty of entire tasks independent of the data. in another attempt, @xcite computes the standard deviation of softmax outputs and average them to quantify the uncertainty of all tasks in mtl. other approaches that leverage uncertainty to reduce the overfitting in mtl framework are presented in @xcite . ## multi-task classification neural text classification has been studied as one of the fundamental nlp problems. some researchers replace hand-crafted features with with word-level and character-level representations obtained by cnns @xcite , others use rnns, convolutional rnns and self-attentive lstms for sequence modeling @xcite @xcite . to highlight the effect of the proposed uncertainty regularization, here we use a simple lstm-based text classifier as in @xcite . ## proposed method let's assume an mtl with k tasks, in which each task k has a dataset d k with n k samples, where to obtain a probability distribution ŷ(k) for labels of task k, the shared feature h (k) t are fed to the final task-specific softmax layer. we train the network by minimizing the cross entropy loss between predicted and true label distributions (ŷ (k) and y (k) ) as follows: ## experiment to evaluate our system, we considered two different settings to investigate the performance of our proposed system in multi-domain and multitask settings. in the former case, a similar task is done on different datasets, while on the latter, several heterogeneous but related tasks are performed. the first experiment comes with extensive analysis on performance, generalization, uncertainty regularization, measuring uncertainty, convergence speed, internal dynamics, and system errors, as well as leveraging auxiliary tasks, unlabeled data, pre-training, and task fine-tuning. ## conclusion in this study, we augment the fully-shared multitask learning framework with a regularization term to improve the shared representation by lowering the classification uncertainty for all tasks while finetuning for each task. the learned representation increased the overall accuracy of the multi-task classifier, achieved competitive results compared to state-of-the-art mtl algorithms, and successfully transferred the knowledge to the unseen tasks.