🦄 How to build a State-of-the-Art Conversational AI with Transfer Learning

In pytorch-pretrained-BERT OpenAI GPT’s model and its tokenizer can be easily created and loaded from the pretrained checkpoint like this:

You probably noticed we’ve loaded a model called OpenAI GPT Double Heads Model which sounds a bit more complex than the language model we’ve just talked about and you’re right!

This is because we need to adapt our model to dialog. Let’s see how this goes!

👻 Adapting a language model to a dialog task

Our language model is trained with a single input: a sequence of words.

But as we saw earlier, in a dialog setting, our model will have to use several types of contexts to generate an output sequence:

• one or several persona sentences,
• the history of the dialog with at least the last utterance from the user,
• the tokens of the output sequence that have already been generated since we generate the output sequence word by word.

How can we build an input for our model from these various contexts?

A simple answer is just to concatenate the context segments in a single sequence, putting the reply at the end. We can then generate a completion of the reply token by token by continuing the sequence:

There are two issues with this simple setup:

• Our transformer is color-blind! The delimiter tokens only give it a weak idea of which segment each word belongs to. For example, the word “NYC” is indicated in blue (persona) in our illustration but our model will have a hard time extracting this information from the delimiters alone: we should add more information about the segments.
• Our transformer is position-blind! Attention is a symmetrical dot-product so we should add position information for each token.

An easy way to add this information is to build three parallel input sequences for word, position, and segments, and fuse them in a single sequence, summing three types of embeddings: word, position, and segments embeddings:

How do we implement this?

First, we’ll add special tokens to our vocabulary for delimiters and segment indicators. These tokens were not part of our model’s pretraining so we will need to create and train new embeddings for them.

Adding special tokens and new embeddings to the vocabulary/model is quite simple with pytorch-pretrained-BERT classes. Let’s add five special tokens to our tokenizer’s vocabulary and model’s embeddings:

These special-tokens methods respectively add our five special tokens to the vocabulary of the tokenizer and create five additional embeddings in the model.

Now we have all we need to build our input sequence from the persona, history, and beginning of reply contexts. Here is a simple example:

👑 Multi-tasks losses

We have now initialized our pretrained model and built our training inputs, all that remains is to choose a loss to optimize during the fine-tuning.

We will use a multi-task loss combining language modeling with a next-sentence prediction objective.

The next-sentence prediction objective is a part of BERT pretraining. It consists in randomly sampling distractors from the dataset and training the model to distinguish whether an input sequence ends with a gold reply or a distractor. It trains the model to look at the global segments meaning besides the local context.

Now you see why we loaded a “Double-Head” model. One head will compute language modeling predictions while the other head will predict next-sentence classification labels. Let’s have a look at how losses are computed:

The total loss will be the weighted sum of the language modeling loss and the next-sentence prediction loss which are computed as follow:

• Language modeling: we project the hidden-state on the word embedding matrix to get logits and apply a cross-entropy loss on the portion of the target corresponding to the gold reply (green labels on the above figure).
• Next-sentence prediction: we pass the hidden-state of the last token (the end-of-sequence token) through a linear layer to get a score and apply a cross-entropy loss to classify correctly a gold answer among distractors.

Let’s see how we can code this:

We now have all the inputs required by our model and we can run a forward pass of the model to get the two losses and the total loss (as a weighted sum):

We are ready to start the training 🎉

🦊 Training on a dialog dataset

The ConvAI2 competition used an interesting dataset released by Facebook last year: PERSONA-CHAT.

It’s a rather large dataset of dialog (10k dialogs) which was created by crowdsourcing personality sentences and asking paired crowd workers to chit-chat while playing the part of a given character (an example is given on the left figure).

This dataset is available in raw tokenized text format in the nice Facebook’s ParlAI library. To bootstrap you, we also uploaded a JSON formatted version that you can download and tokenize using GPT’s tokenizer like this:

