Multi-task dataset mixing

[ghomasHudson]

It seems like many of the best performing models on the GLUE benchmark make some use of multitask learning (simultaneous training on multiple tasks).

The T5 paper highlights multiple ways of mixing the tasks together during finetuning:

• Examples-proportional mixing - sample from tasks proportionally to their dataset size
• Equal mixing - sample uniformly from each task
• Temperature-scaled mixing - The generalized approach used by multilingual BERT which uses a temperature T, where the mixing rate of each task is raised to the power 1/T and renormalized. When T=1 this is equivalent to equal mixing, and becomes closer to equal mixing with increasing T.

Following this discussion huggingface/transformers#4340 in transformers, @enzoampil suggested that the nlp library might be a better place for this functionality.

Some method for combining datasets could be implemented ,e.g.

dataset = nlp.load_multitask(['squad','imdb','cnn_dm'], temperature=2.0, ...)

We would need a few additions:

• Method of identifying the tasks - how can we support adding a string to each task as an identifier: e.g. 'summarisation: '?
• Method of combining the metrics - a standard approach is to use the specific metric for each task and add them together for a combined score.

It would be great to support common use cases such as pretraining on the GLUE benchmark before fine-tuning on each GLUE task in turn.

I'm willing to write bits/most of this I just need some guidance on the interface and other library details so I can integrate it properly.

[patrickvonplaten]

I like this feature! I think the first question we should decide on is how to convert all datasets into the same format. In T5, the authors decided to format every dataset into a text-to-text format. If the dataset had "multiple" inputs like MNLI, the inputs were concatenated. So in MNLI the input:

• Hypothesis: The St. Louis Cardinals have always won.

• Premise: yeah well losing is i mean i’m i’m originally from Saint Louis and Saint Louis Cardinals when they were there were uh a mostly a losing team but

was flattened to a single input:

mnli hypothesis: The St. Louis Cardinals have always won. premise:
yeah well losing is i mean i’m i’m originally from Saint Louis and Saint Louis Cardinals
when they were there were uh a mostly a losing team but.

This flattening is actually a very simple operation in nlp already. You would just need to do the following:

def flatten_inputs(example):
    return {"input": "mnli hypothesis: " + example['hypothesis'] + " premise: " + example['premise']}

t5_ready_mnli_ds = mnli_ds.map(flatten_inputs, remove_columns=[<all columns except output>])

So I guess converting the datasets into the same format can be left to the user for now.
Then the question is how we can merge the datasets. I would probably be in favor of a simple

dataset.add()

function that checks if the dataset is of the same format and if yes merges the two datasets. Finally, how should the sampling be implemented? Examples-proportional mixing corresponds to just merging the datasets and shuffling. For the other two sampling approaches we would need some higher-level features, maybe even a dataset.sample() function for merged datasets.

What are your thoughts on this @thomwolf @lhoestq @ghomasHudson @enzoampil ?

[ghomasHudson]

I agree that we should leave the flattening of the dataset to the user for now. Especially because although the T5 framing seems obvious, there are slight variations on how the T5 authors do it in comparison to other approaches such as gpt-3 and decaNLP.

In terms of sampling, Examples-proportional mixing does seem the simplest to implement so would probably be a good starting point.

Temperature-scaled mixing would probably most useful, offering flexibility as it can simulate the other 2 methods by setting the temperature parameter. There is a relevant part of the T5 repo which should help with implementation.

According to the T5 authors, equal-mixing performs worst. Among the other two methods, tuning the K value (the artificial dataset size limit) has a large impact.

[enzoampil]

I agree with going with temperature-scaled mixing for its flexibility!

For the function that combines the datasets, I also find dataset.add() okay while also considering that users may want it to be easy to combine a list of say 10 data sources in one go.

dataset.sample() should also be good. By the looks of it, we're planning to have as main parameters: temperature, and K.

On converting the datasets to the same format, I agree that we can leave these to the users for now. But, I do imagine it'd be an awesome feature for the future to have this automatically handled, based on a chosen approach to formatting 😄

E.g. T5, GPT-3, decaNLP, original raw formatting, or a contributed way of formatting in text-to-text.

[thomwolf]

