A Little Depth Goes a Long Way: the Expressive Power of Log-Depth Transformers

Thanks for your review! We appreciate your detailed technical comments and questions, and we will incorporate your feedback to improve the paper.

W1: The motivation for log-depth Transformers lacks sufficient conviction. Given the bounded context assumption in line 50, it would be more appropriate to treat all elements as constant. If bounded context and model sizes are to be discussed carefully, as indicated in lines 76-78 regarding the maximum problem size of graph connectivity with fixed depth d, the coefficient is crucial, while the result of 2^O(d) remains unclear.

If we understand correctly, your concern is that when considering the bounded context case, it would be appropriate to treat all architectural parameters (depth, width, number of heads, …) as constants (w.r.t. input length $n$). We agree – given a specific input size bound $n’$, there is indeed a fixed $d’$, a fixed $w’$, a fixed number of heads, etc., that suffice to solve the task.

Our theoretical and newly reported empirical results reveal what values of $(d, w)$ suffice (for the considered reasoning problems) for any desired context bound $n$. In other words, if one is only interested in solving, say, the state tracking task for inputs up to some bounded length $n$, then depth $d(n) = O(\log n)$ suffices with a small constant factor (theoretically 7, empirically less than 5), while width must scale exponentially as $w(n) = \exp(O(n))$.

Does this help address your concern?

W2: The results presented lack clarity. In the main theorems (Theorems 1 and 2), only the model depth of O(log⁡ n) is specified, omitting other critical configurations such as the required model width and the number of attention heads.

We will change the theorem statements to clarify that the model width and number of attention heads are $O(1)$ w.r.t. $n$ in our constructions in Theorem 1 and 2 – depth is the only parameter that can grow (minimally) with $n$.

W3: Additional clarification is needed regarding the results in relation to prior work. Specifically, Lemma 3 of [1] demonstrates that a Transformer with depth 2d maps ⟨C,x⟩ to the circuit value C(x), where C is a threshold circuit with depth d. Since regular language recognition and the graph connectivity problem both belong to TC1, this lemma could already suggest comparable results in this paper.

This is an interesting connection! However, the argument does not quite go through. If the transformer is allowed to receive an input-length dependent prompt of size poly(n) before the input x, then indeed you’re right that, following from [1], there is an input-dependent prompt C(|x|) (namely a log-depth threshold circuit that implements regular language recognition or graph connectivity for inputs of size |x|) such that a log-depth circuit will map ⟨C(|x|),x⟩ to y.

However, there are a couple problems with using this idea to try to prove Theorem 1 or 2. First, the mapping from x to a representation of C(|x|) is actually quite complicated, and it’s unclear whether a transformer can compute it. Further, transformer gets as input just x, and it cannot “prepend” additional tokens C(|x|) to be read by later layers. In other words, while regular language recognition and graph connectivity are reducible in some sense to the circuit value problem (CVP), transformers cannot directly implement the reduction. Thus, it does not directly follow from the CVP construction in [1] that log-depth transformers can express automata or graph connectivity. We thus obtain these results through other means in Theorems 1 and 2.

To the extent that space allows, we will add a discussion along these lines to the paper!

W4: The absence of experimental results to validate the theoretical findings is a significant gap. Empirical verification of the solid lines in Figure 1 is also essential.

Regarding the lack of experimental results, we have added extensive experimental results on the A5 state tracking task that confirm the theoretical predictions around depth and width (see the general response). We will replace Figure 1 with a figure showing these empirical results. We hope these experiments, in addition to the comments above, have addressed your main concerns and believe they have improved the quality of the paper.

Q1: What are the other configurations (e.g., model width, maximum number of attention heads, etc.) for the Transformers in the main theorems, specifically Theorems 1 and 2?

All other configurations in the model architecture are constants; we will clarify this (see response above to W2).

Q2: Can you provide more clarification of your results in relation to prior work, such as [1]? (Refer to weakness 3)

