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

1 Introduction

• The framework ranks retrieved knowledge by contextual relevance, using sentence-embedding weak supervision because gold knowledge labels are unavailable.Top-ranked knowledge sentences guide the conditional generator.Source: §1 · pp. 1–2
• MEGATRON-CNTRL dynamically incorporates external knowledge and controls generation through predicted or manually replaceable keywords.Its framework combines a keyword predictor, knowledge retriever, contextual ranker, and conditional generator.Source: §1 · p. 2, Figure 1 · p. 2
• Compared with prior ROC-story work, the model produces more fluent, consistent, and coherent stories with lower repetition and higher diversity.The paper reports this pattern using both automatic metrics and human evaluations.Source: §1 · p. 2
• Replacing generation keywords enables manual control, with human evaluation finding up to 91.5% of stories successfully controlled by new keywords.The controllability procedure replaces the keywords used to generate stories and reruns generation.Source: §1 · p. 2
• Scaling from 124 million to 8.3 billion parameters improves both generation quality and controllability.The paper demonstrates this trend across model capacities, with larger models producing increasingly coherent, fluent, and logically consistent stories.Source: Table 1 · p. 1, §1 · p. 2

2 Framework

• The framework repeatedly predicts keywords, retrieves and ranks external knowledge, then generates the next sentence from the story context and selected knowledge.The generated sentence is appended to the context, and the four-step process repeats for subsequent sentences.Source: §2 · p. 3
• Dynamic keyword replacement lets users control generation during the story rather than pre-specifying all keywords at the beginning.The framework supports manual external keywords and avoids the difficulty of statically planning knowledge for long generations.Source: §2 · p. 3

2.1 Keyword Predictor Model

• The keyword predictor models the next sentence’s keyword sequence with GPT-2, trained by maximum likelihood on RAKE-extracted keywords from ground-truth sentences.Predicting keywords instead of triples makes control interpretable and reduces the knowledge-triple search space.Source: §2.1 · p. 3

2.2 Knowledge Retrieval

• Keyword retrieval converts matching knowledge-base triples into templated natural-language knowledge sentences, with future work proposed for replacing this simple retriever.The method matches predicted keywords against transformed triples in the external knowledge base.Source: §2.2 · pp. 3–4

2.3 Building Pseudo Label of Ri

• Weak pseudo-labels for latent controlling knowledge are built by matching sentence keywords to triples and ranking the resulting knowledge sentences by USE similarity to the preceding and current sentences.The preceding sentence is included because using only the current sentence can be ambiguous, while additional earlier context may overwhelm its information.Source: §2.3 · p. 4

2.4 Contextual Knowledge Ranker

• A BERT contextual ranker scores retrieved knowledge against the story context with learned representations and margin loss, after which GPT-2 generates the next sentence conditioned on the context and top-ranked knowledge.Positive pseudo-label knowledge and retrieved non-positive knowledge provide the ranker’s training pairs.Source: §2.4 · p. 4, §2.5 · p. 4

3.1 Datasets

• The study uses 98,161 five-sentence ROC stories, with 88,344/4,908/4,909 examples for train/validation/test, and ConceptNet’s 600k knowledge triples.Names and entities are replaced with gendered or neutral placeholders, and the first sentence conditions generation of the remaining story.Source: §3.1 · p. 5

3.2 Models

• MEGATRON-CNTRL initializes its contextual ranker and generator from Megatron-LM BERT and GPT-2 configurations, using top-k sampling with k = 40 and temperature 0.7.The contextual ranker uses a margin of 5.0, ten positive retrieved sentences, forty negatives, and achieves validation accuracy of 0.9.Source: §3.2 · p. 5

3.3 Controllability Experiment Setup

• Controllability experiments replace the first keyword having a WordNet antonym and regenerate the remainder from the preceding story context, shared earlier keywords, and the antonym.The resulting generations are labeled MEGATRON-CNTRL-ANT, with model-size variants such as MEGATRON-CNTRL-355M-ANT.Source: §3.3 · p. 5

3.4 Baselines

• The comparison includes Plan and write, Knowledge enhanced GPT-2, and GPT-2-124M, with Knowledge enhanced GPT-2 identified as the ROC story state of the art at the time.The study omits Fan et al. (2019) because Guan et al. (2020) had already reported significantly stronger performance.Source: §3.4 · pp. 5–6

3.5 Evaluation

• Automatic evaluation measures repeated 4-grams, distinct 4-gram diversity, and conditional-generator perplexity after sequentially sampling knowledge before each story sentence.Human quality evaluation covers fluency, coherence, and consistency through independent AMT tasks.Source: §3.5 · p. 6, §3.5.1 · p. 6, §3.5.2 · p. 6

3.5.2 Human Evaluation on Quality

• Human evaluations use 200 randomly sampled stories, five annotators per story, majority voting, and a separate antonym-based task for controllability.Quality judgments use qualified United States workers, while controllability asks whether the antonym-conditioned story changes consistently with the new keyword meaning.Source: §3.5.2 · p. 6, §3.5.3 · p. 6

4.1 Automatic and Human Evaluations

• The 124M model has higher distinctness and consistency than prior work but worse repetition than Yao et al. (2019).Compared with GPT-2-124M and Guan et al. (2020), it has lower repetition and higher distinctness.Source: §4.1 · p. 7, Table 4 · p. 7
• MEGATRON-CNTRL-124M achieves higher coherence and consistency than Guan et al. (2020), with gains of +7.0% and +7.5%, respectively.Fluency is similar between the two models; the authors attribute the gains to retrieved knowledge that lets generation focus on producing text rather than predicting topics and text jointly.Source: §4.1 · p. 7
• 93% consistency is achieved by the 8B model, while scaling generally improves fluency, coherence, and consistency.The 124M model already exceeds prior work on distinct and consistency, and the 355M model significantly outperforms Guan et al. (2020) across evaluated aspects.Source: §4.1 · pp. 7–8, Table 4 · p. 7

