When Can Transformers Count to n?

We thank the reviewer for the positive feedback.

Weaknesses:

1. Our work is indeed focused on the expressive power of transformers, which is very common in the learning theory community, as there are many papers on expressiveness of neural networks. Studying optimization theoretically is out of scope for our paper, although experimentally we show our results for both transformers trained from scratch and pre-trained models (e.g. Gemini 1.5). This indicates that the issues discovered in our paper persists also for optimization.

2. We agree that it would be interesting to understand how transformers implement counting tasks internally. However, our focus is on when transformers are expected to fail in this task, and our empirical results agree with our theoretical predictions. In these failure cases, the transformer does not carry out the counting task, and therefore mechanistic interpretability becomes less relevant.

Questions:

1. The constructions require log precision. Namely to count until n, having O(log(n)) bits suffices. Note that the lower bound in Theorem 5.1 specifically includes the number of bits, namely having n bits can break this lower bound. This is not surprising as having very large precision can replace the requirement of having a large embedding dimension by encoding the different coordinates into a single large-precision number.

2. As noted above, our focus is on understanding when transformers are expected to fail, so we are less focused on solutions that are implemented in success cases. But we agree of course that this is an interesting question for future study.

3. The lower bound in Lemma 4.4 holds up to depth log(n), although with a smaller width. Showing this limitation requires counting the number of possible linear regions which was done in previous works that study piecewise linear activations. The lower bound in Theorem 5.1 is only for depth 1. It uses a communication complexity argument, and it is currently an open problem to extend this lower bound for more than 1-layer. Both of these lower bounds do not hold for CoT prompting, as we did not study this setting. It was shown that unbounded CoT prompting is a universal approximator (Auto-Regressive Next-Token Predictors are Universal Learners, 2023), so extending our lower bounds to this setting requires some bound on the allowed CoT prompting. This will be an interesting direction for future work.

We thank the reviewer for the thorough review.

We believe that the reviewer missed the main point of the paper. Namely, that transformers are limited when performing simple counting tasks for large vocabulary sizes. Our results show that simple solutions fail when the dictionary size is large (Theorem 4.2), and that any solution to the MFE tasks requires the embedding dimension to depend on the dictionary size (Theorem 5.1). The histogram solution that the reviewer mentions is not the main result of our paper. It just shows that there exists a simple solution when the dictionary size is small.

In particular, papers about the parity problem (e.g. Overcoming a Theoretical Limitation of Self-Attention) only deal with dictionaries of size 2, while our paper focuses on problems where the dictionary size is large. Specifically, our theoretical results predict that counting would become harder as vocabulary size grows. And, our empirical results are in very good agreement with this, and also agree with the phase transition predicted by theory.

To answer the specific issues:

1. The histogram solution that the reviewer mentions is only applicable in unrealistic situations when the embedding dimension is larger than the dictionary size. In realistic settings, we show the limitations of transformers to solve such tasks. Also note that the parity problem is inherently different from counting when the dictionary size is large, thus the works on parity and other bit-related problems are not comparable to our work. With that said, we will add “Overcoming a Theoretical Limitation of Self-Attention” to the related works section.

2. We indeed agree that transformers lack the expressive ability to count, and as far as we know, ours is the first work that shows this limitation theoretically for large vocabulary sizes (compared to vocabulary size of 2 in bit-related problems). One important thing to notice is that transformers DO NOT have expressivity problems when it comes to counting a small vocabulary. This is precisely what our histogram solution shows. Specifically, counting bits is straightforward (e.g. if the input is 0,1 then summing it will produce the count. There is still of course the need to translate the count to the set of right tokens, but in our case the expressivity issue arises even without this requirement). The reviewer should note that counting bits is very different from the parity task, where there is the added challenge of calculating the parity of the count.

3. Our theoretical analysis does in fact make use of positional embeddings and precision (see Theorem 5.1 which has a limitation on the required bit-precision). We also mention layer normalization in lines 159-160 and that it does not affect our theoretical results. We also mention residual connections several times and make use of them (see lines 198, 231-232). We indeed focus on the ReLU activation which is done in many previous works. Also, in the experimental framework we use all the above components. We do not understand this criticism as our work covers all of them, both theoretically and empirically.

4. We specifically make use of finite-bit precision in Section 5, and Theorem 5.1 includes the term “p” which is the precision of the model. Theorem 4.1 and the histogram solution can work with O(1) bit-precision, independent of the context length or dictionary size.

