Multitask
State: closed | Author: ghomasHudson (CONTRIBUTOR)

Following our discussion in #217, I've implemented a first working version of `MultiDataset`.

There's a function `build_multitask()` which takes either individual `nlp.Dataset`s or `dicts` of splits and constructs `MultiDataset`(s). I've added a notebook with example usage.

I've implemented many of the `nlp.Dataset` methods (cache_files, columns, nbytes, num_columns, num_rows, column_names, schema, shape). Some of the other methods are complicated as they change the number of examples. These raise `NotImplementedError`s at the moment.

This will need some tests which I haven't written yet.

There's definitely room for improvements but I think the general approach is sound. 

Comments: It's definitely going in the right direction ! Thanks for giving it a try

I really like the API.
IMO it's fine right now if we don't have all the dataset transforms (map, filter, etc.) as it can be done before building the multitask dataset, but it will be important to have them in the end.
All the formatting methods could easily be added though.

I think there are some parts that will require some work with apache arrow like slicing. I can find a way to do it using pyarrow tables concatenation (I did something similar when implementing `__getitem__` with an input that is a list of indices [here](https://github.com/huggingface/nlp/pull/322/files#diff-73270df8d7f08c62a27e40806e1a5fb0R463-R469)). It is very fast and it allows to have the same output format as a normal Dataset.

Also maybe we should check that not only the columns but also the schemas match ?
And maybe add the `seed` of the shuffling step as an argument ?



Maybe we should remove the methods that are not implemented for now, WDYT @thomwolf ?

That's an interesting first draft, thanks a lot for that and the user facing API is really nice.

I think we should dive more into this and the questions of #217 before merging the first version though.

In particular, the typical way to do multi-tasking is usually to sample a task and then sample a batch within the selected task. I think we should probably stay be closer to this traditional approach, or at least make it very easy to do, rather than go to close to the T5 approach which is very specific to this paper.

In this regard, it seems important to find some way to address the remarks of @zphang. I'm still wondering if we should not adopt more of a sampling approach rather than an iteration approach.

@thomwolf Thanks! I mainly wanted to get something working quickly for my own MTL research. I agree with a lot of the points you made so I'll convert this pull request back to a draft.

For your specific point about 'batch-level' multitask mixing, it would be a pretty trivial change to add a `batch_size` parameter and ensure every `batch_size` examples are from the same task. This would certainly work, but would add a notion of 'batches' to a Dataset, which does feel like a 'Sampler-level' concept and not a Dataset one. There's also the possibility of wanting some specific task-level sampling functionality (e.g. applying `SortishSampler` to each task) which would only work with this kind of 2 step sampling approach. My first proposal in the transformers repo was actually a Sampler https://github.com/huggingface/transformers/issues/4340. I wonder whether functionality at the sampler-level has a place in the vision for the `nlp` repo?

I imagine following a sampling approach you'd have to abandon maintaining the same user-facing API as a standard dataset (A shame because replacing a single dataset seamlessly with a multitask one is a really nice user-experience).

Random half-Idea: You could have a class which accepts a list of any iterables (either a Dataset or a DataLoader which already is doing the batching). Not sure what interface you'd present though. hmmm. 

There's definitely more discussion to have. 


Are there any updates on making multi-task learning more officially supported in the datasets/transformers libraries? 
Given that many papers use more than one task, it would be great to have multi-task learning more officially supported and easier to use. There are a few notebooks/blogs about using HF Transformers for this, but they all mention that it's more of a hack and not really officially supported (e.g. [this notebook](https://colab.research.google.com/github/zphang/zphang.github.io/blob/master/files/notebooks/Multi_task_Training_with_Transformers_NLP.ipynb#scrollTo=xW8bnTgCsx5c), or [this blog](https://medium.com/@shahrukhx01/multi-task-learning-with-transformers-part-1-multi-prediction-heads-b7001cf014bf)). 

[jiant](https://github.com/nyu-mll/jiant) was a framework built on transformers that made multi-task learning a first class feature of the library until recently, but they stopped maintaining their library a month ago ([see here](https://github.com/nyu-mll/jiant)). 
This could be a good reason to increase support from the HF team? @lhoestq @thomwolf 

I'm not advanced enough to contribute on this, but an up-to-date notebook showing how to train a model e.g. on both MLM and next-sentence-prediction would already be very useful!

> Are there any updates on making multi-task learning more officially supported in the datasets/transformers libraries? Given that many papers use more than one task, it would be great to have multi-task learning more officially supported and easier to use. There are a few notebooks/blogs about using HF Transformers for this, but they all mention that it's more of a hack and not officially supported (e.g. [this notebook](https://colab.research.google.com/github/zphang/zphang.github.io/blob/master/files/notebooks/Multi_task_Training_with_Transformers_NLP.ipynb#scrollTo=xW8bnTgCsx5c), or [this blog](https://medium.com/@shahrukhx01/multi-task-learning-with-transformers-part-1-multi-prediction-heads-b7001cf014bf)).
> 
> [jiant](https://github.com/nyu-mll/jiant) was a framework built on transformers that made multi-task learning a first class feature of the library until recently, but they stopped maintaining their library a month ago. This could be a good reason to increase support from the HF team? @lhoestq
> 
> I'm not advanced enough to contribute on this, but an up-to-date notebook showing how to train a model e.g. on both MLM and NSP would already be very useful!

I kinda stopped working on this as I didn't really get any response on an actual workable solution.

The problem that I came up against after initially being redirected here after [proposing this in the transformers repo](https://github.com/huggingface/transformers/issues/4340) ([among](https://github.com/huggingface/transformers/issues/6872) [others](https://github.com/huggingface/transformers/issues/1856)) , was the request be able to do the multitask mixing at the batch level as well as at the level of individual examples. As this repo doesn't really have the concept of 'batches' it would need to be implemented in the transformers repo, rather than here. You could then pick which level to do your multitask learning on.

Work on T5 and as of last week, on [exT5](https://arxiv.org/pdf/2111.10952.pdf), have shown that multitask mixing on the example level works incredibly well (with a big enough batch size), so if you're ok doing that, then this pull request works.

I completely agree that multitask learning is a vital part of modern NLP, nearly every piece of research code I write has at least some aspect of multitask learning (currently using this patch). Many of the top GLUE and SuperGLUE submissions are using some aspect of mutlitask learning. We need to support it.

Fully agree. Batching and data loading is one important thing. The part I'm struggling with right now is the classification head (which is more part of the Transformers repo, but also essential for multi-task learning). @ghomasHudson, how do you tune two classification heads simultaneously? Say, when I want to fine-tune an existing base-model on some classification task (like NLI, or next-sentence-prediction) and at the same time add some MLM for regularisation & domain adaptation. In this case I need two classification heads, but I don't know how to switch them between the batches. 

> Fully agree. Batching and data loading is one important thing. The part I'm struggling with right now is the classification head (which is more part of the Transformers repo, but also essential for multi-task learning). @ghomasHudson, how do you tune two classification heads simultaneously? Say, when I want to fine-tune an existing base-model on some classification task (like NLI, or next-sentence-prediction) and at the same time add some MLM for regularisation & domain adaptation. In this case I need two classification heads, but I don't know how to switch them between the batches.

This pull request is mainly focused on getting the data in the right format, but you're right that there's no easy way to pick between the heads without something like jiant. You could of course replicate this functionality yourself - probably by making a class that implements the functionality of both `ModelNameForSequenceClassification` or `ModelNameForMaskedLM` picking between them depending on some task parameter you add to the forward pass. 

jiant make this approach model agnostic by [ignoring the custom per-model head implementations of huggingface](https://github.com/nyu-mll/jiant/blob/386d4e726a27becda1b03c241f064eb13c54860f/jiant/proj/main/modeling/heads.py#L17-L18), instead making generic versions. Then the jiant code [passes a `task` parameter](https://github.com/nyu-mll/jiant/blob/386d4e726a27becda1b03c241f064eb13c54860f/jiant/proj/main/modeling/primary.py#L107-L109) into their [JiantModel](https://github.com/nyu-mll/jiant/blob/386d4e726a27becda1b03c241f064eb13c54860f/jiant/proj/main/modeling/primary.py#L36-L79) wrapper. To implement this in huggingface transformers would require quite a few modifications to the current approach (potentially interfering with some other project aims e.g. code readability), so you might find it tricky to get a change like that accepted. It would be super cool though.

And there's of course the exT5 way of doing things too where you sidestep this issue entirely by treating both tasks as text-to-text problems so you can end up with 100% shared parameters, e.g.
MLM: `Lorem <mask_0> amet, consectetur <mask_1> do eiusmod tempor incididunt ut labore <mask_2>`
NLI: `Premise: The Old One always comforted Ca'daan, except today. hypothesis: Ca'daan knew the Old One very well.`
This also allows you to do mixed batches of both tasks.

Personally, my research mainly focuses on this last approach, using the structure of the data itself to indicate the task rather than swapping in and out different parts of the network.

Hi! `jiant` maintainer here, don't have much to add to the conversation yet but I'm happy to share my experience/thoughts on working with Multitask models if people have questions.

Hi ! I think it could be easier to simply share as examples in `transformers` some code that uses `jiant` and/or subclass/reimplement some part of `transformers` for multitask ?

> Hi ! I think it could be easier to simply share as examples in `transformers` some code that uses `jiant` and/or subclass/reimplement some part of `transformers` for multitask ?

Well since `jiant` requires new huggingface models to be explicitly added (as there are ["subtle differences in the models that jiant must abstract"](https://github.com/nyu-mll/jiant/blob/master/guides/models/adding_models.md)), and isn't being maintained anymore, then the first option might be out of date quickly.

If `transformers` could move towards making the task-specific heads more generic and as well as [creating a new base model in the `__init__` method](https://github.com/huggingface/transformers/blob/43f953cc2eec804eba04e2a9ae164d1a33fd97a8/src/transformers/models/bert/modeling_bert.py#L1502), allowing it to be passed as an argument (along with other little tweaks to standardize the approach), then this functionality could be moved into `transformers` itself.

It does seem a little redundant to have `jiant` as a library abstracting all the idiosyncrasies of each model type, where this could be done directly in the `transformers` repo in a single place alongside the model.

It's not an easy problem to solve though, especially balanced with the desire to expose models with minimal abstraction. @zphang probably knows more about this than me though.

As mentioned, one of the main obstacles is that HF/T doesn't support generic heads. At first glance, this should be easy, since the interface is quite simple: models output both a token-wise and a sequence representation (e.g. `[CLS]`), and heads use either one and output the corresponding predictions/losses.

However, there are a number of cases where this doesn't work. One of them is multiple-choice tasks like HellaSwag, which is a multiple choice task with 4 text options. The way this is normally formatted is that you encode `context + question + option_X` for X=1..4, and then score all four options based on a scoring head and pick the highest scoring option as the prediction. This requires you to run the encoder on 4 separate inputs, which breaks the above abstraction (the task-specific model might need to call the encoder multiple times).

Another thing is batching. You can imagine with the above that you might want a different batch size for multiple-choice tasks compared to simpler classification tasks. This means you need task-specific batching as well. In addition, [it's been shown](https://arxiv.org/abs/2101.11038) that you really want to mix tasks within a single batch. This also leads into issues like how you want to sample different task examples, early stopping on them, how to mix the validation scores, etc. (`jiant` addressed these, through probably more-complicated-than-necessary configurations.)

None of these are insurmountable problems, but it requires some tweaking of the current code layout to get it to work. I would guess that it wouldn't take much work to get a 90% implementation.

> Another thing is batching. You can imagine with the above that you might want a different batch size for multiple-choice tasks compared to simpler classification tasks. This means you need task-specific batching as well. In addition, [it's been shown](https://arxiv.org/abs/2101.11038) that you really want to mix tasks within a single batch. This also leads into issues like how you want to sample different task examples, early stopping on them, how to mix the validation scores, etc. (`jiant` addressed these, through probably more-complicated-than-necessary configurations.)

That's reassuring. exT5 find the same thing - that mixing tasks together in a batch gives better performance (provided the batch size is big enough that each batch contains a mix of different tasks). Assuming this, we can ignore doing things at the batch-level and just do this at the individual example level - in which case this pull request already does the data mixing part of the problem! Balancing different tasks could easily be added here by implementing temperature-scaled mixing, custom weights, etc...

To make a generic implementation of this using different heads would be hard (impossible?) without doing the sub-batching that Muppet do - in which case we're back at dealing with the 'batch' (sub-batch) level which would need an implementation in `transformers` not here.



Mixing at example should work fine. One issue though is that, as mentioned above, different tasks maybe actually require different amounts of memory, so downstream the user would have to find some way to handle that. But this might be one of those "the last/edge-case 10% is the hardest" to handle kind of deals.

Very true - there's always going to be those cases. I also feel that the way things are going, if we just leave this for a few years no one will be wanting to use task-specific heads anymore - it'll all be task prompts included in the input a-la GPT, T5, etc... which will make this substantially simpler to implement.

It's quite tricky to make a suitably non-opinionated generic version of this at the moment.

> Is there an advantage to varying the proportions of each task in each batch

Some tasks have much less data than others. E.g. SNLI vs. CoLA is almost a 100x difference, so people often sample differently-sized tasks differently.

As a short-term solution, I like @lhoestq's suggestion to create a notebook that shows how to implement multi-task learning by subclassing some transformer & dataset classes in a general way. I've been trying to get @zphang's [great but old notebook](https://colab.research.google.com/github/zphang/zphang.github.io/blob/master/files/notebooks/Multi_task_Training_with_Transformers_NLP.ipynb#scrollTo=CQ39AbTAPAUi) on multi-task learning running today and I didn't get it to work, probably because it was implemented a long time ago with `transformers==2.11`, `torch==1.2`~ etc and installing older versions still caused errors.
There is also this [interesting new repo](https://github.com/shahrukhx01/multitask-learning-transformers), which has a cool way of enabling you to save and load a model with two classification heads ([see model here](https://huggingface.co/shahrukhx01/bert-multitask-query-classifiers) and blog post [here](https://medium.com/@shahrukhx01/multi-task-learning-with-transformers-part-1-multi-prediction-heads-b7001cf014bf)). Haven't tried it yet, but it only uses `BertForSequenceClassification` instead of the more general AutoModelForXYZ

@zphang, would you maybe be up for contributing an updated version of your older notebook with the latest version of `transformers` and `datasets` which runs in today's colabs? I feel like this would be very helpful for the community and if you keep the classes/functions somewhat general, people can easily adapt it to their use cases! 🙏 :) 
Would be a great addition to the [HF notebooks](https://huggingface.co/docs/transformers/notebooks).

In the medium-term, I agree that it would be great to have more native support for this via the HF libraries. I feels weird that you can neither train the old BERT (trained on two tasks) nor any of the newer models, without some hacks. 

@zphang would love to see the newer notebook as suggested by @MoritzLaurer