Please see the answer to W3 above.

Q3: Can you conduct experiments on the two tasks considered in the paper…

Yes! Please see the answer to W4 above.

Thanks for the reminder about the page limit. We've now updated the submission PDF to fit within 10 pages. We've also cited Figure 10 from Li et al in Section 6 and given them credit (while also maintaining our proofs in the appendix for completeness).

While we agree that we focus on the problems of regular language recognition and graph connectivity (and plan to extend to more general results about classes of problems in the future), we note that analyzing specific problems often paves the way towards general results, and also that these are quite important problems. Regular language recognition, in particular, formalizes many types of real-world state tracking, as has been noted in previous work:

• https://arxiv.org/abs/2404.08819
• https://arxiv.org/abs/2210.10749

We appreciate your suggestions and believe our revised manuscript is substantially improved to address your original weaknesses (clarifying definitions, connection to prior work, and adding experiments). Regarding the motivation for log depth, we have also added the following paragraph to Section 2:

An $(s,r,t)$-transformer unrolled to some fixed depth can be viewed as a uniform special case of a fixed-depth transformer. Thus, constructions of dynamic-depth transformers (depth $d(n)$ for inputs of length $n$) imply that, given any bounded context length $N$, there also exists a fixed-depth transformer with depth $d(N)$ for the task at hand. The fact that this can be done with a looped transformer with dynamic depth is, in fact, a stronger condition that shows the construction is uniform, which is formally important as non-uniform models of computation can have very strong and unrealistic power. In this way, our results about looped transformers will provide insights about standard, non-looped transformers with bounded context lengths.

I encourage the authors to engage in a discussion with me, even if you can only address some of my questions at this stage.

As the rebuttal phase approaches its conclusion and considering the various points I have raised, I believe an early discussion would be beneficial.

because it is often called st-connectivity in the computational complexity literature,

Okay, I see this connection. For me, reachability is the natural one. Do as you wish, but I would recommend mention it somewhere as I think I am not the only one whose more used to reachability.

..., which hopefully satisfies your concern. [...] We are happy to hear alternative suggestions. [...] Our main goal is to highlight that this is the first formal analysis that does not consider the depth to be a fixed constant (w.r.t. input size).

The first part sounds like I should actually be named "reviewer pIckY53". Jokes aside, I am aware that this is only an informal statement and should be treated as such. But, especially when concerned with formal research like you are, these informal statements should be handled with extra care. To me, this sentence conveys "most expressiveness research focuses on models with a bounded depth". I see there can be other interpretations, but you probably agree that its problematic as soon as there is more than one. However, this was only a minor point.

This sentence was simply meant as an analogy to help readers better understand our research question and results. If you find it unhelpful or distracting, we would also be happy to delete it.

I appreciate your mission here. But I do not think that this analogy helps very much.

So far, I feel that this submission touches the threshold of acceptance. I see its worth and am willed to increase my rating, but please answer the follow-up questions I raised first, because these are important to me.

Thanks for your review!

W1: Major weakness is that the theoretical results lack experimental validation on the tasks author suggested, limiting real world applicability.

We agree that experiments would improve the original paper, especially because the theory makes testable predictions. We have added extensive experiments that confirm the predicted theoretical relationship between effective context length for state tracking and depth and width (see the general response). We believe that these results strengthen our argument and its applicability and would be curious to hear if they address your concern.

W1: Also, dismissing CoT with a couple of lines seems unfair, given the potential showed by CoT on arithmetic and graph related problems, above all.

We agree that the comparison to chain of thought in section 6 could be made more nuanced. Specifically, we will clarify that the context in which we have found growing depth is more efficient than chain of thought is only for specific reasoning problems. For other sorts of reasoning problems, it is feasible that chain of thought could be a better computational fit than growing depth, and there might be other concerns (simplicity, trainability, etc.) under which chain of thought is more appealing than looping layers. We will include this discussion.

