SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For the first task, if the aggregation function is max or min, we use one attention head to simply copy the embedding  $e_j^{\mathrm{dp}}$  for index  $j$  such that  $f_j^{\mathrm{answer}} = 1$  and  $e_j^{\mathrm{dp}}$  is the largest/smallest, according to Lemma C.7. If the aggregation function is  $\sum$ , we use two attention heads, where one attention head computes the mean of  $e_j^{\mathrm{dp}}$  for index  $j$  such that  $f_j^{\mathrm{answer}} = 1$ , and the other attention head calculates the fraction of elements in the sequence such that  $f_j^{\mathrm{answer}} = 1$ . Finally, the second task is a conditional selection operation and thus can be implemented by an MLP (Lemma C.4).
G.3 Proof of the Theorem 4.8
Theorem G.2. Assume  $\mathsf{TC}^0\neq \mathsf{NC}^1$ . There exists a context-free language such that for any depth  $L$  and any polynomial  $Q$ , there exists a sequence length  $n\in \mathbb{N}$  where no log-precision autoregressive transformer with depth  $L$  and hidden dimension  $d\leq Q(n)$  can generate the correct answer for the CFG Membership Testing problem for all input strings of length  $n$ .
Proof. According to Barrington's theorem [4, 5], given a fixed unsolvable automaton, judging whether the automaton accepts an input is complete in  $\mathsf{NC}^1$ , which is a special case for the CFG Membership Testing problem. Therefore, the CFG Membership Testing problem is  $\mathsf{NC}^1$ -hard. With the assumption that  $\mathsf{TC}^0 \neq \mathsf{NC}^1$ , the CFG Membership Testing problem is out of the capacity of the log-precision autoregressive transformer.
In this section, we present the experimental details.
We set the number field  $p = 11$  in the math experiments. In the LIS experiment, we set the number of different input tokens to 150; in the ED experiment, we set the number of different input tokens to 26. The vocabulary is constructed by including all symbols. For all tasks and settings (direct v.s. CoT), the size of the training and testing dataset is 1M and 0.1M respectively. The constructions of different datasets are introduced below.
Arithmetic Expression. All arithmetic expression problems are generated according to Algorithm 1. In Algorithm 1, we first create a number that serves as the answer to the problem. We then decompose the number using sampled operators sequentially, serving as the problem, until the maximum number of operators is met. The CoT procedure is precisely defined by reversing this problem generation process. For example, a sample in the direct dataset looks like
$$
1 + 5 \times (1 - 2) = 7
$$
while the corresponding sample in the CoT data looks like
$$
1 + 5 \times (1 - 2) = 1 + 5 \times 1 0 = 1 + 6 = 7
$$
Algorithm 1: Arithmetic Expression Problem Generation
Input :Number of Operators n  
Input :Vocabulary of numbers  $V = \{0,1\dots 10\} //$  number field  $p = 11$   
Output :Arithmetic expression s  
1 Sample the first number  $t$  uniformly from  $V$ ;  
2  $s = []$ ;  
3 Append  $t$  to  $s$ ;  
4 for  $i \gets 1$  to  $n$  do  
5 Sample  $p$  uniformly from  $\{0,1,\dots,len(s) - 1\}$ , satisfying  $s[p]$  is a number;  
6 Sample  $o$  uniformly from  $\{+, -, \times, \div\}$ ;  
7 Sample numbers  $t_1, t_2$ , satisfying the result of  $o(t_1, t_2)$  equals  $s[p]$ ;  
8 if  $s[p - 1] = \div$  or  $(o \in \{+, -\}$  and  $s[p - 1] \in \{-, \times\})$  or  $(o \in \{+, -\}$  and  $s[p + 1] \in \{x, \div\})$  then  
9 pop  $s[p]$ ;  
10 insert  $([], [t_1], [o], [t_2], ])$  sequentially into  $s[p]$ ;  
11 else  
12 pop  $s[p]$ ;  
13 insert  $[t_1], [o], [t_2]$  sequentially into  $s[p]$ ;  
14 end  
15 end
Linear Equation. All linear equation problems are generated according to Algorithm 2. In Algorithm 2, we consider the linear systems that only have a unique solution. Given a sampled linear system that satisfies this condition, we "translate" it to a sequence by concatenating all the equations (separated by commas), which serves as the problem. The answer to the problem is also a sequence consisting of variables and the corresponding values. The CoT solution of each problem is
the calculation process of the Gaussian elimination algorithm applied to each variable sequentially. For example, a sample in the direct dataset looks like
$2x_{1} + 3x_{2} + 3x_{3} = 8,1x_{1} + 7x_{2} + 0x_{3} = 0,0x_{1} + 2x_{2} + 1x_{3} = 1,$  [SEP]  $x_{1} = 4,x_{2} = 1,x_{3} = 10$  while the corresponding sample in the CoT dataset looks like
$$
2 x _ {1} + 3 x _ {2} + 3 x _ {3} = 8, 1 x _ {1} + 7 x _ {2} + 0 x _ {3} = 0, 0 x _ {1} + 2 x _ {2} + 1 x _ {3} = 1,
$$
$$
[ \text {S E P} ] x _ {1} + 7 x _ {2} + 7 x _ {3} = 4, 0 x _ {2} + 4 x _ {3} = 7, 2 x _ {2} + 1 x _ {3} = 1,
$$
$$
[ \text {S E P} ] x _ {1} + 9 x _ {3} = 6, x _ {2} + 6 x _ {3} = 6, 4 x _ {3} = 7,
$$
$$
[ \text {S E P} ] x _ {1} = 4, x _ {2} = 1, x _ {3} = 1 0,
$$
Algorithm 2: Linear Equation Data Generation
Input :Number of Variable n  
Input :Vocabulary of numbers  $V = \{0,1\dots 10\} //$  number field  $p = 11$   
Output:Linear Equation s  
1 Sample b uniformly from  $V^{n\times 1}$   
2 do  
3 Sample A uniformly from  $V^{n\times n}$   
4 while A is not invertible;  
5  $s \gets "A_{11}x_1 + \ldots + A_{1n}x_n = b_1, \ldots, A_{n1}x_1 + \ldots + A_{nn}x_n = b_n"$
Longest Increasing Subsequence. All input sequences (i.e., problems) are generated according to Algorithm 3. To make the task challenging enough, we first concatenate several increasing subsequences of given length, and then randomly insert numbers into the whole sequence. The inputs have 150 different tokens, ranging from 101 to 250 to avoid token overlap with DP array. The CoT solution to the problem is the DP array plus the final answer, which is defined in (20). Here, we consider the DP formulation (20) because the CoT output length is much shorter than the one corresponding to formulation (21). This allows us to consider more challenging input sequences with longer length. While this DP formulation does not precisely obey the theoretical assumption given in Assumption 4.3, we found that the Transformer can still learn it easily.
For example, a sample in the direct dataset looks like
103 107 109 112 101 103 105 107 115 109 111 113 102 [SEP] 7
while the corresponding sample in the CoT dataset looks like
103 107 109 112 101 103 105 107 115 109 111 113 102
[SEP] 1234123455672
[SEP] 7
Edit Distance. All input sequences (i.e., problems) are generated according to Algorithm 4. In Algorithm 4, we generate the first string randomly. For the generation of the second string, we use two methods. In the first method, we generate the second string randomly, corresponding to a large edit distance. In the second method, we copy the first string with random corruption, corresponding to a small edit distance. The two strings are concatenated by "l", and the concatenation is used as the model input. For the calculation of edit distance, we assign different costs to different operators. The costs for the ADD and DELETE operators are set to 2, while the REPLACE operator is assigned a cost of 3, since REPLACE should be more costly than ADD/DELETE while less costly than their summation. The CoT procedure is also the DP array, defined in Section 4.1. For example, a sample in the direct dataset looks like
$$
\mathrm {a s} \mid \mathrm {p a s s} [ \mathrm {S E P} ] 4
$$
a s|p a s s
[SEP] 3246
[SEP] 5424
[SEP] 4
Algorithm 3: LIS Data Generation
Input:Sequence Length n  
Input :Vocabulary of numbers  $V = \{101, 101\dots 250\}$   
Output:Sequence s  
1 Sample l uniformly from {3,4...n};  
2 Sample t uniformly from {1,2,3};  
3 a = [];  
4 push 0 to a;  
5 if t = 2 then  
6 Sample j uniformly from {1,2...[l/2] + 1};  
7 push j to a;  
8 else if t = 3 then  
9 Sample j uniformly from {1,2...[l/3] + 1};  
10 Sample k uniformly from {1,2...[(l-j)/2] + 1};  
11 push j to a;  
12 push j + k to a;  
13 push l to a;  
14 s ← Sample l numbers from V;  
15 for i ← 1 to t do  
16 Sort s[a[i-1]:a[i]];// This process makes sure the LIS of the generated sequence s is at least [l/t].  
17 end  
18 r ← Sample n - l numbers from V;  
19 Randomly insert r into s;
Algorithm 4: ED Data Generation
Input:Length of the First String  $n$    
Input:Alphabet  $V = \{a,b\dots z\}$    
Output:Sequence  $s_1,s_2$    
1 Sample t uniformly from  $\{3,4\dots 10\}$  .   
2  $T\gets$  Sample t letters from  $V$  .   
3  $s_1\gets$  Sample n letters uniformly from  $T$  .   
4 Sample p uniformly from [0,1] ;   
5 if  $p <   0.4$  then   
6 Sample l uniformly from  $\{n - 3,n - 2,\dots ,n + 2\}$  .   
7  $s_2\gets$  Sample l letters uniformly from  $T$  .   
8 else   
9 do   
10  $s_2\gets s_1$  .   
for i  $\leftarrow$  1 to n do   
12 Sample p uniformly from  $\{0,1\dots len(s_2) - 1\}$  .   
13 Sample l uniformly from  $T$  .   
14 Randomly conduct one of the followings: pop  $s_2[p]$  , substitute  $s_2[p]$  with  $l$  insert  $l$  into  $s_2[p]$  .   
end   
while  $\mathrm{len}(s_2)$  not in  $[n - 3,n + 2]$  .   
17 end
H.2 Model training
We use the minGPT implementation for all the experiments, where the detailed Transformer layer are listed below.
$$
\boldsymbol {X} ^ {(0)} = \operatorname {L a y e r N o r m} \left(\left[ \boldsymbol {v} _ {1} + \boldsymbol {p} _ {1}, \dots , \boldsymbol {v} _ {n} + \boldsymbol {p} _ {n} \right] ^ {\top}\right) \tag {23}
$$
$$
\operatorname {A t t n} ^ {(l)} (\boldsymbol {X}) = \sum_ {h = 1} ^ {H} \left(\operatorname {s o f t m a x} \left(\boldsymbol {X} \boldsymbol {W} _ {Q} ^ {(l, h)} \left(\boldsymbol {X} \boldsymbol {W} _ {K} ^ {(l, h)}\right) ^ {\top} / \sqrt {d} + \boldsymbol {M}\right)\right) \boldsymbol {X} \boldsymbol {W} _ {V} ^ {(l, h)} \boldsymbol {W} _ {O} ^ {(l, h)} \tag {24}
$$
$$
\operatorname {F F N} ^ {(l)} (\boldsymbol {X}) = \sigma \left(\boldsymbol {X} \boldsymbol {W} _ {1} ^ {(l)}\right) \boldsymbol {W} _ {2} ^ {(l)} \tag {25}
$$
$$
\boldsymbol {Y} ^ {(l - 1)} = \boldsymbol {X} ^ {(l - 1)} + \operatorname {A t t n} ^ {(l)} (\operatorname {L a y e r N o r m} (\boldsymbol {X} ^ {(l - 1)})) \tag {26}
$$
$$
\boldsymbol {X} ^ {(l)} = \boldsymbol {Y} ^ {(l - 1)} + \operatorname {F F N} ^ {(l)} \left(\operatorname {L a y e r N o r m} \left(\boldsymbol {Y} ^ {(l - 1)}\right)\right) \tag {27}
$$
We use sinusoidal positional embedding and use Xavier initialization for all the parameters. The activation function is chosen to be GeLU. The dimension of the embedding is set to 256, and the number of heads is set to 4. The hidden size in the FFN layer is set to 1024.
We use the same hyperparameter configuration for all experiments, i.e., the performance comparison between the models trained on the direct and CoT datasets of Arithmetic, Equation, LIS, and ED tasks, and the additional length extrapolation experiments (which we use relative positional encodings [48] instead of absolute positional encodings). In detail, we use AdamW optimizer with  $\beta_{1},\beta_{2} = 0.9$  0.999. The learning rate is set to 1e-4, and the weight decay is set to 0.01. We set the batch size to 512 during training with a linear learning rate decay scheduler. We use learning rate warm-up, and the warm-up stage is set to 5 epochs. The dropout rate is set to 0.1. The total number of training epochs is set to 100.
I Robustness in Training Data Quality
In the training of language models, the quality of data, especially intermediate steps, is crucial for optimal performance. However, real-world training datasets may contain corrupted or missing intermediate steps. This highlights the importance of evaluating model resilience under imperfect conditions. To investigate this, we conducted a series of experiments focused on arithmetic tasks, introducing varying rates of corruption and omission, encapsulated by the parameter  $\gamma$ . For instance, when  $\gamma = 0.1$ , it denotes that  $10\%$  of the intermediate steps are omitted and  $10\%$  of the remaining steps are corrupted through a single-token random replacement. Our motivation for introducing this metric is to simulate potential imperfections in real-world datasets, thus assessing the robustness of our model more thoroughly. The number of operators in the arithmetic task is set to 10.
Table 1: Experimental results of robustness.
γ	0	0.1	0.2	0.3
Accuracy	100.0%	98.5%	97.6%	95.8%
Table 2: Model performance on Arithmetic task with a 3-layer Transformers.
Our experimental results are presented in Table 1. As the rate of corruption and omission increases, there is a slight decrease in accuracy. These results elucidate the inherent robustness of training CoT demonstrations. It is noteworthy that the accuracy of the model remains commendably high despite significant distortions present in the training data, highlighting its potential utility in real-world scenarios where data quality may be compromised.
$^{6}$ https://github.com/karpathy/minGPT