This is an interesting discussion indeed and it would be nice to make multi-task easier.

Probably the best would be to have a new type of dataset especially designed for that in order to easily combine and sample from the multiple datasets.

This way we could probably handle the combination of datasets with differing schemas as well (unlike T5).

[ghomasHudson]

@thomwolf Are you suggesting making a wrapper class which can take existing datasets as arguments and do all the required sampling/combining, to present the same interface as a normal dataset?

That doesn't seem too complicated to implement.

[ghomasHudson]

I guess we're looking at the end user writing something like:

ds = nlp.load_dataset('multitask-t5',datasets=["squad","cnn_dm",...], k=1000, t=2.0)

Using the t5 method of combining here (or this could be a function passed in as an arg)

Passing kwargs to each 'sub-dataset' might become tricky.

[ghomasHudson]

From thinking upon @thomwolf 's suggestion, I've started experimenting:

class MultitaskDataset(DatasetBuilder):
    def __init__(self, *args, **kwargs):
        super(MultitaskDataset, self).__init__(*args, **kwargs)
        self._datasets = kwargs.get("datasets")

    def _info(self):
        return nlp.DatasetInfo(
            description=_DESCRIPTION,
            features=nlp.Features({
                    "source": nlp.Value("string"),
                    "target": nlp.Sequence(nlp.Value("string"))
                })
        )

    def _get_common_splits(self):
        '''Finds the common splits present in all self._datasets'''
        min_set = None
        for dataset in self._datasets:
            if min_set != None:
                min_set.intersection(set(dataset.keys()))
            else:
                min_set = set(dataset.keys())
        return min_set

....

# Maybe this?:
squad = nlp.load_dataset("squad")
cnn_dm = nlp.load_dataset("cnn_dailymail","3.0.0")
multitask_dataset = nlp.load_dataset(
    'multitask_dataset',
    datasets=[squad,cnn_dailymail], 
    k=1000, 
    t=2.0
)

Does anyone know what methods of MultitaskDataset I would need to implement? Maybe as_dataset and download_and_prepare? Most of these should be just calling the methods of the sub-datasets.

I'm assuming DatasetBuilder is better than the more specific GeneratorBasedBuilder, BeamBasedBuilder, etc....

One of the other problems is that the dataset size is unknown till you construct it (as you can pick the sub-datasets). Am hoping not to need to make changes to nlp.load_dataset just for this class.

I'd appreciate it if anyone more familiar with nlp's internal workings could tell me if I'm on the right track!

[thomwolf]

I think I would probably go for a MultiDataset wrapper around a list of Dataset.

I'm not sure we need to give it k and t parameters at creation, it can maybe be something along the lines of:

squad = nlp.load_dataset("squad")
cnn_dm = nlp.load_dataset("cnn_dailymail","3.0.0")

multitask_dataset = nlp.MultiDataset(squad, cnn_dm)

batch = multitask_dataset.sample(10, temperature=2.0, k=1000)

The first proof-of-concept for multi-task datasets could definitely require that the provided datasets have the same name/type for columns (if needed you easily rename/cast a column prior to instantiating the MultiDataset).

It's good to think about it for some time though and don't overfit too much on the T5 examples (in particular for the ways/kwargs for sampling among datasets).

[ghomasHudson]

The problem with changing k and t per sampling is that you'd have to somehow remember which examples you'd already returned while re-weighting the remaining examples based on the new k and tvalues. It seems possible but complicated (I can't really see a reason why you'd want to change the weighting of datasets after you constructed the multidataset).

Wouldn't it be convenient if it implemented the dataset interface? Then if someone has code using a single nlp dataset, they can replace it with a multitask combination of more datasets without having to change other code. We would at least need to be able to pass it into a DataLoader.

[ghomasHudson]

A very janky (but working) implementation of multitask_dataset.sample() could be something like this:

import nlp
import torch

