Ensembles and Cocktails: Robust Finetuning for Natural Language Generation

1 INTRODUCTION When finetuning a pretrained language model for natural language generation tasks like summarization ( Narayanan et al. , 2018 ) and table-to-text ( Gardent et al. , 2017 ) , one is currently faced with a tradeoff . One can achieve strong in-distribution ( ID ) performance or strong out-of-distribution ( OOD ) performance—but not both—by choosing one of two families of finetuning approaches . Finetuning all parameters of the pretrained language model achieves strong ID performance ( Howard & Ruder , 2018 ) . We call this full finetuning . Alternatively , freezing most of the pretrained parameters during finetuning , in e.g. , adapters ( Rebuffi et al. , 2017 ; Houlsby et al. , 2019a ) , prefix-tuning ( Li & Liang , 2021 ; Lester et al. , 2021 ) and bitfit ( Zaken et al. , 2021 ) , achieves stronger OOD performance than full finetuning ( Li & Liang , 2021 ; Lester et al. , 2021 ) , but worse ID performance . We call these methods lightweight finetuning . In this work , we are the first to demonstrate that this tradeoff is not necessary in natural language generation : we can achieve the best of full and lightweight finetuning both ID and OOD . We test this across three natural language generation tasks in English : summarization ( Narayanan et al. , 2018 ) , table-to-text ( Gardent et al. , 2017 ) , and open-domain question answering ( Roberts et al. , 2020 ; Lee et al. , 2019a ) , and two lightweight finetuning methods : adapters and prefix-tuning . Achieving the best of both worlds is important since across our 6 settings ( 2 lightweight finetuning methods and 3 datasets ) , full finetuning significantly outperforms lightweight finetuning ID in 5 settings , and lightweight finetuning significantly outperforms full finetuning OOD in 6 settings . We present two methods for achieving the best of full and lightweight finetuning . We first show that a weighted ensemble of one full finetuning model and one lightweight finetuning model achieves performance comparable to the best of the two , both ID and OOD . To avoid running two models at once , we propose cocktail finetuning : first train a lightweight finetuning model , and then train a full finetuning model on a mixture of student-teacher distillation loss ( Bucilă et al. , 2006 ; Hinton et al. , 2015 ) ( with the lightweight model as teacher ) and cross-entropy loss , on the original ID data.1 The ensemble weight ( resp . the cocktail mixture loss weight ) is chosen to optimize for ID performance , and we observe that the resultant model achieves comparable OOD performance to lightweight finetuning . Our work contributes to an ongoing realization in machine learning : robust models that deviate minimally from a pretrained model can be gainfully combined with a model adapted more precisely to a specific dataset . Most saliently , in concurrent work , Wortsman et al . ( 2021 ) find that zeroshot image recognition through CLIP ( Radford et al. , 2021 ) can be combined with a finetuned network to perform well ID and OOD . Our work demonstrates that similar results hold true in natural language generation as well . We draw from simple , well-known methods : ensembles , student-teacher distillation ( Abnar et al. , 2020 ; Hinton et al. , 2015 ) —but our application to achieving the best of ID and OOD performance in natural language generation is novel . To better understand how cocktail finetuning can transfer the out-of-distribution behavior of the lightweight finetuning model when distilling only on in-distribution data , we theoretically study multiclass logistic regression with a large number of output classes , viewing language generation as a classification problem over a large output space . We study a distribution shift setting where not all classes are seen in the ID training data ( as is true in language generation due to Zipf ’ s law ( Zipf , 1949 ) ) , and getting good OOD accuracy requires generalizing to unseen outputs . First , we prove that a model trained with gradient descent can not distinguish the unseen output classes , leading to poor OOD accuracy . Second , we prove that distilling a teacher model with good OOD accuracy on only ID training data can transfer the OOD behavior of the teacher to the student , such that the student and teacher make the same prediction on any new input . As a result , the student also has good OOD accuracy . In summary , we present two simple methods—ensembling and cocktail finetuning—for achieving strong ID and OOD natural language generation with only the original ID data , achieving the “ best of both worlds ” of full and lightweight finetuning methods . 2 RELATED WORK . Lightweight finetuning . Recently , as pretrained models scale up in the number of parameters , lightweight finetuning techniques have been developed to avoid updating all of the parameters . These techniques update a subset of the parameters ( Zaken et al. , 2021 ) , introduce a small number of new parameters between layers of the model ( called “ adapters ” ) ( Rebuffi et al. , 2017 ; Houlsby et al. , 2019b ; Pfeiffer et al. , 2020 ) , or learn auxiliary inputs that are prepended to the models inputs ( called “ prompts ” ) ( Shin et al. , 2020 ; Li & Liang , 2021 ; Qin & Eisner , 2021 ; Liu et al. , 2021b ; Logan IV et al. , 2021 ; and see Liu et al. , 2021a for a survey ) . These techniques empirically have the benefit of being able to achieve strong out-of-distribution extrapolation performance at the cost of worse in-distribution performance compared to full finetuning models both in classification and generation settings ( Li & Liang , 2021 ; Lester et al. , 2021 ; Zhou et al. , 2021 ; Utama et al. , 2021 ) . In our work , we aim to overcome this trade-off by proposing techniques for achieving the best of both worlds , with a focus on language generation . Student-teacher distillation . Buciluǎ et al . ( 2006 ) introduce distillation for compressing large ensembles into smaller , faster , more efficient models , and Hinton et al . ( 2015 ) generalize their method and call it knowledge distillation . Kasai et al . ( 2021 ) distill Transformers into RNNs for linear-time inference . Abnar et al . ( 2020 ) show that knowledge distillation can transfer the inductive bias of one architecture to another : for example , distilling an LSTM to a Transformer leads to a Transformer with 1We use the term cocktail as it is a refreshing mixture involving distillation . representations that are more similar to independently trained LSTMs than independently trained Transformers . The efficiacy of our proposed cocktail finetuning is surprising even in context of these results , as it empirically transfers just the strong OOD performance of the teacher lightweight model to the student , while not transferring the weak ID performance . Robust finetuning in vision . In concurrent work , Wortsman et al . ( 2021 ) propose a similar approach to ours , but focus on image classification : they combine a full finetuning model with a lightweight finetuning ( linear probe ) model by ensembling the two models in weight space , and show that their approach maintains both good ID and OOD performance . It is yet surprising that in the structured prediction setting of natural language generation , we further show that we can achieve high OOD performance even when weight-space ensembling is impossible ( for instance , since there are weights not held in common between the full finetuning model and an optimized prompt ) . Moreover , we further this line of work by providing a theoretical study into how distillation can transfer OOD behavior while only using ID data . Self-training for robustness . In adversarial robustness , standard models achieve good accuracy on clean examples but are not robust to adversarial examples , while adversarially trained models are more robust but are less accurate on clean examples . Raghunathan et al . ( 2020 ) show that the benefits of these two models can be combined using robust self-training ( RST ) ( Carmon et al. , 2019 ; Uesato et al. , 2019 ; Najafi et al. , 2019 ) on unlabeled data to improve clean and robust accuracy . Similarly , Khani & Liang ( 2021 ) show that RST can be used to combine models with spurious features removed ( which are more robust ) and without removing spurious features ( which are more accurate ) . In general OOD robustness , In-N-Out ( Xie et al. , 2021 ) uses self-training to transfer the robustness benefits of a model pre-trained with auxiliary self-supervision to a model trained with extra auxiliary features . Berthelot et al . ( 2021 ) use unlabeled data and a consistency regularization loss for domain adaptation . While many of these methods combine the benefits of two models or use a self-training/distillation loss to improve robustness , cocktail finetuning differs in that it performs distillation only on ID training data without any additional ID or OOD unlabeled examples . 3 PROBLEM STATEMENT . 3.1 SETUP . We consider a conditional generation task from an input space X ( e.g. , a news article ) to an output space Y ( e.g. , a summary of the article ) , where both spaces consist of sequences of tokens from a shared vocabulary V . Data . Let Pid and Pood denote the distribution of ( x , y ) pairs in-distribution and out-of-distribution , respectively , where x ∈ X and y ∈ Y . In summarization , Pid may consist of articles about world and business news while Pood consists of health and technology news . The training data Dtrain and validation data Dval consist of n and nval in-distribution ( x , y ) pairs respectively , drawn from Pid . Model and metrics . Let qγ ( y | x ) = ∏len ( y ) j=1 qγ ( yj | x , y < j ) be a pretrained neural autoregressive sequence model with parameters γ. Finetuning ( full or lightweight ) produces a finetuned model pθ ( y | x ) from qγ , where θ denotes the set of trainable parameters . In full finetuning , θ is the same set of parameters as γ ; in lightweight finetuning , θ often refers to the parameters of inserted modules ( e.g. , adapters modules or prefixes ) . The final predictor is fθ ( x ) which approximates argmaxy pθ ( y | x ) through , e.g. , beam search . The goal is to learn parameters θ such that the predictor fθ obtains good ID and OOD performance : Sid ( fθ ) = Ex , y∼Pid [ score ( y , fθ ( x ) ) ] ( 1 ) Sood ( fθ ) = Ex , y∼Pood [ score ( y , fθ ( x ) ) ] ( 2 ) where score may be defined differently for each task ( e.g. , exact match score or BLEU ( Papineni et al. , 2002b ) ) . 3.2 BASELINE FINE-TUNING METHODS . Full finetuning has been the defacto approach to adapt pretrained language models for downstream generation tasks ( Lewis et al. , 2020 ; Kale & Rastogi , 2020 ) . Under this framework , the model pθ is initialized with its pretrained parameters γ , and we optimize the following log-likelihood objective . L ( θ ) = E ( x , y ) ∼Dtrain [ − log pθ ( y | x ) ] ( 3 ) where the set of trainable parameters θ refer to the set of pretrained parameters γ. Lightweight finetuning is another family of adaptation approaches , which freezes most of the pretrained parameters and augments the model with small trainable modules . Lightweight finetuning shares the same training objective as full finetuning , but the trainable parameters θ are different . Within the lightweight finetuning family , we focus on adapter-tuning ( Houlsby et al. , 2019a ) and prefix-tuning ( Li & Liang , 2021 ) . Both approaches keep the pretrained parameters intact . Adaptertuning inserts task-specific MLP layers between each layer of the pretrained language models , and the trainable parameter θ consists of these inserted MLP parameters . Prefix-tuning prepends a sequence of trainable , task-specific prefix vectors to the input , and θ consists of these prefix parameters.2

Summary: The present paper first discusses the trade-off between performance for out-of-domain data and in-domain data with respect to whether the model is fully fine-tuned or lightweight fine-tuned on NLG tasks. Second, it argues that such a trade-off is not necessary if one can make use of both of these two fine-tuning schema in a clever way. To this end, it proposes cocktail fine-tuning, which augments full fine-tuning via distillation from a lightweight model and which achieves equal performance as an ensemble of the two fine-tuning schema. At length, this paper also explains the behavior of the cocktail fine-tuning through a toy model.