Baked-in State Probing (2022)

Abstract: Neural language models have been analyzed for their linguistic and extra-linguistic knowledge via probing. Of particular interest has been the following question: how much can a language model trained only on form learn about meaning? Recent work has demonstrated via probing classifiers that in the setting of simple procedural text, where by “meaning” we mean the underlying world state, language models have a non-trivial performance on world state tracking. However, our proposed evaluation based on model predictions shows differing results, suggesting that these models are either not capturing the world state or not using it. How do these results change if the model has access to the world state? We explore this alternate setting with access to the underlying world state only during training and investigate ways of “baking in” the state knowledge along with the primary task of language modeling. Our proposed approaches allow for state probing during inference simply via text prompts, avoiding any probing classifier machinery. In terms of performance, we show that baking in the state knowledge during training leads to significant improvements in state tracking performance and text generation quality,

## introduction there has been recent interest in the extent to which models trained solely on text (form) capture aspects of the underlying world state (meaning) (bender and koller, 2020; bisk et al., 2020; wu et al., 2021; bender et al., 2021; li et al., 2021; toshniwal et al., 2022). li et al. (2021) show via the use of probing classifiers that language models trained for simple semantic domains learn to represent the state of the world described by the text without any explicit state supervision. however, one of the key limitations of probing classifiers is that the extractability of any information, such as world state, doesn't necessarily mean the model is using this information (see belinkov (2022) for a detailed discussion on limitations of probing classifiers). our proposed evaluation based on model predictions shows results contrary to li et al. (2021) wherein we find that the model performance is close to chance in capturing/utilizing the state knowledge (see how do these results change if the language model has access to the ground truth world state knowledge? we explore this alternate setting where we assume that: (a) the language model has access to the world state only during training, and (b) that the world state (or parts of it) can be translated to text. access to such oracle annotations during training has also been explored in recent work (nye et al., 2021; lampinen et al., 2022). concretely, we build on the setup of li et al. (2021) where we explore simple approaches to explicitly adding the state to the language model training sequences. we show that this baking-in of the world state knowledge allows for state probing during inference simply via prompting. our experiments on the alchemy dataset (long et al., 2016) show that our proposed approaches, particularly a multitask learning approach, significantly improve over a baseline language model both in terms of state prediction and language modeling. more broadly speaking, the proposed approaches could present an alternate way for injecting linguistic knowledge during end task training for nlp tasks (wu et al., 2021). ## task setup in this section, we first formally describe the setup in which we conduct the language modeling and state probing experiments. following that, we briefly describe the probing classifier setup used by li et al. (2021), and finally our proposed lan- w 0 w 1 @xmath0, the second beaker has 2 purple, • • •, the second to last beaker has 1 orange, the last beaker is empty t (w1) → the first beaker has 1 red, the second beaker has 2 purple, • • •, the second to last beaker is empty, the last beaker is empty t (w0, fifth beaker) → fifth beaker has 4 red t (w1, fifth beaker) → fifth beaker has 4 red ## results ## conclusion we show via our proposed evaluation of valid next sentence prediction that a baseline language model is at par with chance performance, suggesting that the model struggles to capture the state knowledge. based on this evidence, we explore language model variants which assume access to the world state during training. the proposed language model variants can be easily probed for world state via text prompt, and significantly improve over a canonical language model for both state tracking and text generation. our results suggest that injecting semantic knowledge could be beneficial more broadly as well, and our proposed approaches can be useful for this end.