C ol D Fusion: Collaborative Descent for Distributed Multitask Finetuning (2023)

Abstract: Pretraining has been shown to scale well with compute, data size and data diversity. Multitask learning trains on a mixture of supervised datasets and produces improved performance compared to self-supervised pretraining. Until now, massively multitask learning required simultaneous access to all datasets in the mixture and heavy compute resources that are only available to well-resourced teams. In this paper, we propose ColD Fusion, a method that provides the benefits of multitask learning but leverages distributed computation and requires limited communication and no sharing of data. Consequentially, ColD Fusion can create a synergistic loop, where finetuned models can be recycled to continually improve the pretrained model they are based on. We show that ColD Fusion yields comparable benefits to multitask training by producing a model that (a) attains strong performance on all of the datasets it was multitask trained on and (b) is a better starting point for finetuning on unseen datasets. We find ColD Fusion outperforms RoBERTa and even previous multitask models. Specifically, when training and testing on 35 diverse datasets, ColD Fusion-based model outperforms RoBERTa by 2.19 points on average without any changes to the architecture.

## introduction over the last few years, pretrained language models are changing the landscape of nlp, where finetuning a pretrained model typically yields state-ofthe-art performance on a diverse set of nlp tasks @xcite . consequently, improving a pretrained model has the potential to boost every model finetuned on it. however, pretraining is often so computationally expensive that practitioners rarely seek to pretrain new models from scratch. in contrast, finetuning is usually dramatically cheaper, allowing a given pretrained model to be finetuned many times; e.g., there are thousands of finetuned bert variants on the hugging face hub 2 . motivated by this, we study if and how finetuned models can be "recycled" to create a better pretrained model (c.f., @xcite . to avoid confusion, henceforth we refer to any starting point for finetuning a base model and only the vanilla model as the pretrained model. to recycle models, we take inspiration from multitask learning ( §2). in multitask learning the pretrained model is finetuned over multiple datasets at once, which was shown to create a better base model than the original pretrained model @xcite @xcite . given the availability of many finetuned models, our aim is to obtain the benefits of multitask learning by mixing multiple models rather than multiple datasets (c.f. §2.3). to achieve that, we suggest the following iterative approach ( §3): in each iteration, contributors finetune the most up-to-date base model (which is presumably also the most performant) on their task, and share the fine-tuned model with the rest of the community. then, those contributed models are fused together, by simply averaging their parameters @xcite , to create the base model for the next iteration. we call this method collaborative descent fusion, or cold fusion. cold fusion fits the common finetuning paradigm, where each contributor finetunes for their own benefit and does not share their data. however, by merely requiring the finetuned model to be shared, the finetuning step can be recast as a training step for the collective's benefit. in doing so, our method allows reusing compute and data consumed by practitioners and researchers to the benefit of the entire community. our experimental results indicate that our approach of combining finetuned models not only produces a better base model but also allows this base model to keep evolving. instead of pretraining or multitasking on a predefined amount of data, we suggest accumulating finetuned models to continuously improve the model. our method is hence limited only by the amount of finetuned models that are shared by the entire community. we discuss limitations in ( §9). we show that cold fusion produces a model that performs well on the finetuned tasks, despite never manipulating more than one task at a time neither by constituent models nor their fusing ( §5). moreover, we show that cold fusion increases the performance of the base model substantially, outperforming the pretrained model by 2.33 points on average on 35 datasets. through additional analysis, we further show that similar improvements are achieved regardless of whether the target tasks were seen or unseen during training ( §5.2) and that accumulating models trained on additional data provides continuous improvement ( §6). ## background we start by motivating the use of further training on diverse data for enhancing the base model abilities ( §2.1). then, we continue with defining our framework's goals ( §2.2) and constraints ( §2.3). ## methodology -cold fusion our proposed method (see fig. 1 ), called cold fusion, is an iterative process that aims to perform multitask learning in the constrained setting outlined above. specifically, cold fusion involves an iterative process where each individual contributor downloads the current base model from the repository, finetunes this base model over their dataset, communicates the resulting model back to the repository, and lastly, the repository fuses @xcite all of the contributors' models into one and sets the new fused model as the new base model for further finetuning. more formally, the repository first initializes the shared model parameters θ 0 using a preexisting pretrained model. then, at each iteration @xmath0, 1, 2, . . .}, each contributor @xmath1. for the purposes of our study, finetuning is any optimization process that aims to minimize the loss over a dataset d. typically, finetuning involves minimizing the loss using a variant of gradient descent. after finetuning, each contributor sends their model's parameters θ c i to the repository. next, the repository fuses the contributor's models by averaging all of the contributor's model's parameters to produce a new shared model as θ i+1 = 1 |c| c θ c i . finally, the process repeats for iteration i + 1. ## experimental setup in this section, we detail the datasets, models, baselines, general experiment setup, and specific experiments settings. ## cold multitask results in this section, we show cold fusion can produce multitask models. we show in §5.1 that cold fusion fulfills both multitask objectives defined in §2. we verify that improvements replicate on datasets that were not seen during training ( §5.2). then we find that base model improvements are even more apparent in few shot settings ( §5.3). finally, we consider the importance of the number of contributors hyperparameter ( §5.4). ## single dataset analysis we now analyze the interacting effects of the core characteristics of cold fusion: additional data across iterations, the amount of training data per iteration, and the number of contributors in each iteration. doing so with multiple datasets would introduce noise. for example, we can not expect additional data coming from different sources (e.g., mnli or twitter) to equally affect the performance. to overcome this, we explore the case where a single dataset is distributed across contributors. using a single dataset allows us to reduce variability due to differences in the datasets (e.g., distribution, task, etc.), and isolate the parameter we wish to control. cold fusion may converge faster with models from a single dataset, but we still expect the general tendencies found to replicate in multiple datasets settings. we chose mnli @xcite for its large size (392k examples). ## related work our work strongly relies on model fusion. model fusion was first introduced as a way to improve pretrained models by @xcite . in parallel, several works such as @xcite and lately @xcite suggested different ways of fusing for other purposes such as improved finetuning. another fusion usage is the stochastic weight averaging, aiming to stabilize the sgd process by averaging multiple points along the sgd trajectory @xcite . unlike the previous, this method utilizes only one model and dataset. low-communication distributed training was proposed in similar settings to ours. @xcite proposed distributed finetuning and model fusing in order to produce better finetuned models. this suggestion is equivalent to one iteration of cold fusion where all models share the same dataset. @xcite ; together (2022) also share the similarity of distributed training, but during pretraining on unlabeled data. understanding why averaging different models improve quality may be related to theoretical works discussing weight and loss spaces. these works state there is a path of minimum loss between models @xcite on which the loss along the path is not increasing. @xcite ; @xcite ; @xcite claimed that under some constraints, this path is linear, which suggests that fusing the weights could produce a model that retains the capabilities of the fused models. although different models on the same task may converge to different locations in the loss space without linear connectivity @xcite , and although the case of multitask is more complex @xcite , we still believe that these works can partially explain why fusing preserves the capabilities gained by the constituent and when it does not that the next iteration fixes it. @xcite further suggests the linear connectivity path is merely a line in a whole connected region, future work may tell whether cold fusion searches in this region or crosses it to find new ones. the literature also includes methods for better aligning models during training @xcite @xcite or after it @xcite to aid in fusing. we did not use those as we wanted to reduce the load on the repository and avoid restricting the contributors' finetuning. however, these methods may improve results in cold multitask. we mention that multitask learning does not optimize the base model objective directly ( §2.3). some works aim to do so @xcite through meta-learning, finding models that can learn a new task well or efficiently @xcite . reptile @xcite meta learns in a way that resembles ours by iteratively using models trained for several batches. ## conclusion and discussion we proposed a scheme for utilizing abundant finetuned models to enhance a pretrained model. our approach does not necessitate the sharing of datasets, but rather assumes each contributor solely finetunes on their own dataset. hence, we believe that applying this scheme as a collaborative pretraining platform is a viable option and that doing so would result in ongoing improvement of base models. to scale this approach, it would be beneficial if the repository was updated asynchronously, perhaps relying on recent fusing techniques @xcite . in the usual finetuning setting, robustness can be improved by tuning batch size and learning rate. in analogy, in cold fusion, one can either increase the number of contributors (batch) and/or restrict the effect of each iteration (learning rate) @xcite to improve the process. following this line, future work may consider regularizing the distance from the pretrained model (learning rate) when a small number of contributors exist (batch) or consider assigning individual weights to each contributor. there are many hyper parameters to optimize which might improve the method substantially. e.g., fusing the contributions with a weighted average, improving fusing itself @xcite , controlling the datasets seen in each iterations (related to; @xcite and backtracking when a harmful update was done to the model. we hope that future work will shed more light on these issues, to further improve the approach proposed in this work.