SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathbb {E} _ {p \left(i _ {1: 2}, v _ {1: 2}\right)} \left[ - \log q \left(v _ {2} \mid i _ {1: 2}, v _ {1}\right) \right] (14) \\ = \mathbb {E} _ {p \left(i _ {1: 2}, v _ {1}\right)} \left[ \mathbb {E} _ {p \left(v _ {2} \mid i _ {1: 2}, v _ {1}\right)} \left[ - \log q \left(v _ {2} \mid i _ {1: 2}, v _ {1}\right) \right] \right] (15) \\ = \sum_ {i _ {1}} \sum_ {v _ {1}} \sum_ {i _ {2}} \mathbb {E} _ {p \left(v _ {2} \mid i _ {1: 2}, v _ {1}\right)} \left[ - \log q \left(v _ {2} \mid i _ {1: 2}, v _ {1}\right) \right] p \left(i _ {1: 2}, v _ {1}\right) (16) \\ = \sum_ {\left(i _ {1}, i _ {2}\right) \in \mathcal {Y} _ {\text {o b s}}} \sum_ {v _ {1}} \mathbb {E} _ {p \left(v _ {2} \mid i _ {1: 2}, v _ {1}\right)} \left[ - \log q \left(v _ {2} \mid i _ {1: 2}, v _ {1}\right) \right] p \left(i _ {1: 2}, v _ {1}\right) (17) \\ \end{array}
$$
Importantly, notice that the outer sum over the variable indices  $i_1, i_2$  is over  $\mathcal{V}_{\mathrm{obs}}$  due to the observation distribution.
On the other hand, consider the right sum.
$$
\begin{array}{l} \mathbb {E} _ {u \left(i _ {1: 2}, v _ {1: 2}\right)} [ - \log q \left(v _ {2} \mid i _ {1: 2}, v _ {1}\right) ] (18) \\ = \mathbb {E} _ {u (i _ {1: 2}, v _ {1})} [ \mathbb {E} _ {u (v _ {2} | k _ {1: 2}, v _ {1})} [ - \log q (v _ {2} | i _ {1: 2}, v _ {1}) ] ] (19) \\ = \sum_ {i _ {1}} \sum_ {v _ {1}} \sum_ {i _ {2}} \mathbb {E} _ {u \left(v _ {2} \mid i _ {1: 2}, v _ {1}\right)} \left[ - \log q \left(v _ {2} \mid i _ {1: 2}, v _ {1}\right) \right] u \left(i _ {1: 2}, v _ {1}\right) (20) \\ = \sum_ {\left(i _ {1}, i _ {2}\right) \in \mathcal {Y} \times \mathcal {Y}} \sum_ {v _ {1}} \mathbb {E} _ {u \left(v _ {2} \mid i _ {1: 2}, v _ {1}\right)} \left[ - \log q \left(v _ {2} \mid i _ {1: 2}, v _ {1}\right) \right] u \left(i _ {1: 2}, v _ {1}\right) (21) \\ \end{array}
$$
We consider two different cases. Suppose the variable indices are adjacent so that  $i_2 = i + 1$  and  $i_1 = i$  for some  $i$ . In addition, fix some value for  $v_1$ . Then, we can take exactly two terms, one from the left and right sum.
$$
\begin{array}{l} \frac {1}{2} \mathbb {E} _ {p (v _ {2} | i _ {1: 2}, v _ {1})} [ - \log q (v _ {2} | i _ {1: 2}, v _ {1}) ] p (i _ {1: 2}, v _ {1}) \\ + \frac {1}{2} \mathbb {E} _ {u \left(v _ {2} \mid i _ {1: 2}, v _ {1}\right)} \left[ - \log q \left(v _ {2} \mid i _ {1: 2}, v _ {1} \right] u \left(i _ {1: 2}, v _ {1}\right) \right. \tag {23} \\ \end{array}
$$
The expectations that appear in Eq 23 are either 1) a cross entropy between either  $q(Y_{i}|Y_{j})$  and  $p_d(Y_i|Y_j)$  or 2) a cross entropy between  $q(Y_{i}|Y_{j})$  and the uniform distribution. By an application of Proposition 1, the sum of these two terms is minimized by taking  $q^{*}(Y_{i}|Y_{j}) = \lambda_{i,j}\frac{1}{|\mathcal{X}|} + (1 - \lambda_{i,j})p(Y_i|Y_j)$  for  $\lambda_{i,j} = \frac{u(i_{1:2},v_1)}{u(i_{1:2},v_1) + p_d(Y_i|Y_j)}$ . The result holds for any adjacent pairs.
Finally, we consider non-adjacent pairs  $Y_{i}$  and  $Y_{j}$ . By construction, we can never have  $i_{2} = i, i_{1} = j$  in the left sum. We will only have one cross entropy term between  $q^{*}(Y_{i}|Y_{j})$  and  $u(Y_{i})$ . Therefore, minimizing the risk is equivalent to minimizing the entropy regularization term and taking  $q^{*}(Y_{i}|Y_{j}) = \frac{1}{|\mathcal{X}|}$ .
In the following corollary, we will show that the squared bias of the scaffolded generation estimator, which explicitly marginalizes out intermediate variables, is less than the squared bias of the direct estimator. For this result, we assume that the true conditional probabilities satisfy a doubly stochastic condition; that is,  $\sum_{y_i}p(Y_i = y_i|Y_j = y_j) = 1 = \sum_{y_j}p(Y_i = y_i|Y_j = y_j)$ . We discuss how to relax this assumption under different assumptions about the risk minimizer.
Theorem A.2. For all  $y_{i},y_{j}\in \mathcal{X}$  with  $|i - j| > 1$
$$
\left| \mathbb {E} [ \hat {q} _ {S} ^ {*} (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) ] - p (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) \right| ^ {2} <   \left| \hat {q} _ {D} ^ {*} (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) - p (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) \right| ^ {2}
$$
Proof. First, we show that  $\mathbb{E}[\hat{q}_S^* (Y_i = y_i|Y_j = y_j)] = \lambda p(Y_i = y_i|Y_j = y_j) + (1 - \lambda)\frac{1}{|\mathcal{X}|}$  for some  $\lambda \in (0,1)$ . Here the expectation is taken with respect to  $y_{i}\sim q^{*}(Y_{i}|Y_{i - 1} = y_{i - 1}),y_{i - 1}\sim$ $q^{*}(Y_{i - 1}|Y_{i - 2} = y_{i - 2})\ldots y_{j + 1}\sim q^{*}(Y_{j + 1}|Y_{j} = y_{j})$ . For conciseness, in our notation, we omit the random variables that the expectation is taken with respect to.
We prove this by induction on  $|i - j|$ , the distance between the nodes corresponding to  $Y_{i}$  and  $Y_{j}$  in the directed chain. We consider the base case  $|i - j| = 2$ .
$$
\begin{array}{l} \mathbb {E} [ \hat {q} _ {S} ^ {*} (Y _ {3} = y _ {3} | Y _ {1} = y _ {1}) ] = \mathbb {E} [ (1 - \lambda_ {3, 2}) p (Y _ {3} = y _ {3} | Y _ {2} = y _ {2}) + \lambda_ {3, 2} \frac {1}{| \mathcal {X} |} ] \\ = \lambda_ {3, 2} \frac {1}{| \mathcal {X} |} + (1 - \lambda_ {3, 2}) \mathbb {E} [ p (Y _ {3} = y _ {3} | Y _ {2} = Y _ {2}) ] \\ = \lambda_ {3, 2} \frac {1}{| \mathcal {X} |} + (1 - \lambda_ {3, 2}) \sum_ {y _ {2}} p \left(Y _ {3} = y _ {3} \mid Y _ {2} = y _ {2}\right) \left[ (1 - \lambda_ {2, 1}) p \left(Y _ {2} = y _ {2} \mid Y _ {1} = y _ {1}\right) + \lambda_ {2, 1} \frac {1}{| \mathcal {X} |} \right] \\ = \lambda_ {3, 2} \frac {1}{| \mathcal {X} |} + (1 - \lambda_ {3, 2}) (1 - \lambda_ {2, 1}) \sum_ {y _ {2}} p \left(Y _ {3} = y _ {3} \mid Y _ {2} = y _ {2}\right) p \left(Y _ {2} = y _ {2} \mid Y _ {1} = y _ {1}\right) \\ + (1 - \lambda_ {3, 2}) \lambda_ {2, 1} \sum_ {y _ {2}} \frac {1}{| \mathcal {X} |} p \left(Y _ {3} = y _ {3} \mid Y _ {2} = y _ {2}\right) \\ = \lambda_ {3, 2} \frac {1}{| \mathcal {X} |} + (1 - \lambda_ {3, 2}) (1 - \lambda_ {2, 1}) p (Y _ {3} = y _ {3} | Y _ {1} = y _ {1}) \\ + (1 - \lambda_ {3, 2}) \lambda_ {2, 1} \frac {1}{| \mathcal {X} |} \\ = (1 - \lambda) \frac {1}{| \mathcal {X} |} + \lambda p (Y _ {3} = y _ {3} | Y _ {1} = y _ {1}) \\ \end{array}
$$
Here  $\lambda = (1 - \lambda_{3,2})(1 - \lambda_{2,1})$ . Importantly, the expected value is a strict convex combination of the uniform and conditional probabilities. For the induction step, we note that the expectations  $\mathbb{E}[\hat{q}_S^* (Y_i = y_i|Y_j = y_j)]$  and  $\mathbb{E}[\hat{q}_S^* (Y_{i - 1} = y_{i - 1}|Y_j = y_j)]$  are related as follows
$$
\mathbb {E} [ \hat {q} _ {S} ^ {*} (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) ] = \lambda \frac {1}{| \mathcal {X} |} + (1 - \lambda) \sum_ {y _ {i - 1}} p (Y _ {i} = y _ {i} | Y _ {i - 1} = y _ {i - 1}) \times \mathbb {E} [ \hat {q} _ {S} ^ {*} (Y _ {i - 1} = y _ {i - 1} | Y _ {j} = y _ {j}) ]
$$
We now use this characterization of the scaffolded estimator as a mixture distribution to prove the desired inequality. The bias of direct estimator can be computed as  $\left|\frac{1}{|\mathcal{X}|} - p(Y_i = y_i | Y_j = y_j)\right|^2$ . By the previous results,  $\mathbb{E}\left[\hat{q}_S^*(Y_i = y_j | Y_i = y_i)\right] = \lambda p(Y_i | Y_j) + (1 - \lambda) \frac{1}{|\mathcal{X}|}$  for some  $\lambda \in (0, 1)$ .
$$
\begin{array}{l} | \mathbb {E} [ \hat {q} _ {S} ^ {*} (Y _ {i} | Y _ {j} = y _ {j}) ] - p (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) | ^ {2} \\ = | \lambda p (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) + (1 - \lambda) \frac {1}{| \mathcal {X} |} - p (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) | ^ {2} \\ = | \lambda p (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) + \frac {1}{| \mathcal {X} |} - \lambda \frac {1}{| \mathcal {X} |} - p (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) | ^ {2} \\ = | (\lambda - 1) p (Y _ {i} = y _ {i} \mid Y _ {j} = y _ {j}) + (1 - \lambda) \frac {1}{| \mathcal {X} |} | ^ {2} \\ = (1 - \lambda) ^ {2} | (\frac {1}{| \mathcal {X} |} - p (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) | ^ {2} \\ <   | \frac {1}{| \mathcal {X} |} - p (Y _ {i} = y _ {i} | Y _ {j} = y _ {j}) | ^ {2} \\ \end{array}
$$
Interestingly, if we assumed that the risk minimizer interpolates between the marginals and the true conditionals, we could relax the doubly stochastic condition.
B Pseudocode for data generation
This section contains pseudocode for the algorithms we use to generate Bayes nets and select variables to include in a given sample according to the observation distribution.
First, we use Algorithm 1 to create Bayes nets. This algorithm takes in a number of nodes and edges to create and defines a directed acyclic graph by randomly adding edges between pairs of nodes. Next, it assigns conditional probability tables to the nodes given the values of their parents to create a Bayes net.
For each sample, we only show a subset of all the variables according to an observation distribution. Algorithm 2 shows the procedure we use to select which variables to display in a given sample. We first sample central variable  $c$  (uniformly) and a distance  $k$  (according to a geometric or Zipfian distribution), then get all the variables within distance  $k$  of  $c$  in the graph. Next, we remove each variable with probability 0.2. If any of the held-out pairs remain, we randomly remove one variable from each pair. In the wrong locality structure training conditions, we use a graph  $G$  that does not correspond to the net from which our samples are drawn but has the same variable names. In the fully observed condition, we start with all variables, but still remove held-out pairs using the final for loop.
C Full sample of training data
The following is an example of a sample from a local neighborhood of a Bayes net in string format, as used in training our models:
target: X5
X17=0
X92=0
X13=0
Algorithm 1 Algorithm for generating a Bayes net
Input: Number of nodes  $N$ , number of edges  $M$ $G = (E, V) \gets$  empty graph  
for  $i \in \{1, \dots, N\}$  do  
 $V \gets V \cup \{Xi\}$   
end for  
for  $i \in \{1, \dots, M\}$  do  
 $v_1, v_2 \gets$  random pair of vertices  $\in V$   
while  $(v_1, v_2) \in E$  or  $(v_1, v_2)$  would create a cycle do  
 $v_1, v_2 \gets$  random pair of vertices  $\in V$   
end while  
 $E \gets E \cup \{(v_1, v_2)\}$   
end for  
 $T \gets$  empty associative array  
for  $v \in$  TOPOLOGICALSORT $(G)$  do  
 $t \gets$  empty conditional probability table  
for  $c \in$  possible configurations of  $v$ 's parents do  
 $p \gets$  SAMPLE(BETA $\left(\frac{1}{5}, \frac{1}{5}\right)$ )  
 $t[c] \gets p$   
end for  
 $T[v] \gets t$   
end for  
 $D \gets$  Bayesian network defined by graph  $G$  and conditional probability tables  $T$   
return  $D$
Algorithm 2 Algorithm for selecting variables to include based on an observation distribution
Input: Bayes net graph  $G = (V,E)$  , set of held-out pairs  $P$  size distribution  $D$ $c\gets$  random variable  $\in V$ $k\gets$  SAMPLE(D)   
 $R\gets \{\}$    
for  $v\in V$  do if  $v$  is within distance  $k$  of  $c$  in  $G$  then  $R\leftarrow R\cup \{v\}$  end if   
end for   
for  $v\in R$  do if SAMPLE(Bernoulli, 0.2) = 1 then  $R\gets R\setminus \{v\}$  end if   
end for   
for  $(v_{1},v_{2})\in P$  do if  $v_{1}\in R$  and  $v_{2}\in R$  then if SAMPLE(Bernoulli, 0.5) = 1 then  $R\gets R\setminus \{v_1\}$  else  $R\gets R\setminus \{v_2\}$  end if   
end if   
end for   
return  $R$
X52=1  
X24=1  
X26=1  
X91=0  
X36=0  
X34=0  
X12=1  
X20=0  
X5=1
The full training set consists of 1,000,000 samples like this concatenated together.
D Formatting estimators as prompts
D.1 Direct prediction
In direct prediction, we use a simple prompt that specifies the label and value for the observed variable then states the name of the target variable. For example, if we wanted to estimate  $p(X_{2}|X_{1} = 0)$  we would use the following prompt:
###
target: X2
X1=0
X2=
We would then take the softmax of the log probabilities the language model assigns to '1' and '0' being the next token and use the resulting probability of '1' as the estimate.
D.2 Scaffolded generation
In scaffolded generation, we pre-compute a sequence of intermediate variables, then use the language model to estimate their values. For example, if we wanted to estimate  $p(X_4|X_1)$  and we knew  $X_2$  and  $X_3$  were scaffold variables, we would start with the following prompt:
```c
target: X4
X1=0
X2=
We would then sample the next token (either 0 or 1) from the language model and append it to the prompt, followed by the name of the next scaffold variable on a new line. For example, if we got a value of 1 from the language model, we would give it the following prompt next:
###
target: X4
X1=0
X2=1
X3=
We repeat this process until all scaffold variables have values, then compute the probability assigned to the target variable at the end, in the same way we do for direct prediction. The final prompt we use to estimate the probability might look like this:
###
target: X4
X1=0
X2=1
X3=0
X4=
We repeat this sampling process 10 times to produce a Monte Carlo estimate over the values of the scaffold variables and average the probabilities assigned to the target variable over samples.
D.3 Free generation
Free generation is similar to scaffolded generation, but we sample from the model to choose which intermediate variables to instantiate, rather than just their values. For example, if we were inferring  $p(X_4|X_1 = 0)$  we would first prompt the model like this:
```c
target: X4
X1 = 0
We then sample the next two tokens from the model, which is exactly enough for one variable name. We add an equals sign, then sample the variable's value in the same way we do in other estimators. For example, our prompt might look like this after generating one intermediate variable and its value
```c
target: X4
X1 = 0
X5 = 1
We repeat this process until the model outputs the name of the target variable. At that point, we have a prompt that might look like this
###
target: X4
X1=0
X5=1
X2=0
X7=0
X3=1
X4=
We then extract the probability of the target variable as in other estimators. Again, we average target variable probabilities over 10 samples of the intermediate variables.
The main architecture we used is a smaller version of the GPT-2 architecture, with 512-dimensional embeddings, 10 layers, and 8 attention heads. It has 32, 573, 440 parameters in total. We fit a Byte Pair Encoding tokenizer to data from our samples, which produced 356 unique tokens.
Our text was grouped into chunks of 1024 tokens and batches of 3, for a total of 3072 tokens per gradient step. We used the Adam optimizer, with an initial learning rate of  $10^{-3}$  and Beta values of 0.9 and 0.999.
All models were trained on Nvidia Titan Xp GPUs. They were trained for 300,000 gradient steps, which took approximately 20 hours each.
F Comparison of reasoning gaps across different architectures
To ensure that our results are robust across different transformer architectures, we compare mean squared error by estimator for four alternative architectures. We focus on training data consisting of geometrically-sized local neighborhoods. Figure 4 shows mean squared error by architecture type trained on geometrically-sized local neighborhoods of variables. Results are very similar across all architectures, except for the tiny architecture which fails to match the distribution at all.
We compare the original architecture to a smaller (4, 473, 600-parameter) architecture, a larger (86, 628, 864-parameter) architecture, and a wider architecture with larger embeddings but fewer
Figure 4: Mean squared error by architecture type and estimator. All architectures were trained on the same data consisting of geometrically-sized local neighborhoods of data. Error bars indicate  $95\%$  confidence intervals.
Table 3: Embedding size, number of layers, and number of attention heads for each of the five architectures used in the architecture comparison
Architecture type	Embedding Size	# Layers	# Attention Heads
original	512	10	8
smaller	256	5	4
larger	768	12	12
wider	1024	3	8
tiny	4	2	2
layers (39, 887, 872 parameters). Finally, we compare against a tiny architecture that is too small to match the distribution of the true Bayes net (8,644 parameters). Table 3 shows the embedding sizes, numbers of layers, and numbers of attention heads for each architecture.
G Mean squared error by number of samples
As mentioned in the main text, the number of Monte Carlo samples taken over intermediate variables influences the variance of the scaffolded and free generation estimators. We demonstrate this computationally by reporting the results of each estimator for different numbers of samples for the model trained on geometrically-sized local neighborhoods of variables. Figure 5 shows MSE by number of samples using each estimator that averages over samples, compared against the baseline of direct prediction. Free and scaffolded generation both have higher MSE than direct prediction with one or two samples due to the increase in variance. However, they achieve lower MSE when averaging over more than three samples. Negative scaffolded generation exhibits very similar performance to direct prediction for all numbers of samples, which suggests that the values of intermediate variables do not influence the final prediction for that estimator.
H Data efficiency of fully observed training with no held-out pairs
To supplement our data efficiency analyses, we ran a version of the fully observed training condition without any held-out pairs. In this setting, the transformer can directly memorize the true conditional probabilities, but it takes a long time for direct prediction to perform as well as free generation with locally-structured data. Figure 6 compares the performance of direct prediction and free generation
Figure 5: Mean squared error by number of Monte Carlo samples of intermediate variables for models trained on geometrically-sized local neighborhoods of variables. The blue line indicates the MSE achieved by direct prediction. Ribbons indicate  $95\%$  confidence intervals.
Figure 6: Learning curves comparing mean squared error on held-out pairs, estimated using free and direct prediction for training data consisting of either geometrically-sized local neighborhoods or the full set of variables. Unlike the version reported in the main text, no pairs of variables are held out in this fully observed condition. Even though the model is trained directly on the held-out pairs in the fully observed condition, there is a substantial data efficiency advantage to using locally structured training data and free generation at inference time.
for locally structured training data with held-out pairs and fully observed training data with no held-out pairs. Direct prediction with fully observed training data can outperform free generation with locally-structured training data. However, it takes approximately 650 million tokens of training to reach the same performance that free generation with geometric locally-structured training data converges to after only about 200 million tokens of training.