class MultiDataset():
    def __init__(self, *args, temperature=2.0, k=1000, maximum=None, scale=1):
        self.datasets = args
        self._dataloaders = {}
        for split in self._get_common_splits():
            split_datasets = [ds[split] for ds in self.datasets]
            mixing_rates = self._calc_mixing_rates(split_datasets,temperature, k, maximum, scale)
            weights = []
            for i in range(len(self.datasets)):
                weights += [mixing_rates[i]]*len(self.datasets[i][split])
            self._dataloaders[split] = torch.utils.data.DataLoader(torch.utils.data.ConcatDataset(split_datasets),
                                                        sampler=torch.utils.data.sampler.WeightedRandomSampler(
                                                            num_samples=len(weights),
                                                            weights = weights,
                                                            replacement=True),
                                                        shuffle=False)

    def _get_common_splits(self):
        '''Finds the common splits present in all self.datasets'''
        min_set = None
        for dataset in self.datasets:
            if min_set != None:
                min_set.intersection(set(dataset.keys()))
            else:
                min_set = set(dataset.keys())
        return min_set


    def _calc_mixing_rates(self,datasets, temperature=2.0, k=1000, maximum=None, scale=1):
       '''Work out the weighting of each dataset based on t and k'''
        mixing_rates = []
        for dataset in datasets:
            rate = len(dataset)
            rate *= scale
            if maximum:
                rate = min(rate, maximum)
            if temperature != 1.0:
                rate = rate ** (1.0/temperature)
            mixing_rates.append(rate)
        return mixing_rates

    def sample(self,n,split):
        batch = []
        for example in self._dataloaders[split]:
            batch.append(example)
            n -= 1
            if n == 0:
                return batch


def flatten(dataset,flatten_fn):
    for k in dataset.keys():
        if isinstance(dataset[k],nlp.Dataset):
            dataset[k] = dataset[k].map(flatten_fn,remove_columns=dataset[k].column_names)

# Squad
def flatten_squad(example):
    return {"source": "squad context: " + example['context'] + " question: " + example['question'],"target":example["answers"]["text"]}
squad = nlp.load_dataset("squad")
flatten(squad,flatten_squad)

# CNN_DM
def flatten_cnn_dm(example):
    return {"source": "cnn_dm: " + example['article'],"target":[example["highlights"]]}
cnn_dm = nlp.load_dataset("cnn_dailymail", "3.0.0")
flatten(cnn_dm,flatten_cnn_dm)

multitask_dataset = MultiDataset(squad, cnn_dm)
batch = multitask_dataset.sample(100,"train")

There's definitely a more sensible way than embedding DataLoaders inside.

[thomwolf]

There is an interesting related investigation by @zphang here https://colab.research.google.com/github/zphang/zphang.github.io/blob/master/files/notebooks/Multi_task_Training_with_Transformers_NLP.ipynb

[ghomasHudson]

Good spot! Here are my thoughts:

• Aside: Adding MultitaskModel to transformers might be a thing to raise - even though having task-specific heads has become unfashionable in recent times in favour of text-to-text type models.

• Adding the task name as an extra field also seems useful for these kind of models which have task-specific heads

• There is some validation of our approach that the user should be expected to map datasets into a common form.

• The size-proportional sampling (also called "Examples-proportional mixing") used here doesn't perform too badly in the T5 paper (it's comparable to temperature-scaled mixing in many cases but less flexible. This is only reasonable with a K maximum size parameter to prevent very large datasets dominating). This might be good for a first prototype using:

  def __iter__(self):
      """
      For each batch, sample a task, and yield a batch from the respective
      task Dataloader.
  
      We use size-proportional sampling, but you could easily modify this
      to sample from some-other distribution.
      """
      task_choice_list = []
      for i, task_name in enumerate(self.task_name_list):
          task_choice_list += [i] * self.num_batches_dict[task_name]
      task_choice_list = np.array(task_choice_list)
      np.random.shuffle(task_choice_list)
  
      dataloader_iter_dict = {
          task_name: iter(dataloader) 
          for task_name, dataloader in self.dataloader_dict.items()
      }
      for task_choice in task_choice_list:
          task_name = self.task_name_list[task_choice]
          yield next(dataloader_iter_dict[task_name])    

  We'd just need to pull samples from the raw datasets and not from DataLoaders for each task. We can assume the user has done dataset.shuffle() if they want to.

  Other sampling methods can later be implemented by changing how the task_choice_list is generated. This should allow more flexibility and not tie us to specific methods for sampling among datasets.