Original text: https://zhuanlan.zhihu.com/
p/138597214
Author: Liu Shikun
Recently, there have been numerous breakthroughs in the research progress of Multi-task Learning (MTL) and many interesting new directions have been explored. This has greatly inspired me to write a new article, attempting to summarize and encapsulate the recent research advancements in MTL and explore the possibilities for future research directions in MTL.
This article will follow the general framework of my master’s thesis from 2018: Universal Representations: Towards Multi-Task Learning & Beyond, supplemented by recent methods from new articles and discussions on future directions.
Disclaimer: The articles mentioned in my master’s thesis are all preliminary results that were unpublished at the time. For anyone interested in the details of the papers, please refer directly to the published conference articles.
Two Major Research Branches of Multi-task Learning
In most cases, MTL research can be categorized into the following two directions: one is MTL Network design; the other is MTL Loss function design. We will provide a detailed interpretation of these two directions below.
Multi-task Learning Network Design / What to share? [Network Design]
Initially, MTL network design can typically be classified into two scenarios: Hard parameter sharing and soft parameter sharing.
Hard-parameter sharing — is now one of the essential baselines for almost all MTL implementations. It involves using the entire backbone network as a shared network to encode task information, which is then split into several task-specific decoders for prediction at the final layer. Hard-parameter sharing serves as a lower bound on the parameter space of network design (not strictly, assuming no network pruning is considered), thus providing an important reference for evaluating the efficiency vs. accuracy balance of new network designs.
Soft-parameter sharing — can be seen as another extreme of hard-parameter sharing and is not commonly found in current MTL network designs. In soft-parameter sharing, each task has its own backbone network of the same size as the parameter space. We can impose specific constraints on this parameter space, such as sparsity, gradient similarity, or LASSO penalty, to softly constrain the representation space of different task networks. If we do not impose any constraints on the parameter space, soft-parameter sharing collapses into single-task learning.
Any MTL network design can be viewed as finding a balance point between hard and soft parameter sharing: 1. How to design the network to be compact and lightweight. 2. How to design the network to maximize the sharing of information among different tasks.
MTL network design is all about sharing.
- Cross-Stitch Network
The Cross-Stitch Network is a classic network design that has been widely used as a baseline in various MTL studies over the past few years. Its core idea is to connect each independent task-specific network using learnable parameters (cross-stitch units) in a linear combination to link the convolutional blocks of different tasks.
For tasks A and B, at the output layer of each convolutional block, we will compute:
Through this computation, the input layer of the next convolutional block will be.
Inspired by the design of Cross-stitch, NDDR-CNN also has a similar approach. However, unlike Cross-stitch, for the information fusion of the intermediate convolutional blocks, they used concatenation and a [1 x 1] convolutional layer to reduce dimensionality. This design allows each task’s channels to intermingle information with channels of different indices, avoiding the limitation of the original Cross-stitch, which could only infuse information from the same channel. When the non-diagonal elements of NDDR’s convolutional layer weights are 0, NDDR-CNN is mathematically equivalent to the Cross-Stitch Network.
The biggest weakness of both Cross-Stitch Network and NDDR-CNN is that each task requires a new network, leading to a linear increase in the entire parameter space with the number of tasks, making it inefficient.
- Multi-task Attention Network
Based on the efficiency drawbacks of the Cross-Stitch Network, I later proposed the Multi-task Attention Network (MTAN) to make the network design more compact and lightweight, with the entire network’s parameter space increasing sub-linearly with the number of tasks.
The core idea of MTAN is to assume that after encoding general representation information in the shared network, we only need a small number of parameters to refine the task-shared representation into task-specific representations to obtain a good representation for any task. Therefore, the entire network only needs to add a small number of task-specific attention modules, two layers of [1 x 1] convolution layers, as additional parameter space to attend to the task-shared representation. The model’s parameters are significantly reduced compared to the Cross-Stitch Network.
- AdaShare
AdaShare takes the efficiency of MTL network design to the extreme. Instead of adding extra convolution layers to refine the task-shared representation, AdaShare treats a single backbone network as the overall representation, using a differentiable task-specific policy to decide whether to update or utilize the representation of a block of this network for any given task.
Since the entire network is applied to the representations of all tasks, the network parameter space is agnostic to the number of tasks, remaining constant, equivalent to hard-parameter sharing. The attached task-specific policy uses Gumbel-softmax for each convolution block to categorical sample “select” or “skip” two policies, allowing the entire MTL network design to change with different tasks, similar to the recently popular Neural Architecture Search concept.
- MTL + NAS
MTL-NAS is another example of combining MTL and NAS. It builds on the core idea of NDDR, extending it to the fusion of any block, thus searching for how to fuse different blocks of different tasks to achieve the best performance.
I personally prefer the search method of AdaShare, which searches layer by layer within a single network. This task-specific representation is already good enough compared to defining each task as a new network. Thus, MTL-NAS cannot escape the characteristic of linear increase in network parameters, but it provides new ideas for MTL network design.
MTL + NAS and traditional single-task NAS will have different requirements and training methods.
1. MTL + NAS is not suitable for the most common two-stage training method in NAS: using validation performance as supervision to update architecture parameters, obtaining a converged architecture, and then re-training the entire network. Because the fused information in MTL has a training-adaptive nature, fixing the network structure will result in the loss of this training-adaptive information, leading to performance lower than that of the one-stage method, which searches and converges simultaneously. In other words, oscillation during training and feature fusion are more important for MTL networks, while there is no concept of feature fusion in single-task learning. This indirectly leads to different training method requirements for NAS.
2. MTL + NAS is task-specific. In NAS training, if the dataset’s complexity is too high, we sometimes use proxy tasks to speed up training. The most common case is using CIFAR-10 as a proxy dataset to search for good network structures, which are then applied to the larger ImageNet dataset. However, this approach is not applicable to MTL because the feature information and task characteristics required for any task or pair of tasks are different, making it impossible to accelerate training through proxy tasks. Each group of tasks has a unique and distinct network.
I believe that in the future research of MTL network design, we will achieve networks with more interpretable/human-understandable characteristics, allowing us to understand the relationships and complexities between tasks. Through the obtained task relationships, we can use them as good prior knowledge to initialize a better starting network, thus obtaining a superior model, creating a virtuous cycle.
A Better Task Relationship A Better Multi-task Architecture
Multi-task Learning Loss Function Design / How to learn? [Loss Function Design and Gradient Optimization]
Parallel to network design, another popular direction is the design of MTL loss functions, or understanding how to better update the task-specific gradients in the network.
For any task i, we have a loss function: where are the task-specific learning parameters. We need to find a good set to optimize the performance of all tasks i. The simplest and most direct way is equal weighting: , which assumes that each task contributes equally to the representation.
- Weight Uncertainty
Weight Uncertainty is one of the earliest studies on MTL loss function design. This paper assumes that there exists a data-agnostic task-dependent uncertainty in each model, referred to as Homoscedastic uncertainty (this term is quite peculiar, and only the Cambridge group seems to use it). By maximizing the log-likelihood of the model prediction uncertainty, we can balance MTL training. Here, the likelihood (usually parameterized by Gaussian) can be seen as the relative confidence between tasks.
For any model prediction y, we define the Gaussian likelihood of the model: where here is a learnable noise scalar (Gaussian variance), and we need to maximize:
,
From this, we can derive a newly defined loss function:
.
The final derived formula is very concise and has thus been used in many MTL benchmarks.
However, this paper has been quite controversial, with one of the most notable points being that the authors have consistently refused to open-source such a simple formula and ignored numerous inquiries from other researchers regarding implementation details, angering many peers (including myself). Additionally, weight uncertainty is highly dependent on the choice of optimizer; in my personal experiments, I found that only the ADAM optimizer allows weight uncertainty to converge correctly, while no convergence trend was observed with other optimizers. This blog further points out that this weight uncertainty formula can directly yield a closed-form solution: when the learnable is minimized, the entire loss function will transform into the geometric mean of task losses, thus raising questions about the feasibility of the uncertainty assumption.
- GradNorm
GradNorm is another early paper on MTL loss function. The implementation of GradNorm is to calculate the inverse training rate as an indicator to balance the gradient updates between different tasks.
We define as W parameters in the weighted multi-task loss calculated at the t step to the L2-norm of the gradient; the mean task gradient is; the relative inverse training rate is . GradNorm updates task-specific weighting through the following objective:
where is the ideal gradient L2-norm (as a constant) to adjust task-specific weighting. As a balancing hyperparameter, the larger it is, the more balanced the task-specific weighting becomes. Since each calculation requires backpropagation for all tasks at each layer, it is computationally expensive. Therefore, the authors approximated by calculating the shared layer of the last layer to speed up training.
- Dynamic Weight Average
Since GradNorm requires two backpropagation calculations for task-specific weighting, it is very complex in implementation. I later proposed a simpler method that calculates task weighting only through the relative descending rate of loss:
Here, is calculated as the loss ratio of two adjacent time steps as the descending rate. Therefore, the smaller it is, the greater the convergence rate, the simpler the task, and the smaller the obtained weight.
- MTL as Multi-objective Optimization
The previously introduced task-weighting methods are based on specific heuristics, making it difficult to guarantee optimum in MTL optimization. In MTL as Multi-objective Optimization, the authors view the MTL problem as a multi-objective optimization problem, aiming to achieve Pareto optimum.
Pareto optimum refers to any improvement in the performance of one task will necessarily lead to a decline in the performance of all other remaining tasks. The authors utilized a method called multiple gradient descent algorithm (MGDA) to find this Pareto stationary point. The general approach is that after calculating task-specific gradients, it obtains a common direction to update shared parameters. If this common direction exists, the entire optimization has not converged to Pareto optimum. This convergence method ensures that shared parameters do not encounter conflicting gradients, allowing each task’s loss to converge more smoothly.
- Other heuristic methods
There are various other heuristic methods for designing MTL loss functions, either based on task difficulty (Dynamic Task Prioritization) or directly mapping computed task gradients to prevent conflicting gradients between tasks (Gradient Surgery).
However, among the various MTL loss function designs, it is rare for one method to outperform all others across all datasets. In fact, in some datasets, the simplest equal task weighting also performs quite well. On one hand, the effectiveness of task weighting heavily depends on the design of the MTL network itself; additionally, the updates of task weighting also depend on the dataset and the optimizer itself. If the core goal is merely to achieve the best MTL performance, I would suggest spending more time researching better networks rather than task weighting. However, it is undeniable that research on task weighting can help people better understand the relationships and complexities between tasks, which in turn can assist in better model design.
Auxiliary Learning — Not all tasks are created equal [Auxiliary Learning]
A closely related direction to MTL is called Auxiliary Learning (AL): its training process is identical to MTL. The only difference is that in Auxiliary Learning, only the performance of some tasks needs to be considered (primary tasks), while other (auxiliary) tasks exist to help the primary tasks learn better representations.
- Supervised Auxiliary Learning
Auxiliary Learning is more universal than we might imagine. For example, MTL is actually a special form of AL, where we can treat any task as a primary task and the remaining tasks as auxiliary tasks. In MTL, we assume that all tasks have a mutually beneficial relationship, so all learning tasks benefit from this correlation.
Auxiliary Learning has also been applied in many fields; for instance, this paper found that training depth and normal prediction simultaneously can effectively help improve the accuracy of object detection. Or this paper found that when reconstructing short sequences, it can help RNN train very long sequences more effectively.
In the setting of supervised auxiliary learning, the entire network and task selection heavily rely on human prior knowledge and do not possess absolute universality.
- Meta Auxiliary Learning
Considering the limitations of supervised auxiliary learning in task selection, I later proposed a meta-learning-based method to automatically generate auxiliary tasks, which I call Meta Auxiliary Learning (MAXL).
In traditional supervised auxiliary learning, I discovered the following two patterns:
1. If the primary and auxiliary tasks are in the same domain, the performance of the primary task will improve if and only if the complexity of the auxiliary task is higher than that of the primary task.
2. If the primary and auxiliary tasks are in the same domain, the final performance of the primary task depends only on the complexity of the highest auxiliary task.
The definition of task complexity here is tricky and not general; currently, I only consider the simplest case of image classification: the complexity of fine-grained tasks is higher than that of coarse-grained tasks.
For example, in the figure below, we see the classification information of cats and dogs, which intuitively is lower than that of fine classification, such as Yorkshire and Persian cats. The information required to classify cats and dogs is a subset of the information needed for fine classification, leading to the conclusion of pattern 2.
Due to pattern 2, we only need to consider the training situation of the simplest two tasks: one primary task and one auxiliary task. Therefore, the problem of “generating” a good auxiliary task can be transformed into finding a good network to generate fine classification labels for any input image.
Thus, in the MAXL framework, we have two networks: one network is a multi-task network similar to hard parameter sharing for multi-task training. The other network is a label-generation network to generate fine classification labels, providing predictions for the multi-task network as auxiliary task labels.
For the Multi-task network, we need to perform the following updates:
For the label-generation network, we need to perform the following updates through meta updates:
Where,
represents the meta update, aiming to maximize the validation performance of the primary task to find the best auxiliary label. The second derivative trick here is similar to the gradient update of the classic meta-learning algorithm MAML.
In the label-generation network, there is also a hyper-parameter representing human-defined dataset hierarchy. Assuming we are doing simple binary classification, then means subdividing the first class into 2 classes and the second class into 3 classes. The label-generation network will generate the corresponding suitable auxiliary labels through this hierarchy using a masked version of softmax. Thus, the multi-task network will perform two classification tasks: the primary task is binary classification, and the auxiliary task is five-class classification.
Through MAXL, we find that it can improve performance on such image classification tasks. We later visualize the generated labels and find that in some simple datasets, they have human-understandable clustering meanings.
In the above figure, the upper part shows the classification of CIFAR-100, while the lower part shows the classification of MNIST. The images in these three auxiliary classes are generated by the label-generation network, showing the five images with the highest confidence in this class.
In the more complex CIFAR-100 dataset, it is challenging to understand what MAXL’s classification is doing. However, in MNIST, we can see that different thicknesses of the digit 3, different orientations of the digit 9, and digits 7 with or without a horizontal bar cluster together. This interesting phenomenon opens up a novel direction for exploring automated auxiliary task generation.
- Self-supervised Auxiliary Learning: Learning to X by Y
Auxiliary learning can also be applied in self-supervised learning, assuming that neither the primary task nor the auxiliary task has any human labels. The self-supervised training here, like supervised auxiliary tasks, effectively utilizes human prior knowledge regarding the characteristics of certain task combinations.
Common combinations include: depth and ego-motion prediction can have temporal consistency. Similarly, time consistency can also be utilized in colorization and tracking, feature learning and ego-motion.
Such consistency can also exist in RL-based robot learning, such as grasping and pushing, or robot manipulation.
Conclusion
Finally, we come to the conclusion! In recent years, there have been many novel and valuable works in MTL research, but there is still significant room for improvement and progress in understanding the tasks themselves and the relationships between tasks. In Taskonomy, the authors attempted to create a relationship graph of different tasks by exploring thousands of combinations (with significant CO2 emissions). However, I believe that in real MTL training, this relationship graph should change over time and depend on the network itself. Therefore, how to better optimize the network structure through the relationships between tasks remains an unsolved mystery, and how to design/generate auxiliary tasks to better assist the primary task is also not fully understood. I hope to see more in-depth explorations of MTL in future research, achieving the ultimate vision of universal representation.
About Deep Deliver
Deep Deliver is a community focused on deep recommendation systems and CTR estimation, sharing first-hand papers, resources, and related technologies in recommendation, advertising, and NLP fields. We welcome your attention! To join the technical exchange group, please add the assistant deepdeliver, and note your name + school/company + direction.
You click on every “Looking”, I take it seriously as a like.