SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chess Piecetype and Chess Move. The chess tasks primarily measure the ability of a model to track the identity and state of variables over a sequence of changes. In the Chess Piecetype task, both the Self-Notes and Scratchpad models outperform the Vanilla model. As with other tasks, this confirms that Vanilla transformers are improved with extra tokens in order to accurately track the state of a set of variables, particularly when the test-time sequence lengths vary from the training length. For Chess Piecetype, Self-Notes is not significantly better than Scratchpad. This is a fairly simple task for Scratchpad since it simply requires copying the piece at each move, assuming it knows where the pieces start. This is different from the Algorithmic and Boolean Variable tasks which not only need to copy the variable, but also increment, decrement, or negate it.
In the Chess Move task, Self-Notes is slightly better than Vanilla, but Scratchpad is significantly worse than both. In this task, the Self-Notes and Scratchpad "note" tokens (pieces) are not the same as the final question (move board position). We hypothesize that Scratchpad cannot learn to simultaneously copy the identity of pieces and predict the chess move.
(a) Toy-Story
(b) Algorithmic
Figure 2: Performance of the Semi-supervised Self-Notes method with varying amounts of Self-Notes supervision on the Toy-Story and Algorithmic tasks.
Table 3: Unsupervised Self-Notes: Toy-Story task accuracy (%)
Method	3-hop	4-hop
Vanilla	79.4	57.9
+ Self-Notes (unsupervised)	79.7	61.8
+ Boost Questions	82.4	68.2
+ Multi-sample	91.3	79.1
+ Finetune	94.2	85.8
4.2 Semi-supervised Self-Notes
Figure 2 shows the performance of the Self-Notes method with varying amounts of Self-Note supervision for the Toy-Story and the Algorithmic tasks. That is, we randomly sample some percentage of the training samples that get Self-Note supervision. For Toy-Story, we find that even Self-Note supervision using as little as  $1\%$  of the training set (100 samples), leads to performance gains over the Vanilla model, and the performance starts to saturate around  $25\%$  supervision. On the other hand, for the Algorithmic task, we observe gains with Self-Note supervision starting at around  $5\%$  supervision, and the performance steadily improves with more Self-Note supervision.
4.3 Unsupervised Self-Notes
In our third set of experiments, we apply Self-Notes to a 1-variable Algorithmic task and Toy-Story task when we have no ground-truth Self-Notes to train on.
First, we conduct experiments in the unsupervised setting for Algorithmic task. We first test on a 1-variable version of the dataset (each sample contains only 1 variable). We train on datasets that contain varying length samples (i.e. varying numbers of algorithmic statements per sample), so the model will generate intermediate Self-Notes on its own in a QA form. In this task, we allow the model to generate Self-Notes, and then conditioning on the previous note to predict the next Self-Note and final answer during training, departing from the standard parallel training procedure. The model therefore has to do two simultaneous tasks during training, write the correct Self-Notes, and predict the final answer given the written Self-Notes. Since we only use 1-variable samples, it makes it straightforward to learn which Self-Note questions to write (it will always be print x, where x is the variable in that sample. We can see from Figure 3, that around 10k samples, the unsupervised Self-Notes method starts to learn how to write and leverage Self-Notes that improve accuracy over the Vanilla method. With 20k samples, the unsupervised Self-Notes method achieves near  $100\%$  accuracy, with a significant increase over the Vanilla model.
The second task we consider for unsupervised Self-Notes is Toy-Story. Here, the training data has 100k samples with 1 and 2 hop questions, but contains no Self-Notes. This task is more difficult since there are many more variables (people, objects, locations) and model needs to ask the right questions in Self-Notes.
Figure 3: Unsupervised Self-Notes vs Vanilla sample comparison on 1-variable Algorithmic.
Figure 4: 10k Supervised Self-Notes samples vs Vanilla sample comparison on Toy-story.
We first train a Vanilla model to generate the final question and answer, with test accuracy shown at the top of Table 3. Next, we test the vanilla model with Self-Notes by allowing it to generate QAs as notes. Here, we only add the answer parts to the context because the model has never seen a question in the context during training. Additionally, because the model is trained on 1-hop questions, it often asks a question whose answer is already in the context. We ignore such duplicate answers and move on to the next context sentence.
Simply adding Self-Notes to the Vanilla model during testing does not improve the performance much because the model is not trained to ask the right questions. We therefore encourage the model to ask more questions by boosting the probability of the question start token "Q:". Boosting Self-Notes by  $B = 5$  does improve the performance over the Vanilla model. Furthermore, generating multiple different Self-Notes by sampling questions and selecting the most confident one also helps. This is likely because when the right question is asked and answered, the model becomes more confident in its final answer. Finally, finetuning the model on a Self-Note version of the original training data improves the performance, as it adapts to longer stories. In summary, we see a significant increase in accuracy over the Vanilla results by allowing the model to generate Self-Notes and finetuning the model on the generations.
4.4 Few-shot Prompted Self-Notes
For the MultiArith task, we test a few-shot prompted GPT-J model [23]. For the Vanilla and Scratchpad baselines, we use the same few-shot prompt demonstrations from [4]. For Self-Notes, we use the same few-shot examples, but additionally annotated with Self-Notes that we wrote (provided in the Appendix). In Table 4, we compare Self-Notes to the Vanilla and Chain-of-Thought (CoT) methods for similar sized models (6.7B-8B). We find that the Self-Notes method on the 6B GPT-J model outperforms all of the Vanilla and CoT results. This result highlights two key findings. First, Self-Notes can write pertinent, yet mathematically accurate notes on arithmetic word problems. Second, we can effectively use Self-Notes on pretrained LLMs using few-shot prompting.
For the more challenging GSM8K task, we employed a GPT-3 model with 175B parameters. As shown in Table 5, Self-Notes outperforms CoT, but by a small margin. Since we prompted the model using the public API and had no access to token output probabilities over the prompt itself, we limit Self-Notes positions to immediately after ".", ".", or "?" to reduce the API calls. See [4] for more comprehensive comparison of CoT against Vanilla prompting (note that they used a more powerful version of GPT-3 than us).
Lastly, we run few-shot prompting experiments using the state-of-the-art open model, Llama 2 (70B) [24]. We compare Vanilla, Chain-of-Thought, and Self-Notes across three datasets: Algorithmic, MultiArith, and GSM8K. For GSM8K and MultiArith, we use the 8-shot arithmetic reasoning prompt from [4]. The CoT and Vanilla prompts are the same as from [4] and the Self-Notes prompt adapts the CoT prompt by putting some of the reasoning steps inside the context. For Algorithmic, we use a 5-shot prompt of randomly selected examples from our training set, ranging from 6-33 statements.
3Here we see better Vanilla performance than Table 2 because the training data is larger.
Table 4: GPT-3, LaMDA, PaLM, and GPT-J few-shot prompting accuracy (\%) on MultiArith.
Model	Size	Vanilla	CoT	Self-Notes
GPT-3 [4]	6.7B	4.5	2.8	-
LaMDA [4]	8B	5.8	1.5	-
PaLM [4]	8B	4.2	15.8	-
GPT-J	6B	6.0	13.5	21.0
Table 5: GPT-3 few-shot prompting accuracy  $(\%)$  on GSM8K.
Model	Size	CoT	Self-Notes
GPT-3	175B	11.4	13.4
Table 6: Llama 2 (70B) few-shot prompting accuracy  $(\%)$ .
Dataset	Vanilla	CoT	Self-Notes
Algorithmic	23.3	37.4	40.3
MultiArith	33.3	94.5	96.2
GSM8K	16.5	55.9	59.9
We test on the 2-100 statement test set from our paper. Results are shown in Table 6. In summary, we find that Self-Notes outperforms Chain-of-Thought and the Vanilla baseline on all three tasks.
4.5 Ablation: Labeled training set size comparison
Each Self-Notes training sample in Toy-Story has intermediate questions and answers, therefore increases the total number of QA pairs in the training set. For example if the final question and answer is "Where is the ball? The ball is in the park", but there was a Self-Note in the middle of the context labeled "Who has the ball? Alice has the ball", then that sample has two QA pairs. We therefore also run a comparison of the total number of labelled QA pairs between Self-Notes and Vanilla. Specifically, the 10k Self-Notes training data for Toy-Story has 10k total samples, and therefore 10k final QA pairs. However, it also includes roughly 70k Self-Note QA pairs which means that the total amount QA pairs is around 80k. Figure 4 shows the effect of increasing the training size for the Vanilla baseline compared to a fixed set of 10k Self-Notes training samples (80k labeled QA pairs) in the Toy-Story task. The Self-Notes model with 10k samples still vastly outperforms the Vanilla model with a  $1500\%$  increase in training samples (and roughly a  $100\%$  increase in the amount of QA pairs to that of Self-Notes). See Appendix for more ablation experiments, including understanding the difference between content and compute in Self-Notes by using "dummy" tokens (Appendix subsection A.3).
Implicit Reasoning. bAbI [22] was a set of synthetic tasks for testing different reasoning capabilities [25] and showed the advantage of attention-based models over recurrent neural networks [26, 27]. More recently, attention-based transformers [1] have become the foundation of language-based reasoning [28]. However, the feedforward nature of transformers makes it unsuitable for state-tracking [11] and several recurrent versions have been proposed [29, 13, 14]. Further, transformer-based large LMs are shown to struggle at multi-step reasoning [7].
Explicit Rationales. Use of rationales has been explored for interpretability [30], and for performing intermediate computations [31, 3, 4]. In particular, the Scratchpad method by Nye et al. [3] is closest to our proposed Self-Notes method which can be interpreted as an online-variant of Scratchpad. Use of rationales for reasoning and arithmetic tasks, referred to as "chain-of-thought", has been shown to be particularly beneficial for zero- and few-shot in-context learning with large language models [4, 32, 7]. Zelikman et al. [9] showed the possibility of bootstrapping from a small set of reasoning labels to a larger unlabeled dataset. Trivedi et al. [33] propose interleaving chain-of-thought reasoning with knowledge retrieval steps (both after context). Other works have generated "inner monologue" tokens as a form of intermediate reasoning after a prompt [10, 34]. However, as with Scratchpad, the "chain-of-thought" and "inner monologue" reasoning is done after reading the entire input context rather than while reading it as in Self-Notes. This emergent capability of chain-of-thought prompting is only possible in large ( $>6\mathrm{B}$ ) models [4]. In concurrent work, Toolformer [35]
introduces a method to finetune LMs to make API calls to external tools in the middle of the context. We propose a more general method to augment the context without API calls.
Length Extrapolation. Length extrapolation, or generalization to longer instances during inference than those seen during training is an important property for an intelligent agent. In the context of transformer-based models, length extrapolation has been explored for language modeling [36], machine translation [37, 38], models trained on artificial datasets [39, 16], and a variety of other tasks. One of the reasons for the limited length generalization capability of transformer models is due to the input position being represented with learnable embeddings [38, 40].
Adaptive Computation. When humans read and write text, we often spend a different amount of time per sentence. Transformers are designed to process and generate each sentence using the same amount of computation regardless of the complexity. Several works have addressed the problem of fixed computation time [41, 42, 43], but require modifying the training procedure or model architecture. Self-Notes can be viewed as a form of adaptive computation because the model decides when to deviate from the context and "think". Unlike previous adaptive computation approaches, Self-Notes can easily be applied to existing LM architectures and training procedures.
Editing of Generated Text. Several works have introduced variants of the transformer architecture to allow insertions, deletions, and revisions to the generated text [44, 45, 46, 47, 48]. Contrary to these methods that revise post-context generations, Self-Notes revises the original prompt context by inserting tokens on the fly.
We proposed a general method that allows language models to explicitly reason and memorize in the form of taking Self-Notes. Unlike scratchpad and chain-of-thought methods that postpone reasoning until all input tokens are processed, our method is more general and can deviate from the input sequence at any time for writing a Self-Note. One advantage of interleaving reasoning with the context in this way is that the reasoning steps can be closer to their relevant context. Another advantage is that it can act as a recurrent memory as the Self-Note answers are fed back to the model. Both these advantages make the method generalize better than previous approaches, as shown in our experiments. We demonstrate this improvement with four different learning paradigms: supervised, semi-supervised, unsupervised, and few-shot prompted, using four different models of varying sizes. We note that scratchpad and few-shot chain-of-thought also require additional human annotations. So, while it is a drawback compared to vanilla training, it requires the same amount of additional annotation as scratchpad/CoT. The goal of this work is to experimentally validate the difference between in-context thoughts (Self-Notes) vs purely post-context thoughts (CoT/Scratchpad). Future work should explore two complementary directions aimed at reducing the amount of supervision: (1) using reinforcement learning to discover the optimal Self-Notes, and (2) whether future models can generate relevant Self-Notes out of the box. Our results suggest that large language models can generate meaningful Self-Notes via few-shot prompting. This excites us about the applications that Self-Notes will enable in the future.
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need. In NeurIPS, 2017.
[2] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language Models are Few-Shot Learners. In NeurIPS, 2020.
[3] Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena. Show Your Work: Scratchpads for Intermediate Computation with Language Models. arXiv, abs/2112.00114, 2021.
[42] Tolga Bolukbasi, Joseph Wang, Ofer Dekel, and Venkatesh Saligrama. Adaptive Neural Networks for Efficient Inference. In ICML, 2017.
[43] Andrea Banino, Jan Balaguer, and Charles Blundell. PonderNet: Learning to Ponder. In ICML Workshop on Automated Machine Learning, 2021.
[44] Timo Schick, Jane Dwivedi-Yu, Zhengbao Jiang, Fabio Petroni, Patrick Lewis, Gautier Izacard, Qingfei You, Christoforos Nalmantis, Edouard Grave, and Sebastian Riedel. PEER: A Collaborative Language Model. In ICLR, 2023.
[45] Jiatao Gu, Changhan Wang, and Junbo Zhao. Levenshtein Transformer. NeurIPS, 2019.
[46] Mitchell Stern, William Chan, Jamie Kiros, and Jakob Uszkoreit. Insertion Transformer: Flexible Sequence Generation via Insertion Operations. In ICML, 2019.
[47] Ahmed Elgohary, Denis Peskov, and Jordan Boyd-Graber. Can You Unpack That? Learning to Rewrite Questions-in-Context. In EMNLP-IJCNLP, 2019.
[48] Gangwoo Kim, Sungdong Kim, Kang Min Yoo, and Jaewoo Kang. Generating Information-Seeking Conversations from Unlabeled Documents. In EMNLP, 2022.
[49] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv:2104.09864, 2021.
[50] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. PaLM: Scaling Language Modeling with Pathways. arXiv:2204.02311, 2022.
A.1 Tasks
Here we explain the tasks in more detail than outlined in Section 3.1. A summary of dataset statistics is provided in Table 8.
Toy-Story. As we read a story, we are often required to infer things that are not explicitly mentioned [15]. For example, reading "Frodo went to Mount Doom. Sam accompanied him.", a reader can infer that "Sam went to Mount Doom". Making such inferences in an online manner as we're reading the story makes it easier to understand the rest of the story. Such forward reasoning is natural for understanding sequential stories like books or movies. It is also more fitting for dialog models as such a model needs to make inferences as conversation happens and respond accordingly. In contrast, backward reasoning starts with a question and tries to find the relevant facts from a given context to answer it, potentially leading to a more narrow understanding of context.
Toy-Story is a synthetic QA task that can be used to test the ability of language models to perform forward reasoning. Each sample consists of a short story followed by a question and its corresponding answer. Sentences in the story state a simple relation between people, items, and places such as "Alice is at the park" or "The ball is in the bag". There are 5 different types of relations. This dataset is inspired by the bAbI tasks [22], with greater controllability of required reasoning steps. Unlike bAbI, our dataset mixes different reasoning steps to create more "hops" in order to answer a question.
We call a question  $k$ -hop if it requires  $k$  observations combined through  $k-1$  reasoning steps (1-hop questions only require repeating of an observed fact). While a backward reasoning model needs to take  $k$  reasoning steps to answer a  $k$ -hop question, a forward reasoning model will infer unseen relations as each sentence is processed. As a result, forward reasoning can uncover all relations by the end of the story and can answer any question with no additional reasoning.
Considering the relevance of forward-reasoning to the Toy-Story task, Self-Notes is therefore a natural fit. Self-Notes should explicitly infer all implied relations. For this dataset, the Self-Note start and end tokens are “SQ:” and“.”, respectively. Following the start token, the model can ask and answer a question, e.g., “SQ: Where is Bob? Bob is at the park.”
If the model correctly learns to infer relations during training, then it can easily answer 3 and 4-hop queries by inferring the intermediate (2-hop) relations. Specifically, by writing a Self-Note, the model can turn a 3-hop query into two separate 2-hop queries.
A sample from the test set is shown in Table 9, where the Vanilla and Scratchpad models fail to predict the correct answer, but the Self-Notes model correctly writes the answer in its notes.
Algorithmic. While the Toy-Story task is designed for testing multi-step reasoning, it doesn't require tracking the state or value of an entity over multiple steps since it assumes the world is static. Algorithmic task from [11] requires printing the state, or value, of an integer variable given a sequence of algorithmic program statements such as increment, decrement, and conditionals. We unify the input and label tokens into a single sequence to fit the language modeling task. The start token is "print" and the end token is ".".
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). So the Self-Notes look like: "print x False ;". Like Algorithmic, the start token is "print" and the end token is " ;". The main difference to the Algorithmic task is that the statements are constrained to boolean logic operations. We show an example of the Self-Notes and Scratchpad samples for this task in Table 7.
Chess Piecetype. This is another state-tracking task designed 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 are never given in the samples, but can be implicitly tracked given the
4To ease tokenization, we split a move in UCI notation from c2c4 to c2 c4. We add all the 64 board squares to the language model's vocabulary.
Table 7: Input-Output pairs for the Vanilla, Scratchpad, and Self-Notes method for the non-prompting tasks that were not listed in Table 1. The input consists of the input context and the question, the answer is to be generated by the model. The highlighted text for Scratchpad and Self-Notes is available only during training, and is to be generated by the model during inference.
Task	Vanilla	Scratchpad	Self-Notes
Boolean Variable	w = False ; v = True ; v = w xor v ; w = v and v ; print w True ;	w = False ; v = True ; v = w xor v ; w = v and v ; print w [w = False ; v = True ; v = w xor v ; print v True ; w = v and v ; print w True ;] True ;	w = False ; v = True ; v = w xor v ; print v True ; w = v and v ; print w True ; print w True ;
Chess Piecetype	c2 c4 e7 e5 g2 g3 b8 c6 f1 g2 g8 f6 b1 c3 f8 b4 c3 PIECE N	c2 c4 e7 e5 g2 g3 b8 c6 f1 g2 g8 f6 b1 c3 f8 b4 c3 PIECE [c2 P c4 e7 P e5 g2 P g3 b8 N c6 f1 B g2 g8 N f6 b1 N c3 f8 B b4 c3] N	c2 P c4 e7 P e5 g2 P g3 b8 N c6 f1 B g2 g8 N f6 b1 N c3 f8 B b4 c3 PIECE N
Chess Move	c2 c4 e7 e5 g2 g3 b8 c6 f1 g2 g8 f6 b1 c3 f8 b4 c3 MOVE c5	c2 c4 e7 e5 g2 g3 b8 c6 f1 g2 g8 f6 b1 c3 f8 b4 c3 MOVE [c2 P c4 e7 P e5 g2 P g3 b8 N c6 f1 B g2 g8 N f6 b1 N c3 f8 B b4 c3] c5	c2 P c4 e7 P e5 g2 P g3 b8 N c6 f1 B g2 g8 N f6 b1 N c3 f8 B b4 c3 MOVE c5
moves. For example, since each game starts with a pawn (P) at board position c2, we know that given the move c2 c4, there is now a pawn at position c4. 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 Self-Notes, the notes are the piecetypes. So the previous example would be written as "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". For the chess tasks, the start token is any of the pieces and there is no end token (since the note is only 1 token). To differentiate the Self-Notes from the final question, which also asks for the piecetype, the final question is asked using the PIECE token. Samples for this task are shown inin Table 7.
We consider a different number of moves during training and testing to test length generalization. We train our models on 200k samples which include up to 80 moves. We evaluate on both up to 80 moves (in-distribution) as well as more than 80 moves (OOD).