SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zeming Chen<sup>1</sup> Gail Weiss<sup>1</sup> Eric Mitchell<sup>2</sup> Asli Celikyilmaz<sup>3</sup> Antoine Bosselut<sup>1</sup> EPFL<sup>1</sup> Stanford University<sup>2</sup> Meta AI Research<sup>3</sup> {zeming.chen, gail.weiss, antoine.bosselut}@epfl.ch eric.mitchell@cs.stanford.edu aslic@meta.com
Recent studies on transformer-based language models show that they can answer questions by reasoning over knowledge provided as part of the context (i.e., in-context reasoning). However, since the available knowledge is often not filtered for a particular question, in-context reasoning can be sensitive to distractor facts, additional content that is irrelevant to a question but that may be relevant for a different question (i.e., not necessarily random noise). In these situations, the model fails to distinguish the necessary knowledge to answer the question, leading to spurious reasoning and degraded performance. This reasoning failure contrasts with the model's apparent ability to distinguish its contextual knowledge from all the knowledge it has memorized during pre-training. Following this observation, we propose teaching the model to reason more robustly by folding the provided contextual knowledge into the model's parameters before presenting it with a question. Our method, RECKONING, is a bi-level learning algorithm that teaches language models to reason by updating their parametric knowledge through back-propagation, allowing them to answer questions using the updated parameters. During training, the inner loop rapidly adapts a copy of the model weights to encode contextual knowledge into its parameters. In the outer loop, the model learns to use the updated weights to reproduce and answer reasoning questions about the memorized knowledge. Our experiments on three diverse multi-hop reasoning datasets show that RECKONING's performance improves over the in-context reasoning baseline (by up to  $4.5\%$ ). We also find that compared to in-context reasoning, RECKONING generalizes better to longer reasoning chains unseen during training, is more robust to distractors in the context, and is computationally more efficient when multiple questions are asked about the same knowledge.
Consider the sentence: "John is David's dad, and Tom is John's dad". Concluding that Tom is David's grandfather involves reasoning about the information in the sentence. Specifically, it requires understanding the direct information, or contextual knowledge, given in the sentence: the stated relationships between John, David, and Tom; and combining it with our existing, commonsense knowledge of the world: someone's dad's dad is their grandfather. Achieving such logical reasoning automatically has long been a goal of AI [17, 53, 74, 82].
The example above demonstrates two necessary abilities required for successful reasoning: first, holding large amounts of commonsense or general knowledge about the world, and second, processing and combining new information with existing knowledge. Transformer-based large language models have shown a remarkable capacity for the first of these abilities, repeatedly being demonstrated to memorize large amounts of data, or parametric knowledge, in their weights [8, 11, 50, 64].
Figure 1: Our algorithm, RECKONING, solves reasoning problems by encoding external contextual knowledge into a model's parameters through gradient updates. At inference time, RECKONING performs a few parameter updates using the gradients of a language modeling loss to encode the relevant facts. Then, the updated model answers the question using only its implicit knowledge.
For the second, recent work showed that transformers fine-tuned to predict answers over a concatenated context ("The cow is big; If something is big then it chases the dog; If the cow chases the dog then the cow sees the rabbit") and question ("Did the cow see the rabbit?") achieve high performance on reasoning tasks where all necessary knowledge is given in the context [17]. We refer to this general setting as in-context reasoning (ICR).
In real-world question-answering settings [16, 22, 40, 42], large amounts of contextual knowledge may be provided at once, and the information may not be perfectly filtered for a specific question. Unfortunately, in-context reasoning is highly sensitive to distractors [70]: additional facts that are not relevant to a question (e.g., "The cow is round" for the above example). Indeed, when fine-tuning and evaluating GPT-2 [60] for ICR, we find that adding distractors to the context drops performance from  $99.4\%$  to only  $70.9\%$  accuracy for the same questions ( $\S 4.2$ ). This sensitivity to distractors in contextual knowledge contrasts with GPT-2's apparent robustness to distractors in parametric knowledge: for any specific example, most of the training data seen by GPT-2—which forms its parameters—is likely to be completely irrelevant to that example. Naturally, we wonder whether presenting contextual knowledge in the same way as memorized knowledge, by encoding it into a model's parameters, will improve the reasoning abilities of transformer-based language models.
In this work, we propose a novel bi-level optimization algorithm, RECKONING, that learns to memorize (and reason) over facts (i.e., knowledge) by performing inference-time parameter updates using gradients computed from a language modeling loss on those facts. The updated model is then used to answer any questions about those facts. Our training framework involves two nested loops: the inner loop performs fast adaptations from a set of initial weights to memorize a set of external knowledge through a few gradient updates, and the outer loop optimizes those same initial weights such that the updated model will solve reasoning problems associated with the memorized knowledge. In other words, the outer loop learns optimal meta-parameters that can rapidly memorize and successfully reason over contextual knowledge, allowing knowledge memorization to be optimized directly for downstream reasoning. At inference time, instead of including external knowledge in the input sequence as the prefix to a question prompt, the model can encode it in its parameters through gradient updates and then reason over its updated parametric knowledge to reach a conclusion.
We evaluate RECKONING on two synthetic multi-hop reasoning datasets: ProofWriter [74] and CLUTRR-Systematic-Generalization (CLUTRR-SG) [29], and one real-world dataset, FOLIO [30], comparing against a fine-tuned ICR (FT-ICR) baseline that uses the same underlying model. Our results show that RECKONING consistently outperforms the FT-ICR baseline on each benchmark, demonstrating that it successfully learns to answer multi-hop reasoning questions as desired. In particular, we find that RECKONING more successfully generalizes to adversarial settings, such as the presence of distractor facts and the introduction of longer reasoning chains at inference time. Finally, while the inference-time gradient updates make RECKONING slower to process new knowledge than a typical ICR forward pass, our run-time analysis shows that RECKONING is more efficient when answering multiple questions about a shared knowledge set. This speedup occurs because RECKONING only needs to encode the knowledge once to answer multiple questions about it. Overall, we demonstrate that RECKONING is an effective algorithm for reasoning through dynamic and controllable knowledge encoding, overcoming an observed weakness in the common reasoning setting and providing multiple additional benefits.
Figure 2: The two-stage training process of RECKONING with an inner and outer loop.
Notation We use  $f: \mathcal{X} \times \theta \to \mathcal{Y}$  to refer to parameterised functions in which  $\mathcal{X}$  is the set of possible inputs and  $\theta$  are their possible weights (parameters). We use  $f_{\theta}: x \mapsto f(x, \theta)$  to easily refer to any  $f$  with a given set of parameters  $\theta$ . We describe reasoning problems using tuples  $(\mathcal{K}, \boldsymbol{x}, y^*, Y)$  such that  $y \in Y$  is the correct answer for the question  $\boldsymbol{x}$  given facts  $\mathcal{K}$ , and use  $\mathcal{D}$  to refer to sets of such problems. When it is clear from context, we drop  $Y$  and use only  $(\mathcal{K}, \boldsymbol{x}, y^*)$ .
Language Modeling and Memorization In the causal language modeling (CLM) objective, a parameterized model  $f_{\theta}$  is trained to estimate the conditional probabilities of each token in a sequence given its predecessors:  $p(x_{t}|x_{<t})$ . Specifically, we train  $f_{\theta}$  to approximate  $p$  using the CLM loss:
$$
\mathcal {L} _ {\mathrm {C L M}} \left(f _ {\boldsymbol {\theta}}, \boldsymbol {x}\right) = - \sum_ {t = 1} ^ {T} \log f _ {\boldsymbol {\theta}} \left(x _ {t} \mid x _ {1}, \dots , x _ {t - 1}\right). \tag {1}
$$
This training objective allows language models to memorize individual training examples [10, 11], and we will exploit this ability to memorize and draw on contextual knowledge in our work.
Transformers as Soft Reasoners In natural language reasoning tasks, we are given reasoning problems  $(\mathcal{K},\boldsymbol {x},y^{*},Y)$  in natural language and attempt to recover the correct answer  $y^{*}$  from the context  $\kappa$  , question  $\pmb{x}$  , and possible answers  $Y$  alone. In in-context reasoning, language models  $f_{\theta}$  trained with a CLM objective are applied to this task by selecting as the response the answer  $y\in Y$  with a maximum probability according to the model's next-token prediction from the concatenated context and question:  $y = \arg \max_{y^{\prime}\in Y}f_{\theta}(y^{\prime}||[\mathcal{K};\boldsymbol {x}])$  . Previous works show that, after relevant supervised fine-tuning, transformer language models can achieve high performance in this setting [17, 29, 74], though this degrades significantly in the presence of irrelevant facts (distractors) [70].
Addressing these challenges, we propose RECKONING (REasoning through dynamiciC KnOwledge eNcodING), which solves reasoning problems by memorizing the provided contextual knowledge, and then using this encoded knowledge when prompted with downstream questions. Specifically, RECKONING uses bi-level optimization to learn a set of meta-parameters primed to encode relevant knowledge in a limited number of gradient steps. The model can then use its updated weights to solve reasoning problems over this knowledge, without further presentation of the knowledge itself.
Overview: Inference Given a reasoning problem  $(\mathcal{K},\boldsymbol {x},y,Y)$ , we initialize our model with weights copied from a set of meta-parameters  $\theta$  and perform a constant number  $N$  of gradient descent steps on these with the goal of minimizing the CLM objective on the knowledge set  $\mathcal{K}$ . This allows the model to memorize  $\mathcal{K}$  in its updated parameters, which we refer to as  $\hat{\theta}_N^{\mathcal{K}}$ . Next, we pass the question  $\boldsymbol{x}$  to the model, using  $f_{\hat{\theta}_N^{\mathcal{K}}}$  to obtain a distribution over  $Y$ , and taking as output the answer  $y\in Y$  with the highest probability. For this method to consistently output the ground truth  $y^{*}$ , we seek a set of optimal meta-parameters  $\theta^{*}$  that can quickly memorize (i.e., learn) the given knowledge in a way that then allows accurate reasoning when queried about the knowledge downstream.
Training RECKONING Similar to the meta-learning algorithm MAML [25], which focuses on optimizing the initial parameters of a model so that the model can be good at learning new tasks with a few examples, we also optimize the initial parameters such that the model learns to encode knowledge for more effective reasoning. Given a distribution  $p(\mathcal{D})$  of reasoning problems, our proposed bi-level optimization framework RECKONING (seen in Figure 2) optimizes the following objective:
$$
\boldsymbol {\theta} ^ {*} \in \underset {\boldsymbol {\theta}} {\arg \min } \mathbb {E} _ {(\mathcal {K}, \boldsymbol {x}, y) \sim p (\mathcal {D})} \left[ \mathcal {L} _ {\mathrm {C E}} \left(f _ {\hat {\boldsymbol {\theta}} _ {N} ^ {\mathcal {K}}} (\boldsymbol {x}), y\right) \right] \tag {2}
$$
where for all  $\mathcal{K}$ ,  $n\in \mathbb{N}$ , and  $\pmb {\theta}:\hat{\pmb{\theta}}_0^{\mathcal{K}} = \pmb {\theta}$ , and
$$
\hat {\boldsymbol {\theta}} _ {n + 1} ^ {\mathcal {K}} = \hat {\boldsymbol {\theta}} _ {n} ^ {\mathcal {K}} - \alpha \nabla \mathcal {L} _ {\mathrm {C L M}} \left(f _ {\hat {\boldsymbol {\theta}} _ {n} ^ {\mathcal {K}}}, \mathcal {K}\right). \tag {3}
$$
Here,  $\mathcal{L}_{\mathrm{CE}}(f(\pmb{x}),y)$  denotes the cross-entropy (CE) loss, which we apply with the relevant parameters for each reasoning question in  $\mathcal{D}$ ,  $\mathcal{L}_{\mathrm{CLM}}(f,\mathcal{K}) = \frac{1}{|\mathcal{K}|}\sum_{k\in \mathcal{K}}\mathcal{L}_{\mathrm{CLM}}(f,k)$  denotes the causal language modeling loss, and  $N$  and  $\alpha$  are pre-defined hyperparameters of the fine-tuning. We seek our actual meta-parameters  $\theta$  through gradient descent. In particular, denoting by  $\theta_0$  our initial meta-parameters, and  $\hat{\theta}_{N,i}^{\mathcal{K}}$  the parameters  $\hat{\theta}_N^{\mathcal{K}}$  obtained when initializing  $\hat{\theta}_0^{\mathcal{K}}$  with  $\theta_{i}$ , we iteratively compute
$$
\boldsymbol {\theta} _ {i + 1} = \boldsymbol {\theta} _ {i} - \eta \nabla \frac {1}{| \mathcal {D} _ {i} |} \sum_ {(\mathcal {K}, \boldsymbol {x}, y) \in \mathcal {D} _ {i}} \mathcal {L} _ {\text {T o t a l}} \left(f _ {\hat {\boldsymbol {\theta}} _ {N, i} ^ {\mathcal {K}}}, \mathcal {K}, \boldsymbol {x}, y\right), \tag {4}
$$
where  $\mathcal{L}_{\mathrm{Total}}(f,\mathcal{K},\boldsymbol {x},y) = \mathcal{L}_{\mathrm{CE}}(f(\boldsymbol {x}),y)$  and for each  $i$ $\mathcal{D}_i$  is randomly sampled from  $p(\mathcal{D})$  . This continues until  $\mathcal{L}_{\mathrm{Total}}$  converges.
The training can be seen as two nested loops: at each iteration, the outer loop (Equation (4)) samples a random batch  $\mathcal{D}_i\subseteq \mathcal{D}$  of reasoning problems for evaluating (in order to update) the current meta-parameters  $\theta_{i}$  after the inner loop (Equation (3)) adapts them to encode the associated knowledge through  $N$  steps of gradient updates.
Multi-Task Objective Through our experiments, we find that adding a knowledge-recovery objective to the outer loop—such that the model must also state all of  $\kappa$  when
Algorithm 1 RECKONING
Require: An example distribution  $p(\mathcal{D})$ , a transformer language model  $f$ , initial meta-parameters  $\theta$ , outer step size  $\eta$ , inner step size  $\alpha$ , inner loop length  $N$ .  
1: while not converged do  
2: Sample  $\mathcal{D}' \sim p(\mathcal{D})$   
3:  $\mathcal{L}_{\mathcal{D}'} \gets 0$   
4: for each  $(\mathcal{K}, \boldsymbol{x}, y) \in \mathcal{D}'$  do  
5: Initialize  $\hat{\theta}_0^{\mathcal{K}} = \theta$   
6: for  $n := 0$  to  $N - 1$  do  
7:  $\hat{\theta}_{n + 1}^{\mathcal{K}} \gets \hat{\theta}_n^{\mathcal{K}} - \alpha \nabla \mathcal{L}_{CLM}(f_{\hat{\theta}_n^{\mathcal{K}}}, \mathcal{K})$   
8: end for  
9:  $\mathcal{L}_{\mathcal{D}'} \gets \mathcal{L}_{\mathcal{D}'} + \mathcal{L}_{\mathrm{Total}}(f_{\hat{\theta}_N^{\mathcal{K}}}, \mathcal{K}, \boldsymbol{x}, y)$   
10: end for  
11:  $\theta \gets \theta - \eta \nabla_{|\mathcal{D}'|} \mathcal{L}_{\mathcal{D}'}$   
12: end while
prompted with  $x$  —improves the model's reasoning performance. We evaluate knowledge recovery with a CLM loss and combine the two losses by simple addition, following prior works [24, 77, 78]. The entire change is achieved by redefining the total loss in our outer loop (Equation (4)) as:
$$
\mathcal {L} _ {\text {T o t a l}} (f, \mathcal {K}, \boldsymbol {x}, y) = \mathcal {L} _ {\mathrm {C E}} (f (\boldsymbol {x}), y) + \mathcal {L} _ {\mathrm {C L M}} (f, \boldsymbol {x}, \mathcal {K}) \tag {5}
$$
where  $\mathcal{L}_{\mathrm{CLM}}(f, \boldsymbol{x}, \mathcal{K})$  is the language modeling loss on  $\mathcal{K}$ , as in Equation (3), but this time conditioned on the question  $\boldsymbol{x}$ . The overall process for training RECKONING is depicted in Algorithm 1 and Figure 2. Additionally, we dynamically learn a per-step-per-layer learning rate to replace the shared constant learning rate in the inner loop. We give more details in Appendix D.
Setup We conduct our experiments on three datasets focusing on multi-hop logical reasoning over natural language knowledge: ProofWriter [74], which measures the model's ability to emulate reasoning over facts and rules expressed in natural language; CLUTRR-SG [29], which is generated from the CLUTRR [72] benchmark, a logical reasoning task that involves reasoning over family relationships between entities grounded in first-order logical proofs; and FOLIO [30], a reasoning benchmark with first-order logical reasoning problems written by expert annotators based on real-world knowledge. Each problem in these datasets requires multiple reasoning hops to answer.<sup>1</sup>
<sup>1</sup>In ProofWriter, the number of reasoning hops is called the proof depth. To unify the presentation of the results, we use the term "hop" to describe the number of reasoning steps for both datasets.
We compare our method against the following baselines: (1) a fine-tuned model that performs a forward pass on only the question without access to the knowledge (No-Facts), (2) a fine-tuned model that performs a forward pass on only the knowledge without access to the question (No-Question), (3) a model trained using RECKONING with random knowledge that is not relevant to the questions (Random-Facts), and (4) an ICR baseline that concatenates the knowledge  $\mathcal{K}$  with the question  $x$  in a single context and is trained using supervised learning to predict the answer (FT-ICR). Our first three baselines sanity-check whether any surface-level patterns in the questions and facts can be exploited to make accurate predictions. The last baseline compares RECKONING to the conventional way of reasoning with language models. Unless stated otherwise, we use the GPT-2-small [60] model ( $\sim 124\mathrm{M}$  parameters) as our initialization and refer by RECKONING to our method trained with the multi-task objective. We compute each score from the average across three different runs. For more details on the implementation, datasets, and examples, see Appendix A and Appendix C.
4.1 Multi-hop Reasoning Performance
Main Results We first evaluate whether RECKONING learns to perform reasoning in the base setting. A model is given a set of supporting facts (without distractors) and a question (or hypothesis) as input and begins by performing a few CLM learning steps on the facts. Then, the updated model reads only the question and generates an answer. To answer correctly, the model must reason over both facts and the question, meaning it must encode the facts during the inner loop such that multi-hop reasoning can be performed over them later.
We train our models and the fine-tuned ICR (FT-ICR) baselines with both the
Method	ProofWriter	CLUTRR-SG
2-h	3-h	5-h	2-h	4-h	6-h
No-Facts	64.1	63.0	64.2	0.0	8.8	8.9
No-Question	66.2	67.0	65.2	35.7	36.4	28.7
Random-Facts	64.1	63.0	64.2	0.0	1.3	2.5
FT-ICRST	98.4	98.8	97.8	97.4	91.3	89.1
FT-ICRM	99.4	99.2	99.6	98.1	96.9	90.3
RECKONINGST	98.3	98.3	99.1	96.0	90.2	91.2
RECKONINGMT	99.5	99.7	99.8	98.3	97.6	94.8
Table 1: Label accuracy of RECKONING on ProofWriter and CLUTRR-SG, compared to FT-ICR baselines where the supporting facts are given as part of the input. MT marks models trained with the multi-task objective, which optimizes both question-answering and knowledge memorization.
single-task  $(\mathcal{L}_{\mathrm{CE}})$  and multi-task  $(\mathcal{L}_{\mathrm{CE}} + \mathcal{L}_{\mathrm{CLM}})$  objectives. For multi-task (MT) training, the model learns to answer the question and generate its relevant knowledge in the outer loop. Table 1 shows the evaluation results on question answering (or hypothesis classification). For all hop numbers in ProofWriter and CLUTRR-SG, multi-task RECKONING outperforms the best result of all baselines (consistently obtained by multi-task FT-ICR) by an average of  $1\%$ . We conclude that RECKONING can effectively solve reasoning problems through its updated parametric knowledge and do so better than existing baselines. The multi-task objective is crucial for this success: not only is RECKONING's performance consistently higher (by an average of  $2.8\%$  over the two datasets and their hop counts) when using the multi-task rather than single-task (ST) objective, but it also under-performs both FT-ICR baselines when trained with only the single-task objective. The multi-task objective also improves FT-ICR consistently (average  $1.8\%$ ), though it is not enough to beat the multi-task RECKONING. In all further experiments, we consider only RECKONING and FT-ICR with a multi-task objective.
Generalizing to Longer Reasoning Chains Our first experiments assume an alignment between the number of reasoning hops in the questions in the training and test set. However, we may not be able to train on all  $n$ -hop reasoning questions we encounter in the wild, and we rarely know the number of reasoning hops in a question a priori. Consequently, we also measure the generalization capacity of our model to questions with hop numbers unseen during training. We compile interpolation (fewer hops than the train set) and extrapolation (more hops than the train set) test sets from the CLUTRR-SG dataset. Again, we train models individually on 2-hop, 4-hop, and 6-hop examples and evaluate these three sets of models on the test sets, which contain 2-10-hop reasoning questions. Figure 3 shows that both RECKONING models and ICR baselines retain high performance on the interpolation test sets but exhibit decreasing performance as the number of hops increases. Importantly, though, RECKONING outperforms FT-ICR on all test sets regardless of the number of training hops, with the highest difference being more than  $10\%$  in every training setting ( $15\%, 30\%, 10\%$ , respectively). These performance gains when testing on extrapolation data suggest that training with RECKONING better generalizes to examples with high OOD hop counts than in-context reasoning (ICR).
Figure 3: System generalization evaluation on CLUTRR-SG. From left to right, the models are trained on 2-hop, 4-hop, and 6-hop CLUTRR-SG data portions. We evaluate the model on 2-10 hop test sets. The higher the hops, the more facts a question has, and the more difficult that question is.
Figure 5: Robustness under distractors for ProofWriter. Each of the three plots corresponds to training and testing on a subset of questions in ProofWriter with a different number of hops (2,3,5-hops). Each bar corresponds to the number of distractors in the knowledge sets for those questions.
Does RECKONING's performance depend on the number of inner loop gradient steps? In RECKONING, the model performs multi-hop reasoning over facts by encoding facts using multiple gradient steps in the inner loop optimization (§3). Naturally, this process prompts the question of whether there is a correlation between the number of reasoning hops and the number of gradient steps needed to reliably encode the knowledge (i.e., problems with more reasoning hops require more gradient steps in the inner loop to encode the facts). In Figure 4, we show for CLUTRR-SG that as the number of inner loop steps increases, the label accuracy of the outer-loop task also increases. Furthermore, when considering the performance gains for reasoning with 6 inner loop steps (i.e., knowledge encoding) as opposed to one, we observe that this gap is much more pronounced for 4-hop (42.3%) and 6-hop (34.7%) reasoning than
it is for 2-hop reasoning  $(5.9\%)$ . These results show that problems requiring more hops of reasoning also greatly benefit from more steps of inner loop knowledge encoding.
Figure 4: Multi-hop reasoning performance as a function of the number of inner loop steps (x-axis), with each line focusing (by training and testing) on CLUTRR-SG with a different number of hops.
4.2 Reasoning with Distractors