The JSON version of PERSONA-CHAT gives quick access to all the relevant inputs for training our model as a nested dictionary of lists:

Using the awesome PyTorch ignite framework and the new API for Automatic Mixed Precision (FP16/32) provided by NVIDIA’s apex, we were able to distill our +3k lines of competition code in less than 250 lines of training code with distributed and FP16 options!

We’ve covered the essential parts of the code in the above gists so I’ll just let you read the commented code to see how it all fits together.

The training (train.py) code is here ➱ 🎮

Training this model on an AWS instance with 8 V100 GPU takes less than an hour (currently less than $25 on the biggest p3.16xlarge AWS instance) and gives results close to the SOTA obtained during the ConvAI2 competition with Hits@1 over 79, perplexity of 20.5 and F1 of 16.5.

A few differences explain the slightly lower scores vs our competition model, they are detailed in the readme of the code repo here and mostly consists in tweaking the position embeddings and using a different decoder.

👻 Talking with the Model — the Decoder

The amazing thing about dialog models is that you can talk with them 🤗

To interact with our model, we need to add one thing: a decoder that will build full sequences from the next token predictions of our model.

Now there have been very interesting developments in decoders over the last few months and I wanted to present them quickly here to get you up-to-date.

The two most common decoders for language generation used to be greedy-decoding and beam-search.

Greedy-decoding is the simplest way to generate a sentence: at each time step, we select the most likely next token according to the model until we reach end-of-sequence tokens. One risk with greedy decoding is that a highly probable token may be hiding after a low-probability token and be missed.

Beam-search try to mitigate this issue by maintaining a beam of several possible sequences that we construct word-by-word. At the end of the process, we select the best sentence among the beams. Over the last few years, beam-search has been the standard decoding algorithm for almost all language generation tasks including dialog (see the recent [1]).

However several developments happened in 2018/early-2019. First, there was growing evidence that beam-search was strongly sensitive to the length of the outputs and best results could be obtained when the output length was predicted before decoding ([2, 3] at EMNLP 2018). While this makes sense for low-entropy tasks like translation where the output sequence length can be roughly predicted from the input, it seems arbitrary for high-entropy tasks like dialog and story generation where outputs of widely different lengths are usually equally valid.

In parallel, at least two influential papers ([4, 5]) on high-entropy generation tasks were published in which greedy/beam-search decoding was replaced by sampling from the next token distribution at each time step. These papers used a variant of sampling called top-k sampling in which the decoder sample only from the top-k most-probable tokens (k is a hyper-parameter).

The last stone in this recent trend of work is the study recently published by Ari Holtzman et al. [6] which showed that the distributions of words in texts generated using beam-search and greedy decoding is very different from the distributions of words in human-generated texts. Clearly, beam-search and greedy decoding fail to reproduce some distributional aspects of human texts as it has also been noted in [7, 8] in the context of dialog systems:

Currently, the two most promising candidates to succeed beam-search/greedy decoding are top-k and nucleus (or top-p) sampling. The general principle of these two methods is to sample from the next-token distribution after having filtered this distribution to keep only the top k tokens (top-k) or the top tokens with a cumulative probability just above a threshold (nucleus/top-p).

Here is how we can decode using top-k and/or nucleus/top-p sampling:

We are now ready to talk with our model 🚀

The interactive script is here (interact.py) and if you don’t want to run the script you can also just play with our live demo which is here 🎮

Here is an example of dialog:

👻 Conclusion

We’ve come to the end of this post describing how you can build a simple state-of-the-art conversational AI using transfer learning and a large-scale language model like OpenAI GPT.

As we learned at Hugging Face, getting your conversational AI up and running quickly is the best recipe for success so we hope it will help some of you do just that!

Be sure to check out the associated demo and code:

• the live demo is here and
• the open-sourced code and pretrained models are here.

As always, if you liked this post, give us a few 👏 to let us know and share the news around you!