Exact Expressive Power of Transformers with Padding

How precisely you define $TC^d$? If as $O(\log^d)$-depth threshold circuit, then definition 11 in the supplementary material does not make sense to me, because any circuit belongs to some sequence of circuits from $TC^d$ (by making constant in $O(\log^d)$ big enough). I guess you should for any constant $C, d$, define the $C,d$-threshold circuit evaluation problem, where the input circuit has depth at most $C \log^d n$ precisely.

Indeed, you’re right that $TC^d$ circuit evaluation is the same as $TC$ circuit evaluation because the constant can be folded in. This is why we have taken the care to define wide-$TC^d$ class of circuits and the wide-$TC^d$ circuit evaluation problem in the appendix.

wide-$TC^d$ is defined as the class of $TC$ circuit families where each family has a fixed $c$ such that depth is at most $c \log n$ and, crucially, size is at least $n^c$. With this additional constraint, wide-$TC^d$ circuits express the same languages as $TC^d$, but the evaluation problem can be solved in $TC^d$ (unlike evaluation of $TC^d$, for the reasons you noted). To see why this constraint helps, note that if we increase $c$ to make a wide-$TC^d$ circuit deep, the length of the circuit serialization must increase, which effectively makes the circuit shallow w.r.t. the total input length for the circuit evaluation problem, upper bounding its effective depth by $c^* \log^d n$ for some fixed constant $c*$ chosen a priori. This circuit then can be evaluated by a $(c^* \log^d n)$-depth circuit family (for a fixed $c^*$), which we build for the wide-$TC^d$ circuit evaluation problem.

Thank you for raising this important point. We will clarify this and also make the wide-$TC^d$ a numbered definition (after Definition 11) as it’s an important concept that addresses the issue you mentioned.

Some minor remarks: Definition 1, $h$ is the number of attention heads, and $k$ is their index?

Yes, we will note this.

Line 176, maybe clarify that ``unique hard-attention transformers''

Will do.

Names of Sections 4 and 5 are misleading because you forget to mention uniformity conditions there.

You are right. We omitted the uniformity qualifier to prevent the section titles from becoming too long, but we agree it would be better to include it and will update in revisions.

Line 126, in the positional encoding here, $n$ is the initial length or padded one?

In discussing this question, we realized that $i/n$ can actually be simplified to $1/i$ embeddings, and all constructions will go through. This works because $1/n$ (where $n$ is the initial input length) can be computed first, and then $\phi(i/n)$ can be computed as $\phi(1/i, 1/n)$. This simplifies the conversion from unmasked to masked transformers because $1/i$ encodings can be directly simulated by a uniform causally masked head. We will update the setup in the paper to use the simpler $1/i$ embeddings and clarify these details. Thanks for your question, which led us to think about this!

Thank you for your review! You raise a good point regarding the* comparison of padding and looping with chain of thought*. In the case of $\log n$ looping vs. CoT, it is now known [1][2]] that transformers with logarithmic CoT steps remain in $TC^0$ whereas log-looped transformers can solve some $NC^1$-complete problems, suggesting an advantage of looping. More generally, the comparison between transformers with $t$ steps of looping vs. $t$ steps of CoT remains open. We will update the paper to discuss this comparison, what is known, as well as the open question.

Regarding learnability, we agree this is very interesting, but is beyond the scope of the current work, which we can acknowledge as a limitation.

[1] https://arxiv.org/abs/2503.03961 [2] https://arxiv.org/abs/2210.10749

Thanks for your review! We are happy to hear that you appreciate the technical quality and clarity of our paper as well as its potential for impact.

We agree that the practical learnability of AHAT constructions (with masked pre-norm) is an interesting open question. We will add some discussion either in the main paper or in a limitations section. The major gap between soft attention and AHAT is indeed the ability to implement hard-attention retrieval. For any fixed maximum context length, it is possible to scale the attention head parameters to arbitrarily approximate AHAT attention, but for unbounded context lengths, it seems difficult (and may likely be impossible) to simulate AHAT or hard attention. For this reason, we view AHAT as a mostly mild assumption (suitable for bounded context lenghts) that abstracts away issues with soft attention for simulating hard attention over very long contexts.

Regarding the practicality of padding, $n^k$ is indeed a practical barrier if $k$ must be large for some tasks. Small values of $k$ will suffice for some natural tasks tasks: for example, in the 3-SUM task considered by Pfau et al. But, in general, there may be tasks in $TC^0$ that require a large value of $k$, which would lead to a large memory overhead to store a long context window (the same problem would plague transformers with long CoT or chain of thought). More work is needed to assess the tasks for which small values of $k$ would suffice or whether other techniques could be developed to reduce the memory overhead of padding -- we will acknowledge this as a limitation.

Regarding your question about $L$ vs. $FO$ uniformity, we have a brief discussion of this in the appendix after the proof of Lemma 4 (that $\log^d$ depth transformers are in $TC^d$).

Indeed, for $d \geq 1$, it can be shown that $\log^d$-looped transformers would be in $FO$-uniform $TC^d$, similar to how you suggest. The idea is that looping an $FO$-uniform $TC^0$ circuit $O(\log^d n)$ times can be implemented in $FO$-uniform $TC^d$.

