Coarse-to-Fine: Hierarchical Multi-task Learning
for Natural Language Understanding

Transformer is a neural structure consisting of multiple stacked self-attention modules Devlin et al. (2019); Liu et al. (2019b); Yang et al. (2019). With the bidirectional attention mechanism, the model can capture contextual information from both sides effectively. BERT Devlin et al. (2019) proposes a masked language modeling objective to pre-train a Transformer encoder, and achieves dramatic performance on several natural language understanding tasks. Since then, the pre-train-then-fine-tuning paradigm has gradually become mainstream. Researchers propose various new pre-training strategies to facilitate the model in several aspects Yang et al. (2019); Lewis et al. (2020); Raffel et al. (2020). For example, in order to obtain a more robust model, RoBERTa Liu et al. (2019b) adapts several tailored training strategies and employs more data. There are also many methods extending the pre-training on Transformer decoders. GPT-2 Radford et al. (2019) is a decoder-only structure that is pre-trained in accordance with the language generation objectives. It also performs exceptionally well on many text generation tasks. Though pre-trained models can be fine-tuned for distinct tasks, it is costly to maintain a separate model for each task, especially when the model is huge.

Multi-task learning is an integrated learning method in which multiple tasks share the same structure for training simultaneously. It can enhance the generalization and performance of each task. Consequently, multi-task learning can also be applied to pre-trained language models. A typical practice is to fine-tune the pre-trained language models by conducting multiple tasks concurrently Liu et al. (2019a). It is reported, however, that not all tasks can boost each other, and that noise may also be introduced Aribandi et al. (2021). Therefore, it is essential to analyze the relationship between tasks before training them jointly. Some recent work proposes addressing this issue by sparsely activating distinct modules for different tasks Tang et al. (2022). This allows the modules to be trained by relevant tasks. Unfortunately, these approaches usually rely on predefined activation paths, and they have to activate multiple modules in order to achieve high performance during the inference stage. This undoubtedly increases the delay and cost of model application.

The main differences between our method and others are: (1) We analyze the task correlation, and use it to group the tasks. Relevant tasks within the same group will be used to fine-tune the same modules, while irrelevant tasks will not interfere with one another. (2) Our model only activates a single module in each layer for each task, therefore it has the same number of parameters as a single model.

The overview of our method is illustrated in the right side of Figure 1. In general, there are three different kinds of layers in our structure: shared layers, task-clustering layers, and task-specific layers. The shared layers are in the bottom and optimized by all tasks. The middle part is the task-clustering layers. Tasks that are classified as relevant (introduced in the next section) will optimize the same group of layers. The top layers are task-specific, meaning that each task has its own module.

Our method adopts a coarse-to-fine paradigm. In this manner, the model can obtain generalized text representations in shared layers (a coarse-grained manner) and effectively interact and learn according to task correlation in the task-clustering layer (a fined-grained manner), then send it to the task-specific layer for specific task training.

Existing multi-task learning framework assumes that all tasks can facilitate each other. However, researchers have also reported the negative effect caused by irrelevant tasks Aribandi et al. (2021). To investigate the influence of each task, we devise three methods from different perspectives to measure the task relevance.

We propose a hierarchical multi-task learning framework (HMNet) based on task relevance. It is built on Transformer Vaswani et al. (2017), which has been widely applied in NLP tasks.¹¹1Other structures, such as convolutional or recurrent networks, may also be compatible with our method. This is left for future work. We omit the details of Transformer and refer readers to the original paper. As shown in the right side of Figure 1, HMNet has three different kinds of Transformer layers from the bottom up. All tasks will first pass the shared layers, then, each task will go through task-clustering layers according to their task group (introduced later). Finally, each task has its own task-specific layer, and the associated head is used to accomplish the task.

In the following paragraphs, we first introduce the details of the layers, followed by the task clustering process.

HMNet is based on a multi-layer Transformer, which is similar to existing pre-trained language models, such as BERT Devlin et al. (2019). Therefore, we use BERT as the backbone model to initialize the parameters in each layer. Then, each task has its own path (as described earlier), and we can use it to fine-tune the corresponding layers. The training process our method is summarized in Algorithm 1. In each epoch, a mini-batch $b_{t}$ is packed, and the HMNet is updated by the path for the dataset $\mathcal{D}_{i}$.

We conduct experiments on 13 NLU tasks and compare the performance of our HMNet with other baselines. These tasks can be grouped into four categories:

(1) Natural Language Inference (NLI): These tasks aim to determine whether a hypothesis is entailed, contradicted, or undetermined given a premise. They require the model to measure the logical coherence between two sentences. We select six datasets: MNLI Williams et al. (2018), QNLI, RTE, WNLI Wang et al. (2019b), CB Wang et al. (2019a), and SNLI Bowman et al. (2015).

