Breakpoint Transformers for Modeling and Tracking Intermediate Beliefs (2022)

Abstract: Can we teach models designed for language understanding tasks to track and improve their beliefs through intermediate points in text? Besides making their inner workings more transparent, this would also help make models more reliable and consistent. To this end, we propose a representation learning framework called breakpoint modeling that allows for efficient and robust learning of this type. Given any text encoder and data marked with intermediate states (breakpoints) along with corresponding textual queries viewed as true/false propositions (i.e., the candidate intermediate beliefs of a model), our approach trains models in an efficient and end-to-end fashion to build intermediate representations that facilitate direct querying and training of beliefs at arbitrary points in text, alongside solving other end-tasks. We evaluate breakpoint modeling on a diverse set of NLU tasks including relation reasoning on Cluttr and narrative understanding on bAbI. Using novel proposition prediction tasks alongside these end-tasks, we show the benefit of our T5-based breakpoint transformer over strong conventional representation learning approaches in terms of processing efficiency, belief accuracy, and belief consistency, all with minimal to no degradation on the end-task. To show the feasibility of incorporating our belief tracker into more complex reasoning pipelines, we also obtain state-of-the-art performance on the three-tiered reasoning challenge for the recent TRIP benchmark (23-32% absolute improvement on Tasks 2-3).