This has the interesting consequence that $FO$-uniform $TC^d = L$-uniform $TC^d$ by our uniformity ceiling lemma (Proposition 2, Appendix C) , just like the $L$-uniform vs. $NL$-uniform collapse noted in the appendix. In answering your question, we looked closely at the uniformity ceiling argument (Prop. 2) and we think it also goes through when $XC = wide-TC^d$ for $d \geq 1$, class $B = L$, and class $A = FO$, showing that $L$-uniformity collapses to $FO$-uniformity for deeper $TC$ classes.

Note that this doesn’t go through for $TC^0 (d=0)$ because $L$ is not known to be a subset of $FO$-uniform $TC^0$. Thus, the uniformity ceiling lemma cannot be applied.

Thank you for bringing this refinement to our attention! We will add a brief section highlighting these novel results about uniformity and also the novel complete problems we developed to obtain these uniformity collapse results, since these might be of independent interest.

Thank you for your review!

Thanks for your careful reading and feedback on clarity. To respond to some of your points:

• Regarding $\phi$, we will standardize it to only refer to the layer-norm hash. We will change Definition 3 to explicitly mention this sense of $\phi$ (we apologize that it seems to have gotten removed at the last minute) We will change the symbol for formulas in Definitions 5 and 8 to $P$ or $Q$. This will make both aspects of the paper clearer to read.

• In Definition 6, “universal” indeed means “looped”. We will change this. By “bos” we mean a beginning-of-sequence symbol $ as defined on page 3. We will update the text to make this clear.

• “Masked unmasked” is a typo and should just be “masked”

Regarding Lemma 2, we will take into account your feedback on clarity, and briefly address here your following question:

Could you explain what "Each token corresponds to a specific assignment of all the variables, which we denote $v$" means?

If a formula has $k$ variables ranging from $1$ to $n$, there are $n^k$ possible assignments to that vector of variables. We can think about a single assignment $v$ as a tuple in $[n]^k$, or we can think about it as a number between $1$ and $n^k$. What we mean is that we identify each variable assignment with a specific number. Then, token $v$ will represent variable assignment $v$.

We will also incorporate all the other detailed feedback you mentioned regarding clarity and informal language - thanks for reading carefully! On the use of informal terms, we did that in certain places to make the argument more accessible to broader audience. However, we see that it can result in lack of precision and we will make sure to make the language more precise -- critically in all proofs.

To my current understanding, Proposition 10.3 in Barrington et al. implies that we can express $FO[M^2]$ by using formulas from $FO[M,Bit]$, but not the other way around…

Your interpretation of Proposition 10.3 is correct, but what we use here is Theorem 10.2, which proves the other direction that $FO[M^2]$ can capture all of $TC^0$. We will update our reference to the Barrington et al. paper to include the relevant theorem.

In lines 352-357 in the proof of Lemma 3: Why do we check if $f(w) \in L’$ instead of $L$?

Indeed, have a minor issue with the notation in this paragraph where we flipped L and L’. We will correct it to say that we are checking whether $f(w) \in L$, which is equivalent to checking $w \in L’$. Thanks for finding this!

In Theorem 1 and Theorem 2, do we assume infinitely many padding tokens (given Definition 6), and thus infinite width?

The number of padding tokens needed in the simulations in Theorem 1, 2, and 3 is naturally a function of the complexity of the underlying $TC^0$, $NL$, and $TC^d$ circuits / algorithms. It’s not infinite, but rather unbounded a priori. E.g., if a $TC^0$ circuit being simulated is equivalent to a $FO[M^2]$ formula with $k$ distinct variables, then we need $n^k$ padding tokens (in Lemma 2 and Theorem 1). Since one can define arbitrarily large circuits, their transformer simulation also needs to support correspondingly more (but still polynomially many) padding tokens, i.e., a larger exponent $k$ in $n^k$. Crucially, for every problem, there is a fixed $k$ that works. The subscript $$ in our notation $AHAT^d_*$, representing the union over all $k$, captures this. We will clarify this, in particular that for simulating any fixed circuit family, we need only a fixed polynomial degree number of padding tokens.

Could you explain what exactly is meant by the phrase "converge to recognizing NC" in the context of your theoretical results?

$NC$ can be defined as $\bigcup_{k = 0}^\infty TC^k$. With $O(\log^k)$ depth, we showed padded transformers can recognize $TC^k$. So, as $k \to \infty$, $O(\log^k)$ depth (i.e., polylog-depth) transformers can recognize any problem in $NC$. In other words, for any problem in $NC$, there is some choice of $k$ such that a $\log^k$ depth padded transformer will recognize it. We will clarify the discussion around this to make it more precise.

How expensive is padding and looping in practice? How can padding be efficiently parallelized?

Padding is efficiently parallelizable in practice: in can be thought of as just growing the input with blank symbols, and then making a single call for transformer inference (as opposed to CoT, which requires sequential inference calls to the transformer).

The practical bottleneck it would run into is the length of the context window, as running a transformer with a long context will run into memory limitations. Some problems like 3SUM considered by Pfau et al. are solvable with a small padding degree $k$, but if a problem requires a larger value of $k$, this could definitely become impractical (the same argument holds for CoT, in addition to the sequentiality of CoT). We will clarify this potential limitation in the text.