Chain of Thought Empowers Transformers to Solve Inherently Serial Problems

Q1: Constant precision is less practical than log precision.

Response: We respectfully disagree.

First, we would like to highlight that constant precision and log precision are only properly defined for a sequence of language models where the input sequence length goes to infinity. For a concrete and practically used language model, whether it is of constant precision or log precision is an ill-defined question.

Still, we can test if the implications or the predictions of certain precision hold in practice, and use it to infer the practicality of that precision. A most notable application of log precision is the ability to perform uniform attention. In detail, given sequence lengths $n$, it requires that the model can express all fractionals, $1,\ldots, 1/n$. Or as a weaker requirement, it requires that the model can express all integers between $1$ and $n$.

Below we shall see one of the most successful open-sourced LLM, Llama 2, doesn't have the ability to express uniform attention, or just all natural numbers smaller than its sequence length, 4096. Llama 2 uses BF16, which has 1 sign bit, 8 exponent bits, and 7 precision bits (plus an implicit leading bit). This means $2^8+1=257$ is not expressible in BF16, which is much smaller than the sequence length of Llama 2, 4096. To get more intuition about BF16, we note that 255+2 = 256 in BF16. Nowadays practitioners are even pursuing more extreme low precision models, like this recent work by [Peng et al, 2023], which uses FP8 to language models. Here FP8 contains 1 sign bit, 4 exponent bit, and 3 bits of mantissa.

That being said, we also agree with the reviewer that our new AC0 upper bound for constant precision transformers does not imply constant-depth transformers in practice cannot learn problems beyond AC0 without COT. At this time point, it is not clear what architecture and precision transformers will converge to in the future, and thus we believe the study of both precision is meaningful and provides useful insight to the community.

Q2: Comparison to Lemma 3 from [Merrill & Sabharwal, 2023]

Response: First the goals are different. We are trying to use a transformer with CoT to simulate an arbitrarily given circuit, while Lemma 3 from [Merrill & Sabharwal, 2023] aims to compute a fixed problem, which is the Circuit Value Problem. Our goal is more general in the sense that the CVP problem itself can be solved by some circuit and we just need to simulate that circuit. Second, the results are different. The steps of CoT we need is the number of the gates of the circuit to be simulated, while their depth is an upper bound for the depth of the circuits where the CVP problem can be solved. The major difference here is because the vocabulary is finite and each time though it is possible to compute more in parallel, one step of CoT can only write constant bits of information, instead of the whole layer of the circuit.

Q3: Is it possible to show $\mathsf T[\mathrm{poly}(n), 1, 1] = \mathsf{CoT}[1, \mathrm{poly}(n), 1, 1]$?

Response: Good question. We do not know if it is possible. But Theorem 3.8 and 3.9 suggest that one step of COT does not help if one has $\mathsf{poly}(n)$ embedding size.

Q4: Do you have any examples or thoughts about how transformers might systematically use rounding to gain expressive power (with nonzero exponent)?

Response: Excellent question. Unfortunately, we do not have a good answer. We leave it for future work.

References:

• Peng, H., Wu, K., Wei, Y., Zhao, G., Yang, Y., Liu, Z., Xiong, Y., Yang, Z., Ni, B., Hu, J. and Li, R., 2023. FP8-LM: Training FP8 Large Language Models. arXiv preprint arXiv:2310.18313.

We thank you for your appreciation and your effort devoted to the review. We will fix the typo you pointed out in the revision.

Q1: Our computation model is non-uniform.

Response: Thanks you for your this critical question and constructive feedbacks. Please refer to our general response on the non-uniformty issue to all the reviewers.

Q2: An analog of Theorem 3.2 remains open for log-precision transformers even with a constant number of bits for the exponent. Why cannot we absorb constant bits of exponent into additional constant bits of the significand?

Response: A short answer is that the existing proof of $\mathsf{TC}^0$ upper bound relies on the fact that the addition among fixed point numbers satisfy the association rule (assuming no overflow happens), and floating point numbers with any non-zero bits for the exponent does not satisfy the association rule. Taking the floating point system defined in our manuscript as an example, let both the number of exponents and precision, namely, $e$ and $s$ be $1$, we have $\mathbb{F} _{1,1} = \{-3/2,-1,-3/4, -1/2,-1/4, 0,1/4,1/2,3/4,1,3/2\}$. One can note that $5/4$ and $-5/4$ does not belong to $\mathbb{F} _{1,1}$ due to the limited precision of $\mathbb{F} _{1,1}$. As a result, even though $a= -1,b=1,c=1/4$ all belong to $\mathbb{F} _{1,1}$, we have that $[[a+b] _{1,1}+c] _{1,1} \neq [a+[b+c] _{1,1}] _{1,1}$. This is because $[[-1+1] _{1,1}+1/4] _{1,1} = [0+1/4] _{1,1} = 1/4\neq 0 = [-1+1] _{1,1} = [-1+[1+1/4] _{1,1}] _{1,1},$ where the last is because we break the tie in rounding by picking the number with smaller absolute value (See Definition 3.2).

