Language Modeling with Latent Situations

TL;DR

Language models often produce incoherent text due to weak world-modeling. SituationSupervision introduces auxiliary situation modeling and latent state inference to teach LMs explicit entity-state representations, applicable to both fine-tuning and prompting. Across TextWorld and TRIP, modest state annotations yield coherent improvements; latent inference enables semi-supervised gains, and prompting with scratchpad-style state descriptions enhances performance relative to text-only baselines. The results demonstrate that semantic supervision of world states can substantially improve LM coherence with limited annotation, suggesting a broad role for state-aware learning in open-ended generation.

Abstract

Language models (LMs) often generate incoherent outputs: they refer to events and entity states that are incompatible with the state of the world described in their inputs. We introduce SituationSupervision, a family of approaches for improving coherence in LMs by training them to construct and condition on explicit representations of entities and their states. SituationSupervision has two components: an auxiliary situation modeling task that trains models to predict state representations in context, and a latent state inference procedure that imputes these states from partially annotated training data. SituationSupervision can be applied to both fine-tuning (by supervising LMs to encode state variables in their hidden representations) and prompting (by inducing LMs to interleave textual descriptions of entity states with output text). In both cases, SituationSupervision requires only a small number of state annotations to produce major coherence improvements (between 4-11%), showing that standard LMs can be sample-efficiently trained to model not just language but the situations it describes.

Figures (3)

• Figure 1: Language modeling with SituationSupervision, which comprises two components: (a) An auxiliary situation modeling task: given contexts annotated with explicit textual representations of the situations they describe, we use them to adapt LMs through either an auxiliary fine-tuning loss or a scratchpad-style prompt. (b) A latent state inference procedure, whereby missing situation annotations are semi-supervisedly inferred, enabling auxiliary situation modeling starting from a small number of seed situation annotations.
• Figure 2: Fine-tuning (top) and prompting (bottom) with SituationSupervision. In each, we show the: Left: normal text-only procedure, where the LM is trained/prompted with text only samples and expected to produce next sentences $T'$ from contexts $T$. Middle: auxiliary situation modeling component of SituationSupervision, where the LM is given state descriptions $S$ during training or in the prompt and expected to learn to encode it in its parameters or infer it in-context. Right: latent state inference component of SituationSupervision, where the LM must infer missing state descriptions in the training data or prompt demonstrations. Finally, in fine-tuning, state reasoning is done implicitly during inference-time, meaning the inference procedure is the same for all forms of training: we use the base LM to infer the next sentence from the context.
• Figure 3: We find that the design of the situation representation is important: in particular, a situation representation must be ideally consist of only the intersection between the known state and the causally relevant state (highlighted in gray). The known state consists of all facts deducible from the prior context $T$ (i.e. in TW only of facts about rooms or objects that the player has seen -- in this case, the kitchen and attic but not the living room). The causally relevant state consists of all facts causally relevant to any plausible next sentence $T'$ (i.e. in TW only facts about the currently accessible items, such as, in this case, the old key, but not the chest).