SlowGuess/NeurIPS_2023_OCR huggingface dataset:
CLUTRR-SG The CLUTRR-SG [29] is an evaluation dataset for inductive reasoning on family relations adapted from the [72] dataset for measuring systematic generalization. Each example in the dataset contains (i) a set of facts representing a family graph  $G = (V, E)$  where nodes  $(V)$  are entities and edges  $(E)$  are the relationships. (ii) a question asking the relationship between two entities  $(v_1, v_n \in V)$ , and (iii) a target relationship  $e^* \in E$  as the answer for the question. The facts are expressed as a list of  $(v_i, e_j, v_k)$  tuples. The two entities in the question are separated by more than one hop in the graph. There are 272 unique entities, 20 relationship types, and nearly 1.5M possible facts in the dataset. Following the authors, we define the difficulty of examples based on the number of family graph edges (i.e., the number of reasoning hops required to determine a relation), in which  $k$  edges ( $k$ -hop) correspond to  $k$  facts. We show an example from the dataset in Table 8.
B In-context Reasoning with Distractors
To motivate the advantage of RECKONING on mitigating interference from distractors, we analyze the performance change of fine-tuned incontext reasoning with and without distractors present in the context of the questions. We define distractors as additional facts or rules present in a question's context that are not directly relevant to the questions. A model should not be able to use only these distractors to answer a question correctly. For an example of distractors in a question's context, please see Table 9. We evaluate the baseline on the ProofWriter dataset since it naturally contains contexts including distractors (Table 9). Recall that we have two training objectives. The single-task objective only trains the model to predict an answer for each question given their contexts. The
Figure 7: Label accuracy of fine-tuned in-context reasoning on questions with and without distractors in the context. With the same questions, adding distractors to contexts significantly lowers the performance of in-context reasoning, both in the single-task and multi-task settings.
multi-task objective (MT) trains the model not only to predict an answer but also to reproduce the correct facts and rules (in contrast to distractors) based on the contexts. We evaluate the baseline on 2, 3, and 5-hop datasets with both training objectives, and we report the average label accuracy across hops in Figure 7. Compared to the baseline's performance without distractors in the context, the performance with distractors decreases significantly. For single-task, the performance drops  $23.2\%$  when adding distractors to the contexts, and the performance with the multi-task objective drops  $28.6\%$ . The results highlight in-context reasoning's high sensitivity to the interference of irrelevant information in the contexts.
We select GPT-2-base [60] as the model for our method and all the baselines. We use the version implemented by the Huggingface Transformers library [79]. All the experiments for RECKONING
Dataset	#Train	#Validation	#Test
CLUTRR-SG (2-hop)	96,012	10,972	3,102
CLUTRR-SG (4-hop)	89,972	10,086	9,946
CLUTRR-SG (6-hop)	90,922	10,290	8,788
ProofWriter (2-hop)	6,996	1,098	2,013
ProofWriter (3-hop)	10,854	1,641	3,057
ProofWriter (5-hop)	18,525	2,553	5,175
Table 6: Dataset splits and statistics for our experiments
Identifier	Content
fact 1	Harry is nice.
fact 2	Fiona is quite Nice.
fact 3	Fiona is round.
fact 4	Fiona is white.
fact 5	Dave is furry.
fact 6	Charlie is white.
rule 1	Furry people are green.
rule 2	Round, green people are red.
rule 3	All red people are white.
rule 4	Nice, round people are furry.
rule 5	If someone is nice, then they are round.
rule 6	If Charlie is round and Charlie is nice, then Charlie is white.
question-answer 1	Harry is red? True
question-answer 2	Harry is not red? False
question-answer 3	Dave is not white? Unknown
Table 7: An example from the dataset ProofWriter. There are 6 facts and 6 rules mapped to three question-answer pairs. Each question can be answered based on the given facts and rules.
are conducted on a cluster with NVIDIA A100 (40GB) GPUs. All the baseline experiments are conducted on a local machine with NVIDIA RTX 3090 GPU (24GB).
Fine-tuned In-context Reasoning We set the train batch size to 16 and train the model for 6 epochs with early stopping based on the validation label accuracy. We set the learning rate to 3e-5 and use the AdamW optimizer with  $\epsilon$  set to 1e-8. We validate the model on the development set for every epoch and select the best checkpoint using the validation accuracy as the metric.
RECKONING In the inner loop, we generally perform 4 gradient steps for lower-hop questions (2, 3, 4-hop) and 5 gradient steps for higher-hop questions (5 and 6-hop). We select the AdamW [47] as the optimizer for the inner loop since the main task is language modeling. The inner-loop learning rate is set to 3e-5 before training, and the algorithm dynamically learns a set of optimal learning rates when converged. In our experiments and analysis, we only report the results from RECKONING with a multi-task objective since its performance is better than the single-task objective. In the outer loop, we also use the AdamW with a learning rate of 3e-5. For both optimizers, we set  $\epsilon$  to 1e-8. We set the train batch size to 2 due to memory limitations. We apply the technique of gradient accumulation and set the accumulation step to 2. We train the model for 6 epochs with early stopping. For each epoch, we validate the model twice: once in the middle and once at the end. We select the best model checkpoint based on the validation label accuracy.
D Adaptive Learning Rate
Prior works [3, 4] show that a fixed learning rate shared across steps and parameters does not benefit the generalization performance of the system. Instead, [3] recommends learning a learning rate for
Identifier	Content
fact 1	C is H's father.
fact 2	Z is J's aunt.
fact 3	J is S's daughter.
fact 4	D is C's father
fact 5	S is B's father.
fact 6	H is Z's son.
question-answer 1	How are D and B related to each other? Grandfather
Table 8: An example of 6-hop reasoning from the CLUTRR-SG dataset.
Identifier	Content
fact 1	Harry is nice.
fact 2	Fiona is quite Nice.
fact 3	Fiona is round.
fact 4	Fiona is white.
fact 5	Dave is furry.
fact 6	Charlie is white.
rule 1	Furry people are green.
rule 2	Round, green people are red.
rule 3	All red people are white.
rule 4	Nice, round people are furry.
rule 5	If someone is nice, then they are round.
rule 6	If Charlie is round and Charlie is nice, then Charlie is white.
question-answer 1	Harry is red? True
Table 9: Example of distractors (black) and relevant knowledge (red) in the ProofWriter dataset.
each network layer and each adaptation step in the inner loop. The layer parameters can learn to adjust the learning rates dynamically at each step. To control the learning rate  $\alpha$  in the inner loop adaptively, we define  $\alpha$  as a set of adjustable variable:  $\alpha = \{\alpha_0,\alpha_1,\dots \alpha_L\}$ , where  $L$  is the number of layers and for every  $l = 0,\ldots ,L$ ,  $\alpha_{l}$  is a vector with  $N$  elements given a pre-defined inner loop step number  $N$ . The inner loop update equation then becomes
$$
\hat {\boldsymbol {\theta}} _ {n + 1, l} ^ {\mathcal {K}} = \hat {\boldsymbol {\theta}} _ {n, l} ^ {\mathcal {K}} - \boldsymbol {\alpha} _ {n} ^ {(l)} \odot \nabla \mathcal {L} _ {\mathrm {C L M}} \left(f _ {\hat {\boldsymbol {\theta}} _ {n} ^ {\mathcal {K}}}, \mathcal {K}\right) \tag {6}
$$
where  $\odot$  is an element-wise product and  $\hat{\theta}_n^{(l)}$  is the parameters for layer  $l$  at the inner step  $n$ . We learn the set of optimal inner loop learning rates  $\alpha^*$  by optimizing the parameters in the outer loop:
$$
\boldsymbol {\alpha} \leftarrow \boldsymbol {\alpha} - \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}} [ \mathcal {K} \boldsymbol {\theta} _ {i} ] N} (\boldsymbol {x}), y\right), \tag {7}
$$
where  $\eta$  is the outer loop learning rate and  $\hat{\theta}$  is the updated parameters from inner loop. Below, we show the final algorithm of RECKONING in Algorithm 2.
Algorithm 2 Dynamic Knowledge Encoding for Reasoning
Require: An example distribution  $p(\mathcal{D})$ , a transformer language model  $f$ , initial meta-parameters  $\theta$ , outer step size  $\eta$ , initial 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 \odot \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:  $\alpha \gets \alpha - \eta \nabla \mathcal{L}_{\mathcal{D}'}$   
12:  $\theta \gets \theta - \eta \nabla \mathcal{L}_{\mathcal{D}'}$   
13: end while
Are dynamic learning rates necessary for RECKONING's performance? Following prior works on meta-learning [3, 4], we dynamically learn a set of per-step-per-layer learning rates for RECKONING. In this ablation study, we analyze whether dynamic learning rates for the inner loop effectively improve the outer loop reasoning performance. Similarly, we fix other experimental settings and set the number of inner loop steps to 4. As Figure 8 shows, when using a static learning rate (i.e., all layers and inner loop steps share a constant learning rate), the performance drops by a large margin (average drop of  $34.2\%$ ). The performance drop becomes more significant on questions requiring more reasoning hops ( $45.5\%$  drop for 4-hop and  $39.5\%$  drop for 6-hop), demonstrating the importance of using a dynamic learning rate in the inner loop of our framework.
Figure 8: We study how much the dynamic learning rate in the inner loop contributes to the outer loop performance. We fix all the hyperparameters except the option of using the dynamic or fixed learning rate. We conduct the analysis using the CLUTRR-SG dataset since it is more complex and difficult (lower random performance).
E Experiments with Large Language Models
Method	ProofWriter	ProofWriter_distractor	CLUTRR-SG
2-h	3-h	5-h	2-h	3-h	5-h	2-h	4-h	6-h
GPT-3.5 0-shot	58.4	56.4	53.7	49.1	47.1	45.3	35.6	16.0	18.5
GPT-3.5 8-shot	78.0	82.4	80.1	58.7	57.2	54.5	39.0	18.5	20.8
RECKONINGMT	99.5	99.7	99.8	79.8	83.7	84.0	98.3	97.6	94.8
Table 10: Label accuracy of RECKONING on ProofWriter and CLUTRR-SG compared against a popular Large Language Model (LLM), GPT-3.5. We prompt GPT-3.5 in the zero-shot setting and also the 8-shot in-context learning setting. Models with MT are trained with the multi-task objective in the outer loop.
Recently, Large Language Models (LLMs) with large parameter sizes learned from human preferences have shown remarkable performance in language understanding and generation. These LLMs are powerful zero-shot and few-shot reasoners. Recent works find that LLMs learn to perform multi-step reasoning by first generating new reasoning chains and then predicting the answers. In this experiment, we benchmark the performance of a popular new LLM, GPT-3.5, on the two multi-hop reasoning datasets we used in our paper. We first evaluate GPT-3.5's zero-shot reasoning performance in predicting the correct answers. As Table 10 shows, zero-shot prompting GPT-3.5 significantly underperforms RECKONING's performance. GPT-3.5's performance improves on ProofWriter without distractors but still is behind the performance of RECKONING. When distractors are present in the context, RECKONING performs much better than zero-shot and few-shot GPT-3.5 prompting. This
highlights RECKONING's strength in disentangling irrelevant information from useful knowledge, an ability that even powerful LLMs like GPT-3.5 lack.