We thank the reviewer for the thorough and constructive feedbacks. We also sincerely appreciate the reviewer's recognition of the significant contribution made by our work. Below we address the concerns of the reviewer one by one and would respectfully ask the reviewer to consider increasing the score if our responses satisfacorily address the concerns.

Q1: Training setup not clear. Is there a risk of overfitting?

Reponse: There is no risk of overfitting because in the experiments, we train in an online setting -- at each step, the training batch is i.i.d. sampled from distribution. We will highlight on this in the revision. Thanks for pointint this out!

Q2: Discrete problems have finite problem space

Reponse: We agree with this point, but we do not agree that this is a weakness. Any language problem with finite vocabulary and finite sequence length has a finite problem space. We hope the reviewer could elaborate more on this critism.

Q3: Sampling method of the data distribution make a huge difference.

Reponse: We agree. But to support our theoretical claim that there is a separation between expressiveness of transformers with and without chain of thought, it suffices to show one data distritbution.

Q4: The model details should not be hidden in appendix. The presentation can be more crips and focused.

Response (updated): Thanks for this helpful suggestion on the presentation. Following your suggestion, we moved the original Definitions 3.1 and 3.2 into the appendix and brought the definition of our transformer architecture into the main paper. See Section 2 in the revision. We also added a clarification in our paper that our model architecture is the same as GPT-2 [Radford et al., 2019], except we use ReLU instead of GeLU. This architecture is also used by GPT-3 with slight modification [Brown et al., 2020].

Q5: Expected background of readers

Reponse: We hope both language model and computational complexity community will find this paper interesting. Given the potential diversity in readers' background and that ICLR has no page limit for the appendix, we choose to make our paper more self-contained by providing those basic definitions in the appendix, including both about decoder-only transformers and complexity.

Q6: The label in Figure 1 for the non CoT example seems wrong

Reponse: It is indeed correct. Here the composition rule for permutation is that, given two permutation $\sigma = (\sigma_1,\ldots,\sigma_5)$, $\pi = (\pi_1,\ldots,\pi_5)$, we define $\sigma\circ \pi \triangleq (\sigma_{\pi_1},\ldots, \sigma_{\pi_5})$. Here in the example of Figure 1, we have $\sigma^{(1)} = (2,3,1,4,5)$, $\sigma^{(2)} = (1,2,4,3,5)$, $\sigma^{(3)} = (5,4,3,2,1)$. Thus for chain of thought, we have $\sigma^{(1)} \circ \sigma^{(2)} = (\sigma^{(1)} _{\sigma^{(2)} _1}, \ldots, \sigma^{(1)} _{\sigma^{(2)} _5}) = (\sigma^{(1)}_1, \sigma^{(1)}_2,\sigma^{(1)}_4, \sigma^{(1)}_3, \sigma^{(1)}_5) = (2,3,4,1,5)$.

Similarly, we have $\sigma^{(1)} \circ \sigma^{(2)} \circ \sigma^{(3)} = \left( \sigma^{(1)} \circ \sigma^{(2)} \right) \circ \sigma^{(3)} = (2,3,4,1,5) \circ (5,4,3,2,1) = (5,1,4,3,2)$, which is the output for the example in Figure 1. We will be more clear about the definition of the composition of permutation in the revision.

Q7: incomplete comments: Additionally, we note that the following sentence in the review "Suggestions: This paper spends a lot of time on the setup and has an extensive collection of supporting material. I like t" is incomplete. We would like to hear the reviewer's suggestion and improve our paper based on it.

References:

• Radford, A., Wu, J., Child, R., Luan, D., Amodei, D. and Sutskever, I., 2019. Language models are unsupervised multitask learners. OpenAI blog, 1(8), p.9.

• Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J.D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A. and Agarwal, S., 2020. Language models are few-shot learners. Advances in neural information processing systems, 33, pp.1877-1901.

We thank the reviewer for the helpful feedback. We have corrected the notations

Q1: The main finding that transformers with chain of thought can do serial computation sounds obvious and is not significant.

