SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jack Lanchantin*
Shubham Toshniwal*
Jason Weston
Arthur Szlam
Sainbayar Sukhbaatar
Large language models have been shown to struggle with multi-step reasoning, and do not retain previous reasoning steps for future use. We propose a simple method for solving both of these problems by allowing the model to take Self-Notes. Unlike recent chain-of-thought or scratchpad approaches, the model can deviate from the input context at any time to explicitly think and write down its thoughts. This allows the model to perform reasoning on the fly as it reads the context and even integrate previous reasoning steps, thus enhancing its memory with useful information and enabling multi-step reasoning. Experiments across a wide variety of tasks demonstrate that our method can outperform chain-of-thought and scratchpad methods by taking Self-Notes that interleave the input text.
Figure 1: [top] Vanilla language models directly generate the answer (A) given the context and the question (Q). [middle] Scratchpad and Chain-of-Thought methods encourage the model to generate reasoning tokens before answering the question, but only after it has read the entire context. [bottom] Self-Notes (ours) allows the model to generate multiple internal reasoning notes that interleave the input context and question.
Transformers [1] and similar variants have reshaped the field of machine learning with impressive results on sequence-based tasks [2]. Notably, large language models (LMs) such as GPT-3 [2] use transformers and are capable of solving various complex natural language tasks such as question answering. However, it's increasingly evident that there are still limitations to these models. Namely, transformers are limited in their ability to perform multi-step computations or store intermediate results due to the lack of an explicit internal dialogue or scratchpad [3, 4, 5].
When an LM is used for question answering (QA), it is typically fed a context prompt containing factual information along with a question, and then the model generates the answer directly, as shown in Fig. 1 (top). However, this autoregressive "one-step" approach struggles with multi-step reasoning
*Equal Contribution. Correspondence to sainbar@meta.com
tasks [6, 7, 8, 4]. We argue that this arises from the fact that vanilla LMs have a fixed computation budget for each token, and do not have the option to "think" more depending on the current context.
Recently, Chain-of-Thought and Scratchpad methods [3, 4, 9, 10] encourage the model to generate reasoning tokens or explain their answer one step at a time, leading to improved reasoning capabilities. These extra tokens are added before answering the question, but after it has read the full context and question, therefore postponing all "thoughts" until the end, as illustrated in Fig. 1 (middle).
In addition to the "one-step" problem, transformers lack a recurrent memory for state-tracking and solving highly nonlinear tasks [11], something that recurrent predecessor models such as the LSTM [12] are well equipped for. Modifications to the feed-forward transformer architecture that use a recurrent mechanism improve state-tracking results [11, 13, 14], but still use a fixed computation amount for a given prompt.
In this paper, we propose an approach that simultaneously makes the challenges in multi-step reasoning and state-tracking memory more tractable. Our method, "Self-Notes", allows the LM to deviate from the context prompts at any time to generate explicit reasoning tokens. Unlike a scratchpad or chain-of-thought, the model can interleave generated tokens with the input context and question, as demonstrated in Fig. 1 (bottom). Such Self-Notes can act as both explicit intermediate reasoning steps and working memory for state-tracking. Overall, we see the ability of a model to think as it reads, storing those thoughts for later use, as an important aspect of an intelligent agent, similar to how humans read [15]. Consider answering a question after reading the contents of a book. One would not want to re-read and re-reason over the whole book again to answer a second question. This issue is not addressed by either vanilla language models or chain-of-thought approaches.
To give an example, when reading a text, at any time there may be a reasoning step that requires combining two facts. The resulting inference can be written into a Self-Note and used for future reasoning, thus acting as an intermediate reasoning step. For example, given "Alice has the box" and "Alice is at the park" one can infer "The box is at the park" and write it to a Self-Note, which can be further combined with a future statement "The key is in the box" to conclude that "The key is at the park". Additionally, the Self-Note can act as a form of working memory because the model can write the latest state of an entity as new tokens while it reads the context. For example, in a programming environment, assume  $x = 5$  initially, and then  $x$  gets incremented by 1, the model can write the new state  $x = 6$  and add it to its working memory.
We test our method on seven text datasets designed to evaluate multi-step reasoning and state-tracking: a proposed synthetic Toy-Story task, two synthetic program evaluation tasks [11, 16], two real-world chess game tasks [17], and two math word problem tasks previously used to test chain-of-thought prompting, MultiArith and GSM8K [18, 19]. Across these tasks, we consider four different paradigms of learning Self-Notes: supervised, semi-supervised, unsupervised, and few-shot prompted. Across all tasks and learning paradigms, our method outperforms both an LM which does not do any explicit reasoning, as well as chain-of-thought or scratchpad baselines.
Let us consider an autoregressive transformer model  $\mathcal{M}$  that predicts the next token in a sequence,
$$
x _ {t + 1} = \mathcal {M} (x _ {1}, \dots , x _ {t}).
$$
Such a model,  $\mathcal{M}$  is the foundation of many tasks like language modeling and question answering. In such tasks, the model is given a context  $C = \{x_{1},\dots,x_{t}\}$  and potentially a question  $Q$  as input and asked to generate  $A$  which is the sequence of next words or an answer to a question.
Our Self-Notes method expands the capability of  $\mathcal{M}$  by allowing it to enrich context  $C$  with "note tokens"  $n_i$  before producing the final output  $A$ . Note tokens share the same vocabulary as input tokens, but they are generated by the model itself. Self-Notes generated in this way can interleave with the context tokens and therefore can be used for writing down a newly inferred fact or tracking variable values.
At inference time, while processing input tokens  $x_{t} \in C$  one by one, the model can start taking a note by generating a token that belongs to a predefined set of start tokens  $N_{\mathrm{sta}}$ . In other words, at any point in the input text, if the next most probable token is in  $N_{\mathrm{sta}}$ , then the model can autoregressively generate itself a note. A note ends when the model generates an end token  $n_{i} \in N_{\mathrm{end}}$ , or after a
Table 1: Input-Output pairs for the Vanilla, Scratchpad, and Self-Notes method on Toy-Story and Algorithmic tasks. The input consists of the input context and the question, the answer is to be generated by the model. The highlighted reasoning steps are not given at test time and is to be generated by the model.
Task	Vanilla	Scratchpad	Self-Notes
Toy-Story	Mary has the ball.The ball is inside the box.The key is inside the box.Q: Who has the key?Mary has the key.	Mary has the ball.The ball is inside the box.The key is inside the box.Q: Who has the key?[Q: Who has the box?Mary has the box.Q: Who has the key?Mary has the key.]Mary has the key.	Mary has the ball.The ball is inside the box.SQ: Who has the box?Mary has the box.The key is inside the box.SQ: Who has the key?Mary has the key.Q: Who has the key?Mary has the key.
Algorithmic	e = 3 ; e ++ ;i = 3 ; if i < e : e ++ ;print e e = 5 ;	e = 3 ; e ++ ;i = 3 ; if i < e : e ++ ;print e[e = 3 ; e ++ ;print e e = 4;i = 3 ; if i < e : e ++ ;print e e = 5;] e = 5 ;	e = 3 ; e ++ ;print e e = 4;i = 3 ; if i < e : e ++ ;print e e = 5;print e e = 5 ;
fixed number of tokens are generated. Once the note ends, the generated note tokens are appended to the context where the start token was generated, and the model continues to process the rest of the input tokens. For example, a context  $C = \{x_{1}, x_{2}, x_{3}, x_{4}\}$  can be enriched to become  $\{x_{1}, x_{2}, n_{1}, n_{2}, n_{3}, x_{3}, x_{4}\}$  if the start token is generated after  $x_{2}$ :
$$
n _ {1} = \mathcal {M} (x _ {1}, x _ {2}) \in N _ {\mathrm {s t a}}, \quad n _ {2} = \mathcal {M} (x _ {1}, x _ {2}, n _ {1}) \notin N _ {\mathrm {e n d}}, \quad n _ {3} = \mathcal {M} (x _ {1}, x _ {2}, n _ {1}, n _ {2},) \in N _ {\mathrm {e n d}}.
$$
By repeating this mechanism, the context  $C$  can be enriched with multiple notes at different locations. An overview of our method is shown in Figure 1 (bottom). A detailed example of how the model takes Self-Notes at inference is shown in Figure 5 (bottom).
The model can use notes as a form of working memory by writing information that might be useful in the future. It can also use a note as an intermediate reasoning step by inferring new facts as it reads. In particular, it can ask a question and answer it within it. This is useful in multi-step reasoning where a final question requires answering multiple sub-questions. Unlike implicit reasoning occurring internally within  $\mathcal{M}$ , Self-Notes are fed back to the model, making it available to future reasoning steps. This feedback loop also allows the model to overcome the limitation of transformers as a feedforward network [11], making it possible to do state-tracking. We explain four different paradigms to encourage the model to write Self-Notes as follows.
Supervised Self-Notes. One way to train  $\mathcal{M}$  to generate useful notes is to use supervised learning on data that is enriched with "ground-truth" Self-Notes interspaced within the context. This training procedure is simple as we just have to train  $\mathcal{M}$  on this enriched data using the standard LM training loss. After training, we can use  $\mathcal{M}$  to generate Self-Notes, so we can apply it to test data that does not contain any Self-Notes or reasoning labels.  $\mathcal{M}$  can generate a Self-Note at test time by predicting the next token in the context to be from  $N_{\mathrm{sta}}$ .
Semi-supervised Self-Notes. We also consider a semi-supervised setting where only a subset of the training samples have ground truth Self-Notes. In this case, we prepend a special token  $s$  to training samples without Self-Notes and train all samples with the standard LM loss:  $C = \{s, x_1, \dots, x_t\}$ . As a result, the model is conditioned to generate Self-Notes during test time because the test context does not contain the special token  $s$  prefix. This signals to the model that it should do extra reasoning and generate Self-Notes.
Unsupervised Self-Notes. We introduce a method for utilizing Self-Notes when no ground truth note is available for training. This method relies on the fact that when the model is trained using the LM loss on all tokens in a QA task, it learns to not only generate answers but also questions. We leverage this property by letting the model generate its own questions and insert their answers as Self-Notes (i.e., interleaved throughout the context) during test time.
2We call this method "unsupervised" or "unlabeled" Self-Notes because there is no supervision or labeling on the notes, but there is full supervision on the output answers.
If we train the model to predict the final question and answer with varying length samples, the model will learn to generate a question after any number of statements. At the same time, we allow the model to write a Self-Note after each intermediate statement. Assuming the model has learned how to answer the shorter samples, it is likely to write the correct value in the intermediate locations. It can then leverage that information on the longer samples. If the relevant intermediate questions are asked and answered, this will make it easier to answer the final question.
We consider two potential problems with approach. The first is that as the context is enriched by Self-Notes, it can become longer than what the model has seen during training, or it can contain new tokens that it didn't see in the context during training. A simple solution is to finetune the model on the Self-Notes enriched samples during training. The training procedure therefore has two simultaneous objectives: learn how to write Self-Note QA pairs after any number of context tokens, and leverage the written Self-Note answers for the final question. The second problem is that the model might not ask enough questions because training samples contain only one final question. We solve this by simply amplifying the probability of generating a Self-Note start token (from  $N_{\mathrm{sta}}$ ), by a multiplicative "boosting" constant  $B > 1$ . Furthermore, we can sample multiple versions of Self-Notes per sample and select the enrichment that leads to the most confident answer.
Self-Note Prompting of Large Language Models. Lastly, we consider prompting a pretrained large language model (LLM) with Self-Notes. In this setting, we show the model few demonstrations of how to write Self-Notes in the prompt. During inference, we allow the LLM to insert a Self-Note every time the start token is predicted as the next token within the current context. This is similar to chain-of-thought prompting [4], except we replace the chain-of-thought reasoning that comes after the question with Self-Notes that can be anywhere in the context. Once all the Self-Notes are generated, we let the LLM generate its final answer.
We compare against two baseline methods: a vanilla transformer language model, and a transformer language model trained to generate a chain-of-thought "scratchpad". For all experiments except few-shot prompting, we use the following LMs. The Vanilla baseline is the pretrained GPT-2 base model [20] from Hugging Face [21] fine-tuned to predict answer tokens given only the context and question. For the Scratchpad (i.e. Chain-of-thought) baseline, we fine-tune the same GPT-2 model to write a scratchpad of reasoning steps after it has seen the context and question, similar to Nye et al. [3]. For the proposed Self-Notes model, we fine-tune GPT-2 to take Self-Notes. During testing, no ground-truth scratchpad or Self-Notes are provided, but both Scratchpad and Self-Notes models are allowed to generate tokens in addition to the answer.
In this section, we explain each task we test our models on. Table 1 shows a sample for two task with different methods: Vanilla, Scratchpad, and Self-Notes. For each task, we evaluate on both an in-distribution and out-of-distribution (OOD) test set. More detailed descriptions of each dataset and a summary of dataset statistics is provided in the Appendix.
Toy-Story. Inspired by the bAbI tasks [22], we introduce a new synthetic QA task for testing the ability of language models to do multi-step reasoning. The task is to answer a question after reading a short story that consists of multiple sentences. The challenge in this dataset is that by applying pragmatic principles, unseen relations can be inferred from observed relations. For example, given the text "Alice is at the park. Bob is with Alice.", we can infer that "Bob is at the park". Furthermore, a newly inferred relation can lead to inference of another unseen relation. In the previous example, if the next sentence is "Bob has the key.", then we can infer that "The key is at the park" using our previous inference about Bob's location. This recursive inference in Toy-Story makes it possible to create questions that require multi-step reasoning.
We use Self-Notes to infer all implied relations. Within a Self-Note, the model can ask and answer a question, e.g., "SQ: Where is Bob? Bob is at the park." The Scratchpad method should infer the same relations, but it will be forced to do it after the question is asked, requiring backwardreasoning. To test generalization, we train the model on 10k 1-hop and 2-hop queries, and test on 3-hop and 4-hop queries.
Table 2: Test Accuracy (in %) for the reasoning and state-tracking tasks in the supervised setup. “*” indicates out-of-distribution harder test settings.
Task	Test Set	Vanilla	Scratchpad	Self-Notes
Toy-Story	1/2-hop	92.4 ±0.7	99.6 ±0.1	99.8 ±0.1
3-hop*	57.0 ±0.3	96.4 ±0.9	98.5 ±0.3
4-hop*	37.4 ±0.8	94.2 ±2.0	97.8 ±0.4
Algorithmic	2-100	44.6 ±1.0	72.2 ±5.7	95.5 ±0.2
101-200*	24.4 ±2.1	11.6 ±2.0	85.0 ±0.6
Boolean Variable	3-8	99.7 ±0.1	100.0 ±0.0	100.0 ±0.0
9-19*	71.3 ±0.8	73.7 ±2.4	75.2 ±2.1
Chess Piecetype	≤80	98.5 ±0.4	98.5 ±0.3	98.8 ±0.2
≥81*	82.9 ±2.3	94.7 ±0.7	94.8 ±0.7
Chess Move	≤80	49.0 ±0.4	37.0 ±0.8	50.8 ±1.1
≥81*	39.8 ±0.2	29.9 ±0.8	41.8 ±0.9
Algorithmic. To evaluate the ability of tracking the state or value of an entity over multiple steps, we adopt the Algorithmic task from [11], where the context is the sequence of statements, e.g. “ $x = 1$ ;  $x++; y = 3$ ; if  $x > 2: y++$ ”, and the final question is to print the last value of one of the variables, e.g. “print  $x$ . While the original task has separate input and label tokens, we unify them into a single sequence to fit the language modeling task.
For the Self-Notes model, the notes are print statements specifying the intermediate value of a certain variable as they are modified. For example, if the previous statement was “ $x++$  ;”, the Self-Note would be to print the value of  $x$ : “print  $x x = 1$  ;”. The Scratchpad method generates the identical print statements, but it also has to copy all of the original statements to the scratchpad and figure out where to insert the prints, thus introducing an additional “alignment” complexity in comparison to the Self-Notes method. We train the models on 2 to 100 statements in each sample. We test on 2-100 (in-distribution) and 101-200 (OOD) statements.
Boolean Variable. In this task, the context consists of a valid Python program where each statement contains a boolean variable assignment operation. The question is to print the final value of an arbitrary variable (True or False). The main difference to the Algorithmic task is that the statements are constrained to boolean logic operations. We use the "chain-like" split from Anil et al. [16], which consists only of operations that compose the values of already defined variables. This results in long chains of dependencies between values of the variable. Similar to the Algorithmic task, Self-Notes prints the value of the variable that was modified in the previous statement, e.g. "print x True"; Following Anil et al. [16], we train on 3-8 statements and test on 3-8 and 9-19 statements.
Chess Piecetype. The goal of this task is to track the state of chess pieces over a sequence of moves in a real chess game [17]. Chess games written in UCI notation consist of a sequence of (start position, end position) pairs, e.g. c2 c4, which denotes a move of the piece from c2 to c4. The piecetypes of the moves are never given explicitly, but since each game begins with pieces in the same positions, the piecetypes can be implicitly tracked given the moves. Given a long sequence of moves, e.g. "c2 c4 e7 e5 g2 g3 b8 c6 f1 g2 g8 f6 b1 c3 f8 b4 c6", the objective is to predict the piece at the last position mentioned ("c6").
For our proposed method, we consider the Self-Notes to be the piecetypes. A Self-Note is inserted after the start position of each move to explicitly remind the model which piecetype is at that position. So the previous example would look like: "c2 P c4 e7 P e5 g2 P g3 b8 N c6 f1 B g2 g8 N f6 b1 N c3 f8 B b4 c6", and it is therefore much easier with Self-Notes to predict the piecetype at "c6", since we know that the last piece moved to "c6" was a knight during the move "b8 N c6". To test length generalization, we consider more moves during testing than seen during training.
Chess Move. This task is to predict the end position of the current move given the start position [17]. For example, given the sequence of moves "c2 c4 e7 e5 g2 g3 b8 c6 f1 g2 g8 f6 b1 c3 f8 b4 c6", the answer is the ground truth final position made in the game: "e5". This task is harder than the Chess Piecetype task as the model needs to learn, state tracking, chess rules, and chess strategy in order to predict the most likely move. The Self-Notes are the same as chess piece, where the model is trained to generate the piece at each starting square as it makes a move, e.g. "c2 P c4".
Math Word Problems. For the few-shot prompting, we consider two real-world math tasks that require solving multi-step elementary-level arithmetic word problems: MultiArith [18] and GSM8K [19]. For example, the following example in the MultiArith format requires language understanding, as well as tracking the value of the entities: "Alice had 4 cupcakes and 2 cookies. Alice then ate 1 cupcake. After Alice gives 1 cookie to Bob, how many treats does she have left?" In the few-shot prompt, we wrote useful intermediate reasoning and calculations in Self-Notes surrounded by parenthesis. For example the Self-Notes for the previous sample should look like: "Alice had 4 cupcakes and 2 cookies.  $(4 + 2 = 6$  total) Alice then ate 1 cupcake.  $(6 - 1 = 5$  left) After Alice gives 1 cookie to Bob, how many treats does she have left?  $(5 - 1 = 4$  left)". For MultiArith, we use the few-shot training samples from [4]. GSM8K contains more challenging problems that requires more diverse reasoning. We selected 5 random samples in the training set to use as the few-shot demonstrations in the prompt. For Chain-of-Thought, we used the original solutions (minus the calculation annotations written inside " $\ll\gg$ " ) given in the dataset. For Self-Notes, we manually annotated the few-shot examples with Self-Notes (shown in the Appendix).
4.1 Supervised Self-Notes
Table 2 shows the results for the five tasks described in Section 3.1.
Toy-story. For both the 3-hop and 4-hop settings, we see that the Self-Notes model substantially outperforms the Vanilla model which has to perform multi-step reasoning in "one-step". We observe a slight improvement of the Self-Notes model over the Scratchpad model. We reason that the drop in Scratchpad's performance has to do with the model having to postpone until after processing the entire input context, which increases the distance between the input context and the reasoning. In comparison, the Self-Notes model writes reasoning tokens on the fly as the relevant facts are stated. We note that for this task, the full context fits into the GPT-2 context window.
Algorithmic. We observe that the Vanilla GPT-2 model struggles to track the state of the variables over many statements, and significantly worsens for OOD sequence lengths. Self-Notes, which allows the model to generate intermediate print statements, achieves high accuracy on both the in-distribution and OOD statement splits. Scratchpad fails at most examples since the context length exceeds the maximum length of GPT-2 (1024 tokens). This leads to a worse performance than the Vanilla model because it tries to write the scratchpad which involves copying the original context, but then runs out of room and can't answer the question. These results show a significant advantage of our method: as long as the model takes a Self-Note about a variable, it will keep it in the memory by pushing its value to the most recent context. The Scratchpad method has to copy the entire context in its scratchpad, often going past the maximum context length, resulting in poor accuracy.
Boolean Variable. Unlike the Algorithmic task, none of the models run out of context length for this task since there are fewer statements. Therefore, Scratchpad is able to perform similarly to Self-Notes. However, we still see a small increase in performance with Self-Notes, likely due to copy alignment errors in the Scratchpad. Both the models improve over Vanilla.