W3: The paper is very intricate and a bit more explanation may be needed …

We plan to clarify the transformer definitions and theorem statements as discussed in responses to other reviewers. This will help simplify the proof structure in the main results (Theorems 1-2). Regarding Theorems 3 and 4, we view these as simple results that help contextualize the main results in Theorems 1-2 (by comparing the effect of depth to width or CoT steps). We will add more discussion to these sections to make the role of these results clearer. Regarding justification, note that these theorems do have proofs that were moved to the appendix for space. Regarding self-containedness, if there are specific technical tools that are not adequately introduced beyond those we have mentioned, please let us know and we would be happy to incorporate additional background in the appendix.

W2: Given the memory management issue, implementing dynamical depth is not clear and also it's not clear if a looped transformer can be indeed trained to solve the two mentioned problems.

The constructions in Theorem 1 and 2 show how, in principle, it is possible to get around the memory management issue with hand-crafted constructions. You are right that a question remains about whether such constructions are learnable in practice. The experiments we have added (see general response) get at this question to an extent: we find that, in practice, transformers with depth ~ $\log n$ can learn to express state tracking over $n$ items. However, our empirical results focused on non-looped transformers: it would be interesting to replicate our analysis with looped transformers but we cannot do this currently due to the significant computational cost. We will acknowledge this as an interesting direction for future work in the paper.

W4: Although theorem 3 and 4 are the core part of the authors work, they are not well justified.

We actually view Theorems 1 and 2 as the main contributions of this paper, while Theorem 3 and 4 are meant to contextualize our main results about depth scaling against other width scaling or chain of thought. Regarding the justification of Theorems 3 and 4, we note that we have included proofs for the theorems, but deferred them to the appendix due to space. We agree with the reviewer that more space could be used to explain the significance of these theorems. In particular, we will add some text explaining the importance of Theorem 3: since regular language recognition and graph connectivity are outside TC0, this result suggests width must scale exponentially w.r.t. $n$ to solve these problems up to bounded length $n$. Our new experiments confirm this theoretical prediction.

Thanks for your review! We appreciate that you found our results about depth scaling insightful.

Regarding the consideration of complementary methods, we agree that the discussion in the original draft could be expanded and made more nuanced. For instance, we will clarify that we have found a previously unknown advantage of scaling depth over width for solving reasoning problems outside TC0 (such as state tracking), but there are likely other kinds of tasks (e.g., those that are knowledge intensive or that involve large amounts of parallel search) where width may be more beneficial than depth. Additionally, scaling depth comes with a cost in terms of training time compared to scaling width, which we will mention explicitly. We will also contextualize our results with empirical papers comparing depth and width scaling (e.g., Petty et al., 2024). Regarding the comparison with chain of thought, we agree that there is also nuance to add. Concretely, we will clarify that we have only found depth scaling to be more efficient than chain of thought on the specific problem of state tracking. Further, we will mention that, while transformers with log precision and polynomial chain of thought steps converge to P (all polynomial time solvable problems; Merrill & Sabharwal, 2024), a result like this could not go through for depth scaling because the memory would be bounded at $n^2 \log n$. We believe clarifying these comparisons and expanding the discussion will improve the value of the paper for readers. Thank you, again, for bringing this up.

We also agree that the original submission could be improved by adding experiments, especially since our theory makes testable claims about the relationship between depth and effective context length for certain reasoning tasks. We have therefore run extensive experiments to empirically verify that the context length at which transformers can perform state tracking does scale as expected with depth and width: see the general response for more details. We believe these results are a strong complement to the theoretical results and enhance the overall story and potential impact of the paper. Thank you for suggesting this!

We have updated the PDF to include the new experimental results discussed above. The new text is (temporarily) highlighted in blue to make it easier for you to find it. The lack of experiments was a point of concern in the original reviews, so we would appreciate hearing whether these experiments, which verify the theorized relationships between depth, width, and context length for regular language recognition, change your assessment of the paper.