MEGATRON - CNTRL : Controllable Story Generation with External Knowledge Using Large-Scale Language Models (2020)

Abstract: Existing pre-trained large language models have shown unparalleled generative capabilities. However, they are not controllable. In this paper, we propose MEGATRON-CNTRL, a novel framework that uses large-scale language models and adds control to text generation by incorporating an external knowledge base. Our framework consists of a keyword predictor, a knowledge retriever, a contextual knowledge ranker, and a conditional text generator. As we do not have access to ground-truth supervision for the knowledge ranker, we make use of weak supervision from sentence embedding. The empirical results show that our model generates more fluent, consistent, and coherent stories with less repetition and higher diversity compared to prior work on the ROC story dataset. We showcase the controllability of our model by replacing the keywords used to generate stories and re-running the generation process. Human evaluation results show that 77.5% of these stories are successfully controlled by the new keywords. Furthermore, by scaling our model from 124 million to 8.3 billion parameters we demonstrate that larger models improve both the quality of generation (from 74.5% to 93.0% for consistency) and controllability (from 77.5% to 91.5%).

## introduction text generation has recently attracted significant attention from the research community as large pretrained language models, such as gpt-2 @xcite demonstrated promising results for generating long, grammatically correct, and fluent text. finetuning these models has shown significant improvements in downstream tasks, such as persona chat @xcite . however, one nonnegligible drawback of these large models is the context @xcite was on a long road trip . @xcite she had just driven into town . she had been driving for hours . she was so tired she decided to take a nap . when she woke up , she still had n't slept well . megatron-cntrl-124m keyword: driving she was on her way to see a strange country. she had never seen anything like that before. finally she saw a small town that was so small. it was a great little town with a nice park ! megatron-cntrl-8b keyword: driving she was driving in the middle of @xcite . all of a sudden the tire pressure light came on. she got out to check her tire. it was flat so she used the roadside assistance. megatron-cntrl-8b-ant keyword: attract she really wanted to see a few attractions. the first one she saw was a giant water park. it was amazing. it ended up being a fun experience. in addition, these large pre-trained language models are hard to control. recently, plug-and-play language models @xcite addressed whole document controllability by adding a linear classifier on top of gpt-2 to predict whether generated text observes a particular style or property. @xcite controlled a 1.2b parameter language model generation via the use of control codes prepended to the model input. @xcite controlled the personality of a dialogue agent by conditioning it on prior conversations of a target actor. however, these controlling conditions are predefined, limited in their capability, and are only used once at the beginning to condition the generation of the rest of the document. they do not provide control granularity at either a sentence or sub-document level. in this work, we address these shortcomings and develop an efficient controllable text generation framework that we apply to the story generation task. in order to provide manual control to users through a set of interpretable keywords, we first develop a keyword predictor model for the next sentence. these keywords are then used to retrieve knowledge sentences from an external knowledge base. not all the retrieved knowledge is relevant to the story context and often it is noisy. to this end, we introduce a novel contextual ranker that ranks knowledge sentences based on the relevance to the context. as we do not have access to ground-truth supervision for this contextual knowledge ranker, we make use of sentence embedding for weak supervision. the top-ranked knowledge sentences from the knowledge ranker are then fed to the conditional text generator to guide generation. by giving the knowledge in addition to the context, we provide rich information for the generator to attend to and help the model better understand the rationale between sentences. ## framework in our problem setup, we complete a story using the first sentence as input, similar to @xcite . we augment the generation process with an external knowledge-base and develop a methodology that can guide and control the story generation. our approach consists of the following four steps connected together as shown in figure 1 : this formulation naturally allows controllability by replacing the keyword prediction process with manual external keywords. this work uses dynamic planning of the keywords and knowledge at each generation step. this allows the users to participate and control the generation on the go. as a result, they don't need to pre-specify the keywords explicitly. we also note that it is challenging to statically plan all the knowledge needed for generation at the beginning. this issue becomes severe for long generations. to formalize this method, we start by introducing notation used throughout the paper and then detail each aforementioned four steps in the following subsections. notation: a knowledge-base, g is defined as a set of knowledge triples @xmath0, relation, object). a knowledge sentence, r is defined as @xmath1. for example, (eiffel tower, at-location, paris) is transformed into eiffel tower is at paris. we should highlight that since our framework transforms the triple knowledge database into natural language sentences, any knowledge base in natural language format can be readily incorporated into our framework. we use superscripts to index story sentences and define a story s of length l as a sequence of individual story sentences s i where to denote the keywords associated with story sentence s i . a keyword k i q is made up of subword tokens from our language model's vocabulary. note that the number of keywords q per sentence varies and can be zero. we define r as the knowledge associated with s i , where r i j denotes the j-th knowledge sentence associated s i . the number of knowledge sentences v varies per sentence and can be zero. note that @xmath2. given this notation, we define the story context the goal of this work is to generate x i given x i-1 , that is to first predict the knowledge r i contained in s i and then predict s i itself. ## results in this section, we first perform automatic and human evaluations with different model sizes and compare our framework to the existing baselines. we then evaluate the controllability of our model and finally show ablation study varying gpt-2 and bert model sizes. the detailed configuration of the model sizes are shown in ## future work short story sentences in roc story dataset limits our exploration from several potential research directions. for example, how long the control signal would propagate for longer generations? investigating this issue using longer story datasets (e.g. writingprompts @xcite ) is a subject for future work. other interesting direction may include incorporating the structure-level controllability by adding it as either an extra input for the conditional generator or a multitask learning supervision for each sequence. we also observed that in some cases during the generation, our model simply mentions the given word in the sentence, and talks about things that are not strictly related to or restricted by the given word. for example, the generated story of megatron-cntrl-8b in storytelling there are several different storytelling tasks described throughout the literature. storytelling can be classified into story completion @xcite , story ending generation @xcite , story generation from prompts @xcite or titles @xcite , and story generation from a given sentence @xcite . different approaches have been developed to model the structure of stories with storylines @xcite , skeletons @xcite , conditional variational autoencoders @xcite and a coarse-to-fine framework @xcite . other works focus on incorporating commonsense knowledge into story generation with attention-based models @xcite . recently, pre-trained language models have been used to finetune on both story completion datasets and commonsense knowledge to further improve the quality of story completion @xcite . however, few works concern the controllability of language model generation, especially for the large pre-trained models that are common in today's literature. ## conclusion in this paper, we proposed a novel framework that adds control to text generation with external knowledge. our model first generates a set of keywords and a knowledge retriever then queries an external knowledge base for triples related to the keywords. based on the relevance to the story context, a contextual knowledge ranker ranks the retrieved knowledge sentences and feeds the top ones to a conditional generator to generate the next story sentence. experimental results on the roc story dataset showed that our model outperforms state-of-the-art models by generating less repetitive, more diverse and logically consistent stories. human evaluation of the controllability of our model shows that 91.5% of generated stories are successfully controlled by changing keywords to their antonym. in line with current trends, we also demonstrate that using larger pre-trained language models consistently improves both the quality of the generated stories and controllability.