4.2 Controllability Evaluation

• 91.5% controllability is reached at 8B parameters, increasing from 77.5% at 124M and 84.5% at 355M when keywords are replaced by antonyms.Controlled generations also have lower repetition and higher distinctness than original generations, while distinctness improves as model size increases.Source: §4.2 · p. 8, Table 5 · p. 8, Table 6 · p. 8

4.3.1 Planning Strategy

• Dynamic planning yields +0.7% distinctness and -3.8% repetition compared with static planning.Dynamic planning retrieves knowledge for each sentence, whereas static planning predicts all knowledge from the starting sentence, producing weaker and noisier supervision.Source: §4.3.1 · p. 8, Table 7 · p. 8

4.3.2 External Knowledge

• Removing external knowledge decreases distinctness by 0.8% and increases repetition by 3.6%, highlighting its contribution to generation quality.This effect is observed for dynamic planning; under static planning, adding or removing knowledge produces similar quality, reinforcing the value of dynamic knowledge planning.Source: §4.3.2 · p. 8, Table 7 · p. 8

5 Future Work

• Short ROC story sentences limit exploration of how long control signals propagate during longer generations.The authors propose evaluating this question on longer story datasets such as WRITINGPROMPTS.Source: §5 · pp. 8–9
• Future work could add structure-level controllability as an extra generator input or multitask supervision.Source: §5 · p. 9
• Some generations mention the supplied keyword without centering the story on its meaning.The paper attributes this behavior to RAKE keywords that do not always represent the sentence well.Source: §5 · p. 9

Knowledge Incorporation

• Knowledge has been incorporated into language models for factual generation, commonsense graph construction, entity typing, question answering, and dialogue modeling.Source: §6 · p. 9

Storytelling

• Storytelling research includes story completion, ending generation, prompt- or title-based generation, and generation from a given sentence.Source: §6 · p. 9
• Story-generation methods model structure with storylines, skeletons, conditional variational autoencoders, coarse-to-fine frameworks, and attention-based commonsense incorporation.Source: §6 · p. 9

Controllable Generation

• Few works address controllability of generation, particularly for the large pretrained models common in current research.MEGATRON-CNTRL differs by controlling generation through keywords backed by external knowledge.Source: §6 · p. 9
• Controllable-generation research targets persona, politeness, templates, pivot words, domains, links, and whole-document style or properties.Source: §6 · p. 9

7 Conclusion

• 91.5% of generated stories are successfully controlled when keywords are changed to their antonyms.Source: §7 · p. 9
• The framework predicts keywords, retrieves related knowledge triples, ranks knowledge sentences by story-context relevance, and conditions the generator on the top-ranked sentences.Source: §7 · p. 9
• On the ROC story dataset, the model generates less repetitive, more diverse, and more logically consistent stories than state-of-the-art models.Source: §7 · p. 9
• Larger pretrained language models consistently improve both generated-story quality and controllability.Source: §7 · p. 9

A.1 GPT-2 Hyperparameters:

• The experiments use the BPE subword tokenizer with a maximum sequence length of 1024 learned positional embeddings.Source: §A.1 · p. 12
• Training uses Adam with learning rate 0.0001, dropout probability 0.1, and global gradient-norm clipping at 5.0.Source: §A.1 · p. 12
• GPT-2 models are finetuned for 5 epochs, selecting the checkpoint with the lowest validation perplexity.Source: §A.1 · p. 12

• 512 learned positional embeddings define the maximum sequence length, and tokenization uses the bert-large-uncased WordPiece vocabulary.Source: §A.2 · p. 12
• Adam optimization uses a 0.0001 learning rate, 0.01 weight decay, and global gradient clipping at norm 1.0.Source: §A.2 · p. 12
• Dropout with probability 0.1 is applied to the embedding and multi-head attention layers.Source: §A.2 · p. 12
• A margin of 5.0 performs best among the tested values 0.1, 0.5, 1.0, and 5.0.Source: §A.2 · p. 12

A.3 Model Size

• Increasing the keyword predictor from 774M to 2B reduces repetition while keeping other scores similar, whereas scaling the ranker can reduce diversity and may overfit weakly supervised noise.The authors therefore use the smaller 336M ranker in all experiments; Table 8 summarizes the ablation.Source: §A.3 · p. 12, Table 8 · p. 12

A.4 Datasets Used for pre-trained Models

• Pretrained GPT-2 models use a 174GB corpus containing Wikipedia, OpenWebText, RealNews, and CC-Stories, while BERT models additionally use BooksCorpus.Source: §A.4 · p. 12

A.5 Input Format

• Knowledge sentences are separated with SEP, and EOK marks the end of the knowledge input.Source: §A.5 · p. 12
• The conditional generator concatenates the preceding story context, sentence-end markers, an end-of-text token, and the current knowledge input.Source: §A.5 · p. 12

A.6 Preprocessing

• Preprocessing follows Guan et al. (2020) and replaces “ .” with “. ”.Source: §A.6 · p. 12

A.6.1 Computation Resources

• Experiments use up to 20 DGX-1V servers with 160 Tesla V100 GPUs; GPT-2 training usually takes under three hours and BERT training eight hours.Source: §A.6.1 · p. 12

A.7.1 Single Sentence Context

• The paper presents five generation examples with the context fixed to one sentence and three generation examples that vary the number of sentences in the context.Source: §A.7.1 · p. 13, §A.7.2 · p. 14

A.8 Human Evaluation Interface for Annotators

• Human-evaluation interfaces assess coherence, fluency, logical consistency, and controllability.Source: §A.8 · p. 14