SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Variable dropout Even within a local subset of the world, we may not see everything at once. Certain variables may be missing or go unnoticed. We formalize this intuition with variable dropout. With some probability (0.2 in our experiments), variables are dropped from a local neighborhood and the learner does not see them. Variable dropout may also help a model generalize to pairs of variables that were unseen in training as more subsets of variables appear in the training set.
Held-out pairs Finally, some pairs of variables are held out across all training data. If a local neighborhood, after variable dropout, would include a pair of variables we decided to hold out, we randomly remove one of the two variables in the pair from the sample. We use the ability to match conditional probabilities for held-out pairs of variables, measured via mean squared error, as our main performance metric.
We also create two control conditions to compare against local observation distributions. As one control, we consider a training dataset in which the values of variables are sampled from one Bayes net, but the local neighborhoods are constructed with respect to the structure of a different Bayes net. This condition still has a local co-occurrence structure, but the co-occurrences do not reflect
Figure 2: Mean squared error between estimated and true conditional probabilities for held-out pairs with high mutual information. Error bars denote  $95\%$  confidence intervals. Both free and scaffolded generation significantly outperform direct prediction when the training data is locally structured.
which variables influence each other. As another control, we use a fully observed condition where each sample contains almost all of the variables in the Bayes net. One of the two variables in each held-out pair is randomly dropped, but all other variables are included. These controls enable us to test whether local structure in the training data drives the value of reasoning.
4.2 Estimation
Each of the estimators in Section 2.2 is implemented via sampling from and scoring with the language model. For instance, in free generation, the model is prompted with the target variable name and initial observation, then sampled from until the target variable name is generated. The probabilities for possible values of the target variable are then examined. Figure 1C illustrates how we implement direct prediction and free generation as prompts and Appendix D contains examples of each estimator. For estimators that rely on Monte Carlo estimation, we use 10 samples.
We also introduce negative scaffolded generation as a control estimator that generates irrelevant intermediate variables. For each pair of variables, we select a random set of variables equal in size to the scaffold, but which does not include any of the scaffold variables. We prompt the language model to generate values for the negative scaffolds in the same way as in scaffolded generation.
4.3 Model architecture
We use a smaller version of the GPT-2 autoregressive transformer architecture [20], as implemented in the HuggingFace transformers library [21]. Our model has 512-dimensional embeddings, 10 layers, and 8 attention heads. Data is tokenized via a Byte Pair Encoding tokenizer fit to data in our format [22]. We trained this architecture with randomly initialized weights for 300,000 gradient steps on batches containing 3,072 tokens each, for a total of 921,600,000 tokens of training. We trained models using the Adam optimizer [23]. Each model's training set consisted of 1,000,000 samples from a single Bayes net. The models achieved near-perfect performance on the training task. Further training details are provided in Appendix E. We also compare multiple different architectures in Appendix F, demonstrating that our results are robust across them.
Each held-out pair gives four data points: if we hold out the pair  $X1, X2$  we can estimate  $p(X1|X2 = 0), p(X1|X2 = 1), p(X2|X1 = 0)$ , and  $p(X2|X1 = 1)$ . Pooling across the 10 Bayes nets we trained models on gives us 1,000 data points per training condition and estimator. Figure 1D shows our main result: free generation performs better (lower MSE) than direct prediction and performs
Figure 3: Number of intermediate variables generated in free generation vs. absolute distance between true and estimated conditional probabilities for each training condition. Lower dots indicate more accurate estimation, while dots further to the right indicate more intermediate variables generated.
Table 1: Mean squared errors with  $95\%$  confidence intervals between direct prediction probabilities and either true or marginal probabilities for non-held-out pairs with high mutual information. In the local training conditions, direct prediction is close to perfect at matching conditional probabilities.
fully observed	local (geom)	local (Zipf)	wrong local (geom)	wrong local (Zipf)
true	.013 [.010, .016]	.004 [.003, .005]	.003 [.003, .004]	.052 [.045, .059]	.056 [.049, .063]
marginal	.080 [.073, .086]	.114 [.106, .121]	.112 [.104, .121]	.047 [.041, .052]	.033 [.029, .038]
well with less training when the data is structured locally. The MSEs for all estimators and training conditions after 921.6 million tokens of training are shown in Figure 2.
Appendix G shows how MSE varies with the number of Monte Carlo samples taken for each estimator. Estimation error is high for free and scaffolded generation for low numbers of samples due to high variance, but quickly decreases with the number of samples. Re-sampling the intermediate variables is akin to self-consistency prompting methods which involve re-sampling the chain of thought and taking the majority or average answer [14].
5.1 When reasoning helps
We find a benefit to reasoning when the observation distribution has the correct locality structure. When the training data is structured locally with respect to strong dependencies, both scaffolded and free generation perform significantly better than direct prediction. We call this difference the reasoning gap. Scaffolded and free generation also perform significantly better than negative scaffolded generation, indicating that relevant intermediate variables help in predicting the target variable, but irrelevant intermediate variables do not. Furthermore, the success of free generation demonstrates that a model trained on local subsets of a Bayes net can self-generate helpful intermediate variables.
In free generation, the model generates variables until it produces the target variable. This results in reasoning traces of varying lengths, including many long traces. One might expect long traces to distract the model, but this did not occur in practice: the length of the reasoning trace does not relate strongly to accuracy within training conditions, as is shown in Figure 3. However, free generation does not outperform direct prediction substantially when the training data is fully observed.
It is also of note that training on locally structured data reduces the average length of reasoning traces (Fig. 3). In local training conditions, free generation usually produces a set of variables that  $d$ -separates the observed variable from the target variable  $-70\%$  of the time in the local neighborhood conditions. In contrast, the training conditions with the wrong locality structure only lead to  $d$ -separating reasoning traces  $34\%$  of the time. The fully observed training condition also usually creates  $d$ -separating reasoning traces  $(69\%)$ , likely because it generates a large number of variables.  $d$ -separation is important because it ensures that the intermediate variables capture all the influence that the observed variable exerts on the target variable. These results suggest that training on local clusters of variables is valuable in part because it makes it easy for an autoregressive model to
Table 2: Mean squared errors between estimated conditional probabilities and marginal probabilities of target variables with  $95\%$  confidence intervals. Lower values indicate closer matches.
fully observed	local (geom)	local (Zipf)	wrong local (geom)	wrong local (Zipf)
direct prediction	.086 [.079, .093]	.115 [.108, .124]	.116 [.107, .124]	.029 [.025, .033]	.022 [.019, .025]
negative scaffolded	.089 [.082, .096]	.116 [.108, .124]	.115 [.108, .123]	.036 [.031, .040]	.026 [.022, .030]
free generation	.123 [.115, .132]	.124 [.117, .132]	.127 [.120, .135]	.049 [.044, .055]	.042 [.037, .048]
scaffolded	.110 [.102, .117]	.134 [.126, .142]	.134 [.126, .142]	.051 [.046, .057]	.043 [.038, .049]
naturally output useful intermediate reasoning steps. Still, chain-of-thought reasoning can be helpful even if the reasoning trace does not completely  $d$ -separate the two variables. In the local training conditions, the model does not generate a  $d$ -separating trace  $30\%$  of the time, yet there is virtually no performance difference between free and scaffolded generation.
5.2 Data complexity and reasoning
The combination of locally structured training data and step-by-step reasoning can also lead to accurate conditional probability estimates with much less training than would be required if a model were trained on fully observed data. To demonstrate this difference, we run both direct prediction and free generation on checkpoints from our models that are saved after every 5000 gradient steps (15.3 million tokens) of training. Part D of Figure 1 shows the results of this analysis.
When data is fully observed, direct prediction and free generation performances improve slowly but consistently. There is no reasoning gap, as both estimators perform almost identically. When the model is trained on data from geometrically sized local neighborhoods, direct prediction performance improves quickly, then plateaus. There is a substantial reasoning gap where free generation outperforms direct prediction. Free generation matches the true conditional probabilities very closely after about 120 million training tokens, achieving the best performance of any combination of training condition and estimator. Still, a sufficiently large transformer can eventually memorize the relationships in its training corpus with enough training. We evaluate the data complexity of direct learning by training a language model of the same architecture on fully observed data without holding out any pairs. Direct prediction in this case takes over 3 times as much training to achieve the same performance as local training with free generation, despite the model seeing the relevant pairs in training (see Appendix H for detailed results). Training language models on datasets consisting of local neighborhoods with strong dependencies and performing chain-of-thought reasoning at inference time can therefore be more data-efficient than training on more complete datasets.
5.3 When reasoning is unnecessary
Table 1 shows MSEs between direct prediction estimates and true probabilities for the 25 high-mutual-information variable pairs that were not held out from the training set. In the local training conditions (and to a lesser extent fully observed), direct prediction matches the true conditional probabilities almost perfectly. This result indicates that our language models learn to match conditional probabilities directly when the observed and target variables co-occur in the training distribution. Observing these variables together frequently obviates the need for step-by-step reasoning.
5.4 When reasoning fails
We test the hypothesis that language models match the marginal probabilities of the target variables when they fail to predict the true conditional probabilities by comparing models' estimated conditional probabilities against the true marginal probabilities. The MSEs between the estimated conditional probabilities and marginal probabilities of the relevant target variables are shown in Table 2. We can see the opposite of the trend for true probabilities: the worse a training condition does at matching the true conditional probability, the better it matches the marginal. The language models trained on data with the wrong locality structure generated estimates that were particularly close to the marginal probabilities. When the variables that co-occur with each other frequently are not local in the Bayes net, they often have very little influence on each other. This means that the joint distribution over co-occurring variables is usually very close to the product of the marginal probabilities, i.e.
$P(X_{1},X_{2},X_{3})\approx P(X_{1})P(X_{2})P(X_{3})$  for non-local  $X_{1},X_{2},X_{3}$ . Without the ability to estimate conditional probabilities accurately, there are no reliable 'steps' for step-by-step reasoning to use.
Our theoretical and experimental results demonstrate that estimating conditional probabilities by reasoning through intermediate variables can outperform direct prediction when the training data has local structure. When the training data includes all the variables, free generation does not lead to better performance than direct prediction. When the training data has the wrong locality structure, the models do not reliably learn conditional probabilities that can be chained together.
Our results provide a minimal case in which chain-of-thought reasoning is helpful and suggest conditions under which it is likely to be helpful in more naturalistic settings: we can expect chain-of-thought reasoning to help when a model is tasked with making inferences that span different topics or concepts that do not co-occur often in its training data, but can be connected through topics or concepts that do. Re-sampling the chain of thought multiple times might be especially valuable in domains with moderately strong dependencies. In these settings, introducing intermediate steps increases the variance of estimation and re-sampling multiple times reduces it. We have also shown that chain-of-thought reasoning enables greater data efficiency: a model trained on locally structured data with chain-of-thought reasoning at inference time can accurately match conditional probabilities with less training than a model trained on all the data. This result can provide insight into why humans are more data-efficient than language models. Since humans experience the world from a first-person perspective, the information we encounter is structured in small clusters of entities that are tightly coupled with each other. This idea could also be relevant to data curation for the training of language models. Constructing datasets with tightly correlated observation neighborhoods that collectively cover the space of relevant concepts may amplify language models' ability to perform chain-of-thought reasoning while reducing their data needs.
Future work should explore the structure of the observation distribution for human learners to understand how the information we observe facilitates reasoning. Using modern language models as toy models for the study of reasoning is a promising direction for cognitive science to address long-standing problems in reasoning and problem-solving, such as the value of thought experiments.
The role of abstraction in reasoning is another important direction for future work. Humans often use abstract principles to reason rather than exclusively considering possible concrete states of the world (settings of observable variables). Likewise, the language corpora that language models are trained on contain abstract concepts expressed in natural language. Understanding how abstractions can arise and the effects they have on reasoning would be helpful in connecting our setup to more naturalistic reasoning contexts. The possibility of taking abstract steps makes the problem of choosing useful reasoning steps much harder, so we plan to explore questions of how language models might learn to reason more effectively through fine-tuning, in-context learning, and reinforcement learning.
Limitations Of the many forms of chain-of-thought prompting that exist in the literature, our findings are most relevant to zero-shot prompting [e.g. 9] where a model produces steps on the way to an answer without any other examples of reasoning traces in its context window. It is less applicable to methods that involve giving the model examples of reasoning [e.g. 10]. Our results are also in the context of simple propositional worlds. Reasoning in rich, structured worlds may require more expressive languages with which to reason and specify hypotheticals. Our theoretical results do not explain the phenomenon observed in practice that transformers revert to the marginal distributions for held-out pairs. We leave it to future work to explore how the inductive biases of transformers might induce this behavior on unseen data.
We thank members of the Computation and Cognition Lab, the Causality in Cognition Lab, and the Stanford NLP Group for feedback on earlier versions of this work. In particular, we thank Tobias Gerstenberg for advice on the importance of strong dependencies for reasoning and Eric Zelikman for feedback on the design of our figures. This work is supported by an NSF Expeditions Grant, Award Number (FAIN) 1918771.
[18] D. M. Blei and J. D. Lafferty, “A correlated topic model of Science,” The Annals of Applied Statistics, vol. 1, no. 1, pp. 17 – 35, 2007.
[19] S. C. Chan, A. Santoro, A. K. Lampinen, J. X. Wang, A. K. Singh, P. H. Richemond, J. McClelland, and F. Hill, "Data distributional properties drive emergent in-context learning in transformers," in Advances in Neural Information Processing Systems (A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, eds.), 2022.
[20] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever, "Language models are unsupervised multitask learners," 2019.
[21] T. Wolf, L. Debut, V. Sanh, J. Chaumont, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, et al., "Transformers: State-of-the-art natural language processing," in Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations, pp. 38-45, 2020.
[22] R. Sennrich, B. Haddow, and A. Birch, “Neural machine translation of rare words with subword units,” in Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1715–1725, 2016.
[23] D. P. Kingma and J. Ba, "Adam: A method for stochastic optimization," in Proceedings of the 3rd International Conference on Learning Representations (ICLR), 2015.
A Theoretical analysis
In this section, we theoretically analyze chain-of-thought reasoning when the underlying data distribution can be represented as a directed chain. We prove that any minimizer of a risk consisting of a cross-entropy term and an entropy regularization term will exhibit a "reasoning gap" when the data has a particular locality structure.
A.1 Problem Setup
We assume that we observe a sequence of variable indices  $i_t \in \{1, \dots, N\}$  and variable values  $v_t \in \mathcal{X}$  for some finite set  $\mathcal{X}$ . In the sequence, variable values always follow variable indices. Since our results concern estimators that are only ever conditioned on a single variable, it suffices to consider sequences of length two. However, our setup can be interpreted as assuming we observe independent local neighborhoods of size 2 with token separators.
Factorization of data distribution We assume that the joint distribution  $p_d$  over  $Y_1, Y_2, \ldots, Y_N$  factorizes as
$$
p _ {d} (Y _ {1}, \dots Y _ {N}) = p _ {d} (Y _ {1}) \prod_ {j = 1} ^ {N} P _ {d} (Y _ {j + 1} | Y _ {j})
$$
and that  $Y_{1}, Y_{2}, \ldots, Y_{N}$  take support in some finite set  $\mathcal{X}$ . Note that the variable indices index particular random variables in  $p_d$ .
Local observation distribution We will consider a setting where only adjacent variables in the graphical model can appear adjacent in the sequence. More formally, let  $\mathcal{V}$  refer to the set of all variable identifiers. Let  $\mathcal{V}_{\mathrm{obs}}$  be the set of all allowed pairs of variable indices. Then,  $\mathcal{V}_{\mathrm{obs}} \subseteq \mathcal{P}(\{1, \dots, N\})$  and  $\mathcal{V}_{\mathrm{obs}}$  consists of pairs of the form  $(i, i + 1)$ . We will define the distribution over these valid subsets of variable indices as  $p_{\mathrm{obs}}$ ; in particular,  $p_{\mathrm{obs}}$  takes support only on  $\mathcal{V}_{\mathrm{obs}}$  and assigns equal probability to all of its pairs. The two distributions  $p_d$  and  $p_{\mathrm{obs}}$  define a distribution  $p$  over the sequence of variable indices and variable values.
That is, for any non-adjacent pairs  $Y_{i}$  and  $Y_{j}$  and for any variable values  $y_{i}$  and  $y_{j}$ ,
$$
p \left(i _ {1} = i, v _ {1} = y _ {i}, i _ {2} = j, v _ {1} = y _ {j}\right) = 0 \tag {5}
$$
For adjacent pairs  $Y_{i + 1}$  and  $Y_{i}$  and for variable values  $y_{i + 1}$  and  $y_{i}$ ,
$$
p \left(i _ {1} = i, v _ {1} = y _ {i}, i _ {2} = i + 1, v _ {1} = y _ {i + 1}\right) \propto p _ {d} \left(Y _ {i} = y _ {i}, Y _ {j} = y _ {j}\right) \tag {6}
$$
Intuitively, this property will imply that the risk minimizer only gets to see a small subset of local statistical dependencies in the graphical model. This key property of the observation distribution will induce a reasoning gap.
We begin with a proposition that will be useful in the main theorem. The proposition concerns the minimizer of the sum of two cross-entropy terms.
Proposition 1. Let  $\alpha_{1} \geq 0, \alpha_{2} \geq 0$ . Let  $R(q) = \alpha_{1}\mathbb{E}_{p_{1}(x)}[-\log q(x)] + \alpha_{2}\mathbb{E}_{p_{2}(x)}[-\log q(x)]$ .
Then  $q^{*} = \arg \min R(q) = \frac{\alpha_{1}}{\alpha_{1} + \alpha_{2}} p_{1}(x) + \frac{\alpha_{2}}{\alpha_{1} + \alpha_{2}} p_{2}(x)$
Proof. We assume the probability distributions are discrete. The Lagrangian is given by
$$
\mathcal {L} (q, \lambda_ {0}) = - \alpha_ {1} \sum_ {x} p _ {1} (x) \log q (x) + - \alpha_ {2} \sum_ {x} p _ {2} (x) \log q (x) + \lambda_ {0} (\sum_ {x} q (x) - 1) \tag {7}
$$
The first-order conditions are
$$
\frac {\partial \mathcal {L}}{\partial q (x)} = - \alpha_ {1} \frac {p _ {1} (x)}{q (x)} + - \alpha_ {2} \frac {p _ {2} (x)}{q (x)} + \lambda_ {0} = 0 \tag {8}
$$
$$
\frac {\partial \mathcal {L}}{\partial \lambda_ {0}} = \sum_ {x} q (x) - 1 = 0 \tag {9}
$$
The first condition allows us to write  $q(x) = \frac{\alpha_1p_1(x) + \alpha_2p_2(x)}{\lambda_0}$ . We can substitute this for  $q$  into the other first order condition (i.e.,  $q$  normalizes) to obtain
$$
\sum_ {x} \frac {\alpha_ {1} p _ {1} (x) + \alpha_ {2} p _ {2} (x)}{\lambda_ {0}} = 1
$$
This implies  $\lambda_0 = \alpha_1 + \alpha_2$  and the desired result follows immediately.
A.3 Main Theorem
Our main result is that the minimizer of the risk  $q^{*}$  under the distribution  $p$  defined by  $p_{\mathrm{obs}}$  and  $p_d$  and a cross-entropy loss with entropy regularization will exhibit a "reasoning gap". We begin with a characterization of  $q^{*}$ . For simplicity, in the result below, we'll use  $x_{t}$  to denote elements of the sequence that correspond either to variable values  $v_{t}$  or variable indices  $i_t$ .
Theorem A.1. Let  $u$  be the uniform distribution. Let  $p$  be the distribution over variable indices and variable values defined by  $p_{obs}$  and  $p_d$ . Let  $H(p, q)$  denote the cross entropy between distributions  $p$  and  $q$ . We consider the following risk:
$$
R (q) = H (p, q) + H (u, q) \tag {11}
$$
Then  $q^{*} = \arg \min_{q} R(q)$  satisfies the following properties.
- For all pairs of adjacent random variables  $Y_{i}$  and  $Y_{i+1}$ ,  $q^{*}(Y_{i}|Y_{j}) = \lambda p_{d}(Y_{i+1}|Y_{i}) + (1 - \lambda)\frac{1}{|\mathcal{X}|}$  for some  $\lambda \in (0,1)$ .
- For all pairs of non-adjacent random variables,  $q^{*}(Y_{i}|Y_{j}) = \frac{1}{|\mathcal{X}|}$ .
Proof. We can write the risk as a sum across timesteps.
$$
R (q) = \sum_ {t = 1} ^ {T} \mathbb {E} _ {p \left(x _ {1: t}\right)} \left[ - \log q \left(x _ {t} \mid x _ {1: t - 1}\right) \right] + \sum_ {t = 1} ^ {T} \mathbb {E} _ {u \left(x _ {1: t}\right)} \left[ - \log q \left(x _ {t} \mid x _ {1: t - 1}\right) \right] \tag {13}
$$
Let us consider the left sum. By the law of iterated expectations, we can decompose the fourth term in the sum where  $x_4$  is a variable value as