Response: While we agree that the ability to perform serial computation might sound natural for computation models with large computation depth, a rigorous construction for a specific architecture, transformers with chain of thought, is novel and not trivial, especially given the various restrictions, like constant precision and $\log(n)$ embedding size. For example, previous work (Perez et al., 2021) also shows that encoder-decoder transformers can perform inherently serial computation by showing that transformer can use one step of chain-of-thought to simulate one step execution of Turing Machine. But this construction uses infinite precision and does not apply in our setting, where only constant precision is allowed.

References:

• Pérez, J., Barceló, P. and Marinkovic, J., 2021. Attention is turing complete. The Journal of Machine Learning Research, 22(1), pp.3463-3497.

Q1: Symbols are not clearly described. The authors don't make clear instructions when using symbols like NC^1, AC^0, etc. It is recommended that these symbols be explained when they are first mentioned.:

Response: Due to the space limit, we defer the definition of complexity class to Appendix B, as noted in our section 2, notations and preliminaries.

Q2: The conclusions of Theorem 3.8 and Theorem 3.9 presented in the paper seem to be contradictory. Theorem 3.8 indicates that enlarging the embedding size to poly(n) doesn't improve the model's expressiveness of T(n) = poly(n) step CoT. On the other hand, Theorem 3.9 demonstrates that broadening the embedding width strengthens the model's power for any particular polynomial T(n) = n^k step CoT.

Reponse: There is no contradiction. We would like to clarify that Theorem 3.8 does not say that "enlarging the embedding size to poly(n) doesn't improve the model's expressiveness of T(n) step CoT for any fixed polynomial T(n)". Instead, Theorem 3.8 should be interpreted as "for every polynomial T(n) and every function expressible by poly embedding size transformers with T(n) embedding size, there is a different and potentially larger polynomial function T'(n) than T(n), and another transformer with log-embedding size and T'(n) steps of COT, which can also express the function".

Q3: In Figure 2, it is observed that when the depth is set to 1, the CoT performs less effectively compared to the decoder-only Transformer without CoT. This difference is particularly noticeable when the label represents the sum of the inputs modulo 2 and exceeds a threshold of 20%.

Response: Thanks for pointing this out. We think this is because the modular addition experiments are the only set of experiments which we do not exactly follow the setting of our theorems. Because it is so easy to achieve perfect accuracy, we use a more challenging training task there, that is, the data distribution contains the input of different sequence length. The idea is that if the models can achieve perfect accuracy on such data distribution, then it shows it can express the ground truth of modular addition for each sequence length. The benefit of doing this is that we can save space for the plot.

However, the above deviation from the theoretical setting could indeed make cot not as powerful as hint when the depth is low, as it is beyond the theoretical guarantee. To avoid confusion, we rerun the experiments for the modular addition using data distribution of fixed sequence length, and the results are presented in Figure 2, Figure 5, and Figure 7 in the revision. In the new experiments, we can clearly see cot is always better than hint, and hint is always better than base.

Due to the time constraints, we only rerun the experiments for depth equal to 1, 2, and 3 and we will add the remaining depth up to 12 in the future revision. Thank you again for this critical and constructive feedback!

4. Can transformer architectures with more trainable parameters when increasing input sequence length achieve length generalization? That is, train on shorter lengths, and generalize to unseen longer lengths.

Response: Honestly we don't believe the answer to the above question is positive. It is well-known that absolute position embedding can lead to the failure of length generalization. However, this does not stop the success of GPT architectures, which use trainable absolute position encoding. The explanation behind this is that in the real world, language models still have a limit for input length. Moreover, for longer input sequences, people also tend to use models with large embedding sizes. In other words, there is some disconnection between practice and the theoretical value of uniformity. We do not think we can decide which regime is correct without further advancement in practice. That being said, we do agree it sounds impossible to achieve length generalization with additional trainable parameters for longer input length.

References:

• Hao, Y., Angluin, D. and Frank, R., 2022. Formal language recognition by hard attention transformers: Perspectives from circuit complexity. Transactions of the Association for Computational Linguistics, 10, pp.800-810.

• Merrill, W., Sabharwal, A. and Smith, N.A., 2022. Saturated transformers are constant-depth threshold circuits. Transactions of the Association for Computational Linguistics, 10, pp.843-856.

• Liu, B., Ash, J.T., Goel, S., Krishnamurthy, A. and Zhang, C., 2022. Transformers learn shortcuts to automata. arXiv preprint arXiv:2210.10749.

• Merrill, W. and Sabharwal, A., 2023. The parallelism tradeoff: Limitations of log-precision transformers. Transactions of the Association for Computational Linguistics, 11, pp.531-545.

• Defour, D., De Dinechin, F. and Muller, J.M., 2001, November. Correctly rounded exponential function in double-precision arithmetic. In Advanced Signal Processing Algorithms, Architectures, and Implementations XI (Vol. 4474, pp. 156-167). SPIE.