Ensembles and Cocktails: Robust Finetuning for Natural Language Generation
Abstract
When ﬁnetuning a pretrained language model for natural language generation tasks, one is currently faced with a tradeoff.
Lightweight ﬁnetuning (e.g., preﬁxtuning, adapters), which freezes all or most of the parameters of the pretrained model, has been shown to achieve stronger out-of-distribution (OOD) performance than full ﬁnetuning , which tunes all of the parameters.
However, lightweight ﬁnetuning can underperform full ﬁnetuning in-distribution (ID).
In this work, we present methods to combine the beneﬁts of full and lightweight ﬁnetuning, achieving strong performance both ID and OOD.
First, we show that an ensemble of the lightweight and full ﬁnetuning models achieves the best of both worlds: performance matching the better of full and lightweight ﬁnetuning, both ID and
OOD.
Second, we show that we can achieve similar improvements using a single model instead of two with our proposed cocktail ﬁnetuning , which augments full ﬁnetuning via distillation from a lightweight model.
Finally, we provide some explanatory theory in a multiclass logistic regression setting with a large number of classes, describing how distillation on ID data can transfer the OOD behavior of one model to another.
Introduction
When ﬁnetuning a pretrained language model for natural language generation tasks like summarization (Narayan 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 ﬁnetuning approaches.
Finetuning all parameters of the pretrained language model achieves strong ID performance (Howard & Ruder, 2018).
We call this full ﬁnetuning .
Alternatively, freezing most of the pretrained parameters during ﬁnetuning, in e.g., adapters (Rebufﬁ et al., 2017; Houlsby et al., 2019a), preﬁx-tuning (Li & Liang, 2021; Lester et al., 2021) and bitﬁt (Zaken et al., 2021), achieves stronger OOD performance than full ﬁnetuning (Li & Liang, 2021;
Lester et al., 2021), but worse ID performance.
We call these methods lightweight ﬁnetuning .
In this work, we are the ﬁrst to demonstrate that this tradeoff is not necessary in natural language generation: we can achieve the best of full and lightweight ﬁnetuning both ID and OOD.
We test this across three natural language generation tasks in English: summarization (Narayan 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 ﬁnetuning methods: adapters and preﬁx-tuning.
Achieving the best of both worlds is important since on average across our datasets and lightweight ﬁnetuning methods, full ﬁnetuning achieves only 84% of the OOD performance of lightweight ﬁnetuning, and lightweight ﬁnetuing achieves only 81% of the ID performance of full ﬁnetuning.
on
O u t - o f - d i s t
r i b u t i o n p e r f o r m a n c e `
In-distribution performance
Lightweight Finetuning
Full Finetuning
[Figure]
Ensembles and Cocktail Finetuning:
Best of both worlds λ=1Weight λ=
We present two methods for achieving the best of full and lightweight ﬁnetuning.
We ﬁrst show that a weighted ensemble of one full ﬁnetuning model and one lightweight ﬁnetuning model achieves performance comparable to the best of the two, both ID and OOD.
To avoid running two models at once, we propose cocktail ﬁnetuning : ﬁrst train a lightweight ﬁnetuning model, and then train a full ﬁnetuning model on a mixture of student-teacher distillation loss (Bucil˘a et al., 2006; Hinton et al., 2015) (with the lightweight model as teacher) and cross-entropy loss, on the original ID data.
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 ﬁnetuning.
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 speciﬁc dataset.
Most saliently, in concurrent work, Wortsman et al.
ﬁnd that zeroshot image recognition through CLIP (Radford et al., 2021) can be combined with a ﬁnetuned 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 ﬁnetuning can transfer the out-of-distribution behavior of the lightweight ﬁnetuning 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 classiﬁcation problem over a large output space in Section D in the appendix.
In summary, we present two simple methods—ensembling and cocktail ﬁnetuning—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 ﬁnetuning methods.
Problem Statement 2.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 P id and P ood denote the distribution of ( x, y ) pairs in-distribution and out-of-distribution, respectively, where x ∈ X and y ∈ Y .
In summarization, P id may consist of articles about world and business news while P ood consists of health and technology news.
The training data D train and validation data D val consist of n and n val in-distribution ( x, y ) pairs respectively, drawn from P id .
Model.
Let q γ ( y | x ) = (cid:81) len ( y )
j =1 q γ ( y j | x, y <j ) be a pretrained neural autoregressive sequence model with parameters γ .
Finetuning (full or lightweight) produces a ﬁnetuned model p θ ( y | x ) from q γ , where θ denotes the set of trainable parameters.
In full ﬁnetuning, θ is the same set of parameters as γ ; in lightweight ﬁnetuning, θ often refers to the parameters of inserted modules (e.g., adapters modules or preﬁxes).
The ﬁnal predictor is f θ ( x ) which approximates arg max y 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.
2.2 Datasets and metrics.
We conduct experiments across three natural language generation tasks in English: WebNLG for tableto-text (Gardent et al., 2017), XSUM for summarization (Narayan et al., 2018) and open-domain QA
[Table]
(Kwiatkowski et al., 2019; Berant et al., 2013), and report BLEU (Papineni et al., 2002b), ROUGE- (Lin, 2004), and exact match, respectively (dataset details are in the appendix Section A ).
We deﬁne the scores S id ( f θ ) and S ood ( f θ ) as the average of the metric on P id and P ood , respectively.
2.3 Baseline ﬁnetuning methods
Full ﬁnetuning 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.
[Equation]
where the set of trainable parameters θ refer to the set of pretrained parameters γ .
Lightweight ﬁnetuning is another family of adaptation approaches, which freezes most of the pretrained parameters and augments the model with small trainable modules.
Lightweight ﬁnetuning shares the same training objective as full ﬁnetuning, but the trainable parameters θ are different.
Within the lightweight ﬁnetuning family, we focus on adapter-tuning (Houlsby et al., 2019a) and preﬁx-tuning (Li & Liang, 2021).
Both approaches keep the pretrained parameters intact.
Adaptertuning inserts task-speciﬁc MLP layers between each layer of the pretrained language models, and the trainable parameter θ consists of these inserted MLP parameters.
Preﬁx-tuning prepends a sequence of trainable, task-speciﬁc preﬁx vectors to the input, and θ consists of these preﬁx parameters.
Combining models with a Weighted Ensemble
As we observe in Table 1 (detailed discussed in B), full ﬁnetuning models attain better ID performance but worse OOD performance, whereas lightweight ﬁnetuning models attain better OOD performance at the cost of good ID performance.
Is it possible to achieve both good ID and OOD performance simultaneously?
In this section, we show that a simple ensemble of the two methods accomplishes this.
First, we train a lightweight ﬁnetuning model p ( (cid:96) ) and a full ﬁnetuning model p ( f) on D train .
At decoding time, we combine the prediction of the p ( (cid:96) ) and p ( f) by interpolating them at each token.
[Equation]
At each time step t , both models output a next token distribution conditioned on the common history y <t , and the two distributions are interpolated with a mixture weight λ .
The output is generated by beam searching the ensemble distribution p (Ens) ( y t | x, y <t ) .
Choosing the mixture weight λ .
We start with a set of { λ i } Ki =1 where λ i ∈ [0 , 1] .
For example we use λ ∈ { 0 ., 0 . 1 , 0 .
9 , 1 . 0 } in this experiment.
We choose the mixture weight
λ i that attains the best score on D val .
This method of choosing λ is useful because (1) it crucially does not use OOD data, which is typically unaccessible at model selection time, and (2) it guarantees that
ID performance is not sacriﬁced, because we can match ID performance by setting λ = 0 .
Results.
As shown in Table 1, the ensemble matches and slightly outperforms both the ID performance skyline (achieved by full ﬁnetuning) and the OOD performance skyline (achieved by lightweight ﬁnetuning).
Comparing with the weaker candidates ID and OOD, and averaging acrosssettings (3 datasets, two lightweight ﬁnetuning methods,) the ensemble achieves a 37% improvement over the weaker ID candidate (i.e, lightweight ﬁnetuning) and a 24% relative improvement over the weaker OOD candidate (i.e., full ﬁnetuning).
Therefore, the ensemble—unlike lightweight and full ﬁnetuning—achieves the best of both worlds: performance matching the better of full and lightweight ﬁnetuning, for both ID and OOD.
Runtime and memory.
Despite its ability to combine the best of both worlds, ensembling is expensive at inference time because two models are run simultaneously, doubling the amount of GPU memory required.
Furthermore, two forward passes are required, doubling the runtime for decoding.
For deployment efﬁciency, it’s better to have a single model at inference time.
Cocktail Finetuning: Combining with Distillation
The goal of cocktail ﬁnetuning is to approximate the ensemble in Section 3 with a single model.
Method.
Train a lightweight ﬁnetuning model p ( (cid:96) ) .
Fix mixture weight λ ∈ [0 , 1] .
Train all parameters of the pretrained model q γ ( y | x ) on the following objective.
First, deﬁne a token-level student-teacher distillation loss at token j :
[Equation]
The objective is the sum over the length of y of the λ -weighted mixture of distillation and loglikelihood loss:
[Equation]
Let f θ,λ be the predictor resulting from optimizing on L λ ( θ ) and estimate the ID score S id ( f θ,λ ) on D val .
Perform this across a set of mixture weights { λ i } i .
The cocktail ﬁnetuning model is the predictor f θ,λ i that maximizes the in-distribution score S id ( f θ , λ i ) across the mixture weights λ i .
Results.
As shown in Table 1, cocktail ﬁnetuning matches and slightly outperforms the ID performance skyline (achieved by full ﬁnetuning) and achieves comparable to or slightly worse than the OOD performance skyline (achieved by lightweight ﬁnetuning).
Comparing with the weaker candidates ID and OOD, on average cocktail ﬁnetuning achieves a 39% improvement over the weaker
ID candidate (i.e, lightweight ﬁnetuning) and a 19% relative improvement over the weaker OOD candidate (i.e., full ﬁnetuning).
We speculate that the slight decay in OOD performance relative to the ensembles may be due to imperfect transfer of OOD behavior when ﬁnetuning only on ID data; that it works at all is an interesting phenomenon we explore in Section D of the Appendix.
Analysis of the Weighting of Full and Lightweight Finetuning
Trade-off curves.
Our ensembling and cocktail ﬁnetuning methods are useful in practice because one can choose the mixture coefﬁcient λ without OOD data.
Now, we explore how ID and OOD performance vary as a function of λ .
Figure 2 shows results across ﬁve random seeds of ensembling and cocktail ﬁnetuning on our three datasets, each dot colored according to its λ value.
In each plot, we see a smooth tradeoff curve, in which a range of λ improve on the OOD, ID, or both of each constituent model.
Interestingly, in two of three tasks, there is effectively only one pareto-optimal λ .
This means the λ is best for both for ID and OOD performance, which is why optimizing for ID performance also leads to good OOD performance.
[Figure]
and XSUM, ensembles and cocktails both achieve the best of both preﬁxtuning and full ﬁnetuning for some λ ; for OpenQA, one can achieve approximately the best of both.
[Figure]
Ablations.
One hypothesis to explain the strong performance of ensembles is the known beneﬁts of ensembling two models trained with different random seeds.
Likewise for cocktail ﬁnetuning, one could point to the effectiveness of born-again networks (Furlanello et al., 2018): distilling a teacher model to an identically parameterized student with the same model architecture.
As an ablation to test this hypothesis, we distill from a full ﬁnetuning model into another full ﬁnetuning model, again using λ to mix the distillation loss with the label loss.
The results, in Figure 3, show that this distillation does not improve over a single full ﬁnetuning model, for any of the λ .
Conclusion
Considered together, our work and the work of Wortsman et al.
(2021) provide strong cross-modal evidence that simple ensembles of models that have minimally deviated from pretraining (for them, zero-shot; for us, lightweight ﬁnetuning models) with full ﬁnetuning models provides strong ID and OOD accuracy.
The reasons for this “best of both worlds” behavior are still poorly understood, motivating future work.
Since ensemble models effectively weight by model conﬁdence (though modulated by λ ), one explanation may be that lightweight ﬁnetuning models are less conﬁdent than full ﬁnetuning models in-distribution, yet more conﬁdent out-of-distribution.
In summary, we provide simple, effective methods for achieving the best of lightweight and full ﬁnetuning—strong OOD and ID performance—in language generation.
Our cocktail ﬁnetuning provides system builders with the opportunity to optimize for in-distribution performance (in choosing lambda), often beating full ﬁnetuning, while resting assured that the out-of-distribution performance (i.e., what they observe with real users) will likely be higher than had they performed full ﬁnetuning.
Ethics statement
The methods presented in this work have the potential to improve a wide range of natural language generation tasks.
Natural language generation has known dual-use issues, like scaling misinformation generation and generating naturalistic text in scams.
The WebNLG XSUM, Natural Questions, and
WebQuestions datasets we evaluate on contain personal information (though public and published) information about real people and organizations.
Reproducibility Statement
All code (including experiment scripts and plot creation scripts) will be released upon publication.
All datasets used in this paper are publicly available and will be linked in our released code.
All reported numbers represent the average of results from 5 random seeds.
A Experiment Setup A.1 Datasets and Metrics
Table-to-text.
We evaluate on the WebNLG dataset (Gardent et al., 2017), where the input x is a sequence of (entity, predicate, entity) triples and the output y is a natural language description that covers the input information.
Each example in this dataset is labeled with one of 14 topic (e.g., airports, comic characters).
We select 9 topics as in-distribution and the remaining 5 topics as out-of-distribution.
We evaluate the performance using the ofﬁcial evaluation scripts and report the
BLEU scores (Papineni et al., 2002a).
Summarization.
We use the XSUM (Narayan et al., 2018) dataset, which is an abstractive summarization dataset on BBC news articles.
Each example is labeled with a news category.
We select world, UK and business categories to be in-distribution and let the remaining news categories such as health and technology to be out-of-distribution.
We report the ROUGE-2 scores (Lin, 2004).
Open-domain, closed-book QA (Open-QA).
We use the open-domain variant (Lee et al., 2019b) of the Natural Questions dataset (Kwiatkowski et al., 2019) as the in-distribution data and the Web
Questions dataset (Berant et al., 2013) as the out-of-distribution data.
The distribution shift is induced by different data collection process and their year of collection (roughly 2018 versus 2013).
We report the exact match accuracy.
Note that this task, unlike most QA (e.g., SQuAD (Rajpurkar et al., 2016)), does not provide systems with a text from which to extract answers.
Roberts et al.
(2020) demonstrated that one could ﬁnetune an LM to generate answers from knowledge stored during pretraining.
A.2 Model Architectures and Hyperparameters
We use GPT-2-medium (Radford et al., 2019) as the pretrained model for table-to-text and opendomain QA tasks.
GPT-2 is an autoregressive language model; therefore, we parametrize p θ ( y | x ) by concatenating x , a separator token, and y (e.g., [x; SEP; y]).
We deﬁne the loss function to sum over the token losses that correspond to y .
We use BART-large (Lewis et al., 2020) as the pretrained model for the summarization task.
BART uses an encoder-decoder architecture and parametrizes p θ ( y | x ) by ﬁrst feeding x to the encoder, and decoding y conditioned on the encoder representation.
For open-domain closed-book QA, this is to our knowledge the ﬁrst time GPT-2 has been ﬁnetuned for the task.
We used separate development OOD data to test a handful ( < 10) of hyperparameter settings to determine whether we’d observe a similar ID/OOD tradeoff to the other tasks.
We note this because, in general, one cannot do hyperparameter optimization on OOD data; we did this to provide another testbed for our methods.
Our implementation is based on the Hugging Face Transformers package (Wolf et al., 2020).
We use most of the hyperparameters suggested by the Hugging Face default setup.
The hyperparameters we tune include the number of epochs, batch size, learning rate, and preﬁx length (for preﬁx-tuning), all reported in Table 3.
B Results of full and lightweight ﬁnetuning
There is a trade-off in ID/OOD performance when choosing between full ﬁnetuning and lightweight ﬁnetuning: full ﬁnetuning consistently outperform in-distribution, whereas lightweight ﬁnetuning consistently outperform out-of-distribution (Li & Liang, 2021; Lester et al., 2021).
To quantify this trade-off, we examine the generation performance of both methods on ID and OOD data.
As shown in Table 1, when evaluating in-distribution, full ﬁnetuning achieves the stronger results (denoting this ID skyline as 100%), whereas lightweight ﬁnetuning achieves only 81% of the skyline, on average.
When evaluating out-of-distribution, lightweight ﬁnetuning achieves stronger result (denoting this OOD skyline as 100%), whereas full ﬁnetuning achieves only 84% of the skyline.
Lightweight ﬁnetuning tends to produce models that are more robust at the cost of good ID performance; Full ﬁnetuning produces models that attain good ID performance at the cost of OOD performance.
C Related Work
Lightweight ﬁnetuning.
Recently, as pretrained models scale up in the number of parameters, lightweight ﬁnetuning 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”)
(Rebufﬁ 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 beneﬁt of being able to achieve strong out-of-distribution extrapolation performance at the cost of worse in-distribution performance compared to full ﬁnetuning models both in classiﬁcation 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.
Buciluˇa et al.
(2006) introduce distillation for compressing large ensembles into smaller, faster, more efﬁcient 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 representations that are more similar to independently trained LSTMs than independently trained
Transformers.
The efﬁciacy of our proposed cocktail ﬁnetuning 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 ﬁnetuning in vision.
In concurrent work, Wortsman et al.
(2021) propose a similar approach to ours, but focus on image classiﬁcation: they combine a full ﬁnetuning model with a lightweight ﬁnetuning (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 ﬁnetuning 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 beneﬁts of these two models can be combined using robust self-training (RST)
(Carmon et al., 2019;
Uesato et al., 2019; Najaﬁ 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 beneﬁts 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 beneﬁts of two models or use a self-training/distillation loss to improve robustness, cocktail ﬁnetuning differs in that it performs distillation only on ID training data without any additional ID or OOD unlabeled examples.
D Why does distillation on ID data improve OOD?
Prior works have used unlabeled OOD data and self-training (similar to distillation) to transfer good
OOD performance from another model (Raghunathan et al., 2020; Xie et al., 2021).
In contrast, cocktail ﬁnetuning performs distillation only on ID data.
In this section, we provide an explanation for why this can occur in multiclass logistic regression with a large number of classes.
Prediction problem and distribution shift.
Motivated by the large output space in language generation, we consider a prediction problem from inputs X ⊆ R d to labels Y = { 1 , . . .
, k } where k is large.
The labeled training data is { ( x i , y i ) } ni =1 .
For convenience, we also let y i be the one-hot vector of length k denoting the correct class.
The labels in the training data do not cover all the possible labels, and we let the subset of seen labels be S .
In this problem, the class marginals of the ID and OOD distributions P id and P ood are different.
To have good OOD accuracy, the model must generalize to unseen classes.
Model and losses.
We consider training a multi-class logistic regression model using full batch gradient descent where the learned parameters are θ ∈ R k × d .
In full ﬁne-tuning, we run gradient descent on the negative log-likelihood.
In cocktail ﬁnetuning, we also add a distillation loss using soft labels (a probability vector) from the preﬁx-tuned model, which we denote p x ˜ θ = p ˜ θ ( · | x ) ∈ R k with parameters ˜ θ , given an input x .
Similarly, the predicted class probability vector of the cocktail model given an input x is denoted p xθ = p θ ( · | x ) ∈ R k .
In this setting, the cocktail loss (Equation 4) with parameter λ ∈ [0 , 1] is
[Equation]
[Table]
- 5 ∗
[Table]
− - 5 ∗ 10 −
[Table]
5 ∗ 10 − d=  8 ∗ 10 − d=
[Table]
5 ∗ 10 − d=
[Table] [Table]
∗
[Table]
l=  5 ∗
[Table] [Table]
=
[Table] [Table]
The cocktail loss encourages the model to learn to output a mixture of the hard label and the soft label on the training data.
The ﬁnal score for the classiﬁcation task is computed with respect to the accuracy metric, where score ( y, ˆ y )
= 1 − (cid:96)
− 1 ( y, ˆ y ) and (cid:96) 0 − 1 is the 0-1 loss.
Poor OOD performance with full ﬁnetuning.
Our ﬁrst result shows that when training a logistic regression model using gradient descent (initialized at 0) on the only labeled training set (as in full ﬁnetuning), the learned parameters θ j for an unseen label j / ∈ S are all the same.
Proposition 1.
Let θ ( t ) be the parameters at iteration t of gradient descent and let θ (0)
= 0 .
At any iteration t and for any unseen label index j / ∈ S , θ j = g t where g t is the same for all j / ∈ S .
As a result, the model cannot distinguish between any of the unseen classes, resulting in poor
OOD performance.
Suppose that the model outputs a ﬁxed, arbitrary unseen class j ∈ S when the probabilities of the unseen classes are the largest.
Then, assuming that the OOD distribution P ood is such that the probability of drawing an example from an unseen class (excluding j ) P ood ( y / ∈ S ∪{ j } ) is at least c ∈ [0 , 1] , then the OOD accuracy is bounded as
[Equation]
This OOD accuracy can be low when c is large, which can occur with very large output spaces.
Disillation transfers OOD behavior through soft logits.
We investigate how distillation on ID training data can affect OOD behavior.
In the following, let X ∈ R n × d be the training data matrix.
Proposition 2.
Assume the data matrix X has full column rank.
Let θ be the minimizer of the cocktail loss when λ = 1 , subject to the constraint that 1 (cid:62) exp( θX ) =
α 1 (cid:62) exp(˜ θX ) for some constant α > 0 , such that the model learns the same normalizing constants (up to scaling) as the preﬁx-tuned model on the training data only.
Then
[Equation]
for any new input x , such that the cocktail model has the same OOD accuracy as the preﬁx model.
Therefore, if the preﬁx-tuning model has high OOD accuracy, then S ood is high.
The proof proceeds by showing that the extra information in the soft logits can help the model recover the parameters of the preﬁx-tuned model (up to scaling), which in turn inherits its OOD behavior.
Thus, in generation problems where the output space is large, distillation on soft logits can transfer information about
OOD behavior even through ID data.
E Hyperparameters and Dataset Details E.1 Hyperparameters
Relevant hyperparameters beyond the defaults of Huggingface Transformers Wolf et al.
(2020) are reported in Table 3.
[Table] [Table]
o r m a n
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
r
[Table]
o r m a n c
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
E.2 Dataset Details
Table 5 shows the number of examples in each dataset split.
For the WebNLG dataset, the ID topics include { Astronaut, University, Monument, Building, ComicsCharacter, Food, Airport, SportsTeam, City, and WrittenWork } and the OOD topics include { Athlete,
Artist, MeanOfTransportation, CelestialBody, Politician } .
For the XSUM dataset, the ID topics include { uk, business, world } news and the OOD topics includes { entertainment, election, technology, science, health, education, explainers, live, ukpolitics, world-us-canada, world-europe, uk-england, uk-scotland, world-asia, science-environment, uk-scotland-scotland-politics, disability } .
F Additional Results
Due to space constraints, we did not include in Figure 2 the tradeoff plots using adapters lightweight ﬁnetuning method (instead reporting only those for preﬁx-tuning.)
The plots for adapters are provided in Figure 4.
G Analysis Proofs
We restate and expand upon some of the setup for the analysis here.
Prediction problem.
We consider a prediction problem from inputs X ⊆ R d to labels Y = { 1 , . . .
, k } where k is large.
The labeled training data is { ( x i , y i ) } ni =1 .
For convenience, we also let y i be the one-hot vector of length k denoting the correct class, and will make this clear.
The labels in the training data do not cover all the possible labels, and we let the subset of seen labels be S .
Model.
We consider training a multi-class logistic regression model using full batch gradient descent.
The learned parameters are θ ∈ R k × d and the model prediction given input x is
[Equation]
where θ j is the j -th row of θ .
In standard ﬁne-tuning, we run gradient descent on the negative log-likelihood (cid:96)
[Equation]
where y i ∈
R k is a one-hot vector of the class label.
Distillation loss.
In the Cocktail model, we also add a distillation loss using soft labels from the preﬁx-tuned model, which we denote p ˜ θ with parameters ˜ θ .
The Cocktail loss with parameter λ ∈ [0 , 1] is
[Equation]
The Cocktail loss encourages the model to learn to output a mixture of the hard label and the soft label on the training data.
G.1 Proof of Proposition
Proof.
By induction on t .
By zero initialization, the statement holds with g 0 = 0 .
Suppose for iterations up to t , we have θ ( t ) j = g t for all j / ∈ S .
We focus on some j / ∈ S WLOG.
The gradient for θ ( t ) j is
[Equation]
since y ij = 0 for all i as j is unseen.
This gradient does not vary with j , which shows the result.
G.2 Proof of Proposition
Proof.
Let Y θ ∈ R n × k denote the predicted class probabilities of the model on the training data, and
Y ˜ θ be the class probabilities of the preﬁx-tuned model on training data.
With λ = 1 , the gradient of the loss is
[Equation] [Equation]
[Equation]
Since θ is the minimizer, this gradient is 0 and thus
[Equation] [Equation]
for X with full column rank (right-invertible).
We let the right-inverse of X be X † .
This implies that for all i ,
[Equation]
In matrix form
[Equation]
this is equivalent to
[Equation]
On a new input x
[Equation]
by using the assumption on normalizing constants.
Therefore the probability of the j -th class under the model is
[Equation]