5. The reviewer writes: “The counting discussed in this paper is non-inductive, while most modern large language models (LLMs) rely on inductive counting, requiring O(N) depth complexity.” We ask that the reviewer define exactly what they view counts as “inductive counting” and in what sense “most modern large language models (LLMs) rely on inductive counting”. Our paper provides a clear definition of the counting problem we address, and we see no clear argument for why it is not an interesting counting problem. Specifically, we do not agree that counting to N requires O(N) depth, because our histogram construction shows that one can always solve the query-count task for vocabulary size m, as long as the model dimension d satisfies d>2m, with one transformer layer. This refutes the claim by the reviewer that O(N) depth is needed.
   The reviewer states that RNNs have a natural mechanism for counting (via their state). This is true, but we don’t see how it is relevant when one studies how transformers (which currently underlie all successful LLMs) count. We also note that the reviewer mentions two papers in this comment where one (“Language Models Need Inductive Biases to Count Inductively”) is under review at this ICLR and the other (“Counting Ability of Large Language Models and Impact of Tokenization”) was put on arxiv AFTER the ICLR deadline. Our paper is in any case very different from these, because our focus is on the interplay between vocabulary size and model dimension. However, in any case, these would not count as prior work.

Thanks for your comment, my concern are all addressed and I have changed my scores, the only thing I am still worried about Is experiments, as mentioned in my review, as theoretical side now makes sense to me.

We thank the reviewer for the thorough review. We address the comments below.

Weaknesses:

1. Our work focuses on simple counting tasks, and shows that even for such simple tasks there are inherent limitations of the transformers architecture for large vocabulary sizes. We believe this is an important stepping stone in the study of transformers. For comparison, there are many works on neural networks studying specific tasks that greatly helped the community in understanding the strengths and limitations of such architectures. Some examples include depth separation for specific target functions [Telgarsky, 2016], [Eldan & Shamir, 2017], learning parity under different distributions [Malach & Daniely 2020], learning single neurons using kernels [Yehudai & Shamir 2019], [Kamath et al. 2020] and many more. Additionally, counting is a core capability that underlies more elaborate tasks (e.g., summarization, data analysis). Thus, we believe it is very important to understand when transformers can achieve this task.

2. Indeed our theoretical results are limited to 1-layer of attention, although we do allow an arbitrary sized MLP (for example in Theorem 5.1). Our result that 1-layer transformers cannot count (e.g., for MFE) when vocabulary size grows is novel and surprising, because it shows the special role of vocabulary size in this task. We use a communication complexity argument, and it is an open problem on how to extend it to multi-layer transformers. We are currently not aware of any theoretical work showing expressive limitation of transformers for more than 2-layers, and showing such a result would have great impact.

Questions:

1. We believe that if transformers are limited in solving such simple counting tasks, then solving even more complex tasks would be even more limited. Our paper focuses on such tasks since these can be analyzed theoretically, and we are the first work, as far as we know, to show theoretically the limitations for solving such simple counting tasks when the vocabulary size is big. It would be interesting to formalize and theoretically analyze other tasks related to code and math, although we suspect it would be hard.

2. Empirically, the limitations are also manifested in multi-layer transformers, and also large pretrained models (e.g. Gemini 1.5), as depicted in the experiments section. Theoretically, it would be very interesting to show expressiveness limitations for multi-layer transformers. This issue is also discussed in Representational strengths and limitations of transformers (Sanford et al. 2023).

Thanks for your response. My concerns have been addressed, I will increase my score.

We thank the reviewer for the positive feedback. We address the comments below.

Weaknesses:

1. Indeed our paper focuses on the expressive power of transformers from a theoretical point of view, as was done in several previous works. Clearly it is possible to also pursue a mechanistic interpretability study of counting, and our results indeed suggest what these mechanisms might be. However, we believe this would make sense as a separate contribution, and our results here are of sufficient interest, as they point to core limitations of transformers for this important task.

