The Expressive Power of Transformers with Chain of Thought
Authors: William Merrill
Ashish Sabharwal
Venue: ICLR 2024 poster (2024)
Keywords: The Expressive Power of Transformers with Chain of Thought
Primary area: general machine learning (i.e., none of the above)
TL;DR: We consider the expressive power of chain-of-thought transformers with different length reasoning chains, characterizing the additional reasoning power gained as a function of the length of the chain.

Abstract: Recent theoretical work has identified surprisingly simple reasoning problems, such as checking if two nodes in a graph are connected or simulating finite-state machines, that are provably unsolvable by standard transformers that answer immediately after reading their input. However, in practice, transformers' reasoning can be improved by allowing them to use a "chain of thought" or "scratchpad", i.e., generate and condition on a sequence of intermediate tokens before answering. Motivated by this, we ask: *Does such intermediate generation fundamentally extend the computational power of a decoder-only transformer?* We show that the answer is *yes*, but the amount of increase depends crucially on the amount of intermediate generation. For instance, we find that transformer decoders with a logarithmic number of decoding steps (w.r.t. the input length) push the limits of standard transformers only slightly, while a linear number of decoding steps, assuming projected pre-norm (a slight generalization of standard pre-norm), adds a clear new ability (under standard complexity conjectures): recognizing all regular languages. Our results also imply that linear steps keep transformer decoders within context-sensitive languages, and polynomial steps with generalized pre-norm make them recognize exactly the class of polynomial-time solvable problems—the first exact characterization of a type of transformers in terms of standard complexity classes. Together, this provides a nuanced framework for understanding how the length of a transformer’s chain of thought or scratchpad impacts its reasoning power.

Published as a conference paper at ICLR 2024
THE EXPRESSIVE POWER OF TRANSFORMERS WITH CHAIN OF THOUGHT
William Merrill New York University wi 1 1allmmmmuueeee edu
Ashish Sabharwal Allen Institute for Al ashishs@allenai. org
ABSTRACT
Recent theoretical work has identified surprisingly simple reasoning problems, such as checking if two nodes in graph are connected or simulating finite-state machines, that are provably unsolvable by standard transformers that answer im- practice, transformers' reason- mediately after reading their input. However, "chain of thought" or "scratch- ing can be improved by allowing them to use a pad", i.e., generate and condition on a sequence of intermediate tokens before answering. Motivated by this, we ask: Does such intermediate generation fin- damentally extend the computational power of a decoder-only transformer? We show that the answer is yes, but the amount of increase depends crucially the amount of intermediate generation. For instance, we find that transformer de- coders with logarithmic number of decoding steps (w.r.t. the input length) push the limits of standard transformers only slightly, while a linear number of decoding steps, assuming projected pre-norm (a slight generalization of standard pre-norm). adds a clear new ability (under standard complexity conjectures): recognizing all regular languages. Our results also imply that linear steps keep transformer decoders within context-sensitive languages, and polynomial steps with general- ized pre-norm make them recognize exactly the class of polynomial-time solvable problems-the first exact characterization of type of transformers in terms of standard complexity classes. Together, this provides a nuanced framework for understanding how the length of transformer's chain of thought scratchpad impacts its reasoning power.
INTRODUCTION
series of recent theoretical results (Merrill & Sabharwal, 2023b;a; Merrill et al., 2022: Liu al., 2023; Chiang et al., 2023; Hao et al., 2022) has unveiled surprising limits realistic formal models of transformers. They have shown that standard transformers, even with ideal parameters, cannot perfectly solve many sequential reasoning problems at scale, such simulating finite-state ma- chines, deciding whether nodes in graph are connected, or solving matrix equalities. The intuition here is that the transformer lacks recurrent connections, and recurrence is required solve these sequential reasoning problems. Empirically, reasoning problems inspired by these results cannot be solved by cutting-edge transformer language models such as ChaiGPT and GPT-4 (Zhang al., 2023), and the reasoning performance of GPT-4 negatively correlates with the depth of the problem's 's computation graph (Dziri al., 2023). These results show certain kinds of sequential reasoning pose a challenge for the transformer and motivate extensions to address this issue.
One method that has been empirically successful for improving sequential reasoning with transform- ers is adding so-called chain of thought (Wei et al., 2022) or scratchpad (Nye et al., 2021). These methods allow the transformer to output sequence of intermediate tokens before answering, rather than answering right away after reading the input. Intuitively, such methods could unlock greater expressive power on sequential reasoning problems because the model can use each intermediate token as kind of recurrent state. Feng et al. (2023) recently showed how chain of thought lets transformers solve a specific modular arithmetic problem that they likely cannot solve without one. Yet there is no general characterization of the class of problems transformers thought. Thus, the extent to which chain of thought alleviates transformers' weaknesses is unclear, well the number of chain of thought steps required gain reasoning power.
solve with chain of

---

