Learn What Is Possible, Then Choose What Is Best: Disentangling One-To-Many Relations in Language Through Text-based Games (2022)

Abstract: Language models pre-trained on large self-supervised corpora, followed by task-specific fine-tuning has become the dominant paradigm in NLP. These pre-training datasets often have a one-to-many structure—e.g. in dialogue there are many valid responses for a given context. However, only some of these responses will be desirable in our downstream task. This raises the question of how we should train the model such that it can emulate the desirable behaviours, but not the undesirable ones. Current approaches train in a one-to-one setup—only a single target response is given for a single dialogue context—leading to models only learning to predict the average response, while ignoring the full range of possible responses. Using text-based games as a testbed, our approach, PASA, uses discrete latent variables to capture the range of different behaviours represented in our larger pre-training dataset. We then use knowledge distillation to distil the posterior probability distribution into a student model. This probability distribution is far richer than learning from only the hard targets of the dataset, and thus allows the student model to benefit from the richer range of actions the teacher model has learned. Results show up to 49% empirical improvement over the previous state-of-the-art model on the Jericho Walkthroughs dataset.

## introduction language models pre-trained on large selfsupervised corpora, followed by downstream finetuning on the target domain has become a popular recipe in nlp @xcite . these pre-training datasets often have a oneto-many structure-e.g. in dialogue there are many valid responses for a given context. however, only some of these responses will be desirable in our downstream task, while others will be counterproductive. for instance, dialogue models intended to engage the user in friendly chitchat are pre-trained on reddit @xcite , which contains a mixture of users who are thoughtful and interesting as well as hostile and argumentative @xcite . this raises the question of how we should train the model such that it can emulate the desirable behaviours on the downstream task, but not the undesirable ones. the current predominant approach is to train in a one-to-one setup-only a single target response is given for a single dialogue context @xcite . however, this leads to models only learning to predict the average response, while ignoring the full range of possible responses @xcite . this is particularly problematic when the desired responses are only contained in a small subset of the dataset. to explore solutions we leverage the domain of text-based games @xcite ). text-based games involve interacting in a textual world by issuing commands ("actions") in natural language-which are interpreted by the game's hard-coded parser-and receiving textual observations from the game's interface (see figure 1 ). the objective is to maximise the integer score by completing various in-game tasks such as exploring new areas, collecting valuable items or solving puzzles. text-based games are especially well-suited to the task for two key reasons: our proposed approach, pasa (posterior action selection from discrete latent variables), uses discrete latent variables to capture the range of different behaviours represented in our larger pretraining dataset. particularly, we learn to predict a set of possible actions, where each action reflects a different value for the latent variable. we assign our latent variable using the ground-truth action itself, via a posterior recognition network. unlike previous approaches, we do not then learn to predict the latent variable from the context. this is because our goal is not to learn to output diverse responses as is the goal of previous approaches @xcite , but rather to model the full set of possible responses. then, the second novelty of our work lies in our method for re-transforming what our model has learned into a one-to-one setup. given our teacher model which predicts sets of actions, we use knowledge distillation @xcite to distil the posterior probability distribution into a student model which predicts the gold action given the context, when fine-tuning on the jericho walkthroughs dataset. this probability distribution is far richer than learning from only the hard targets of the dataset, and thus allows the student model to benefit from the richer range of actions the teacher model has learned. as shown in figure 1 , during pre-training we learn the set of possible actions from a given context, which ranges from poor actions ('eat coal'), to average actions ('west') to high-quality actions ('put coal in furnace'). then during fine-tuning we specialise the model to only predict the high-quality actions. we compare our approach to calm (yao et al., 2020)-currently a state-of-the-art approach for action prediction in text-based games-which represents the paradigmatic one-to-one method by training gpt-2 @xcite to generate the next action given the context using the clubfloyd dataset of human transcripts. figure 1 demonstrates that selecting the gold action requires choosing a highly contextually relevant action-i.e. 'put coal in furnace'-which might not be well represented in the dataset, whereas we found calm preferred the more common output 'west'. although we use text-based games as the testbed for our approach, our method could be extended to generative dialogue tasks-e.g. reddit @xcite )-if paired with a suitable generative model. many social media-based dialogue datasets often form a tree-like structure, in which there are many comments for any given post @xcite . this format would naturally allow for the introduction of a latent variable to first learn to predict all of the comments, then to learn to select which comment is best. in summary, our key contributions are: (1) a novel discrete latent variable neural retrieval architecture, (2) a novel distillation algorithm for disentangling one-to-many and one-to-one tasks; language priors in text-base games several previous works have explored learning linguistic priors for text-based games. the overall idea is that instead of relying on reinforcement learning to learn how to play the games through repeated interactions with the same game @xcite @xcite , the model can rely on general knowledge it has learned about language to give a good sense for which actions are possible and/or may be useful. @xcite explore the use of intrinsic motivation @xcite -i.e. additional loss functions encouraging the model to explore more-to obtain more semantically relevant representations. @xcite use glove word embeddings @xcite for discovering new potentially useful unexplored actions. @xcite use the albert pre-trained model @xcite and perform question answering for intrinsic motivation and to build information about the state. @xcite use the distilbert model @xcite as a base for q-learning. there has also been work towards using text-based games as a testbed for ethical value alignment in agents @xcite . most relevant to our work is @xcite who train gpt-2 @xcite to generate the next action given the context using the clubfloyd dataset of human transcripts; ammanabrolu and riedl (2021a) extend this by multi-tasking on world-modeling tasks, using the jericho worlds dataset which extends the jericho walkthroughs dataset @xcite . however, unlike our work, these works have largely focused on generating valid actions, rather than the optimal gold action. one-to-many modelling in nlp the predominant method for one-to-many modelling is to in-troduce some latent information into the context. often this is in the form of a learned latent variable: variational autoencoders have been used for generating language from continuous spaces @xcite ; subsequently, they were extended to dialogue models as conditional variational autoencoders to model the diversity of responses @xcite ; @xcite jointly optimise over diversity and relevance within the latent space of an rnn; bao et al. (2020) train a large transformer with a discrete latent variable to produce more interesting outputs. different from our approach however, the focus of these approaches is largely on generating more diverse outputs, rather than being able to explicitly generate a set of outputs. thus, it is difficult to steer the outputs of these models towards a particular type of desirable outputs. other one-to-many methods involve conditioning the outputs on more interpretable information such as the speaker id @xcite , or consider the highlevel intent of the response @xcite . these are also highly relevant to our paper, as we also explore such interpretable latent variables, although again our purpose differs in attempting to capture the full set of possible actions. knowledge distillation due to the rapid upscaling of model size in nlp, knowledge distillation (kd) is a popular technique for encoding the knowledge of a larger teacher model into a smaller student model @xcite . kd is most typically used for compressing models of similar types, such as generative models @xcite , or language understanding models @xcite . also it can be used for transferring knowledge between heterogeneous models, such as from a reader to a retriever in open-domain qa @xcite , or from a slow re-ranker to a fast bi-encoder in dialogue response selection @xcite . to the best of our knowledge however, our work is the first to use knowledge distillation to transfer between one-to-many and one-to-one tasks. ## conclusion and future work overall we have demonstrated an approach that disentangles one-to-many and one-to-one prediction by first training a discrete latent variable model with posterior-guided latent variable selection and then distilling this into a student model. results showed up to 49% empirical improvement over the current state-of-the-art approach. our proposed approach currently assumes we have access to a ground-truth set of 'valid actions'. this holds in the case of many dialogue datasets from social media @xcite , as there are often many replies (sc. valid actions) to a given post. in other tasks however such as summarisation or translation we may only have a single exemplum to learn from, and therefore additional techniques would be required to utilise pasa in those cases. we chose text-based games as a starting point in this paper because of the availability of the set of possible actions in the jericho dataset, which allowed us to isolate the task of selecting the best action in evaluation. however, expanding these techniques to other nlp domains such as dialogue would be a key next step for future work. as an example, our approach could be extended to reddit @xcite by treating the number of upvotes to a given comment as the reward, and then treating all of the comments to a particular post as the set of possible actions. then, the evaluation could measure the model's ability to select the comment with the most upvotes given all the comments to a particular post. @xcite . learning discourse-level diversity for neural dialog models using conditional variational autoencoders. in acl.