Leveraging Pre-trained Language Model Checkpoints for Encoder-Decoder Models
patrickvonplaten
<a target="_blank" href="https://colab.research.google.com/github/patrickvonplaten/notebooks/blob/master/Leveraging_Pre_trained_Checkpoints_for_Encoder_Decoder_Models.ipynb">
    <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
</a>


Transformer-based encoder-decoder models were proposed in [Vaswani et
al. (2017)](https://arxiv.org/pdf/1706.03762.pdf) and have recently
experienced a surge of interest, *e.g.* [Lewis et al.
(2019)](https://arxiv.org/abs/1910.13461), [Raffel et al.
(2019)](https://arxiv.org/abs/1910.10683), [Zhang et al.
(2020)](https://arxiv.org/abs/1912.08777), [Zaheer et al.
(2020)](https://arxiv.org/abs/2007.14062), [Yan et al.
(2020)](https://arxiv.org/pdf/2001.04063.pdf).

Similar to BERT and GPT2, massive pre-trained encoder-decoder models
have shown to significantly boost performance on a variety of
*sequence-to-sequence* tasks [Lewis et al.
(2019)](https://arxiv.org/abs/1910.13461), [Raffel et al.
(2019)](https://arxiv.org/abs/1910.10683). However, due to the enormous
computational cost attached to pre-training encoder-decoder models, the
development of such models is mainly limited to large companies and
institutes.

In [Leveraging Pre-trained Checkpoints for Sequence Generation Tasks
(2020)](https://arxiv.org/pdf/1907.12461.pdf), Sascha Rothe, Shashi
Narayan and Aliaksei Severyn initialize encoder-decoder model with
pre-trained *encoder and/or decoder-only* checkpoints (*e.g.* BERT,
GPT2) to skip the costly pre-training. The authors show that such
*warm-started* encoder-decoder models yield competitive results to large
pre-trained encoder-decoder models, such as
[*T5*](https://arxiv.org/abs/1910.10683), and
[*Pegasus*](https://arxiv.org/abs/1912.08777) on multiple
*sequence-to-sequence* tasks at a fraction of the training cost.

In this notebook, we will explain in detail how encoder-decoder models
can be warm-started, give practical tips based on [Rothe et al.
(2020)](https://arxiv.org/pdf/1907.12461.pdf), and finally go over a
complete code example showing how to warm-start encoder-decoder models
with 🤗Transformers.

This notebook is divided into 4 parts:

-   **Introduction** - *Short summary of pre-trained language models in
    NLP and the need for warm-starting encoder-decoder models.*
-   **Warm-starting encoder-decoder models (Theory)** - *Illustrative
    explanation on how encoder-decoder models are warm-started?*
-   **Warm-starting encoder-decoder models (Analysis)** - *Summary of
    [Leveraging Pre-trained Checkpoints for Sequence Generation
    Tasks (2020)](https://arxiv.org/pdf/1907.12461.pdf) - What model
    combinations are effective to warm-start encoder-decoder models; How
    does it differ from task to task?*
-   **Warm-starting encoder-decoder models with 🤗Transformers
    (Practice)** - *Complete code example showcasing in-detail how to
    use the* `EncoderDecoderModel` *framework to warm-start
    transformer-based encoder-decoder models.*

It is highly recommended (probably even necessary) to have read [this
blog
post](https://colab.research.google.com/github/patrickvonplaten/notebooks/blob/master/Encoder_Decoder_Model.ipynb)
about transformer-based encoder-decoder models.

Let\'s start by giving some back-ground on warm-starting encoder-decoder
models.

## **Introduction**

Recently, pre-trained language models \\({}^1\\) have revolutionized the
field of natural language processing (NLP).

The first pre-trained language models were based on recurrent neural
networks (RNN) as proposed [Dai et al.
(2015)](https://arxiv.org/pdf/1511.01432.pdf). *Dai et. al* showed that
pre-training an RNN-based model on unlabelled data and subsequently
fine-tuning \\({}^2\\) it on a specific task yields better results than
training a randomly initialized model directly on such a task. However,
it was only in 2018, when pre-trained language models become widely
accepted in NLP. [ELMO by Peters et
al.](https://arxiv.org/abs/1802.05365) and [ULMFit by Howard et
al.](https://arxiv.org/pdf/1801.06146.pdf) were the first pre-trained
language model to significantly improve the state-of-the-art on an array
of natural language understanding (NLU) tasks. Just a couple of months
later, OpenAI and Google published *transformer-based* pre-trained
language models, called [GPT by Radford et
al.](https://s3-us-west-2.amazonaws.com/openai-assets/research-covers/language-unsupervised/language_understanding_paper.pdf)
and [BERT by Devlin et al.](https://arxiv.org/abs/1810.04805)
respectively. The improved efficiency of *transformer-based* language
models over RNNs allowed GPT2 and BERT to be pre-trained on massive
amounts of unlabeled text data. Once pre-trained, BERT and GPT were
shown to require very little fine-tuning to shatter state-of-art results
on more than a dozen NLU tasks \\({}^3\\).

The capability of pre-trained language models to effectively transfer
*task-agnostic* knowledge to *task-specific* knowledge turned out to be
a great catalyst for NLU. Whereas engineers and researchers previously
had to train a language model from scratch, now publicly available
checkpoints of large pre-trained language models can be fine-tuned at a
fraction of the cost and time. This can save millions in industry and
allows for faster prototyping and better benchmarks in research.

Pre-trained language models have established a new level of performance
on NLU tasks and more and more research has been built upon leveraging
such pre-trained language models for improved NLU systems. However,
standalone BERT and GPT models have been less successful for
*sequence-to-sequence* tasks, *e.g.* *text-summarization*, *machine
translation*, *sentence-rephrasing*, etc.

Sequence-to-sequence tasks are defined as a mapping from an input
sequence \\(\mathbf{X}_{1:n}\\) to an output sequence \\(\mathbf{Y}_{1:m}\\) of
*a-priori* unknown output length \\(m\\). Hence, a sequence-to-sequence
model should define the conditional probability distribution of the
output sequence \\(\mathbf{Y}_{1:m}\\) conditioned on the input sequence
\\(\mathbf{X}_{1:n}\\):

$$ p_{\theta_{\text{model}}}(\mathbf{Y}_{1:m} | \mathbf{X}_{1:n}). $$

Without loss of generality, an input word sequence of \\(n\\) words is
hereby represented by the vector sequnece
\\(\mathbf{X}_{1:n} = \mathbf{x}_1, \ldots, \mathbf{x}_n\\) and an output
sequence of \\(m\\) words as
\\(\mathbf{Y}_{1:m} = \mathbf{y}_1, \ldots, \mathbf{y}_m\\).

Let\'s see how BERT and GPT2 would be fit to model sequence-to-sequence
tasks.

### **BERT**

BERT is an *encoder-only* model, which maps an input sequence
\\(\mathbf{X}_{1:n}\\) to a *contextualized* encoded sequence
\\(\mathbf{\overline{X}}_{1:n}\\):

$$ f_{\theta_{\text{BERT}}}: \mathbf{X}_{1:n} \to \mathbf{\overline{X}}_{1:n}. $$

BERT\'s contextualized encoded sequence \\(\mathbf{\overline{X}}_{1:n}\\)
can then further be processed by a classification layer for NLU
classification tasks, such as *sentiment analysis*, *natural language
inference*, etc. To do so, the classification layer, *i.e.* typically a
pooling layer followed by a feed-forward layer, is added as a final
layer on top of BERT to map the contextualized encoded sequence
\\(\mathbf{\overline{X}}_{1:n}\\) to a class \\(c\\):

$$
  f_{\theta{\text{p,c}}}: \mathbf{\overline{X}}_{1:n} \to c.
 $$

It has been shown that adding a pooling- and classification layer,
defined as \\(\theta_{\text{p,c}}\\), on top of a pre-trained BERT model
\\(\theta_{\text{BERT}}\\) and subsequently fine-tuning the complete model
\\(\{\theta_{\text{p,c}}, \theta_{\text{BERT}}\}\\) can yield
state-of-the-art performances on a variety of NLU tasks, *cf.* to [BERT
by Devlin et al.](https://arxiv.org/abs/1810.04805).

Let\'s visualize BERT.

![texte du
lien](https://raw.githubusercontent.com/patrickvonplaten/scientific_images/master/bert.png)

The BERT model is shown in grey. The model stacks multiple *BERT
blocks*, each of which is composed of *bi-directional* self-attention
layers (shown in the lower part of the red box) and two feed-forward
layers (short in the upper part of the red box).

Each BERT block makes use of **bi-directional** self-attention to
process an input sequence \\(\mathbf{x'}_1, \ldots, \mathbf{x'}_n\\) (shown
in light grey) to a more \"refined\" contextualized output sequence
\\(\mathbf{x''}_1, \ldots, \mathbf{x''}_n\\) (shown in slightly darker grey)
\\({}^4\\). The contextualized output sequence of the final BERT block,
*i.e.* \\(\mathbf{\overline{X}}_{1:n}\\), can then be mapped to a single
output class \\(c\\) by adding a *task-specific* classification layer (shown
in orange) as explained above.

*Encoder-only* models can only map an input sequence to an output
sequence of *a priori* known output length. In conclusion, the output
dimension does not depend on the input sequence, which makes it
disadvantageous and impractical to use encoder-only models for
sequence-to-sequence tasks.

As for all *encoder-only* models, BERT\'s architecture corresponds
exactly to the architecture of the encoder part of *transformer-based*
encoder-decoder models as shown in the \"Encoder\" section in the
[Encoder-Decoder
notebook](https://colab.research.google.com/drive/19wkOLQIjBBXQ-j3WWTEiud6nGBEw4MdF?usp=sharing).

### **GPT2**

GPT2 is a *decoder-only* model, which makes use of *uni-directional*
(*i.e.* \"causal\") self-attention to define a mapping from an input
sequence \\(\mathbf{Y}_{0: m - 1}\\) \\({}^1\\) to a \"next-word\" logit vector
sequence \\(\mathbf{L}_{1:m}\\):

$$ f_{\theta_{\text{GPT2}}}: \mathbf{Y}_{0: m - 1} \to \mathbf{L}_{1:m}. $$

By processing the logit vectors \\(\mathbf{L}_{1:m}\\) with the *softmax*
operation, the model can define the probability distribution of the word
sequence \\(\mathbf{Y}_{1:m}\\). To be exact, the probability distribution
of the word sequence \\(\mathbf{Y}_{1:m}\\) can be factorized into \\(m-1\\)
conditional \"next word\" distributions:

$$ p_{\theta_{\text{GPT2}}}(\mathbf{Y}_{1:m}) = \prod_{i=1}^{m} p_{\theta_{\text{GPT2}}}(\mathbf{y}_i | \mathbf{Y}_{0:i-1}). $$

\\(p_{\theta_{\text{GPT2}}}(\mathbf{y}_i | \mathbf{Y}_{0:i-1})\\) hereby
presents the probability distribution of the next word \\(\mathbf{y}_i\\)
given all previous words \\(\mathbf{y}_0, \ldots, \mathbf{y}_{i-1}\\) \\({}^3\\)
and is defined as the softmax operation applied on the logit vector
\\(\mathbf{l}_i\\). To summarize, the following equations hold true.

$$ p_{\theta_{\text{gpt2}}}(\mathbf{y}_i | \mathbf{Y}_{0:i-1}) = \textbf{Softmax}(\mathbf{l}_i) = \textbf{Softmax}(f_{\theta_{\text{GPT2}}}(\mathbf{Y}_{0: i - 1})).$$

For more detail, please refer to the
[decoder](https://huggingface.co/blog/encoder-decoder#decoder) section
of the encoder-decoder blog post.

Let\'s visualize GPT2 now as well.

![texte du
lien](https://raw.githubusercontent.com/patrickvonplaten/scientific_images/master/gpt2.png)

Analogous to BERT, GPT2 is composed of a stack of *GPT2 blocks*. In
contrast to BERT block, GPT2 block makes use of **uni-directional**
self-attention to process some input vectors
\\(\mathbf{y'}_0, \ldots, \mathbf{y'}_{m-1}\\) (shown in light blue on the
bottom right) to an output vector sequence
\\(\mathbf{y''}_0, \ldots, \mathbf{y''}_{m-1}\\) (shown in darker blue on
the top right). In addition to the GPT2 block stack, the model also has
a linear layer, called *LM Head*, which maps the output vectors of the
final GPT2 block to the logit vectors
\\(\mathbf{l}_1, \ldots, \mathbf{l}_m\\). As mentioned earlier, a logit
vector \\(\mathbf{l}_i\\) can then be used to sample of new input vector
\\(\mathbf{y}_i\\) \\({}^5\\).

GPT2 is mainly used for *open-domain* text generation. First, an input
prompt \\(\mathbf{Y}_{0:i-1}\\) is fed to the model to yield the conditional
distribution
\\(p_{\theta_{\text{gpt2}}}(\mathbf{y} | \mathbf{Y}_{0:i-1})\\). Then the
next word \\(\mathbf{y}_i\\) is sampled from the distribution (represented
by the grey arrows in the graph above) and consequently append to the
input. In an auto-regressive fashion the word \\(\mathbf{y}_{i+1}\\) can
then be sampled from
\\(p_{\theta_{\text{gpt2}}}(\mathbf{y} | \mathbf{Y}_{0:i})\\) and so on.

GPT2 is therefore well-suited for *language generation*, but less so for
*conditional* generation. By setting the input prompt
\\(\mathbf{Y}_{0: i-1}\\) equal to the sequence input \\(\mathbf{X}_{1:n}\\),
GPT2 can very well be used for conditional generation. However, the
model architecture has a fundamental drawback compared to the
encoder-decoder architecture as explained in [Raffel et al.
(2019)](https://arxiv.org/abs/1910.10683) on page 17. In short,
uni-directional self-attention forces the model\'s representation of the
sequence input \\(\mathbf{X}_{1:n}\\) to be unnecessarily limited since
\\(\mathbf{x}_i\\) cannot depend on
\\(\mathbf{x}_{i+1}, \forall i \in \{1,\ldots, n\}\\).

### **Encoder-Decoder**

Because *encoder-only* models require to know the output length *a
priori*, they seem unfit for sequence-to-sequence tasks. *Decoder-only*
models can function well for sequence-to-sequence tasks, but also have
certain architectural limitations as explained above.

The current predominant approach to tackle *sequence-to-sequence* tasks
are *transformer-based* **encoder-decoder** models - often also called
*seq2seq transformer* models. Encoder-decoder models were introduced in
[Vaswani et al. (2017)](https://arxiv.org/abs/1706.03762) and since then
have been shown to perform better on *sequence-to-sequence* tasks than
stand-alone language models (*i.e.* decoder-only models), *e.g.* [Raffel
et al. (2020)](https://arxiv.org/pdf/1910.10683.pdf). In essence, an
encoder-decoder model is the combination of a *stand-alone* encoder,
such as BERT, and a *stand-alone* decoder model, such as GPT2. For more
details on the exact architecture of transformer-based encoder-decoder
models, please refer to [this blog
post](https://huggingface.co/blog/encoder-decoder).

Now, we know that freely available checkpoints of large pre-trained
*stand-alone* encoder and decoder models, such as *BERT* and *GPT*, can
boost performance and reduce training cost for many NLU tasks, We also
know that encoder-decoder models are essentially the combination of
*stand-alone* encoder and decoder models. This naturally brings up the
question of how one can leverage stand-alone model checkpoints for
encoder-decoder models and which model combinations are most performant
on certain *sequence-to-sequence* tasks.

In 2020, Sascha Rothe, Shashi Narayan, and Aliaksei Severyn investigated
exactly this question in their paper [**Leveraging Pre-trained
Checkpoints for Sequence Generation
Tasks**](https://arxiv.org/abs/1907.12461). The paper offers a great
analysis of different encoder-decoder model combinations and fine-tuning
techniques, which we will study in more detail later.

Composing an encoder-decoder model of pre-trained stand-alone model
checkpoints is defined as *warm-starting* the encoder-decoder model. The
following sections show how warm-starting an encoder-decoder model works
in theory, how one can put the theory into practice with 🤗Transformers,
and also gives practical tips for better performance.