2. The width bottleneck in the CountAttend construction is indeed presented only for 1-layer MLP. It can be readily extended to deeper MLPs at the cost of having a smaller width constraint, and we will add a remark about this in the final version. Namely, the bottleneck means that the construction requires $n$ linear pieces, which can be done using an $L$-layer MLP with width $m$ as long as $2^L\cdot m \geq \Omega(n)$ (this is a standard result about counting the number of linear pieces for a ReLU MLP. See e.g., “On the Number of Linear Regions of Deep Neural Networks, 2014”). This means that if $L=1$ (as we have in the result) the width needs to be $\Omega(n)$, but it can also be done with depth $\log(n)$ and width $O(1)$. We chose to focus on small and fixed depth, as that is the common implementation in transformers.

3. This is a good point, we will make it clear in the final version.

Questions:

1. Thanks for the great question. In the experimental section we do use a standard architecture with normalization layers, and the effect is evident there, so LayerNorm does not seem to help in solving the task. In the theory part, the MFE lower bound still holds with LayerNorm because the lower bound does not put any restriction on the MLP size, and an MLP can be used to implement LayerNorm (e.g. by adding neurons that compute the square of their input, which can be done with $\log(eps^{-1})$ layers, using the construction from Telgarsky 2016).

2. This sounds like a good suggestion for additional experiments to gain insight into internal implementation. Namely, training on one distribution of tokens (e.g. i.i.d), and testing on a different one (e.g. when different tokens have different probabilities). We chose to focus on the simplest empirical settings, as they already demonstrate our key result (i.e., that increasing vocabulary size makes counting harder for transformers). We will consider adding such an experiment in the final version.

We thank the reviewer for the thorough review. We address the comments below.

We first emphasize that the focus of our paper is the interplay between vocabulary size and embedding dimensionality in counting tasks. Our theoretical results suggest that counting will be harder for transformers as the vocabulary size grows, and this is supported by our empirical results. Second, our lower bounds are general and apply to any construction, not specifically the ones we provide. For example, in Theorem 5.1 we prove that any one-layer transformer cannot solve the MFE task unless the embedding dimension is large enough.

Regarding the questions:

1. Our theoretical constructions and lower bounds are designed to count tokens rather than words or letters. This applies to both the QC and MFE tasks. To simplify the presentation of the task we use a different letter for each possible token. In our experiments, for the transformers trained from scratch, we used a similar setting as in the theoretical section where each letter is presented as a token. For the pre-trained transformer (i.e. Gemini 1.5) we used the standard tokenization used by the model.

2. Throughout the paper, we use an MLP for all of our constructions. This is also mentioned in the theorem statements. Our constructions and lower bounds can be readily extended to include skip connections, we haven’t used them as they are unnecessary for the upper bounds and don’t affect the lower bounds. We will add them if the reviewer thinks it will strengthen the results. Regarding the number of layers, indeed our lower bound in Theorem 5.1 is restricted to 1-layer of attention (although with an MLP that has any number of layers). We use a communication complexity argument for the proof, and it is an open problem on how to extend this proof method for multi-layer transformers, which is also discussed in previous works on the theory of transformers (see e.g. Representational strengths and limitations of transformers, 2023)

3. This is a simple and straightforward fact that we mention informally and was also made in a recent work (Barbero et al. 2024). We can add more formal proof to it. It does happen in practice as it is a limitation of the architecture itself. However, note that this limitation only happens when the model has no positional embedding. Adding even the simplest positional embedding (e.g., by marking the last token) solves this issue.

4. This assumption is only part of the expressiveness result presented in Theorem 4.1 when the input dimension is larger than the dictionary size. This often does not hold in practice, because the dictionary size is several times larger than the embedding dimension. This is also not a main assumption in our paper and is only relevant for this specific result. For this reason, in Section 4.3 we describe how in a more realistic setting the histogram solution fails and a more complex solution is needed (i.e. the CountAttend).

5. For the experiments, all the training details are presented in Section 6, including data generation, and implementation details. In Section 6.2 we specifically discuss pretrained transformers, which still struggle with counting tasks when the dictionary size is large. In Section 6.1 we discuss transformers trained from scratch on this specific task, which is stronger than just fine-tuning, and the problem with counting over large vocabulary persists. Indeed, in this training, we provide supervision of (sequence, count) pairs (corresponding to x1,...,xn and y described in the text, in Tasks section of Section 6.1). Our results in all these settings demonstrate a clear effect of vocabulary size on the ability to count, as predicted by our theoretical analysis.

The reviewer mentioned loose ends in the proofs, we would be very happy if the reviewer could point to specific such points so that we could address these in our response.