Published as a conference paper at ICLR 2024
(n)(n)
(IME
CoT(n)
CSL
(log))
CFL
CIT(1) (no int. decoding)
TC
Reg
Summary of results: transformers with intermediate generation against various classes of formal languages. A logarithmic number of chain-of-thought steps remains in log-space (L). A linear number of steps adds more power, enabling recognizing all regular languages (Reg), but contained within context-sensitive languages (CSL). We assume context-free languages (CFL) require (n²) time to recognize. Some regions with area in the plot not known to be non-empty.
Figure 1:
In this work, we address these open questions by characterizing the reasoning power of transformer decoders that can take intermediate steps before generating an answer and comparing them against transformers without intermediate steps. A transformer with a chain of thought constitutes a special case of a transformer decoder with intermediate steps. Our fine- greiggaaidd results give upper and lower bounds on transformers' power depending on t(n)): the number of allowed intermediate steps function of the input size n. We focus mainly on understanding three regimes: logarithmic steps (when ((n) Ollogn)) linear steps (when ((z) (n))), and polynomial steps:
No Intermediate Steps. Recent work has shown transformer decoders with- out any intermediate steps can only solve problems that lie inside the fairly small circuit complexity class TC (Merrill & Sabharwal, 2023b) and related logical classes (Merrill & Sabharrwa, 2023a; Chiang al., 2023). This implies basic transformers are far from Turing-comples they cannot even solve problems complete for classes larger than TC such simulating automata (NCp-com))ete), deciding directed graph connectivity (NL- complete), or solving linear equalities (P-complete).
Prior Work:
2. Logarithmic Steps. With a logarithmic number of intermediate steps, we show that the upper bound for transformers expands slightly from TC" L. This means transformers with logarithmic number of intermediate steps might gain power, but they still cannot solve NL-complete problems like directed graph connectivity or P-complete problems like solving linear equalities.
3. Linear Steps. Linear intermediate steps allow transformers with projected pre-norm" to simulate automata (NC*-complets complete), which cannot be done without intermediate steps un- less TC NC Polynomial Steps. With polynomial number of decoding steps, show that transformers with strict causal attention and projected pre-norm are equivalent the class This. our best knowledge, is the first equivalence between class of transformers and standard complexity class.
Together, our results provide framework for understanding how the length of a transformer's 's chain of thought affects its reasoning power. We find logarithmic chain does not add much, while linear chain affords more power on inherently sequential reasoning problems.
1.1 MAIN RESULTS: POWER OF TRANSFORMERS WITH INTERMEDIATE DECODING
Let (I)) be the class of languages L for which there exists Turing machine that runs in time O(((n)) and accepts L. Let TIME(t(n)) be the class of problems in TIME(t(n) log" n) for some k. which meaningful for t(n) 2 n. Let SPACE (s(n) be the class of languages L for which there exists Turing machine with tape size bounded by ((s(n)) that accepts L. We show the following
Assuming NC NL, and P do not collapse to TC respectivly..
Assuming NL and P do not collapse L, respectively.
i... standard pre-norm but applied only to a linear projection of a sublayer's 's input; cf. Definition 4 As will define later, this is non-random-acces, multitape Turing machine.
2

---

Published as a conference paper at ICLR 2024
relationship between transformers with t(n) steps and standard time/space complexity classes:
S (P) + logn)
TIME(((1)) C CoT(((n))
(I)
C TIME(t(n)) + n²)
Here (tT()(n)) denotes the set of languages recognized by some transformer using t(n) decoding steps. Our lower bound (left side of Equation (1)) assumes strict causal saturated attention and
projected pre-norm, while upper bounds hold both with and without these architectural assumptions. Both our time lower bound and space upper bound are fairly tight: improving either by factor larger than log ((2) would result in a fundamental complexity theory advance (Hopcroft et al., 1977).
Capabilities of Transformers with COT. The left side of Equation (1) implies that transformer decoders with ((n) steps can simulate real-time models of computation like automata or counter machines (Merrill, 2020). Under standard assumptions in complexity theory, transformers with no decoding steps cannor simulate all automata (Merrill & Sabharwal, 2023b; Merrill, 2023; Liu et al., 2023). Thus, a linear number of decoding steps makes transformers strictly more powerful. Simi- larly, the left side of Equation (I) implies transformers with a quadratic number of steps can express a linear-time algorithm (for a random access Turing machine) to solve directed graph connectivity (Wigderson, 1992), again a problem known be beyond the limits of standard transformers. In the same vein, with a polynomial number of decoding steps, transformers can solve linear equalities, Horn-clause satisfiability, and universal context-free recognition, all of which are P-complete and thus known to be inexpressible by standard transformers (Merrill & Sabharwal, 2023b).
The left side of Equation (1) is proven by showing transformer decoders can simulate & Turing machine steps with t intermediate steps. Similar prior results have assumed a transformer with external memory (Schuurmans, 2023) or an encoder-decoder model with nonstandard-posinn positional decodings (Pérez et al., 2021). Our construction adapts these ideas to work for decoder-on -nlly model without external memory or extra positional encodings, but with strict causal masking and projected pre-norm (cf. Section 2.1).5 The key idea behind our more general construction is the layer-norm hash (Section 3.1): simple module for effectively storing memory in decoder-only transformers, We believe the layer-norm hash could be broadly useful for building algorithms in transformers. For example, Yao et al. (2021) used a related idea to construct transformers that recognize hounded-depth Dyck languages, although in more ad hoc way.
Limitations of Transformers with CoT. The right side of Equation (1) establishes two upper bounds transformer decoders with t(n) intermediate steps that depend on both t(n) and We turn to the implications of this general result in different regimes for t(n):
1. Log Steps: Transformer decoders with O(log n) intermediate steps can only recognize lan- guages in L = SPACE(logn). This implies that transformers with O(log n) intermediate steps cannot solve NL- or P-complete problems like directed graph connectivity, just like transformers with no intermediate decoding (Merrill & Sabharwal, 2023b).
2. Linear Steps: Transformer decoders with O(n) intermediate steps can only recognize
languages that are in both TIME(n?) and SPACE(n). Since SPACE(n) falls within the context-sensitive languages (Kuroda, 1964), transformers with linear steps can recognize at most context-sensitive languages. Alongside our lower bound, this shows transformer decoders with ((n) steps fall somewhere between regular and context-sensitive languages in the Chomsky hierarchy. Further, transformers with O(n) steps cannot recognize all context-free languages unless context-free languages can be parsed in soft quadratic time. Poly
3. Polynomial Steps: If t(n) O(n)) for some c, we get an upper bound of U TIME(n"). Combined with our lower bound, this shows that transformer decoders with a polynomial number of steps recognize exactly the class P. Thus, a polynomial num- ber of steps turns transformers into strong reasoners, though running polynomial number of forward passes with a large transformer is likely intractable in practice.
Together, these results show that intermediate generation like chain of thought or scratchpad can add reasoning power to transformers and that the number of steps matters as a computational resource
"ur construction (Theorem 2) can be casily modified to work with an encoder-decnder model as well. 'The best known algorithms for context-free recognition run in time (nn)), where is the matrix multipli- cation constant (Valiant, 1975); the best lower bounds for context-free parsing sub-quadratic (Lee, 2002).
3

---

Published as a conference paper at ICLR 2024
akin to time space. Some of the limitations identified in prior work (Merrill & Sabharwal, 2023b; Chiang et al., 2023, etc.) can be overcome with linear or quadratic number of steps, and a poly- nomial number of steps covers all problems in P. On the other hand, have not identified any concrete reasoning problem where a logarithmic number steps would help. These results provide unified understanding of the power of transformer decoders across decoding lengths and problems.
PRELIMINARIES
We study the power of decoder-only transformers that can generate intermediate tokens between reading the input and generating an answer. On input E Σ", the transformer consumes tokens .1..... Zn for the first steps, and then, for t(n) intermediate steps, consumes the token gener- ated by the previous step. At each step, the transformer can attend over all previous hidden states. This standard method of generating text from a decoder-only model can described formally as follows. Let Σ be a finite alphabet and f E - E be function mapping prefix to next token (parameterized by a transformer). Let be concatenation. We define the k-step extension of f as
(xx)= =x,
p (x) (((
We say we have run f on æ with t(n) (additional) decoding steps if we compute the func- tion (x)) We consider f with t(n) steps to recognize the language of strings such that (x))=))(x) = 1, where 1 E Σ is a special "accept" symbol. We denote by CoT(t(n)) the set of languages that are recognized by (nn) decoding steps for some transformer f.
2.1 TRANSFORMERS
A transformer is neural network parameterizing a function Σ' + E. Let Dp be the datatype of p-precision floats and define p-truncated addition (+, Σ), multiplication (.), and division (/) over Dp as in Merrill & Sabharwal (2023b). We now define the high-level structure of the transformer in terms of its core components, with the details of those components Appendix
Definition (Merrill Sabharwal 2023a). A p-precision decoder-only transformer with h heads, d layers, model dimension m (divisible by h), and feedforward width w is specified by:
1. An embedding function e Σ x N- + DØ whose form is defined in Appendix A.2;
2. For each v S d and 1 < & < h, a head similarity function s D. x 1. + Dp whose form is defined in Appendix A.3 (and includes projected layer-norm);
3. For each 1 < e < d and 1 < & < h, a head value function of. D + DT//h whose form is defined in Appendix A.3 (and includes projected layer-norm);
4. For each S S d, an activation function f (Dm x D + D whose form defined in Appendix A.4 and implicitly uses the feedforward dimension (and includes projected layer-norm);
5. An output function 7 D. + E parameterized as a linear transformation.
Definition 2. We define one decoding step n + Σ. with a decoder-only transformer follows:
Embeddings: For S <is n, h e( 1).
2. Multihead Self Attention: For each layer 1 < l d, we compute h attention heads:
<(i) $ (h
c(i)
a k
o (h where z = Σ 0
h2-1)
and c(i) for standard causal attention and for strict causal attention.
3. Activation Block: For 1 - { < d, activation block & maps the head outputs to h:;
h = (a ................................ .............. h(-1).
................................
4. Classifier Head: The transformer output is (hg)).
These definitions use l-indexing, but when the input contains a beginning of- orems 1 and 2), will use O-indexing starting at $ in the natural way.
begiininggof--sequeeeee of-seuuence token $ (The-
4

---

Published as a conference paper at ICLR 2024
Transformer Precision. We consider log- log-precision transformers (Merrill & Sabharwal, 2023b), ie., we allow the transformer at most clog m precision for m decoding steps. As transformer with intermediate generation rums for input steps and t(n) intermediate decoding steps, this means we have precision at most e elog(n 1 + f(n)). ) Log precision has been analyzed in prior work (Pérez al., 2021; Merrill & Sabharwal, 2023b;a) because it gives the transformer juste precision represent indexes and sums across different positions. This means it naturally formalizes a bounded- precision transformer that capable of representing position and computing uniform attention, two important capabilities for constructing algorithms with transformers.
Our lower bound constructions (Theorems and 2) assume the following:
1. Saturated Attention. A saturated transformer (Merrill et al., 2021) is an idcalized trans- former with "averaging hard attention" (Strobl al., 2024): per head, all attention scores are either 0 or 1/v for some e. This includes uniform attention (1/n over n tokens) or hard attention as special cases. Following common practice (Pérez al., 2021; Merrill & Sabharwal, 2023b), we use saturated attention for our lower bound constructions.
2. Striet Causal Masking. The formulation attention Definition makes the slightly nonstandard assumption that causally masked attention at position i can view tokens at all positions up to but not the current token i. This is required in Theorem
3. Projected Pre-Norm. Our lower bound constructions require and fe in Definition 2 to allow generalization of standard pre-norm. Normally, a layer-norm is applied to the entire input each sublayer. We generalize this, allowing each sublayer to apply linear projec- tion before layer-norm. Crucially, in particular, this enables each layer to pick out a subset of the previous hidden state apply layer-norm to (cf. Definition 4 in Appendix A.
For convenience, our proofs with projected pre-norm use an even more general notion of pre-norm, namely aulti-pre-norm. which allows each sublayer take k different projections of its input, apply layer-norm to each, and concatenate (cf. Definition 5 in Appendix A.1). Multi-pre-nort can, however, be simulated by multiple layers of projected pre-norm (see Appendix A.1 for proof):
Proposition 1 (Chiang, 2024). Multi-pre-norm with norms can be simulated by k + projected pre-norm layers.
2.2 MODELS OF COMPUTATION
Automata. A deterministic finite-state automaton is a tuple A (,, q0, 0, F3 where E is finite input vocabulary, Q a finite set of states containing initial state To, a transition function Q x Σ + Q, and F C Q is set of final states. A processes starts with state go and reads a one token at time, updating 9; = 8(q.1, σ.) until i = n. A accepts if qu E F and rejects it otherwise. The language recognized by A is the set of strings it accepts.
input string
" follows. A
Turing Machines. Adapting the notation of Hopcroft et al. (2001), multitape Turing machine is tuple (E, D, k. Q, 90, F) where:
1. E is a finite input vocabulary
2. I is finite tape vocabulary with Σ r
3. & is the number of work tapes
4. b is blank symbol such that b r and b g Σ
5. Q is a finite set of states containing initial state
6. ó a transition function (Q) p Q x pk+1 {#1]k+2
7. F C Q is a set of halting states
We define Turing machine computation in the standard way (cf. Appendix B).
LOWER BOUNDS FOR TRANSFORMER DECODERS
Prior work (Merrill Sabharwal, 2023a) has established strong upper bounds on the reasoning problems transformers solve. Specifically, under standard conjectures in complexity, transform- ers without intermediate decoding cannot recognize all regular languages. this section, we show some of these shortcomings can be overcome with suitable number of intermediate decoding steps
5

---

Published as a conference paper at ICLR 2024
(and projected pre-norm). Specifically, a linear number of steps enables simulating automaton. We also show this can be extended to simulate t(n) Turing machine steps with t(n) decoding steps.
3.1 INTRODUCING LAYER-NORM HASII
We first introduce useful building block for our results that we call the layer-norm hash. The layer-norm hash is a mechanism that enables retrieval across different columns in the transformer based on query-key matching of numerical values. Exact-match retrieval is trivial when the query ga and keys ki,... k; are items in a finite set: just one-hot encode and k- and the inner product will be maximized when qi and ky match. But this does not work when the keys and values are counts produced by uniform attention, which many transformer algorithms use (Weiss et al., 2021), as the key q// and query k have different denominators. The layer-norm hash helps by transforming q:/i and k j so hard attention retrieves s.t. qi = k.- Let layer_norm
x' |x
where x' = x-
Definition (Layer-norm hash). For z, y E R, let (xz, y) A layer_ y, -x, -y)
((x, v) is a unit vector in R4. A key feature is scale invariance, and, in particular, that (xx/i, 1/2) is invariant w.r.t. i in the sense that it is only function of z, independent of i. Let dx 6(z 1). Then we have the following properties, whose proof may be found in Appendix C. Lemma (Scale invariance). For any E I and i E Rz,, ((x/i, Φs.
Lemma 2 (Equality check via layer-norm hash). For any q, k R. 0g Øk and only if k.
In other words, the inner product of these representations of two scalars q and k. even if computed at different positions i and j. respectively, allows us to check for the equality of g and k. We can look up key 9 in a sequence of keys k 11/1.... k at position i and key o(k, 1/j) at each j i. By Lemmas 1 and 2 this averages the values at all j such that = ky. The layer-norm hash can also be used to directly compare two values qi: kj without removing the denominator by computing p(q,, 1) and (kk,, 1).
- 1) by attending with query p(q:/t, 1/i)
3.2 SIMULATING AUTOMATA
We can use the layer-norm hash to simulate models of computation like automata or Turing machines transformers. To warm up, we first show how use the layer-norm hash to simulate an automaton (i.e., recogmize regular language) and then extend it in Section 3.3 to show how a transformer can simulate Turing machine for bounded number of steps.
with intermediate-gennraaie
Theorem 1 (Regular language recognition). For any regular language L. there is decoder-only projected pre-norm transformer with strict causal saturated attention (with or without positional encodings) that, on input Sx," 7.8 checks whether E L with [x] + 1 decoding steps.
Proof. Let be state automaton recognizing L. We will simulate one step of A with one transformer decoding step (after first reading input tokens). We refer tokens with 0 indexing: $ is token O, Ii is token 1, etc. At step i, n, S i S 2n, we will output a token Qi-n encoding the next one additional step to output iff qn E F, the set of final states of A. At each token > O, we compute 1/i by attending uniformly over the strict left context with value 11 [rr = $). We show by induction that at step 2 n,,
state of A. After printing the final state 9n,
output gi-n-
Base Case: i = n. For i S n, we output go- Crucially, at i = n, this becomes the next input.
Inductive Case: i n. We already have a sequence of intermediate tokens ... 9i-n-1- Our goal to compute qi- = 5(g--n-1, d.n),, which first involves retrieving -i--1 and 0; q--n-1 is the input to the current column of the transformer. We will use hard attention to retrieve the current input symbol ai-n- To do this, we attend uniformly over the prior decoding tokens and S. with a value of 1 at $ and 0 elsewhere. At tokens i > (i.e., decoding tokens), this yields بيب simulate multi-pre-norm (Proposition 1). We now leverage the
Recall that projected pre-norm multi-pre-nome norm architecture to pass two layer-norms to a feedforward network:
d! (1/i, 1),
Φ e (1/( n),
"Theorem goes through without strict causal masking (but Theorem 2 will require strict masking). *Theorems and 2 both go through without 8 long token j can compute value 21 1 0).

---

Published as a conference paper at ICLR 2024
Let d, A 1[x, Q). where Q is the set of states of A. Based on d,, select between of and oP:
Φ - ReLU(- d!) ReLU I o.)
We attend with query layer_norm(oo) = oi. key layer_norm(dg) = Φ if d, = 0 and Ô otherwise, and value oj if dj = 0 and 0 otherwise. By Lemma 2, at the current step i, the attention score is maximized when j i-n, token -- We conclude by computing qi-n = 8(9--n-1, with feedforward network.
n, thus retrieving d--n- We now have the previous state q--n- and current
Theorem 1 shows that a linear number of decoding steps gives additional reasoning power to log- precision transformers with projected pre-norm (assuming TC 7 NC'). This follows because log- precision transformers with no decoding steps are contained in uniform TC" (Merrill & Sabharwal, 2023b), which means they cannot recognize all regular languages. In contrast, Theorem says linear number of steps is sufficient for recognizing all regular languages, establishing a conditional separation. This is example of simple and familiar additional computational power granted by additional decoding steps. The core challenge in simulating an automaton is recurrence, which cannot be done without decoding steps (Merrill & Sabharwal, 2023b). A linear number of decoding steps allows simulating recurrence, which is where the additional power comes from. However, this added power does not stop with finite-state machines: the layer-norm hash can be used to simulate more complex models of computation such as Turing machines, which we will turn to next.
3.3 SIMULATING TURING MACHINES
We now show how a transformer decoder can simulate a Turing machine in real time using the layer-norm hash. Our decoder-only construction resembles the encoder- Pérez al. (2021). However, avoids simplifying assumptions from Pérez et al. (2021). In addition to assuming non-standard attention and no layer-norm, they required 1/2, 1/12, and i in the posi- tional embeddings, which is problematic because transformers cannot represent unbounded scalars like i due to layer-norm. In contrast, our construction works with or without positional encodings. However, it assumes strict causal masking and projected pre-norm (Section 2.1).
construction of
Theorem 2 (Turing machine simulation). Let M be Turing machine that, on input length 1 + n, runs for at most t(n) steps (at most polynomial). There is decoder-only projected pre-norm transformer with strict causal saturated attention (with or without positional encodings) that, on inpul Sx," 8 lakes (n) decoding steps and then, with (+)) additional steps, oulputs M (2r).
Proof. We construct a transformer decoder that uses a single decoding step to simulate each Turing machine step. The main difficulty is representing a Turing machine tape a sequence of transformer state vectors so that the transformer can always correctly reconstruct the value on the tape at the current head position. The key idea will be store "diffs" to the tape at each step and use the layer- norm hash to dynamically reconstruct the contents at the head position at future steps. Coneretely, let A be a finite vocabulary representing elements of Q x I x (0, +1] The deterministic Turing machine run induces a diff' sequence do...., de(n) written, and directions moved after each token. Following the proof of Theorem I, we use O-indexing starting at the $ token and compute i at each token i > 0 as representation of position. We show by induction that at step i Z n, we can output 6,-n-
A capturing the state entered, tokens
Base Case: i= At every input token (crucially, the last one), we output 8n = {yo, U++1, 0
Inductive Case: i >> n. We first reconstruct hT, the current position on each tape т. For each T, a head attends with query 1, key [ & 2.. and value being the move direction of at j. Since we assume strict causal attention (for reasons that will become clear later), head t thus computes h/-/í.. Since we need hT, we write both (A²_ + 1)/i to the residual stream. When need hī /i going forward, we use a linear layer to select either (h + 1)/i or (h 1)/2 depending on the current input n contains a +1 move
11 move for T, respectively.
We now use two layers to compute the contents at Nς on the input tape. Similar to Theorem 1, we use feedforward network to implement the following piecewise comparison:
0f a (1,1 1(i,
if x; E (hh//,, = ((h,1)) otherwise.
7

---

Published as a conference paper at ICLR 2024
With some abuse of notation, let (.) denote vector concatenation. We attend with query (0f. -1), key (0), 1l [Ir E E]},, and value "y)) Let (0) ,7) be the head output. We show in Appendix D that two properties hold. First, by Lemma 3, क0 of iff 1 ng < Second, by Lemma 4, if 1 S / < n., then ō = Th,- Based on this, we compute the value read from the input tape as 7 ō if nD = of and as 7. -b otherwise.
We now use a single attention layer to compute 7 the contents at hi on each non-input tape т. The layer uses two layer-norm hashes, again taking advantage of the multi-pre-norm architecture that projected pre-norm can simulate (Proposition 1):
of a @h? /i,1/i) o(h;,1)
o7 a (()),,
where f(i) is defined in Definition 7 in Appendix E. Crucially, ((i) is computable with a single transformer layer and monotonically decreasing with i. With strict causal masking, we attend with query (07, e1), key (03. -2)). and value (07, 8)-_))-)-1). Let (Ø 8) be the head output. We show Appendix E that two properties hold. First, by Lemma 6, o = = iff there is some j < s.t. hg h.. Second, by Lemma 7, if there is some j < i s.t. h hj, then the head retrieves (67.8) for the greatest such Based on this, compute the last-writted value on tape Rf [6] 2+ if क- of and of b otherwise. Having obtained all arguments for the transition function, we now compute on = 6(1---1: np, ...... +++11 with a feedforward net.
h as
Finally, we use M(xx steps to write the Turing machine output. We detect we are at an output step if either some 5, token to the left the current input encodes halting state. At each such token i, we compute h/+1// /i as before (recall that tape k + 1 is the output tape) via attention, except the value now d + if a; E and +1 otherwise. We attend as before using 1/+1// /i to retrieve (and output) 7 k+1 Thus, the [x)] tokens generated after a final state are precisely (xx).
LI
The critical role projected pre-norm or multi-pre-norm play in Theorems and 2 suggest it could be interesting to investigate incorporating these generalized pre-norms into transformers in practice.
Theorem 2 gives us a general result connecting the power of transformer decoders with t(n) steps to Turing machines with the same number of intermediate steps:
Corollary 2.1. TIME(t(n)) S CtT(t(n))
Thus, simulating an automaton (cf. Theorem 1) is not the only new capability unlocked with O(n) decoding steps: rather, such transformers can solve any problem Turing machine can solve in O(n) time, such as simulating real-time counter machines (Weiss et al., 2018). With O(n²) steps, we can solve directed graph connectivity using standard graph traversal algorithms like depth-first search. Depth-first search runs in (n)) time on random access Turing machine (Wigderson, 1992). which can be simulated in Onn)) time without random access. Possibly, transformers can solve directed graph connectivity with fewer than ((n)) steps, as results from Zhang et al. (2023) hint at.
4 UPPER BOUNDS FOR TRANSFORMER DECODERS
Having shown lower bounds on transformers with t(n) steps, we present two different upper bounds: one that relates transformer decoders to time complexity classes, and one that relates them to space complexity classes. The relative strength of the two different bounds will vary depending A simple upper bound on transformers with chain of thought can be obtained based on the fact that transformers can be simulated using quadratic number of arithmetic operations.
t(n).
Theorem 3. CoT(((n)) C TIME(n² + ((n))).
Proof. We sketch a multitape Turing machine that will simulate the transformer. Each forward pass i appends key i onto work tape and value onto another work tape. To simulate the forward pass at time it suffices show how to simulate computing self-attention at time i, To compute self
9 As in Theorem second layer norm gets applicd to Φf the start of the layer but has no effect.
8

---

Published as a conference paper at ICLR 2024
attention, the Turing machine first computes the query time i. It then iterates over pairs on the key and value work tapes. For each pair J. we compute the attention score between query and key j and then multiply it by value using additional work tapes. We then add this value to a running sum tape. We treat the final sum at the output of the attention mechanism.
This runs + ((n)) forward passes, and each forward pass loops over + f(n) key-value pairs. This means we run at most O(n +t(n)² inner loop calls. remains to be shown that one inner loop runs in polylogarithmic time. An inner loop just involves adding and multiplying O(log n)-bit numbers. p-bit numbers can be added in time O(p) = Ollog nt). Similarly, p-bit numbers can be multiplied in time O(p log p) < (p²), which comes out to O(log" (n + t(n))) with log precision. Thus, one inner loop can be run in polylogarithmi time. We conclude that a transformer decoder with t(n) intermediate steps can be simulated by a multitape Turing machine in time O(n² + t(n)²).
Our second upper bound relies on the TC upper bound for transformers without intermediate steps. Theorem 4. (tT()(m)) S SPACE(t(n) + logn).
Proof. Since log- reg-precision transformers can be simulated in uniform TC (Merrill & Sabharwal, 2023b), they can be simulated L, i.e., with most clog space overhead on inputs of size n. To compute ((n) intermediate decoding steps of a transformer, we store buffer of at most ((n) generated tokens, which has size (()))). To compute the next token, we call the transformer with an input of size O(++t(n)) +(()) using most c (t)))))))))))))))))))))))))))))) space overhead. We then clear the memory used and append the finite token generated to the input buffer. It follows from this algorithm that
CoT(t(m)) S SPACE(+(n) + clog(n + t(n)))
C (PACE(t(n) + n).
With at least Q(log n) steps, this upper bound can be simplified to SPACE(t(n)). The t(n) = (n) case establishes the context-sensitive languages as upper bound for transformers with linear steps. Given our (I)) lower hound (Theorem 2), the tightest possible space upper bound without making fundamental complexity advances would be (PACE(t(m)) log t(n)) (Hoperoft al., 1977). Conversely, our lower bound can only be tightened to TIME(t(n) log t(r))..
On the other hand, with only O(logn) decoding steps, intermediate decoding does not increase expressive power much beyond TC because the upper bound simplifies to SPACE(!(n) L. Thus, under standard assumptions, transformers with logarithmic number of decoding steps cannot solve directed graph connectivity, Horn formula satisfiability, or other NL- or P-complete problems. Yet, they may be able to solve L-complete problems, unlike transformers without decoding steps.
5 CONCLUSION
We have shown that intermediate decoding steps extend the formal power of transformers well be- yond previously known upper bounds, such as TC circuits and FO(M) logic, on transformers with- out intermediate decoding. Further, the amount of additional power is closely related to the number of decoding steps. In particular, transformers with a linear number of decoding steps have the capac- ity to recognize regular languages, but cannot recognize languages beyond context-sensitive With a log number of decoding steps, such transformers can only recognize languages in L, which is a complexity class relatively close to TC Thus, it appears that linear number of intermediate decoding steps may be required to overcome the limitations of transformers on many sequential rea- soning problems of interest. In future work, it may be possible to derive a strict separation between transformers with a log and a linear number of decoding steps and show that certain problems that currently have quadratic bound can in fact be solved with a roughly linear number of steps.
We have focused on expressive power, rather than analyzing learnability. Whereas our upper bounds directly reveal limitations on what transformers with intermediate generation can learn, one caveat is that our lower bounds do not directly imply transformers can learn to use intermediate steps effectively. It would be interesting to formally investigate transformers with CoT from a learning- theoretic lens, possibly along the lines of Malach (2023), and how different kinds of fine-tuning, such reinforcement learning, might better allow models use the power of chain of thought.

---

Published as a conference paper at ICLR 2024
ACKNOW
We thank David Chiang for the valuable feedback and for identifying mismatch between the trans- former definition in earlier version of this paper and standard pre-norm transformers. We also appreciate helpful comments from Gabriel Faria, Ofrr Press, Abulhair Saparov, Jason Wei, and Avi Wigderson, as well researchers in ML2 at NYU and A12. WM was supported by NSF Award 1922658, an NSF Graduate Research Fellowship, and A12.
REFERENCES
David Chiang. Personal communication, March 2024.
David Chiang. Peter Cholak, and Anand Pillay. Tighter bounds on the expressivity of transformer encoders. In ICML, 2023.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jian, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D. Hwang, Soumya Sanyal, Sean Welleck, Xiang Ren, Allyson Ettinger, Zaïd Harchaoui, and Yejin Choi. Faith and fate: Limits of transformers on compositionality. In NeurIPS, 2023.
Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: A theoretical perspective. In NeurlPS, 2023.
Sophie Hao, Dana Angluin, and Roberta Frank. Formal language recognition by hard attention transformers: Perspectives from circuit complexity. TACL, 10:800-811 810, 2022.
John E. Hopcroft, Wolfgang J. Paul, and Leslie G. Valiant. On time versus space. J. ACM, 24: 332-337, 1977.
John E Hoperoft, Rajeev Motwani, and Jeffrey D Ullman. Introduction to automata theory, lan- guages, and computation. ACM Sigact News, 32(1):60-655 2001.
S-Y Kuroda. Classes of languages and linear-bounded automata. Information and control, 7(2): 207-223, 1964.
Lillian Lee. Fast context-free grammar parsing requires fast boolean matrix multiplication. J. ACM, 49(1):1-15, Jan 2002.
Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts automata. In ICLR, 2023.
Eran Malach. Auto-regressive next-token predictors are universal learners. arXiv, abs/2309.06999. 2023.
William Merrill, On the linguistic capacity of real-time counter automata. arXiv, abs/2004.06866, 2020.
William Merrill, Formal languages and neural models for learning on sequences. In François Coste, Faissal Ouardi, and Guillaume Rabusseau (eds.). ICGI, volume 217 of PMLR, Jul 2023.
William Merrill and Ashish Sabharwal. A logic for expressing log-precision transformers. In NeurIPS, 2023a.
William Merrill and Ashish Sabharwal. The parallelism tradeoff: Limitations of oog-prccision trans- formers. TACL, 11:531-545, 2023b..
William Merrill, Vivek Ramanujan, Yoav Goldberg, Roy Schwartz, and Noah A. Smith. Effects of parameter norm growth during transformer training: Inductive bias from gradient descent. In EMNLP, 2021.
William Merrill, Ashish Sabharwal, and Noah A. Smith. Saturated transformers are constant-depth threshold circuits. TACI, 10:843-856, 2022.
10

---

Published as a conference paper at ICLR 2024
Maxwell Nye, Anders Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena. Show your work: Scratchpads for intermediate computation with language models. arXiv, fos/2112.001 14, 2021.
Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is Turing complete. JMLR, 22(1), January 2021.
Dale Schuurmans. Memory augmented large language models are computationally universal. ArXiv, 2023.
Lena Strobl, William Merrill, Gail Weiss, David Chiang, and Dana Angluin. Transformers as rec- ognizers of formal languages: A survey on expressivity. TACL, 2024.
Leslic G. Valiant. General context-free recognition in less than cubic time. Journal of Computer and System Sciences, 10(2):388-33 -355, 1975.
Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed H. Chi, Quoe V Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), NeurlPS, 2022.
Gail Weiss, Yoav Goldberg, and Eran Yahav. On the practical computational power of finite preci- sion RNNs for language recognition. In ACL, July 2018.
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In ICMI., 2021.
Avi Wigderson. The complexity of graph connectivity. In International Symposium on Mathematical Foundations of Computer Science, pp. 112-132. Springer, 1992.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie- Yan Liu. On layer normalization in the transformer architec- ture. In ICML, 2020.
Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention net- works process bounded hierarchical languages. In 2021.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. NeurIPS, 2019.
Muru Zhang, Ofir Press, William Merrill, Alisa Liu, and Noah A. Smith. How language model hallucinations snowball. arXiv, 2023.
A TRANSFORMER COMPONENTS
This section formally defines our generalizations of pre-norm and then recalls the definition from Merrill & Sabharwal (2023a) for the components of the transformer layer.
A.1 GENERALIZED PRE-NORM
We assume pre-norm (Xiong et al., 2020) parameterization of the transformer for concreteness, as this is more standard in newer transformers. As stated in Section 2.1, we allow transformer sublayers to apply linear projection before layer-norm. Concretely, define lojelayer_nome (v) as follows: Definition 4 (Projected pre-norm). Let M mmm + nm be parameter matrix that projects dimensional vectors to m-dimensioonll vectors.
proj layer norm(v; M) = layer norm(Mv).
We will omit the parameter M for convenience, instead writing layer -nomm_norrmv)......
Our proofs also allow multiple pre-norms different projections of the hidden state for our lower- bound constructions. Concretely, define multi layer norm(v) as follows.
11

