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

Thanks for your review! We take your point that experiments with looped transformers would have been interesting to explore, but we have already invested a significant computational budget (about 1000 GPU hours) into the current regular language recognition experiments and are not in a position to rerun them.

It is not entirely clear to me how realistic the setting of average-hard attention is.

Regarding the motivation and realism of the average-hard attention assumption (i.e., $\tau \to 0$), a central motivation of studying the AHAT model is that, without AHAT, it is unclear whether soft-attention transformers could even implement hard-attention retrieval over long contexts. Over any max context length bound, however, soft-attention transformers can arbitrarily approximate AHAT by multiplying all the weights in a head by a specific $\tau$ (which shrinks as context length bound increases).

Some versions of soft-attention transformers (e.g., if the weights can change with $n$) could thus directly simulate AHAT, but if the weights are fixed wrt $n$, it’s unclear if this simulation is possible (and seems intuitively unlikely). Thus, a central motivation in the AHAT model is to abstract away difficulties transformers might have in implementing hard attention. There is definitely opportunity for future work clarifying the relationship of the AHAT model to practice.

When $z$ is stored as $\psi(z) = \langle z, 1, -z, -1 \rangle$, …

You are right that there is an additional $\sqrt{2}$ factor with layer-norm – we will be sure to double-check that we have not omitted this anywhere. However, the constant factor is not significant, since it can be hardcoded into the weights and thus essentially ignored.

The graph connectivity construction requires a cubic amount of padding tokens. Do you think that it might be possible to reduce this amount of tokens?

We don’t quite see how to do this, but that doesn’t mean it’s not possible!

Thank you for your review!

depth-recurrence models in practice … training can be unstable, and … it's hard to collect intermediate supervision

We agree that both of these details (stability and supervision in practice) going beyond expressivity are definitely important. Since we focused on the theoretical expressivity side here, we don’t have much to say about this, besides noting that there is some empirical evidence that looped transformers can be effective on algorithmic tasks. See for example [1]. We will add some discussion highlighting the learning aspects of looping as a potential limitation of this work.

The experiments are performed with non-weight-tied models. Do you have comments on the difference in inductive biases these non-weight-tied models have compared to the weight-tied models?

We don’t have any novel insights on this. However, [1] makes several interesting claims about this, suggesting weight tying might imbue an inductive bias towards “reasoning”:

Claim 2: For language modeling, looped models have an inductive bias towards good reasoning despite having worse perplexity and memorization to an iso-flop non-looped model

Claim 4: Looping-inspired regularization can leverage this inductive bias towards better reasoning

[1] https://arxiv.org/abs/2502.17416

Thank you for your review!

The *abstract might give the impression that scaling depth is universally more efficient than scaling width or chain-of-thought steps, even though the results are specific to the particular context of regular language recognition.

We are unsure which part of the abstract you found to be potentially misleading. We tried to scope any depth vs. width vs. CoT comparisons to the two reasoning problems we consider here. In any case, we are happy to change any specific phrases to address this concern.

The use of the term uniform is not fully clarified. In circuit complexity, "uniform" means that there's a systematic way to generate the circuit parameters for any input size. In this paper, "uniform" seems to refer to having fixed parameters across input lengths / layers, and the reviewer is unsure if this matches the computational theory definition. This could create some confusion about whether previous constructions of [1] are truly non-uniform in the traditional sense. It is also unclear to the reviewer that why the authors believe Transformer can't learn a context length dependent weight as the training is performed on a fixed context length.

Thank you for raising this point - indeed, the sense of uniformity that we mean is stronger than standard uniformity, and we will clarify it appropriately. In circuit complexity, uniform indeed means there is a systematic way to generate the circuit for a certain size. Applying this to transformers, we can say that the parameters are $C$-uniform if the function mapping $n$ to the parameters for a network that processes strings of length $n$ is computable in the class $C$. The sense that we mean with fully uniform for a transformer is very strong– that $C$ is the class of constant functions, i.e., the parameters of the network do not depend on $n$ at all. We will make sure to clarify this appropriately.

relationship with the closely related or possibly concurrent work [2] (reasoning with latent thoughts)

Thank you for bringing the concurrent Saunshi et al. paper on looped transformers to our attention -- we will make sure to cite them and add some comparison. Their Theorem 5.1 addresses the regular language recognition problem, but it is not fully parameter-uniform because the parameters (indeed the embedding dimension) can depend on $n$. Related to this, they also use nonstandard position embeddings in the proof, where each position is represented by its binary encoding, which has $O(\log n)$ length. In contrast, our construction is fully parameter uniform and does not use any nonstandard positional embeddings.

Given that it has been shown in some previous works that linear RNN models or some variants of Transformer models can solve some $NC^1$ problem with constant depth, will logarithmic depth of such modules potentially be in a larger computational class than softmax Transformer?

Good question! The short answer is that our constructions don’t straightforwardly generalize to linear RNNs because they require using attention to retrieve tokens from arbitrarily far back in the context. It’s unclear whether there is some way that linear RNNs could somehow take advantage of logarithmic depth to gain expressivity via a different construction.

Thanks for your review!

width … is truly independent of the input size

Yes! In fact all parameters (their shape and values) in our construction are fully uniform, meaning they have no dependence at all on the sequence length. Thus, the embedding dimension, feedforward dimension, and number of heads do not depend at all on $n$. We will update the theorem statements to indicate this. For Theorem 2, the number of padding tokens does depend on $n$, but, crucially, this is an inference-time dependence and does not result in different parameters for different values of $n$ (unlike prior constructions of Liu et al [1] and Sanford et al [2]; see below).