(2) Sentiment Classification: These tasks ask the model to classify the sentiment polarity of a sentence (i.e., positive or negative). We use IMDB Maas et al. (2011) and SST-2 Socher et al. (2013) datasets.

(3) Similarity and Paraphrase: These tasks aim at determining whether two sentences have the same or similar meaning. The semantic relationship between two sentences is vital in these tasks. We choose three commonly used datasets: QQP Wang et al. (2019b), STS-B Cer et al. (2017), and MRPC Dolan and Brockett (2005).

(4) Question Answering: This task requires the model to answer a question by reasoning on a given paragraph. BoolQ Clark et al. (2019a) and MultiRC Khashabi et al. (2018) datasets are used.

Detailed statistics of benchmark datasets are mentioned in Appendix B.

We compare with our HMNet with three other training methods:

Single-Task fine-tuning: We fine-tune a pre-trained BERT for each task independently. As a result, 13 models are obtained in total.

Multi-Task fine-tuning: Following MT-DNN Liu et al. (2019a), we use the pre-trained BERT as a shared encoder and add a task-special head for each task. During fine-tuning, all parameters are optimized by all tasks jointly. With multi-task learning, only one model needs to be trained and saved.

SkillNet-style fine-tuning: This is a sparsely activated multi-task learning structure proposed by Tang et al. (2022). They design seven basic skills, such as getting the semantic meaning of a sequence and understanding a question. Then, they divide typical NLU tasks into five categories, each requiring a unique skill combination. The skill module is implemented by adding multiple FFN layers into the original Transformer structure. This model is also initialized by a pre-trained BERT. However, as each task activates multiple skill modules, compared with our method and other baselines, on average 1.0$\times$ more parameters (i.e., 223M vs. 110M) are used in inference. More detail about our reproduce reported in Appendix C.

For a fair comparison, all methods are initialized by the bert-base-uncased checkpoint. The batch size is 32, the max length of sequence is 512, and the initial learning rate is 2e-5, which is linearly decayed. AdamW Loshchilov and Hutter (2019) optimizer is applied. We train all methods for three epochs.

For the NLI and text similarity & paraphrase tasks, similar to the vanilla BERT, we concatenate the sentence pair by adding a separator token [SEP] and a head token [CLS]. For the sentiment classification task, the input is a single sentence, so we only add a head token [CLS]. For the QA task, we concatenate the question, passage, and answer, then separate them by two [SEP] tokens. A head token [CLS] is also added at the beginning of the sequence. All tasks are performed in similar ways, namely processing the embedding of the [CLS] token by a linear layer with a softmax activation function, and output the probability of each category. The default setting of HMNet has eight shared layers, two task-clustering layers, and two task-specific layers.

Table 1 shows the results on all datasets. In general, our HMNet performs better than other baselines on most datasets and achieves the best result in terms of the average score. This clearly demonstrates the superiority of our proposed HMNet. We further have the following observations:

(1) Comparing the performance between single-task and multi-task learning, it is evident that the latter can bring improvement for most tasks. Furthermore, this enhancement is related to the size of the datasets. Specifically, for the large dataset, such as MNLI and QQP (both of which have more than 300k training samples), the single-task learning can perform slightly better than multi-task learning. This implies that not all tasks can complement each other, and for those tasks with sufficient training data, adding extra tasks may not improve or even degrade performance. The potential reason is that the data in other tasks are collected from different domains and require models with distinct capabilities. Simply combining them will result in noise. For QA tasks, multi-task learning often leads to better performance. Indeed, QA is more complicated than other tasks, thus training on other tasks can be beneficial (e.g., better capture the relationship between question and answer).

(2) Sparsely activated (SkillNet-style) methods can achieve better performance than single-/multi-task learning. Its advantages stem from two perspectives: First, this method groups tasks according to their requiring skills, so the tasks rely on similar skills can enhance one another (e.g., both NLI and semantic similarity judgment rely on the skill of understanding how two text segments interact), while avoiding noise from other used skills. On the other hand, since multiple skill modules are activated, more parameters improve the model’s capacity.

(3) Our HMNet with three different task relevance measurements can consistently outperform all baseline methods. Specifically, HMNet brings more than 0.7% absolute improvement over multi-task learning in terms of average score. We attribute this improvement to our architecture of task-clustering and task-specific layers. Instead of sharing all layers, HMNet gradually separates the tasks in higher layers, so that the general language knowledge can be accumulated while the noise can be filtered out. Different from SkillNet-style methods, our HMNet only activates one channel at each layer, so its number of parameters is identical to a vanilla BERT. Surprisingly, HMNet with fewer parameters can even perform better. This demonstrates again the effectiveness of our proposed task clustering and the hierarchical sharing structure.