---

Published as a conference paper at ICLR 2024
Definition 5 Multi-pre--orm) Let be divisible by k. For < <i k. let M, R - Rm/h be a parameter matrix that projects m-dimensional vectors to m//k-dimensional vectors. Let (1)/ denote iterated concatenation. The multi multi-pre-nont pre-norm of v 1 is defined as
multilaayer.norm(;MM... M1, M&) (proj MM
As for projected pre-norm, we will omit the parameters M M2 for multi-pre-norm, instead writing multi caattillyye_noorm(()
As noted earlier, projected pre-norm can simulate multi multi-pre-norm:
Proposition 1 (Chiang, 2024). Multi-pre-norm with k norms can be simulated k projected pre-norm layers.
Proof. We will simulate a multi-pre- layer that takes as input
layer_norm (M,v), v), layer_norm( (M
using only projected pre-norm layers. The idea use & different projected pre-norm layers (one for each input layer norm). At layer i, the layer takes as input aover_normMMe) and writes this to the residual stream. Then, after these k layers, final layer reads as imput
layer norm (layer norm v))
Since each vector in the concatenation has unit norm, this equivalent
(layer (o)yer.orm(MM.v))....
follows that this layer receives essentially the same input as the original multi-pre-norm layer, up to a constant factor. The original weights for the layer can be multiplied by v/: to implement the same computation as the original layer. To make sure that vk: is exactly representable, we can pad the number of entries so that k perfect square.
A.2 TRANSFORMER EMBEDDINGS
For each position 1 < i S the transformer embedding function represents token Σ and its position with a vector. Let be an embedding matrix size |>]] m where each row represents the embedding for some Let f ::N. + Di be computable in time Ologgn) Then,
e(o,, Vo. f()).
A.3 SELF ATTENTION
The two components of the self attention block are s, the similarity function, and the value function. Let h; be the hidden state at the previous layer and h; = multi layer norm(h;). We define similarity of keys and queries follows:
sh; h,) exp (
vm/h/ (1
9 K;
)
where q: W h + b
k; W + by
Then the value function is defined v(h,) W + b
A.4 ACTIVATION BLOCK
The activation function f encapsulates the aggregation of the attention head outputs and the feedfor- ward subnetwork of the transformer. f takes as input attention head outputs a,,,...., aü, E I and the previous layer value h,
The first part of the activation block simulates the pooling part of the self atlfnattention sublayer. The head outputs first concatenated to form vector a,, which is then passed through affine
12

---

Published as a conference paper at ICLR 2024
transformation (Wo. bo) D > D followed by residual connections to form the sublayer output o, E -
o, = Wa + b, + h,
The second part of the activation block first applies multi-layer- and then simulates the feed- forward subnetwork compute the next layer vector hy. Let o, multi_laeer lar.llayrreomoo))) Let be nonlinearity computable in linear time on its input (in the most standard transformer, ReLU). Then,
for affine transformations (W1x 1, b1) D subnetwork can be defined as:
D and (W2, b2) 1 Dp the feedforward
h W (W, + b1) + b2 to
B TURING MACHINES
Turing machine takes as imput string E E" configuration of Turing machine is a finite state q along with the contents of an input tape c², k work tapes c'....,e"" and an output tape c++1. Finally, for each tape т, a configuration specifies head position h². We start with the initial state Yo and the input tape of containing starting at position 0) with infinite b's on each side, and hg = 0. All other tapes start containing all b's and with their head O. At each time step i, q; E F. we recurrently update the configuration by first computing:
(q++1. ....... Y/ .d......, d = ò(g;, c. 14
( (
We then update tape 7 by setting c. 1441 7f and keeping all other tape cells the same. We update the head position on tape according to h[+1 h dy. On the other hand, if q: E F, the Turing machine halts and outputs the string of tokens on the output tape from the current head position on the left up to (but not including) the first b on the right. A Turing machine can also be viewed as language recognizer if we set Σ {0, 1} and check if the first output token 0 or
C LAYER-NORM HASH
Lemma (Scale invariance). For any E IR and E Rso, ((z/i, ox.
Proof. Let Vr = (x/i,,1/, 1/// -//,,/ va constructed with mean 0, so layer-norm reduces to RMS-norm (Zhang & Sennrich, 2019). Thus,
0xx/i, 1/1) vz/l[vall
= V 2rr
(x, 1, 1,, x3,
22 +2
=(x 1)
which, by definition, is dz.
Lemma (Equality check via layer-norm hash). For any k IR, og db if and only if k.
Proof. By the definition of layer-norm hash, have
((g.1) o(k,
229² +2 1, -4, -1)
(k, 1, -k, 1)) 22 +
2qk +
V(2g² + 2)(2 + 2)
gk +1
V(9² 1)(k2 11
13

---

Published as a conference paper at ICLR 2024
The inner product of unit-norm vectors is maximized at 1. In this case, show that it achieves only when k, meaning that the unique maximum:
ok +
1 Vlg² + 1)(k2 + +1)
(gk 1)2 (4 1) +1)
(gk)² + 29k + 1 (gk)² + q² + + +1
2xh 9² + h:2
= (q - k:)²
We conclude that th - du maximized (to 1) and only if (1 = k.
As the layer-norm hash is constructed to have mean 0, does not require fully general layer-norm implementation and can, in fact, be implemented with simplified RMS norm (Zhang & Sennrich, 2019).
D INPUT TAPE RETRIEVAL VIA THE LAYER-NORM HASH
We now describe an attention head that uses the layer-norm hash to read from the input tape in Theorem Define sequence hi,..., h;-1, which represents Turing machine tape position in Theorem 2.
We define the following layer-norm hash based quantity, which is instantiated in Theorem particular way:
(0(,,)) (i,, if iflsisn
Φ. = ( (0h,, 1) otherwise.
The attention head we construct can then described as follows:
Query: (0., -1}
Key: (Øj, in < 1
Value: V, (Ø).
Let v (ō,0) o) be the head output. This head obeys the following properties:
Lemma 3. Let i > Then, 1 < h; S n if and only if o Φ;.
Proof. We proceed by considering the two directions.
Forward Direction, The query-key inner product has two terms kity + xaz By Lemma 2, xiy is maximized either when b. = j (and 1 < j I n) or when h2 = hy (and n < f). However, if n < j. the second term N² = 1. Thus, the attention score is maximized uniquely when j = h;, so the head retrieves v = (o(h,, 1), o..). Thus, ó Φ(h1,
6..
Backward Direction, We establish bidirectionality by proving the contrapositive. Assume that either
h; < 1 or h; > n. The head retrieves 6 ক. ΣjEM Φ(j. 1) for some M C {I,..., n).. It holds that, for all < n. h4 j. the other way around (i.c., for all 1 < j < n. h > j). Thus, by Lemma 5, & oh,, 1) = Φ1-
The following property also emerges from the proof of the forward direction in Lemma 3: Lemma 4. Let > n. Then, if $ h4 S n, δ = Tb.-
The backward direction in Lemma relies on the following lemma:
Lemma 5. Let q Z and K2 E for 1 S $ m. Let >. for all j, q Ky then
m
Φq
m1
14

---

Published as a conference paper at ICLR 2024
Proof. Recall that @- 0(x, I$ has first element v/v2n2 which will denote 2x- Observe that 2z is monotonically increasing function of x E R. Thus, 2g > 2k, for 1 S j S m,
which implies zq > + Σe
ش== Zky from which the lemma conclusion follows.
E RIGHTMOST RETRIEVAL VIA THE LAYER-NORM HASH
We now describe an attention head that can attend to the rightmost token satisfying certain property. capturing the construction in Theorem retrieve the most recent write a Turing machine work tape. Define sequence fi..... f which represents Turing machine tape position in Theorem 2. As is natural for Turing machine tapes, we assume that if * h; for all i, then must that < h2 for all i, where < is fixed as either > or <.
Let (() be a bie-breaking term that we will define later in Appendix E.I. We define two layer-norm hash quantities:
o - oht/it,1/1 4 - (i()),
Recall that 61 0), Construct attention head as follows:
Query: (d., c1)
. Key: (Ø), -
Value: Vj (0), 0z)
Let A (ō, 8) be the head output. The following properties hold for such head:
Lemma 6. There exists such that h2 if and only if Φ,-
Proof. We proceed by considering the two directions.
Forward Direction. The query-key inner product has two terms Kly + x.g-- By Lemma 2, the first term wly maximized at for each such that h4 = hy. For h4 + hy. xly The second component N², monotonically increases with j and satisfies nit f() < 1/(2i*) by such that h h3. Thus,
1/(2i*) by Lemma 8.
Lemma 10. Thus, the attention score is maximized for the largest = and σ; for this j. which means Φ;-
Backward Direction. We establish bidirectionality by proving the contrapositive. Assume there is no such that h; hy. Then assumption (top of Appendix E), we have h, < h, for all j < i. It follows from Lemma that + @he, 1) = di. completing the proof.
M Σyem o(hy, 1) for some M S (1....,m). By
The following property also emerges from the forward direction of the proof above:
Lemma 7. If there exists < i such that fęç = h2, then ō = ôg for the greatest such j.
E. TIE-BREAKING TERM
The construction above tie-breaker that favors retrieving tokens further to the right. We will justify the construction of such a tie- brebbrea term here. To begin, we will establish a bound on the layer-norm hash inner product similarity for inexact matches.
Lemma 8. For any 2 2, $(t,1) 1) @ii 1,1) 1/(2i)).
15

---

Published as a conference paper at ICLR 2024
Proof. Consider the squared dot product:
-i, -1,----,)))---,,,1-,11 1). -
(Ø(i, 1) 0(i 1.1))²=
(2i? 2)(2(i 1)2 +
(i(i (1(i-1)+4 11
1))+((1-1)++ +1
((0--))+2((i1)))
((i- + 2i(i - 1)+1
A P(i-1)) 1)3 + 2:(i 1) +2
20 - 1)2 2i(1 +2
2/3 4 3/ +2
Since (1 w/2)3 - for any we have -1 S y/2 for any y S 1. Applying this to the right hand side of the above equation, obtain:
o(1, (1,1))-((- - 1,1-1,11
213 + 3/2 81
2i +
$
2/4
4 + 21
41 +4
$ 2/4
for i 12
which completes the proof.
To break ties in attention, we aim to construct function of that is computable in the transformer, monotonically decreasing with and smaller than 1/(2i*). The following definition will accomplish this:
Definition 6. We define the following inductively:
f(i,0) 1/2 f(i f(i 1, - f(i
By construction, f(i, Rt) monotonically increasing and a linear combination of 1/i, 1/(i "h....,1(ff - k). The latter property means is computable by single multihead self-attention layer. To do this, we construct & heads in parallel, where head h attends to all tokens besides the first h and puts value 1 at token h and elsewhere. Head thus computes 1/(i h). We use the linear transformation at the end of the layer to compute f(,, k) via a linear combination of the head outputs.
Lemma 9. For any and > k we have
k!
f(i,k)
11, - 5)
Proof. By induction over k. Base Case: = 0. We have f(i,0) 0l/i.
10This head can be implemented by setting flag in the previous layer at each i for whether i S by hardcoding comparison hetween (61, 1/i) and ((h, 1).
16

---

Published as a conference paper at ICLR 2024
Inductive Case. We analyze the form of f(i, & 1):
(i.k+L) = f(i k) f(i,k) k)
k!
h!
(Inductive assumption)
of (i 1)
I oo - j]
T (i -) 1T k j= -)
= -!!
(Form common denominator)
'II i j) - k-
j= (i - j) - (i- - k - 1111]1 -1][[==10 - j)
k!
(Pull out factors)
III - (i 5) - ---1)
(k 1) I[ (i - j)
(Distributive property)
(i j)
1)
(I: + 1)!
(Simplify)
II (1 - -i) k - 1)
(k + 1)!
(1 (i
remains be shown that f(i, k) can be made smaller than 1/(21²). To handle edge cases around small values of i, we define:
Definition 7. Let = 10-10 10 For i 2 1, let
f(i) = 1/10000 - éi fi1,3)/100
if i i ifi
Lemma 10. For i 2 1, f(i) 11/(2t)).
Proof. When i < 4, we have f(i) < 1/1000 < 1/(2it).
When 2 5, by Lemma 9, we have:
((i
3!
100
100
1) - 2) 3)
can be verified that the values of i for which this expression equals 1/(2it) are all in the interval
[0, 5), and that for i 2 5, (100/6)ε(i 1) 2) 3) > 2/a
17


Reviews:
[{"review_id": "PmMBpjLp2U", "reviewer": "Reviewer_DD5G", "summary": "This paper investigates how the formal reasoning power of decoder-only transformers change if they are allowed to generate intermediate outputs that can be read in subsequent steps, akin to chain-of-thought or scratchpad reasoning.  The authors show that for an input of size n, adding log n intermediate steps doesn't add much to the power of transformers without intermediate steps.  However, adding a linear number of intermediate steps results in a significant increase in reasoning power -- allowing transformers to recognize languages somewhere between regular and contex-sensitive languages in the Chomsky hierarchy.  The authors also show that using t(n) intermediate steps allows a transformer to mimic t(n) steps of a Turing machine.  Finally, they also show that allowing a transformer poly(n) intermediate steps precisely captures the class of problems solvable in deterministic polynomial time.  The paper presents complexity-theoretic lower and upper bounds of the computational power of transformers with the ability to use t(n) intermediate steps, for various functions t(n).", "questions": "1. Please explain the intuition behind norm-hash using examples.  Otherwise, its introduction appears too abrupt.  This is an important concept for the proofs, and hence it should have been explained more clearly with examples.\n2. Why is it necessary to consider a multi-tape Turing machine when you want to show that a move of a Turing machine can be simulated by one move of a transformer with intermediate steps?  Wouldn't showing this for a single tape Turing machine suffice?\n3. The result concerning characterization of P is not adequately explained in the paper.  For example, there is no subsection devoted to it. It is only mentioned as a bullet on page 3.  Please elaborate on this a bit more, like the other results.\n4.  The paper presumes a lot about prior knowledge of the reader in the formal language theoretic modeling of transformers.  A more gentle introduction would help the readability of the paper.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper establishes both lower and upper bounds on the computational power of transformers with intermediate steps (akin to chain-of-thought or scratchpad reasoning).  The lower and upper bounds are almost tight.  Although I'm not an expert in this area, I find that this work significantly extends state-of-the-art in relating the computational power of transformers to classical complexity classes.  The notion of norm-hash is likely useful in other contexts as well.", "weaknesses": "The paper is technically intricate, and may not be accessible to the general ICLR audience.  The core idea of using norm-hash is not illustrated well.  Using examples to explain the intuition will help in improving the readability of the paper.  The paper assumes a lot in terms of prior knowledge of the reader in the formal language theoretic modeling of transformers.  The notation is at times unnecessarily complicated -- I understand this was done for the sake of rigour, but this does hurt readability of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tLJdtHOVQu", "reviewer": "Reviewer_XLkW", "summary": "A series of recent theoretical work has analysed the limitations of transformers when it comes to solving certain simple sequential problems. In this work, the authors analyse whether these negative results remain when transformers are enriched with a chain-of-thoght. The main results show that these extensions extend the power of transformers to beyond TC^0 and firs-order logic.", "questions": "No questions.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The theoretical results are interesting. In particular, the authors have proved that transformers with a linear number of decoding steps have the capacity to regocnize regular languages. Nevertheless they cannot recognised all context sensitive languages.", "weaknesses": "The paper focus on the analysis of the expressive power of transformers enriched with chain of thought. For the audience of ICLR, it would be interesting to have some results analysing aspects related to learnability of such models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nFSBke6Yr3", "reviewer": "Reviewer_a7kJ", "summary": "This paper studies the reasoning ability of transformer models, which underlie some state-of-the-art deep neural networks such as ChatGPT or GPT-4, in solving sequential reasoning problems such as simulating finite-state machines, deciding connectivity of two nodes in a graph, and solving matrix equalities, which are formalized as classical language classes (context free grammars or regular languages in the Chomsky hierarchy) or traditional complexity classes (bounded-depth threshold circuits $TC^0$, logarithmic space $L$, polynomial time $P$).\n\nThe paper shows that the number of intermediate steps (length of Chain-of-Thought, or amount of scratchpad space) is closely related to the sequential reasoning ability of a decoder-only architecture transformer. In particular, Equation 1 shows that the class of problems solvable with $t(n)$ intermediate steps are those solvable by multitape Turing machines using similar time or space.", "questions": "While some future questions are mentioned (for example to study how different kinds of fine-tuning such as reinforcement learning or improving a model's ability to use chain of thought), how could such learnability problems be approached at a high level? The current paper studies expressive power by simulation results (between chain-of-thought transformers and Turing machines in both directions), which are somewhat standard in the computer science literature and appears unable to answer learnability questions by far.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The connection to traditional complexity classes (Equation 1) effectively refactors out all recurrence/use of scalable memory into the intermediate steps in chain of thought/scrathpad using transformers. This is as expected, and such intuition shows up also in studying traditional complexity classes such as $L$ and $P$ in a natural way, contrasting the need of memory to sequential reasoning.\n\nThe proof of this connection (Equation 1) is also natural. In particular, the simulation of Turing machines by decoder-only transformers using an encoding based on layer-norm hash is also natural.", "weaknesses": "The study of log-precision transformer models with $h$ heads, $d$ layers, dimension $m$, and feedforward width $w$, while capturing some of the trends in Machine Learning in recent years, may not apply when the trend changes to other non-transformer architectures in a decade or two.\n\nThe first equivalence between a transformer-defined classes (the chain of thought class) with a standard complexity class ($P$) is not too unexpected, given that both sides allow for a more lenient notion of equivalence/reductions, namely closed under polynomial time reductions instead of more strict notion of reduction such as logarithmic space reductions. Equating under polynomial time reductions are easier targets to hit.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FJO5BStRyP", "reviewer": "Reviewer_7d6w", "summary": "This submission studies theoretical expressive power of transformer models equipped with chain of thought reasoning. The authors provide theoretical lower and upper bounds for different decoding step sizez, log (L), linear (Time n^2; Space n), and polynomial (P). Their proof for the lower bounds relies on a construction that enables retrieval across different columns in the transformer (“Layer-norm hash”), which is used to simulate automata and Turing machines in transformer forward passes. The upper bounds are proven by straightforwardly simulating attention via a TM.", "questions": "- For the upper bounds: Doesn’t the same reasoning as for the lower bounds apply? For fixed data distributions, complexity classes are irrelevant. How does this mismatch fit into the theory work?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Interesting theoretical work on the expressive power of transformers that sheds light on the limitations (and possible capabilities) of transformers for researchers and practicioners alike\n- The theory is well developed and well presented", "weaknesses": "- The paper is not very well self-contained and relies on the reader being familiar with two previous papers\n- The last paragraph in the conclusion should be clearly marked as limitation of this theoretical bounds (the lower bounds and upper bounds might not have any relecance in practice).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)