Experiments

We took the max over runs (i.e., the best trained model) because we wanted the results to more carefully target expressivity (e.g., ignoring learning difficulties, zooming on what is possible for a transformer to represent). We focused our experiments on the regular language task because it has the nice property that every prefix forms a valid subinstance of the problem: thus, to evaluate the maximum length $n^*$ up to which the model succeeds, we can simply find the first index where the prediction was wrong. This property was very handy in our experimental methodology, and is not satisfied with graph connectivity.

informal language

There are some points in the paper where we indeed use slightly informal language to make the paper more readable (e.g., “likely” to mean assuming $TC^0 \neq NC^1$). We see your point that this sacrifices precision in certain cases and will address this – most crucially in the proofs. In particular, we will clarify Lemma 7 and the other specific points you brought up (thanks for reading carefully).

Depths vs. width scaling: Width can be more easily parallelized in comparison to depth; have you considered the practical implications of depth vs. width scaling…

Indeed, scaling width is easier than scaling depth. Our claim is in the context of inherently sequential problems such as regular language recognition and graph connectivity, which we focus on. Our theoretical results show that scaling width, even if easier, fundamentally cannot enable solving these problems (Theorem 3) whereas scaling depth logarithmically via looping layers can (Theorems 1 & 2).

We also note that depth scaling in a looped transformer is an inference-time method: the same trained transformer can be looped more, depending on $n$. Width scaling, on the other hand, is a training-time method: one must train a different transformer for each desired width.

For Theorem 1 and following corollaries, are there any assumptions on the required width for transformers, or is this completely independent of, e.g., the size of the alphabet, the number of states in the finite automaton that captures the language, etc.? On a similar note, what about the width of the feedforward layers, is this also problem independent? And in Lemma 1, are there any assumptions on how many attention heads are needed?

Good question. The embedding dimension, while independent of input length $n$, does need to be large enough to distinguish different tokens in $\Sigma$ and to encode all the possible states. A conservative way to satisfy this would be to construct it as \max \{ |\Sigma|, |Q} \}.

The feedforward network is implementing a finite lookup table from $\Sigma \times Q \to Q$. Thus, its width will have to grow with both the number of states of the minimal DFA and the size of the vocabulary. We have not closely analyzed how many attention heads are needed, but it is fixed independent of the language (i.e., there is some fixed number of heads that would work for any regular language, no matter the size of the minimal DFA).

We will clarify this.

If we consider a (2,1,1)-transformer where $r$ is repeated only one time, then, given definition of $L^\ell$, $L^1$ and $L_2$ are $s$-layers 1 and 2, but $L^3$ is $r$-layer $(3 - 2) \mod 1 = 0$; shouldn't the indexing of layers start at 1 as well?

Thanks for reading carefully and catching this. We’ll update the notation so that the three layer lists (s,r,t) are either all 0-indexed or all 1-indexed.

Comparison to Liu et al. [1]: The main difference to the transformer architecture used in Liu et al. seems to be the universal part, i.e., sharing weights in a dedicated block. However, in Section 4 it is said that Liu et al. strongly simplify their transformer architecture by assuming (i) "specific, nonstandard" positional encodings, and (ii) removing residual connections. In [1], the authors state that they use "fairly standard positional encodings" and use residual connections (both in p. 4, second paragraph). Could you further address these discrepancies?

• Related to what you said, a major difference is that our construction is fully parameter uniform (the parameters are the same for any context length), whereas the Liu et al. construction requires the parameters to change with input length $n$.

• As mentioned on page 44, the Liu et al. construction assumes each index $t$ gets two sinusoidal position embeddings $\cos(\pi * t / T) and \sin(\pi * t / T)$. While this somewhat resembles the absolute positional embeddings used in the original transformer, our construction is much more general, since it holds without any positional embeddings (as long as a beginning of sequence symbol is used or the first token can be uniquely attended to). We will clarify this difference in revisions and remove the current vague language.

• Another difference between our construction and theirs is the presence of $layer norm$, which is standard in all practical implementations. As they mention on page 25, Liu et al. omit layer norm from their construction, whereas our construction uses (and in fact highly leverages) layer norm.

• Regarding residual connections, it’s straightforward to adapt the Liu et al. construction to work with residual connections in a non-parameter-uniform setting. However, with full parameter uniformity (like we assume), it becomes technically challenging because looped layer blocks need to reuse the same cells in the residual stream. Dealing with this appropriately (residual stream management) is a novel contribution of this paper. We will revise the comparison with Liu et al. to better clarify these points.

Comparison to Sanford et al. [2]: Could you please put Theorem 2 in context with the results for parallelizable tasks in [2], i.e., Theorem 11 and Theorem 18, which also state that log depth is required for transformers that are able to solve st-connectivity.

A key difference between our construction and that of Sanford et al. is, again, full parameter uniformity. Whereas our construction uses fixed embedding dimension (and fixed parameters) for all $n$, Theorem 18 in Sanford et al. uses an embedding dimension that grows as a sublinear polynomial in $n$.

Interestingly, their construction uses a similar amount of “padding” tokens ($n^3$ to $n^4$ tokens, depending on the embedding dimension) to ours ($n^3$), unless the embedding dimension is allowed to grow as at least $\sqrt(n)$, in which case their construction can work without padding. We will add some discussion of this comparison, focusing on the central difference of full parameter uniformity for our construction.