SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ben Prystawski
benpry@stanford.edu
Michael Y. Li
michaellyli@stanford.edu
Noah D. Goodman
ngoodman@stanford.edu
Humans have a powerful and mysterious capacity to reason. Working through a set of mental steps enables us to make inferences we would not be capable of making directly even though we get no additional data from the world. Similarly, when large language models generate intermediate steps (a chain of thought) before answering a question, they often produce better answers than they would directly. We investigate why and how chain-of-thought reasoning is useful in language models, testing the hypothesis that reasoning is effective when training data consists of overlapping local clusters of variables that influence each other strongly. These training conditions enable the chaining of accurate local inferences to estimate relationships between variables that were not seen together in training. We prove that there will exist a "reasoning gap", where reasoning through intermediate variables reduces bias, for the simple case of an autoregressive density estimator trained on local samples from a chain-structured probabilistic model. We then test our hypothesis experimentally in more complex models, training an autoregressive language model on samples from Bayes nets but only including a subset of variables in each sample. We test language models' ability to match conditional probabilities with and without intermediate reasoning steps, finding that intermediate steps are only helpful when the training data is locally structured with respect to dependencies between variables. The combination of locally structured observations and reasoning is much more data-efficient than training on all variables. Our results illustrate how the effectiveness of reasoning step by step is rooted in the local statistical structure of the training data.
The human mind is a ship – the immediate inferences we make from instinct keep us afloat, but reason is the compass that brings us to the shore of wisdom. Many tasks that we find hard to do immediately – solving math problems, planning vacations, understanding our relatives – become much easier when we talk ourselves through a reflective reasoning process. Likewise, by considering thought experiments or “intuition pumps” in science we can form strong beliefs – such as that the rate at which an object falls should not depend on its mass – purely by thinking through a set of steps
[1, 2]. It is not a priori obvious that step-by-step reasoning should be helpful. Reasoning does not give us any new data from the world, yet it can still improve our inferences. In investigating the origins of reasoning, we must thus ask, why does reasoning help at all?
Large language models have been shown capable of performing a wide variety of tasks by immediately answering a question [3-6]. However, they struggle with some complex tasks like math word problems [7]. A recent line of work has demonstrated that inducing language models to produce a "chain of thought" consisting of intermediate steps toward a solution, before giving an answer, leads to better performance than prompting them to produce an answer directly [8-11]. Other work has built on these findings, showing that providing worked solutions in context is helpful across a broad array of tasks [12, 13] and proposing methods to augment language models' reasoning performance [14-16]. These findings raise an important question: why is chain-of-thought reasoning useful? Exploring this question in language models may also provide insight into the origins of human reasoning.
We hypothesize that chain-of-thought reasoning is useful in language models due to local structure in the training data. Human experience is governed by our first-person perspective: we see and hear aspects of the world that are near to us in time and space. Yet our reasoning transcends these local bounds, supporting plans and conclusions that span distance and years. Similarly, language models are trained on documents in natural language, which are usually about a few closely interconnected topics [17, 18]. When concepts co-occur frequently in experience or training data, estimating the effect of one on the other is easy to do directly with simple statistical estimators. However, when we need to infer the effect of one piece of information on another but have not encountered them together, we must make a series of inferences that jump between pairs of concepts to connect what we know with what we want to infer. We posit that chain-of-thought reasoning becomes useful exactly when the training data is structured locally, in the sense that observations tend to occur in partially overlapping neighborhoods of related concepts.
To illustrate, we may know the value of some variable  $A$  and want to know about another variable  $C$ , so we try to estimate  $P(C|A)$ . However, if we need to estimate probabilities using observed samples from joint distributions and we have not often seen  $A$  and  $C$  together, we would struggle to estimate  $P(C|A)$  directly. Instead, we might estimate it by reasoning through intermediate variables. If conditioning on an intermediate variable  $B$  renders  $A$  and  $C$  independent of each other, we can compute the conditional probability by marginalizing over  $B$ , using the fact that  $P(C|A) = \sum_{B} P(C|B)P(B|A)$ . For an example in natural language, suppose that we want to answer a question like "What is the climate of France's capital?" using a model trained on text from Wikipedia. Wikipedia pages for cities usually have information about the city's climate and Wikipedia pages for countries have information about capital cities, but pages for countries typically do not directly mention the climate of the capital city. It is possible that our model could succeed in answering the question directly, but it would likely have more success if it took the intermediate step of stating the capital of France. By introducing the step "the capital of France is Paris" before answering "Paris has an oceanic climate," an autoregressive model can better leverage the dependencies that are well-represented in its training data.
Our hypothesis is similar to that advanced by Chan et al. [19] to explain why in-context learning occurs (generalizing from examples shown in context). They show that "burstiness," the tendency of instances of a class to occur close together, in the training data is important for language models to learn to infer tasks based on examples in the context window. Burstiness and locality are two different statistical properties of training data that could arise from topic structures in language models or from a first-person perspective in humans. While burstiness concerns how a class is distributed over time, locality concerns which classes co-occur together. In datasets with local structure, concepts that strongly influence each other are also seen together frequently in the training set, whereas concepts that do not are seen together less frequently.
In this work, we investigate the properties of a model's training data that make generating intermediate variables useful for inference. We first formalize the problem mathematically and prove that reasoning through the right intermediate variables reduces bias in a model trained on local samples from a chain. We then train autoregressive transformers from scratch on synthetic data with different structures and evaluate the models' ability to match conditional probabilities for pairs of variables that were not observed together in training. Generating values for relevant intermediate variables enables models to match conditional probabilities more accurately, but only when the training data is locally structured. Allowing the model to select intermediate variables to reason through is similarly helpful to giving it
Figure 1: Overview of our training and estimation setup. A: visualization of a Bayes net. The pink variable is an example observed variable and the yellow variable is an example target variable. Gray variables are examples of useful intermediate variables for reasoning. Lines show examples of local neighborhoods from which training samples are drawn. B: format of the training samples. C: illustration of direct prediction and free generation estimators as prompts. We prompt the model to either immediately predict the target variable (direct prediction) or do so after generating intermediate variables and their values (free generation). We then compute mean squared errors between the estimated and true conditional probabilities. D: mean squared error by number of training tokens for each training condition and estimator. Ribbons indicate  $95\%$  confidence intervals.
the ideal intermediate variables to reason through. Our results suggest that chain-of-thought reasoning is useful for language models because 1) direct prediction is inaccurate for some inferences because the relevant variables are rarely seen together in training and 2) chain-of-thought reasoning improves estimation by incrementally chaining local statistical dependencies that are observed frequently in training. We also find that the combination of locally structured training data and reasoning with self-generated intermediate variables yields much greater data efficiency than training on data containing all variables.
2 Task setup
Our theory and experiments take place in the context of conditional inference in joint distributions, as represented by Bayesian networks. In our framing, a learner sees samples from a Bayes net and needs to accurately estimate conditional probabilities. However, the learner may not see all variables together, having access only to locally-structured observations instead.
We describe our framework formally. First, we assume there is some underlying Bayes net with random variables  $\{Y_i\}_{i=1}^N$  taking support on a finite set  $\mathcal{X}$ . We denote the distribution defined by the Bayes net as  $p_d$ , the data distribution. Our training data is a sequence of variable indices  $i \in \{1, \dots, N\}$ , and variable values  $v_i \in \mathcal{X}$ . Variable values always follow variable indices.
2.1 Observation distribution
We create local structure in our task via an observation distribution  $p_{obs}$ , which is a distribution over subsets of the variable indices. Observation distributions take support on a set  $\mathcal{Y}_{\mathrm{obs}} \subseteq \mathcal{P}(\{1, \dots, N\})$ . As a consequence, our learner will only see samples from the joint distribution over certain subsets of random variables in the Bayes net. The generative process for our sequence consists of two steps. First, we sample a set of variable indices  $\{i_t\}_{t=1}^K$  according to  $p_{\mathrm{obs}}$ . These indices correspond to some set of random variables  $\{Y_{i_t}\}_{t=1}^K$  in the Bayes net, which appear in random order in training. We
then sample the values  $\{v_{i_t}\}_{t = 1}^K$  corresponding to the variable indices according to the distribution  $p_d(Y_{i_1},\ldots ,Y_{i_K})$ . Together,  $p_d$  and  $p_{\mathrm{obs}}$  define a distribution  $p$  on training sequences.
We are particularly interested in local observation distributions as they enable models to learn strong dependencies and reflect some of the statistical structure of human experience. In local observation distributions, the learner observes a subset of variables that are close together in the Bayes net graph. Figure 1A shows several local neighborhoods in an example Bayes net.
2.2 Estimators
Given an autoregressive conditional probability estimation model  $q$  that is trained to predict variable indices and values, we can use one of three different estimators for the conditional probabilities. The estimators differ in how a sequence of variables and values is produced before estimating the target variable. If we want to compute the conditional probability of a target variable  $Y_{i}$  taking some value  $y_{i}$  given that we have observed variable  $Y_{j}$  to be  $y_{j}$ , we can use one of the following estimators.
Direct prediction The direct prediction estimator is the baseline without any reasoning: the probability directly output by  $q$ .
$$
\hat {q} _ {D} \left(Y _ {i} = y _ {i} \mid Y _ {j} = y _ {j}\right) = q \left(Y _ {i} = y _ {i} \mid Y _ {j} = y _ {j}\right). \tag {1}
$$
We simply use the model to directly estimate the conditional probability of the target variable.
Scaffolded generation The scaffolded generation estimator represents ideal reasoning if we know the best set of steps to work through. A scaffold is an ordered set  $S$  consisting of variables that were each observed together with the observed or target variable (or another scaffold variable) and collectively  $d$ -separate the observed variable from the target variable. In the case of a chain, the scaffold consists of all variables between  $Y_{i}$  and  $Y_{j}$ . Variables are ordered by their distance from the observed variable in the Bayes net, increasing. We estimate each variable given the observed variable and previously generated scaffold variables using  $q$  before estimating the target probability. We approximately marginalize over the scaffold variables' values using  $M$  Monte Carlo samples from the conditional distributions.
$$
\hat {q} _ {S} \left(Y _ {i} = y _ {i} \mid Y _ {j} = y _ {j}\right) = \frac {1}{M} \sum_ {k = 1} ^ {M} q \left(Y _ {i} = y _ {i} \mid \left\{Y _ {s} = y _ {s} ^ {(k)} \right\} _ {s \in S}, Y _ {j} = y _ {j}\right) \tag {2}
$$
$$
\text {w h e r e} y _ {s} ^ {(k)} \sim q \left(Y _ {s} \mid \left\{Y _ {t} = y _ {t} ^ {(k)} \right\} _ {t \in S \mid t \prec s}, Y _ {j} = y _ {j}\right) \tag {3}
$$
We will denote the distribution defined by the sampling procedure in 2 by  $\hat{q}_S$ .
Free generation The free generation estimator tests whether trained models spontaneously generate useful intermediate variables. It is similar to scaffolded generation in that it generates intermediate variables before the target, but free generation uses the model to choose which variables to instantiate in addition to estimating their values. We simply sample variable indices and values from  $q$  until it generates the index of the target variable. At that point, we compute the probability of the target variable. We average over  $M$  such samples.
In order to understand the range of situations in which reasoning might be useful, we study the conditions under which a risk-minimizing sequence model  $q$  benefits from reasoning through intermediate variables. We analyze a simplified version of our task in which a model is trained on pairs of adjacent variables from a directed chain. In this setting, we prove that the minimizer of a risk consisting of cross entropy with entropy regularization exhibits a reasoning gap: the bias of direct conditional probability estimates between pairs of random variables that do not appear together is higher than indirect estimates that chain together conditional probabilities of intermediate random variables.
We assume that a risk minimizer  $q$  is trained on a sequence of alternating variable indices  $i_t \in \{1, \dots, N\}$  and variable values  $Y_i \in \mathcal{X}$  where the  $Y_i$ 's take support in some finite set  $\mathcal{X}$ . Variable values always follow variable indices. We assume that the joint distribution  $p_d$  over  $Y_1, Y_2, \ldots, Y_N$
factorizes as  $p_d(Y_1, \ldots, Y_N) = p_d(Y_1) \prod_{i=1}^N p_d(Y_{i+1}|Y_i)$ . The observation distribution  $p_{\mathrm{obs}}$  only assigns non-zero probability to adjacent variable pairs, i.e.  $p_{\mathrm{obs}}(\{i,j\}) = 0$  if  $|i-j| > 1$  and  $p_{\mathrm{obs}}(X) = 0$  if  $|X| \neq 2$ . The training set consists of i.i.d. samples from  $p$ , which is the distribution over complete sequences of variable indices and values defined by  $p_{\mathrm{obs}}$  and  $p_d$ .
We show that when our data has the locality structure as described, marginalizing over intermediate random variables using the learned conditional probabilities between adjacent variables leads to lower-bias estimates than using the learned estimate directly:
Theorem 3.1. Let  $S$  be the space of possible sequences consisting of variable indices followed by variable values. Let  $u$  be the uniform distribution over  $S$ . 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 {4}
$$
Let  $q^{*} = \arg \min_{q}R(q)$  be a minimizer of the risk over all possible probability distributions. Then, for all non-adjacent random variables  $Y_{i}$  and  $Y_{j}$ , reasoning through the intermediate variables has lower bias than direct prediction. That is, for any  $y_{i},y_{j}\in \mathcal{X}$ :
$$
\begin{array}{l} \left| \mathbb {E} _ {q _ {S} ^ {*}} \big [ \hat {q} _ {S} (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) \big ] - p _ {d} (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) \right| ^ {2} \\ <   | \hat {q} _ {D} (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) - p _ {d} (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) | ^ {2} \\ \end{array}
$$
Here, the expectation is over the randomness in the Monte Carlo samples of intermediate variables.
Proof sketch. We characterize  $q^*$  in terms of  $p_d$  and  $u$  by analyzing the risk (Eq 4), which decomposes into a sum of cross-entropy terms, either between  $q(\cdot)$  and  $p(\cdot)$  or between  $q(\cdot)$  and  $u$ .
We consider two cases. For any adjacent pairs  $Y_{i}$  and  $Y_{i + 1}$ ,  $q^{*}(Y_{i + 1}|Y_{i})$  will interpolate between the true conditional probability  $p_d(Y_{i + 1}|Y_i)$  and the uniform distribution. This is because the risk will contain a term of the form  $-\mathbb{E}_{p_d(Y_{i + 1}|Y_i)}[\log q(Y_{i + 1}|Y_i)]$  and a term of the form  $-\mathbb{E}_u[\log q(Y_{i + 1}|Y_i)]$ . The minimizer of the sum of these terms is a mixture between  $p_d(Y_{i + 1}|Y_i)$  and the uniform distribution. On the other hand, for non-adjacent pairs  $Y_{i}$  and  $Y_{j}$ ,  $q^{*}(Y_{i}|Y_{j}) = \frac{1}{|\mathcal{X}|}$ . As a consequence of the observation distribution  $p_{\mathrm{obs}}$ ,  $Y_{i}$  and  $Y_{j}$  will never appear together. There will only be a term of the form  $-\mathbb{E}_u[\log q(Y_i|Y_j)]$  in the risk. Therefore, minimizing the risk for  $Y_{i}$  and  $Y_{j}$  is equivalent to minimizing the entropy regularization term since the risk does not penalize  $q^{*}(Y_{i}|Y_{j})$  for not matching  $p_d(Y_i|Y_j)$ .
Given  $q^{*}$ , we can calculate the bias of the scaffolded estimator  $\hat{q}_S$ , where we approximately marginalize out the intermediate random variables by sampling from the conditional probability distributions  $q^{*}(Y_{i + 1}|Y_{i})$ . Since these conditional distributions are mixture distributions, we can also show that the scaffolded estimator is a mixture between  $p(Y_i|Y_j)$  and  $u$  with mixture weight  $\lambda \in (0,1)$ . Therefore, its bias can be expressed as  $(1 - \lambda)^2 |p(Y_i|Y_j) - \frac{1}{|\mathcal{X}|}|$ . The second term in the product is exactly the bias of the direct estimator and  $(1 - \lambda)^2 < 1$ . The desired inequality immediately follows. Intuitively, the bias of the scaffolded estimator is lower because it leverages the dependency structure of the data, chaining together local conditional probabilities of variables that appeared together in training to estimate conditional probabilities for unseen pairs.
In practice, we find that density estimators tend to predict the marginal distributions for held-out pairs, i.e.,  $q^{*}(X_{i}|X_{j}) \approx p(X_{i})$ . If we instead assume that the risk minimizer is a mixture between the true conditional distributions and the marginal distributions, we can relax some assumptions we make about the underlying data distribution. For further discussion, full details of the theorem, and proofs see Appendix A. In the next section, we explore whether the reasoning gap occurs in practice for transformer language models in situations with more complex dependency and observation structures.
4 Experimental methods
We test the hypothesis that locally structured training data leads to a reasoning gap by training a language model (i.e., an autoregressive density estimator with a transformer architecture) and testing conditional inference accuracy. We manipulate the structure of the training data and measure the effectiveness of different estimators. The key question is whether there are training conditions under
which scaffolded or free generation outperforms direct prediction. Code and data are available at https://github.com/benpry/why-think-step-by-step.
As in our theoretical analysis, we are interested in understanding the bias of different estimators. In practice, we evaluate our estimators using Mean Squared Error (MSE), averaging over 10 Monte Carlo samples for estimators that include intermediate variables. Averaging over samples reduces variance, meaning that the remaining estimation error is primarily driven by bias.
4.1 Training data
The first step in generating our training data is to create Bayes nets. We generate a random topology for the networks, creating 100 variables and incrementally adding 100 random edges between pairs of variables. After generating a topology, we create conditional probability tables for each variable. We design the conditional probability tables and select among generated Bayes nets to create non-adjacent pairs of variables with high mutual information. The probability of a variable being 1 for each setting of its parents is randomly sampled from  $\mathrm{Beta}(\frac{1}{5},\frac{1}{5})$ , which favors moderately strong dependencies. If we did not favor strong dependencies, conditional probabilities would be very close to marginal probabilities for non-adjacent pairs of variables and a language model could come very close to matching true conditional probabilities without learning the relationships between variables. Strong dependencies are therefore an important precondition to the effectiveness of reasoning. We generate 100 Bayes nets according to this process, then compute the mutual information between each pair of variables. We select the 50 pairs of non-adjacent variables that have the highest mutual information, then randomly sample 25 of those pairs to hold out from the training set. Finally, we select the 10 Bayes nets, out of the initial 100, for which the 25 held-out pairs of variables have the highest mean mutual information. Pseudocode for Bayes net generation is shown in Algorithm 1 of Appendix B.
For each of the 10 selected Bayes nets, we generate a training set consisting of 1 million samples formatted as strings. The strings consist of the name of each variable followed by its value; variables are ordered randomly. Variable names are the letter 'X' followed by a number from 0 to 99. For example, if the variable 'X42' has value 1, the string would include the line "X42=1". We mention the name of the last variable in the sample at the beginning of the string and mark it with the word "target." We also include '#' before the sample to indicate where it begins. A schematic example of a formatted sample string is shown in Figure 1B and a full example is shown in Appendix C.
Each sample includes only a subset of all the variables in the Bayes net, chosen according to the observation distribution, which is a distribution over subsets of variables. Pseudocode for selecting a subset of variables from an observation distribution is shown in Algorithm 2 in Appendix B. At a high level, observation distributions have three key properties:
Locality Observed samples contain only variables from a local neighborhood, consisting of a central variable along with all variables of distance at most  $k$  from it. To sample from the observation distribution, we sample the central variable uniformly randomly and then sample  $k$  from some distribution that controls the local neighborhood size. In our experiments, we draw  $k$  either from a geometric distribution with a parameter of 0.5 or a Zipfian distribution with a parameter of 2.