SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In cases where multiple questions must be answered about the same knowledge set, some knowledge that is relevant to one question will likely be irrelevant to another question. For example, in Table 7, the fact "Charlie is White." is not needed to answer the question "Harry is red?". Thus, it is important to evaluate the robustness of RECKONING when there exists irrelevant information (i.e., distractors) in the knowledge set. In this experiment, we analyze RECKONING's ability to focus on the correct knowledge and ignore distractors when answering questions. We use ProofWriter as the evaluation dataset since it already has a setting with distractors included in the knowledge. For systematic analysis, we gradually add distractors to the context (starting from 2 and finishing at all possible
distractors, of which there are an average of 7 per question). We train RECKONING and the baseline using the multi-task objective, where the model must (1) recall all of the facts and rules relevant to the question and (2) predict the conclusion based on the correct knowledge. In this case, we adapt training such that for each question  $x$ , the outer-loop (Equation (5)) CLM loss is only computed with respect to the relevant facts from  $\mathcal{K}$ , thereby learning to recall only relevant facts during training.
In Figure 5, we see that RECKONING's performance is consistently more robust under distractors than the FT-ICR baseline. When we include all of the distractors in the context, RECKONING achieves a significantly higher average label accuracy (82.5%) across hops than the baseline (70.9%), as computed by the average of the 3 considered hop depths. Additionally, compared to performance with no distractors, RECKONING's performance only drops 17.1% while the baseline performance drops 28.6%, thereby exhibiting a better ability to disentangle the correct knowledge from the distractors.
Finally, we also explore RECKONING's generalizability to models with a larger parameter size. We scale up the language model we used, GPT-2-small (124M), to GPT-2-XL (1.5B) by adopting a parameter efficient finetuning method LoRA [34]. For simplicity, we only evaluate the models on the most difficult settings, i.e., ProofWriter-5-hop with all the distractors. With GPT-2-XL-LoRA, in-context reasoning achieves  $65\%$  accuracy on the test set, while our RECKONING model achieves  $70.2\%$  accuracy, a  $5\%$  performance gain. This result suggests that RECKONING's advantages in the presence of distractors hold even as models scale in size.
4.3 Generalization to Real-World knowledge
To investigate how generalizable our method is to real-world knowledge beyond the synthetic setting, we evaluate RECKONING on a more real-world multi-hop logical reasoning task, FOLIO [30], and report the result in Table 2. The dataset has a rich vocabulary, diverse logic patterns, and abundant
language variations. It has been shown to challenge LLMs in both supervised fine-tuning and in-context learning settings. We fine-tune the GPT-2 model following the in-context reasoning setting as the baseline. As before, we train the GPT-2 model and RECKONING using the multi-task objective. We also compare to more advanced baselines, including GPT-3.5 (text-davinci-003 [56]) and ChatGPT(gpt-3.5-turbo 2), two popular large language models with around 175B parameters. For these two large models, we evaluate both in the zero-shot and few-shot settings. In the few-shot setting, we prompt the model with 8 single-task examples randomly sampled from the training set to perform in-context learning. We find that RECKONING's performance (which is initiated
Model	Acc.
Random	33.3
GPT-2	53.3
GPT-3.5 (text-davinci-003) 0-shot	45.1
GPT-3.5 (text-davinci-003) 8-shot	52.9
ChatGPT (gpt-3.5-turbo) 0-shot	40.0
ChatGPT (gpt-3.5-turbo) 8-shot	42.6
RECKONING	54.9
Table 2: Evaluation results on FOLIO. We compare RECKONING against the FT-ICR baseline with GPT-2 and two popular large language models.
here from GPT-2) is better than the GPT-2 in-context reasoning baseline. Compared to the two advanced large language models, RECKONING outperforms them by a significant margin (12% 0-shot and 7% 8-shot). We conclude that RECKONING is effective and significantly benefits reasoning tasks using real-world knowledge.
4.4 Run-time Analysis
One of the advantages of RECKONING is the ability to memorize a large set of knowledge  $\kappa$  and answer multiple related questions about that knowledge at a little extra cost per question. Specifically, in contrast to ICR, RECKONING can encode  $\kappa$  once and answer multiple questions without needing to reprocess it for each question asked. To test whether RECKONING could be a more efficient method for inference in this setting, we measure the wall-clock time (in seconds) of the complete inference pipeline of RECKONING vs. ICR. For this experiment, we use a synthetic reasoning dataset in which  $\kappa$  is a sequence of random letters, and the question  $x$  asks for the most frequent letter in the context. The total number of tokens in each example is 1024: 7 for  $x$ , 1 for the answer, and the remaining 1016 for  $\kappa$ , broken into 8 "facts".
The FT-ICR baseline receives a sequence including all 8 facts and the question. In contrast, RECKONING receives the 8 facts as a batch of eight segments of 127 tokens and encodes them in parallel in the inner loop. In the outer loop, the model only receives the question or a batch of questions. We focus on two settings: (1) inference time for a single question and (2) inference time when answering multiple questions. In the multiple-question setting, we set the number of questions to 18 (the same as in ProofWriter). For RECKONING, the inference process includes the inner-loop knowledge encoding and the final forward pass to encode the question. We set the number of inner loop gradient steps to 1 and 4. In Table 3, we see that when answering a single question, RECKONING does not perform inference faster than in-context reasoning. However, RECKONING shows significant advantages under a multi
question setting. Both the 1-step inner loop and the 4-step inner loop are faster than the baseline. Since RECKONING encodes the knowledge in model parameters, it does not need to reprocess the knowledge for a related question and is more efficient. We run this experiment on 1 RTX 3090 GPU. $^{3}$
Model	Wall-clock Time (s)
Single question
FT-ICR	0.1887
RECKONING1step	0.2532
RECKONING4step	0.9664
Multiple questions (18)
FT-ICR	2.0436
RECKONING1step	0.6228
RECKONING4step	1.4839
4.5 Memorizing Knowledge
In Table 1, we saw that training RECKONING with a multi-task (MT) outer loop objective improved over training with the single-task (ST) objective, potentially because the MT objective improves the model's ability to memorize the knowledge in the inner loop. To validate our hypothesis, we analyze RECKONING's performance in reproducing memorized knowledge. First, we show in Table 4 the inner loop average loss
$(\mathcal{L}_{\mathrm{CLM}})$  and average change  $(\Delta \mathcal{L}_{\mathrm{CLM}})$  (from first inner loop evaluation to last) on validation examples from the 5-hop ProofWriter data. We see that the average inner loop loss for RECKONINGST is much higher than RECKONINGMT, and indeed starts out much higher as well. This shows that the ST outer loop objective, which optimizes the model only for question answering, does not learn to encode the knowledge in the inner loop by memorizing it. In contrast, the MT objective forces the model to learn to memorize the knowledge, too: we observe that RECKONINGMT minimizes the inner loop loss as it processes the knowledge. This pattern is also shown in the average inner-loss difference  $(\Delta \mathcal{L}_{\mathrm{CLM}})$ : the inner loop loss decreases more after the gradient updates when trained with the MT objective. Next, we report in Table 5 the model's ability to reproduce memorized facts correctly under a multi-task setting, as measured by an exact match score between the reproduced facts and the gold facts.4 We evaluate on the ProofWriter dataset both with and without distractors in the context and compare the results to the FT-ICR baseline.
Table 3: Wall clock run-time, in seconds, of the fine-tuned ICR baseline and RECKONING.
Method	CLM	ΔCLM
RECKONINGST	10.74	9.55
RECKONINGMT	0.167	12.61
Table 4: Average inner loop validation loss: final  $(\mathcal{L}_{CLM})$  and difference from start to finish  $(\Delta \mathcal{L}_{CLM})$
Method	ProofWriter	ProofWriterdistractor
2-h	3-h	5-h	2-h	3-h	5-h
FT-ICRMT	99.8	99.0	98.7	42.3	50.3	55.6
RECKONINGMT	98.9	98.6	98.2	71.2	74.4	75.1
Table 5: Exact match score for reproducing memorized knowledge. In contrast to in-context reasoning, RECKONING does not have direct access to the knowledge.
RECKONING<sub>MT</sub> no longer has direct access to them. When the context includes distractors, both RECKONING and FT-ICR struggle to identify and reproduce only the relevant facts. However, the per
The results show that RECKONING<sub>MT</sub> can successfully (average exact match score of 99.3%) recover the relevant facts from its model parameters when the context does not include any distractors. Note that this is comparable to the FT-ICR baseline, for which the task is much easier as it can directly attend to and copy the facts from input, while
<sup>3</sup>We perform this experiment in a limited setting and do not handle the case where hidden states could be cached for the forward pass of in-context reasoning, likely speeding up multi-question inference [12].
This is done by prompting the model with the question and comparing its output (after its answer to the question) to the concatenation of all facts. The model is able to produce these facts in the expected order due to an implementation detail: they are numbered and labeled when given to the inner loop.
formance for FT-ICR (average  $49.4\%$  ) drops far below that of RECKONING  $(73.6\%)$  , demonstrating that RECKONING is much better at disentangling the relevant knowledge from the distractors.
Finally, we show that RECKONING with a multi-task objective is also more robust to distractors as it trains the model to reproduce only the facts that would be relevant to a particular question we ask in the outer loop. As in Section 4.2, we use the ProofWriter dataset and, for each question, add all the distractors to the context. We train the model using the multi-task objective and report the label accuracy. While in Table 1, we originally saw a  $\sim 1\%$  improvement from training with a multi-task objective on ProofWriter with no distractors, we see a much more significant performance gap in Figure 6 ( $\sim 18.2\%$ ) when distractors are available. We also note that the performance of the single-task model is essentially random (see the Random-Facts baseline from Table 1). By learning how to memorize knowledge in the inner loop to recall relevant facts in the outer loop, the model also learns how
Figure 6: RECKONING performance when trained with a single-task and a multi-task objective under distractors. When trained with the multi-task objective, the model learns to memorize and reason over the relevant facts.
to encode facts more robustly over them.
Logical Reasoning Datasets and Benchmarks As a central building block of human cognition and intelligence [28], logical reasoning has been a long-pursued topic in the field of AI [2, 9, 13, 17, 46, 48, 55, 73]. Logical reasoning, in general, can be categorized in a trichotomy of deductive, inductive, and abductive reasoning [26]. Multiple datasets have been published that evaluate neural models' ability on these three types of logical reasoning [6, 17, 72]. Initially, logical reasoning tasks focused on hypothesis classification, where, given a theory consisting of multiple facts and rules, a model would determine whether the hypothesis was correct. Recently, transformer-based language models have been directly used to solve this task in synthetic [17, 66], real-world [30], and adversarial [27, 62, 68] settings. However, simply predicting whether the hypothesis is valid does not elucidate whether the model correctly reasons over the provided knowledge. To better analyze and interpret the reasoning process of language models, new tasks focus on generating the valid proof that explains the model's decision [20, 74]. Our proposed method, RECKONING, is optimized for the hypothesis classification reasoning task and evaluates on many of these datasets [29, 74, 30].
Logical Reasoning over Natural Language Historically, automatic logical reasoners used symbolic systems and formal languages as a knowledge representation [1, 43, 49, 52, 55, 81]. However, these systems were hard to scale up due to the knowledge-acquisition bottleneck and the brittleness of formal representation [35, 84]. With recent advances in transformer-based language modeling [76] and self-supervised pre-training [21, 60, 61], a novel paradigm for logical reasoning emerged, where pre-trained language models (PLMs) could be used as soft reasoners over knowledge expressed in natural language. Natural language as a knowledge representation allowed PLMs to handle raw input with diverse formats [15, 32], resulting in PLMs being applied to various types of deductive [17], abductive [6], and inductive [29] reasoning tasks. However, language models as soft reasoners also showed structural weaknesses, as their performance dropped on complex logical operations [13, 80], and their reasoning process was not interpretable [45, 65]. Consequently, a new line of work uses neuro-symbolic methods to combine the best of both language models and symbolic reasoning [7, 14, 36, 41, 44]. Specifically, the interpretability gap motivated modular and step-wise reasoning systems that use PLMs as intermediate modules [33, 59, 67, 69, 75, 83] to generate reasoning steps (e.g., proofs). In contrast to these works, our method RECKONING dynamically encodes natural language knowledge into the model parameters, thereby reasoning by mixing contextual knowledge with pre-encoded parametric knowledge and allowing the model to determine a conclusion based on its updated parametric knowledge.
Model Editing While our motivations are grounded in research on machine reasoning, our methods are more often used in the area of model editing. Model editing is a method to edit a model's parameters to correct its errors or update the model. Several works propose hypernetwork-based methods to edit knowledge in a model by predicting updates conditioned on new factual statements [31] or transforming the gradients from new provided facts [54] to make local edits to a model. Other
approaches focus on more direct edits of model behavior, such as directly modifying neuron outputs [19, 85], localizing distinct feed-forward layers that are responsible for factual recall, and modifying these weights [50], and performing weight updates across multiple layers to perform simultaneous edits [51]. Similarly, our method also rapidly edits the model parameters to add knowledge. However, our bi-level framework optimizes model edits for the reasoning task in the outer loop, allowing the model to learn to quickly memorize knowledge that can support the model's reasoning ability.
Language Models as Knowledge Bases Our work learns to reason by dynamically encoding contextual knowledge in the parameters of language models before answering questions about them. Previous studies have found that LLMs can store real-world facts learned during pre-training [5, 10, 11, 50, 64]. Learning these facts during pre-training allows language models to be prompted [39, 58, 71, 86] or adapted [8, 37, 38, 63] to produce these facts on-demand. However, LLM knowledge is latent and hard to identify or control. The model generation is sensitive to specific words or phrases. LLMs emit knowledge encoded in the parameters only when prompted appropriately [10, 18, 23, 57]. It is also difficult to inject or update knowledge for LLMs [50], and the memorization of knowledge in LLMs is not optimized toward their reasoning ability. In our work, we seek to find a way to add knowledge to LLMs in a controllable and adaptive way that can benefit downstream reasoning applications.
We present RECKONING, a bi-level learning framework for multi-hop reasoning that encodes knowledge verbalized using natural language into a model's parameters through gradient updates. During training, the inner loop encodes the contextual knowledge into the model parameters by backpropagating a language modeling loss. In the outer loop, given only the question as input, the model solves reasoning problems using the memorized knowledge. Through bi-level optimization, RECKONING finds a set of meta-parameters that allows it to perform quick knowledge-based updates for reasoning. Our experiments show that RECKONING learns to reason only by relying on its parametric knowledge after the external knowledge has been encoded. Using a multi-task objective that jointly optimizes reasoning and knowledge memorization in the outer loop, RECKONING outperforms ICR baselines that are trained to encode external knowledge as part of the context. Through our analysis, we show that RECKONING is more generalizable to problems with longer reasoning chains, less susceptible to irrelevant distractor knowledge, and that RECKONING is more efficient than the baseline when answering multiple questions that require common knowledge.
We thank Shikhar Murty and Christopher Manning for helpful discussions in crafting ideas for this project. We also gratefully acknowledge the support of the Swiss National Science Foundation (No. 215390), Innosuisse (PFFS-21-29), the EPFL Science Seed Fund, the EPFL Center for Imaging, Sony Group Corporation, and the Allen Institute for AI.
[1] Lasha Abzianidze. A tableau prover for natural logic and language. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pages 2492-2502, Lisbon, Portugal, September 2015. Association for Computational Linguistics.
[2] Lasha Abzianidze. LangPro: Natural language theorem prover. In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 115–120, Copenhagen, Denmark, September 2017. Association for Computational Linguistics.
[3] Antreas Antoniou, Harrison Edwards, and Amos Storkey. How to train your MAML. In International Conference on Learning Representations, 2019.
[4] Sungyong Baik, Myungsub Choi, Janghoon Choi, Heewon Kim, and Kyoung Mu Lee. Meta-learning with adaptive hyperparameters. ArXiv, abs/2011.00209, 2020.
[74] Oyvind Tafjord, Bhavana Dalvi, and Peter Clark. ProofWriter: Generating implications, proofs, and abductive statements over natural language. In Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021, pages 3621-3634, Online, August 2021. Association for Computational Linguistics.
[75] Oyvind Tafjord, Bhavana Dalvi Mishra, and Peter Clark. Entailer: Answering questions with faithful and truthful chains of reasoning. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 2078-2093, Abu Dhabi, United Arab Emirates,
December 2022. Association for Computational Linguistics.
[76] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[77] Xinyi Wang, Yulia Tsvetkov, and Graham Neubig. Balancing training for multilingual neural machine translation. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 8526-8537, Online, July 2020. Association for Computational Linguistics.
[78] Zirui Wang, Zachary C. Lipton, and Yulia Tsvetkov. On negative interference in multilingual models: Findings and a meta-learning treatment. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 4438-4450, Online, November 2020. Association for Computational Linguistics.
[79] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumont, Clement Delangue, Anthony Moi, Pierrick Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Huggingface's transformers: State-of-the-art natural language processing, 2020.
[80] Hitomi Yanaka, Koji Mineshima, Daisuke Bekki, and Kentaro Inui. Do neural models learn systematicity of monotonicity inference in natural language? In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 6105-6117, Online, July 2020. Association for Computational Linguistics.
[81] Hitomi Yanaka, Koji Mineshima, Pascual Martínez-Gómez, and Daisuke Bekki. Acquisition of phrase correspondences using natural deduction proofs. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers), pages 756-766, New Orleans, Louisiana, June 2018. Association for Computational Linguistics.
[82] Kaiyu Yang, Jia Deng, and Danqi Chen. Generating natural language proofs with verifier-guided search. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 89-105, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.
[83] Kaiyu Yang, Jia Deng, and Danqi Chen. Generating natural language proofs with verifier-guided search, 2022.
[84] Zonglin Yang, Xinya Du, Rui Mao, Jinjie Ni, and Erik Cambria. Logical reasoning over natural language as knowledge representation: A survey, 2023.
[85] Yunzhi Yao, Shaohan Huang, Li Dong, Furu Wei, Huajun Chen, and Ningyu Zhang. Kformer: Knowledge injection in transformer feed-forward layers, 2022.
[86] Wenhao Yu, Dan Iter, Shuohang Wang, Yichong Xu, Mingxuan Ju, Soumya Sanyal, Chenguang Zhu, Michael Zeng, and Meng Jiang. Generate rather than retrieve: Large language models are strong context generators. In The Eleventh International Conference on Learning Representations, 2023.
ProofWriter The ProofWriter [74] dataset has 500k pairs of questions, answers, and proofs over natural-language rule bases. Each example in the dataset contains a set of facts, a set of rules, a hypothesis, and a label indicating whether the hypothesis is true, false, or unknown. The dataset comprise five datasets named D0, D1, D2, D3, D5, each with 100k examples. Each dataset's questions require reasoning up to depths  $D$  ( $D = 0, 1, 2, 3, 5$ ) to determine their answers. In our experiments, we only focus on the datasets that require more reasoning depths (D2, D3, D5). We show an example from the dataset in Table 7. In these datasets, a set of facts and rules are mapped to 18 questions, where the questions can be answered based on a subset of the facts and rules. Thus, some of the facts or rules can be irrelevant to some questions, and we call them distractors in Section 4.2. In the experiment for knowledge encoding with distractors, we encode all the facts in the model parameters and evaluate its ability to reproduce and reason over the correct facts. We show an example of distractor and relevant knowledge of a question in Table 9. For detailed statistics on the two datasets, please see Table 6.