## introduction despite considerable progress made recently in natural language understanding (nlu), driven largely by advances in language model pre-training @xcite and the development of large-scale nlu benchmarks @xcite , understanding the behavior of models remains a formidable and highly consequential challenge for model safety. such a challenge is particularly acute in tasks such as narrative understanding, where one must piece together many individual (possibly implicit) facts through time in order to solve problems. for example, in the story in figure 1 , answering the question where is the apple? requires knowing how to track objects through time (e.g., knowing the location of the john and mary and their interaction) and how to compartmentalize other types of knowledge across the story. in such a setting, where models are trained to narrowly answer questions, a natural question arises: do models acquire the kind of requisite background knowledge and world tracking abilities, and ultimately learn representations that give rise to correct beliefs foot_0 about intermediate states? john went to the store and bought an apple @xcite breakpoint token he then returned home to the kitchen john is in the store john is in the kitchen john has the apple john has the ball john has the apple john is in the kitchen the apple in in kitchen he refers to john queries proposition encoding t f t u ... t t t t ... belief state belief state proposition encoding propositions story input s encoder end-task (e.g., qa) [b] for a given story and a set of textual queries corresponding to intermediate points in the story (breakpoints), truth assignments are assigned to queries to form belief states based on a projection over encodings of breakpoints and individual proposition encodings using a single task-specific encoder. using a state-of-the-art pretrained model, t5 @xcite , we develop and investigate a breakpoint transformer to do belief prediction on three categories of tasks: narrative understanding on babi @xcite , relational reasoning on clutrr @xcite and physical commonsense reasoning over human authored stories on trip @xcite . in the former two cases, we focus on training and evaluating models on a novel belief prediction task. we report improvements over a conventional transformer-based representation learning approach @xcite both in terms of prediction accuracy (4% to 8% absolute improvement on clutrr dev) and belief consistency, all with significantly improved processing efficiency (i.e., minimal forward calls to the full transformer) and minimal effect on endtask performance when jointly trained with qa. in the latter case for trip, we show how to integrate our modeling approach into a more complex transformer pipeline and report state-of-the-art results on the three-tiered reasoning task (with 23-32% absolute improvement on two component tasks) over existing task-specific architectures. taken together, our results show the viability of building an end-to-end trainable belief tracking mechanism and integrating it within existing transformer-based reasoning systems. to our knowledge, our work is among the first to look at at general-purpose sentence representation learning for intermediate states in text as a way to facilitate complex situation reasoning. ## related work our work brings together two recent areas that aim to understand model behavior (broadly model prob- task example stories breakpoint propositions relational reasoning (clutrr) john is the brother of susan [b]1 susan's mother is janice [b]2, ... p1 : { 'susan is the sister of john' true, 'susan is the sister-in-law of janice' false,'janice is the mother of john' unk } p2 : { 'janice is the mother of john' true, 'john is the father of janice' false, ...} story understanding (babi) john moved to the kitchen [b]1 he picked up an apple [b]2 john then gave the apple to mary [b]3 ... p1 :{ 'john has the apple' false, 'john is in the kitchen' true,...} p2 : { 'john has the apple' true,'john is in the kitchen'...}... p3 : { 'john has the apple' false, 'mary has the apple' true } commonsense (trip) tom dropped his radio ...carpet. [b]1 the radio broke .. [b]2 tom turned on the radio ... [b]3 ... p2 : { 'radio is in pieces' true, 'radio is powered' false}, ... p3 : { 'radio was powered' true } figure 3 : three tasks rendered as stories with special breakpoint tokens [b] j (for convenience, marked with an index j). each intermediate breakpoint is aligned to a set of propositions p j marked with truth conditions (i.e., true, false, unknown) corresponding to the truth value of each proposition at that breakpoint. our breakpoint model closely relates to lateinteraction encoder architectures that tease apart the encoding of problems and solutions. this includes the sentence transformer from @xcite , which we compare against in our experiments, as well as read once transformers @xcite , colbert @xcite and others. given that the types of narrative tasks we focus on require modeling many intermediate points, we follow this work in putting an emphasis on representation and encoding efficiency. in contrast to this, and other related work on sentence representation learning @xcite , we uniquely focus on learning representations of intermediate states in text for complex situational reasoning. we are also inspired by the situation modeling literature in cognitive science @xcite @xcite , and proposals for their integration with nlp research @xcite . these works also studied neural models of narrative comprehension in carefully controlled micro-worlds, but typically focused on relatively short sentence-level inputs. our work also relates to efforts on building interpretable models by making the underlying reasoning processes transparent, either through explicit decomposition @xcite @xcite or generation of rationales @xcite and other reasoning structures @xcite @xcite . in contrast, we focus on belief representations that are ultimately faithful @xcite to end-tasks by training knowledge directly into a model's task-specific representations. ## breakpoint modeling the goal of breakpoint modeling is to capture the intermediate states and beliefs of models at arbitrary positions in text. our models take stories as inputs, or pieces of text containing one or more intermediate positions (breakpoints), as well as sets of text propositions that align to certain intermediate points (see figure 3 ). such propositions play the role of auxiliary supervision if provided at training time or as queries to the model for performing probing; when coupled with predictions they constitute the beliefs of the model. while breakpoint models can technically take different forms, their basic function is to assign encodings to intermediate states in text and their corresponding propositions ( § 3.1) and to make predictions about the truth/falsity of each proposition ( § 3.2). learning ( § 3.3) reduces to the problem of teaching a model to have a correct and consistent set of beliefs for each target task given a set of representative intermediate propositions and beliefs provided at training time ( § 3.4). ## proposition prediction tasks we focus on three categories of tasks: text-based relational reasoning, story understanding and commonsense reasoning, each considered in turn. in the former two cases, we devise new proposition and belief prediction tasks that involve training on intermediate belief state annotations. we also include out-of-distribution (o.o.d) generalization qiana is lisa's mother qiana is derick's wife derrick is lisa's father story lisa is jerry's granddaughter [b] derrick is lisa's father [b] qiana is derrick's wife [b] atomic belief annotations: the basic facts that should be predicted e(1,'jerry is the grandfather of lisa') e(1,'derick is the father of lisa'), c(1,'lisa ..father of derick'), e(3,'qiana..wife of derick'), e(3,qiana is the mother of lisa) knowledge: constraints to satisfy (implies(anddsde(2,'derick...father of lisa') dsde(3,'qiana is the wife of derick')) e(3,'qiana is the mother of lisa')) figure 5: how are intermediate propositions collected? an illustration of constructing intermediate propositions from clutrr proof trees (above) in gontier et al. (2020). bottom: an example ground constraint, which we use for analyzing consistency. ## modeling details and metrics here we detail our main breakpoint transformer ( §5.1) following the framework in § 3 and all metrics used in our experiments ( § 5.2). ## results and discussion we focus on the following questions: 1. can our main model effectively and efficiently solve our new belief proposition prediction tasks (introduced in § 4) and model intermediate state? 2. can we effectively integrate our breakpoint model into joint models for solving more complex tasks? proposition prediction we found breakpoint models to be effective at our proposition prediction tasks, most notably improving on the transformer perform worse than our bilstm baselines. multi-pass baselines for clutrr prediction from 81.9 to 85.2 (top of model i.i.d hard qa prop.% qa% qa% majority 65.87 --ft-t5-base (qa) -97.29 (±0.14) 69.09 (±0.79) ft-bart-base -97.57 (±0.31) 67.21 (±0.80) bilstm (multi) 80.2 (±0.16) --t5-base (multi) 99.1 (±0.21) --bpt-base 98.5 (±0.10) --bpt-base + qa 98.5 (±0.10) 94.9 (±0.60) 70.51 (±0.29) for babi ( our model's proposition prediction consistency is 7.5% higher than that of the baseline, in terms of the ρ metric reported in through our results on trip ( ## conclusion being able to track the beliefs of models remains a formidable challenge at the forefront of model interpretability. in this paper, we presented a new representation learning framework, breakpoint modeling, that facilitates end-to-end learning and tracking of beliefs at intermediate states in narrative text. on a diverse set of nlu tasks, we show the benefit of our approach (based on t5) over conventional learning approaches in terms of improved belief prediction performance on new belief tracking tasks and processing efficiency. we also show the feasibility of recasting existing tasks into our framework and integrating our approach into existing transformer-based nlu pipelines, which we believe can help to improve the interpretability of these models as part of this larger challenge.