When Can Transformers Count to n?

## Abstract

Abstract Large language models based on the transformer architecture can solve highly complex tasks. But are there simple tasks that such models cannot solve? Here we focus on very simple counting tasks, that involve counting how many times tokens in the vocabulary appeared in a string. We show that if the dimension of the transformer state is linear in the context length, this task can be solved. However, the solution we propose does not scale beyond this limit, and we
provide theoretical arguments for why it is likely impossible for a size-limited transformer to implement this task. Our empirical results demonstrate the same phase-transition in performance, as anticipated by the theoretical argument. Our results demonstrate the importance of understanding how transformers can solve simple tasks.

## 1 Introduction

Large language models (LLMs) have demonstrated striking performance on a wide array of tasks, from creative writing to solving complex math problems. Given these successes, a key question arises: what can these models do, and just as importantly what can they not do. There are multiple ways to address this question of expressiveness of LLMs. First, it can be studied empirically by probing LLMs and seeking relatively simple tasks that they cannot perform successfully. Indeed recent work has found several such tasks, including “needle in haystack” [, ] as well as extrapolating tasks to longer sequences []. A second, complementary, approach is theoretical studies which chart the computational capabilities of transformers [, ].

In the current work, we focus on a simple task that transformers often struggle with, and analyze it theoretically. Specifically, we consider a very simple “Query Counting” task, defined as follows. The model is presented with a sequence of tokens, and is then asked how many times a given query token appears in the sequence. For example:

> Consider the sequence a a b b a c c d a. How many times does the letter “a” appear in the sequence?

Our interest in this problem is that it seems like a basic module needed for working with long input streams. Indeed, there is a long line of work on sketching and streaming, that studies similar tasks [, ]. Furthermore, this task can be viewed as an extension of the “needle in haystack” problem which has recently garnered much attention, since it was shown that most LLMs struggle with it for long context. In the needle in haystack problem, the goal is to find a specific string in a long text. In our counting problem, the model is tasked with counting how many times a given string has appeared, which is a harder task than finding one appearance.

There is, however, a key difference between the needle in haystack problem and the counting problem we consider above. The needle in haystack problem is clearly solvable by transformers, regardless of the context length. This is because detecting a similar token and extracting it (or nearby tokens) is a simple task for single attention head [e.g., using induction heads ]. On the other hand, for the counting problem, as we argue here, it is much less clear that transformers can solve it for arbitrary context length. Note that when we write that transformers cannot perform a task, we are referring to transformers whose number of parameters is not dependent on context size.

Modern LLMs indeed struggle with counting tasks (e.g., see Sec. 6 and []). Of course if these models can use code, the task becomes easy, but our focus is on understanding the capabilities of the transformer architecture itself. Specifically, we ask whether transformers have an architectural limitation related to the counting task.

We next turn to ask when transformers can count. We show a construction that works as long as the transformer embedding size $d$ is greater than the vocabulary size $m$. Our construction uses one-hot embeddings, or more generally orthogonal embeddings, of the vocabulary, which allow the model to maintain a histogram of counts of tokens previously observed. When $d<m$ this orthogonal construction is no longer possible. A natural approach is to consider the “most orthogonal possible” embedding (a notion formalized in Welch bounds []), and try to use it in a similar scheme. However, we show that this does not allow implementing the histogram solution if the dimension is smaller than the vocabulary.

The above discussion suggests that in the $d<m$ regime,(^1^11Some of our results involve regimes such $d<cm$ for a small constant $c$. We leave $c$ out for simplicity.) the naive histogram approach does not seem to work. Our study of this regime reveals both a positive and a negative result. On the positive side, we show that there does exist a construction that allows counting, which can be done with a single transformer layer. On the negative side, we prove that this construction requires an MLP width that grows with context size, meaning it is not applicable to arbitrarily long contexts. Indeed, when training transformers on these counting tasks, we find that they fail in the $d<m$ regime.

We next study a somewhat more complex counting task, which we refer to as “Most Frequent Element”. Here we present the model with a sequence of tokens, and ask for the count of the most frequent token. This is the same as taking the maximum of the histogram of the counts. Similarly to the Query Counting task, we show that in this case, a solution exists for $d>m$ based on an orthogonal construction. However, for $d<m$ we show, using a communication complexity argument, that no solution exists for a one layer transformer. Thus again we obtain a phase transition for counting at $d=m$.

Taken together, our results reveal an interesting landscape for simple counting task, where the $d=m$ threshold separates between transformers that can count and those that cannot.

Our results highlight the importance of studying basic counting problems, and their dependence on vocabulary size.
They also point to limits of solving seemingly simple problems using transformers, and further emphasize the advantages of using code as a tool to sidestep these issues.

## 2 Related Work

Since the introduction of transformer architectures [] and the success of LLMs, there has been much work on evaluating their performance on various tasks [e.g., see ]. In particular, much recent work has explored performance on long context tasks, the dependence of accuracy on context length [], and the ability of models to extrapolate to lengths beyond those seen at training [].

The fact that models often do not perform well on these evaluations has prompted works that try to pinpoint the inherent limitations of transformer models. One line of work is to use computational complexity lower bound approaches to show that transformers of a certain size and architecture cannot implement a given function of interest. For example, Sanford et al. [] show that certain functions cannot be implemented without the transformer size scaling with input size. A similar limitation is also shown in [] for certain compositions of functions.

A related line of work is to relate transformers to known complexity classes. For example it has been shown that transformers can efficiently approximate Turing machines [], and that transformers with bounded percision can only solve problems in uniform $TC^{0}$ [].
Chain-of-Thought [] has also been analyzed from an expressiveness viewpoint, demonstrating it can substantially improve the expressive power of transformers [].

Our focus here is not on the general capabilities of transformers but rather on a specific, seemingly simple problem, and on the ability of transformers to solve it.

Figure: (a)
Refer to caption: extracted/5907230/figures/histogram_method.png

## 3 Problem Setup

We consider inputs that consist of a sequence of $n$ tokens: $x_{1},\ldots,x_{n}$. The desired output for these is denoted by $y$. We consider the following two counting tasks: Query Count (QC) and Most Frequent Element (MFE) where $y$ is defined as follows:

- •
For the QC task: $y$ is the count of the token $x_{n}$ in the set $x_{1},\ldots,x_{n}$ (i.e., $y\geq 1$ always).
- •
For the MFE task: $y$ is the count of the most frequent token in $x_{1},\ldots,x_{n}$.

We denote the dictionary size by $m$, namely $x_{i}\in\{1,\ldots,m\}$.
Furthermore, we use the following notations for model-related quantities:

- •
$d$: the key-size (i.e. embedding dimension of each head).
- •
$h$: the number of attention heads.
- •
$L$: the number of layers
- •
$p$: the numerical precision of computations: We assume that all arithmetic operations (including the softmax) are performed exactly using registers of $p$ bits.
- •
$D$: the overall embedding dimension, where $D=d\times h$
- •
The embedding of token $i$ is $\boldsymbol{v}_{i}\in\mathbb{R}^{D}$.
- •
The query, key, value matrices for layer $i$, attention head $j$ are denoted by $Q_{i,j},K_{i,j},V_{i,j}$. All are matrices in $\mathbb{R}^{dh\times d}$.
- •
$\boldsymbol{p}_{i}$: the positional embedding for location $i$.
- •
$\boldsymbol{u}_{i,j}$: the output of head $i$ in layer $j$.

Most of our solutions work with an architecture consisting of a single layer and a single head.
When this is the case we omit the indices $i$ and $j$ from our notation.
Also note that if $h=1$ then $D=d$.
In our theoretical results we do not use normalization layers, although we can easily add degenerate normalization that doesn’t alter the input, as was done in []. In our empirical results we use standard transformer layers that do include normalization layers.

## 4 Analyzing Query Count (QC)

In this section we focus on the QC problem, and ask which transformer architectures can implement it successfully.
We first remark on a general limitation of transformers without positional embeddings for the counting problems we consider.(^2^22We note that a similar observation was made in [], we provide it here for completeness.)
We next show in Sec. 4.2 that if $d>2m$ a one-head one-layer transformer can implement QC. We refer to this as the histogram solution. We then show that the histogram solution stops working if $d<m$. The natural question is then whether there are other solutions for the $d<m$ case. We argue that in this case, solutions are likely to require calculating the function $1/x$, and show that this function would require an MLP layer of width $n$. This means we cannot expect the transformer to extrapolate to long context sizes, and therefore a one-layer transformer is unlikely to be able to implement QC, at least using the two most natural solutions.

### 4.1 The Need for Positional Embeddings

Transformers use self attention to average over previous representations. The fact that they average, rather than sum, leads to an interesting limitation on their ability to count. Specifically, it is easy to see that for variable context size, they cannot perform any counting task without the use of positional embedding.
Consider the QC task and an input sequence $S_{1}=x_{1},\ldots,x_{n}$, where the goal is to return the count of $x_{n}$ in the sequence. Now consider the length $2n$ sequence $S_{2}=x_{1},\ldots,x_{n},x_{1},\ldots,x_{n}$. The correct output for this sequence is twice the correct output for $S_{1}$. However, a transformer without positional embeddings that is applied to $S_{1}$ will have exactly the same output as the one for $S_{2}$. This follows because averaging is invariant to input duplication.

The above restriction no longer holds when positional embeddings are used, and it is easy to see that it can be rectified with even a simple positional embedding that just signifies the last position (see our construction in Sec. 4.4).
This implies that if a transformer has access to the legnth of the sequence, it may make it easier to count. Another thing to note is that while the above difficulty arises for counting, it does not arise if we are interested in calculating proportions (e.g., what is the fraction of the items of the sequence that are equal to $x_{n}$).

### 4.2 A “Histogram” Solution for d > 2 ⁢ m 𝑑 2 𝑚 d>2m italic_d > 2 italic_m

We begin by providing a solution for the case where the model dimension is larger than the vocabulary size.

###### Theorem 4.1 .

For the Query Count problem and any context length $n>0$, if $d>2m$, there exists a transformer that solves it, which has one layer, one head, and an MLP with $d$ neurons.

We provide the construction below (see also Fig. 1(a)). We begin by describing it as a two head solution, but a one-head implementation is also simple using a residual connection.
The idea is to construct a histogram of all previous tokens (i.e. the number of times each token appears) and then given the query token $x_{n}$ extract the count for this particular token from the histogram.

First, we assume that the embeddings are orthonormal. Namely:

$$ $\boldsymbol{v}_{i}\cdot\boldsymbol{v}_{j}=\delta_{ij}\quad\quad\forall i,j\in \{1,\ldots,m\}$ (1) $$

where $\boldsymbol{v}_{i}$ is the embedding into $\mathbb{R}^{d}$ of the dictionary item $i$.
This is possible because of the assumption $d>m$. For simplicity, we assume that $\boldsymbol{v}_{i}=\boldsymbol{e}_{i}$ where $\boldsymbol{e}_{i}$ is the standard basis in $\mathbb{R}^{d}$.

Next, we construct an attention head whose output at position $n$ is the histogram of the tokens up to and including this token. Let $Q_{1}=0$ (the zero matrix) and $V_{1}=I_{d}$ ($I_{d}$ is the identity matrix in $\mathbb{R}^{d}$). Then the output of this attention head is

$$ $\boldsymbol{u}_{1}=\sum_{i=1}^{n}\boldsymbol{e}_{x_{i}}=\sum_{j=1}^{m}c_{j} \boldsymbol{e}_{j}$ (2) $$

where $c_{j}$ is the number of occurrences of item $j$ in the context normalized by $n$, the length of the context.
That is $c_{j}=|\{i\in[n]\mid x_{i}=j\}|/n$.
In words, $\boldsymbol{u}_{1}$ is a vector in $\mathbb{R}^{d}$ whose $i^{th}$ entry is the number of times that token $i$ appeared in the input $x_{1},\ldots,x_{n}$.

The second head is set to simply copy the input embedding. This is done by setting $Q_{2}$ and $K_{2}$ such that $K_{2}^{\top}Q_{2}=TI_{d}$ where $T$ is sufficiently large and $V=I_{d}$.
After this we have:

$$ $\boldsymbol{u}_{2}=\boldsymbol{e}_{x_{n}}$ (3) $$

The outputs of the two heads consist of the histogram and a one-hot identifier for the query token. Recall that our desired output is the count $c_{x_{n}}$ of the query token.
We can extract this count using an MLP with $d$ ReLU gates in the hidden layer. Gate $i$ computes ReLU of $n\cdot u_{1}[i]-B\cdot(1-u_{2}[i])$ for some sufficiently large constant $B$. It is easy to see that the output of gate $i$ is $c_{x_{n}}$ if $x_{n}=i$ and 0 0 otherwise.

Remarks:
1) Note that the above solution will only work for input of a fixed length $n$. As noted in 4.1, a transformer without positional embeddings (as the one in our construction) cannot possibly count inputs of variable lengths. Our construction here can be extended to variable lengths by using positional embeddings, together with an MLP for computing $1/x$. We elaborate on such an approach in the next section.

2) We can also implement the above histogram scheme with one head, by taking advantage of a residual connection.
The idea is to use half of the coordinates in the embedding dimension to store the result of the attention module,
and pass the original token to the MLP using the residual connection.

3) In the construction above, we assumed that $\boldsymbol{v}_{i}$ are the standard basis. However, a similar construction is possible when $\boldsymbol{v}_{1},\ldots,\boldsymbol{v}_{m}$ are orthonormal, but not one-hot. In this case, the MLP will have to take the dot product of $n*\boldsymbol{u}_{1}$ and $\boldsymbol{u}_{2}$ to extract the count. This is less natural to do with RELU gates. One could, however, first apply to $\boldsymbol{u}_{1}$ and $\boldsymbol{u}_{2}$ a linear transformation (i.e., a rotation), that changes the basis to the standard basis and then extract the count as before.

###### Theorem 4.1 .

### 4.3 The Histogram Solution Breaks for d < m 𝑑 𝑚 d<m italic_d < italic_m

Our solution in the previous section
uses the fact that if $d>m$ we can embed the dictionary into orthogonal vectors in $\mathbb{R}^{d}$.
When $d<m$ this is not possible.
One may try to extend this solution by embedding the dictionary into a collection of “almost” orthogonal vectors.
However any collection of $m$ $m\geq 2d$ vectors in
$\mathbb{R}^{d}$ contains a pair of vectors whose
inner product is at least $\Omega(1/\sqrt{d})$ in absolute value. This is a result of the Welch bounds [] which provide upper bounds on the minimum dot product between $m$ unit-vectors in $\mathbb{R}^{d}$.
This implies the following
lower bound, which states that counting will fail in this regime.

###### Theorem 4.2 .

Consider the “Histogram” solution for the counting problem presented in Sec. 4.2, and embedding vectors $\boldsymbol{v}_{1},\dots,\boldsymbol{v}_{m}$. For an input $\bar{\mathbf{x}}=(x_{1},\dots,x_{n})$ to the counting problem, denote by $c_{x_{n}}$ the correct solution and by $\texttt{hist}(\bar{\mathbf{x}})$ the output of the “Histogram” solution.(^3^33The Histogram solution in this case is $\boldsymbol{v}_{x_{n}}\cdot\sum_{j}c_{j}\boldsymbol{v}_{j}$, which is the natural generalization to the non-orthogonal case. As noted in Remark 3, it is generally more elaborate to implement because of the dot product, but as Theorem 4.2 shows, this solution has an inherent limitation, irrespective of this implementation difficulty.) If $m\geq 2d$, then for any embedding vectors $\boldsymbol{v}_{i}$’s there are inputs to the counting problem for which: $|\texttt{hist}(\bar{\mathbf{x}})-c_{x_{n}}|\geq 0.25\sqrt{n}$.

###### Proof.

Let $\boldsymbol{v}_{1},\dots,\boldsymbol{v}_{m}\in\mathbb{R}^{d}$ with $m\geq 2d$, and let $A=\max_{i\neq j}|\langle\boldsymbol{v}_{i}\cdot\boldsymbol{v}_{j}\rangle|$.
Assume without loss of generality that $A=\boldsymbol{v}_{1}\cdot\boldsymbol{v}_{2}$.
By the Welch bounds [] for $k=1$ we have that $A\geq\frac{1}{\sqrt{2d-1}}$. Consider the input $x_{1},\dots,x_{n}$ to the counting problem where $x_{1},\dots,x_{n-c}$ are equal to the same token which is different from $x_{n}$ and mapped to the embedding $\boldsymbol{v}_{1}$, and $x_{n-c},\dots,x_{n}$ are all equal to $x_{n}$ which is mapped to embedding $\boldsymbol{v}_{2}$.
Then the output for the histogram solution is:

$$ $\displaystyle|\texttt{hist}(\bar{\mathbf{x}})|=$ $\displaystyle|\langle(n-c)v_{1}+cv_{2},v_{2}\rangle|$ $\displaystyle\geq c+\frac{n-c}{\sqrt{2d-1}}~{}.$ $$

By choosing $c=0.5n$ and $n=d$ we have the desired result.
∎

The theorem implies that even if the dictionary size is only linear in the embedding dimension, any solution will incur an error that grows with the context size. In practice, the dictionary can contain millions of tokens, while the embedding dimension is at most a few thousands, thus this error can be quite large. Note that picking the embedding vectors at random (e.g. i.i.d Gaussians) will result in an even greater error than what is stated in the theorem, since the inner product between each two vectors will be larger (with high probability) than the lower bound we used in the proof.

###### Theorem 4.2 .

###### Proof.

### 4.4 The CountAttend solution for all d 𝑑 d italic_d and m 𝑚 m italic_m

In the previous section we considered a histogram based solution, which required $d>m$. Here we provide an alternative approach to solve the counting problem which works for any $d$ and $m$. However, as we shall see, this solution requires a large MLP, that must scale with the length of the input $n$. As a result, a transformer with a fixed MLP size will not be able to learn such a solution that will work with arbitrary $n$ values.

We first present a high level description of this construction, which uses a 1-layer transformer with a single head.
The idea explicitly uses attention to count (hence the name CountAttend) as follows (see also Fig. 1(b)). Assume that the query token is $x_{n}=4$, so that we are seeking the number of elements in $x_{1},\ldots,x_{n}$ that are equal to $4$, and assume that the number of these elements is $7$. Then the token $x_{n}$ can attend to all other tokens, such that attention weight is high for all $i$ such that $x_{i}=4$ (and same for all these) and near-zero otherwise. Thus, the attention weight will be $\frac{1}{7}$ for all the $x_{i}=4$ tokens, including $x_{n}$. We next need to extract this number and invert it, in order to get the answer $7$.

Extracting the attention weight can be done by using a single coordinate in the positional embedding that is one for position $n$ and zero otherwise. The value aggregation of self-attention can then extract the number $\frac{1}{7}$ in this coordinate. Finally, to invert this number we need an MLP to implement the function $1/x$. If the smallest number we need to invert is $1/n$ then this can be done with an MLP with 4n neurons.

The following proposition, proved in Appendix A, summarizes the properties of this construction.

###### Proposition 4.3 .

For any $d,m,n$ there exists a transformer that solves the corresponding QC problem. The transformer has one layer, one attention head, dimension $d$, and an MLP of size $O(n)$. Furthermore, its matrix $K^{\top}Q$ is diagonal with elements of magnitude $O(\log{n})$.

###### Proposition 4.3 .

#### 4.4.1 Limitations of the CountAttend solution

The advantage of Proposition 4.3 is that it works for any dimension $d$, and does not restrict $d$ to be larger than $m$ as in the histogram solution.
However, the solution in Proposition 4.3 has two major limitations compared to the histogram solution presented in Sec. 4.2. We discuss these below.

First, Proposition 4.3 has an $O(n)$ sized MLP, which is a result of its internal implementation of the function $x\mapsto\frac{1}{x}$ using an MLP, where $x\in\left[\frac{1}{n},1\right]$, and the desired precision is $0.5$ (because we are counting, and can round the result). In the proof of 4.3 we used a naive implementation of this function. It is natural to ask if a smaller implementation exists. The following result shows this is not possible.

###### Lemma 4.4 .

Any $2$-layer MLP with ReLU activations
that approximates the function
$f(x)=1/x$ in the interval
$x\in\left[\frac{1}{n},1\right]$ to within
$L_{\infty}$ error of less than $1/2$
has $\Omega(n)$ neurons.

###### Proof of Lemma 4.4 .

Let $g$ be a piecewise linear approximation of $f(x)$. Then for
for $x=1/k$, $k=1,\ldots,n$, we must have $k-1/2\leq g(x)\leq k+1/2$.

Consider the line $\ell(x_{1},x_{2})$ between
$(1/x_{1},x_{1})$ and
$(1/x_{2},x_{2})$ for some integers $x_{1}$ and $x_{2}$,
$1\leq x_{1},x_{2},\leq n$.
The equation of this line is
$y=(-x_{1}x_{2})x+x_{1}+x_{2}$.
Let $x_{1}=k$ and $x_{2}=k-c$
for some constant $c$ that we determine below.
Then the equation of $\ell(k,k-c)$ is
$y=-k(k-c)x+2k-c$.
Let $\ell^{\prime}(k,k-c)$ be the line $y=-k(k-c)x+2k-c-0.5$
which is parallel and below
$\ell(k,k-c)$.
We claim that the point $A=(1/(k-c/2),k-(c/2)+0.5)$ lies below $\ell^{\prime}(k,k-c)$.
By convexity this implies that $g$ must have a breakpoint between $1/k$
and $1/(k-c)$.

To prove the claim we have to show that

$$ $k-(c/2)+0.5\leq-k(k-c)\frac{1}{k-(c/2)}+2k-c-0.5$ $$

It is easy to check that this holds for $c=3$ and any $k$.
This shows that $g$ must have at least $\Omega(n)$ linear pieces. Note that any $2$-layer MLP with ReLU activations with $\ell$ neurons is a piecewise linear function with at most $2\ell$ pieces. This is because each ReLU neuron is a piecewise linear function with at most $2$ pieces, and the MLP is just the sum of those neurons.
∎

Note that although the lemma focuses on $2$-layer MLPs, it can be readily generalized to deep MLPs, e.g. using the lower bound on the number of linear pieces for deep ReLU networks from Telgarsky []. Although deeper networks can have more linear pieces using fewer neurons than shallow network, the depth would still need to scale with $\log(n)$ which is infeasible in practical implementations.

The second limitation of Proposition 4.3 is that the magnitude of its attention matrices scales logarithmically with the context size $n$. Since the temperature is inside the exponent, it means that the magnitude of the gradient should scale polynomially with the context size. This is possible given high-precision computational resources. However, transformers are trained with limited precision (e.g. $8$- or $16$-bit) which can make the optimization of such large weights infeasible.

Taken together the two observations in this section suggest that the despite the fact that QC has a transformer based implementation with one layer, this representation is too large to be applicable to arbirary $n$, and also potentially suffers from optimization difficulties.

###### Lemma 4.4 .

###### Proof of Lemma 4.4 .

## 5 Analyzing Most Frequent Element

In this section we consider the task of finding the most frequent element (MFE) in a given sequence of tokens. This problem is very much related to the counting problem, as intuitively it requires to count every token separately and compute the token that appears the maximum number of times.
We show that there are strict bounds on the size of the embedding compared to the size of the dictionary,
in order for transformers to be able to perform this task.

### 5.1 MFE Models Must have d ≥ Ω ⁢ ( m ) 𝑑 Ω 𝑚 d\geq\Omega(m) italic_d ≥ roman_Ω ( italic_m )

We begin by showing a lower bound on the required size of a $1$-layer transformer solving the MFE task. The following result
establishes that MFE can be implemented only when $dhp=\Omega(\min\{m,n\})$. This means that at a fixed precision $p$ and if $n>m$, the dimension $d$ must grow linearly in the vocabolary size in order to implement MFE. This is the same trend that we saw for the QC problem.

###### Theorem 5.1 .

Suppose that there is a $1$-layer transformer with $h$ heads, embedding dimension $d$, and $p$ bits of precision, followed by an MLP of arbitrary width and depth, that solves the MFE task for sequences of length $n$. Then, we must have that $dhp\geq\Omega(\min\{m,n\})$, where $m$ is the vocabulary size.

The full proof can be found in Appendix B. The proof uses a communication complexity argument that is inspired by a lower bound of Sanford et al. []. This lower bound implies that transformers which solve the MFE task need to have an embedding size that scales with the size of the dictionary, or have many attention heads. Note that increasing the size of the MLP which follows the attention cannot break the lower bound.

###### Theorem 5.1 .

### 5.2 MFE Can be Implemented when d = O ⁢ ( m ) 𝑑 𝑂 𝑚 d=O(m) italic_d = italic_O ( italic_m )

The previous result showed that MFE cannot be implemented with a one layer transformer when $d$ is smaller than $m$.
Here we show that it is possible to implement MFE when $d=O(m)$. This implies that $d=O(m)$ is tight for the MFE problem.
The result is described below.

###### Theorem 5.2 .

There exists a $1$-layer transformer that solves the MFE task for sequences of size $n$ and dictionary size $m$, where the parameters $d,h,p$ are equal to: $d=m,~{}h=1,~{}p=\log(n)$, and the MLP has $d^{2}$ neurons.

The construction is again based on the histogram approach. Because $d>m$, one can compute the histogram of counts in the last position (as for QC). The only part left to be done
is to extract the maximum from the histogram, which can
be done via a one layer MLP with $m^{2}$ units (each unit performs a maximum between two distinct elements in the histogram).

One limitation of the above result is that it requires an MLP that grows with $m$. This can be avoided if using two layers of attention. A two layer implementation is simple: use the first layer to calculate the “Query Count” for each element, and then use softmax to calculate the maximum over tokens. This construction does not need an MLP at all. Another option is to use a depth $\log(m)$ MLP with $m$ neurons at each layer to calculate the maximum from the histogram [see ], however having depth which relies even logarithmically on the dictionary size is infeasible for practical implementations.

To summarize the above results, we have shown that MFE cannot be implemented by a one layer transformer if $d<m$, and that if $d>m$ MFE can be implemented either with a one layer transformer with wide MLP, or a two layer transformer without an MLP.

###### Theorem 5.2 .

## 6 Experiments

Our analysis considers the dependence between the transformer model size $d$, and its ability to perform counting tasks. Specifically, we show that for vocabulary size $m$ that exceeds $d$, exact counting is likely to become impossible. In this section we perform experiments that support this observation. We begin with results for training a model from scratch and then also consider results with a pretrained LLM (Gemini 1.5).

### 6.1 Training from Scratch

##### Tasks:

We consider the two counting tasks described in the text: Most Frequent Element (MFE) and Query Count (QC). We generate instances of these by sampling sequences of length $n$ uniformly from a set of $m$ tokens. Denote each such sequence by $x_{1},\ldots,x_{n}$. The expected output $y$ for these is as follows:

- •
For the QC task: $y$ is the count of the token $x_{n}$ in the set $x_{1},\ldots,x_{n}$ (i.e., $y\geq 1$ always).
- •
For the MFE task: $y$ is the count of the most frequent token in $x_{1},\ldots,x_{n}$.

During training and evaluation we sample batches from the above distribution. Evaluation used $1600$ examples in all cases.

##### Model:

We train a transformer model with the standard architectural components (self attention, MLP, layer norm, etc.). We use two layers and four heads (theoretically we could have used less, but optimization was faster with this architecture). Training uses Adam for optimization, with batch size $16$ and step size $10^{-4}$. Training is run for $100K$ steps. Positional embeddings were optimized. For predicting the count $y$, we use a linear projection on top of the embedding of the last token in the last layer (i.e., we do this instead of the vocabulary prediction). Training was done via colab and takes about 15 minutes per model with the standard GPU provided therein.

##### Parameter Settings:

We experimented with several values of $d$ (between 8 and 128). For each, we varied $m$ in order to test dependence on vocabulary size (we use $20$ values between $m=5$ and $m=150$). In order to keep the average count at a constant value of $c$, we set $n=cm$. We used $c=10$ in all experiments.

##### Results:

Our focus is on understanding the dependence between $d$ and $m$ and the ability to count. Thus, we report results as follows. For each value of $d$, we find the value of $m$ at which counting begins to fails. Specifically, we consider $m$ at which counting accuracy falls below 80% . We refer to this as $m_{thr}(d)$. Fig. 2(a) shows these for the two counting tasks. It can be seen that in both cases, the threshold indeed increases roughly linearly with $d$, agreeing with our theoretical analysis.

Figure: (a)
Refer to caption: extracted/5907230/figures/mthr.png

### 6.2 Evaluation of a pretrained LLM

Our theoretical results highlight the role of vocabulary size in counting problems. Here we provide an exploration of this role in a trained LLM, Gemini 1.5. We provide the model with the query count task.(^4^44Specifically we use the prompt: “consider the following array [1,1,2,2,3] of length 5. How many times does the word 3 appear in the array? Respond in one number.”.) We then vary $m$, the number of different tokens used in the sequences (e.g., for $m=5$ we’ll use a sequence with just numbers $1,\ldots,5$), while keeping the expected counts of all elements at a constant $c=10$. Namely, for each $m$ we use context length $mc$. As a baseline for these, we also use the same sequence length, but with binary sequences matched to have expected count $c$ for the query token (which we set to “10” throughout). This allows us to estimate the error attributable to just the vocabulary size and not the sequence length and count . Results are shown in Fig. 2(b) and it can be seen that increasing vocabulary size indeed has a negative effect on performance. Furthermore, this effect cannot be explained just by increasing the sequence size, since the binary curve is lower.

## 7 Conclusion

We focus on the basic task of counting using a transformer architecture. When the dimension of the model is sufficiently large, we show that this task can be easily implemented by letting the transformer calculate the histogram of the input sequence. For smaller dimensions, we provide theoretical support suggesting that a one layer transformer cannot implement this function. Our empirical results support this phase transition.

Understanding such limitations of transformers are key to developing new architectures. For example, our results show that in a sense it would be impossible to have transformers count arbitrarily well and for long contexts, without increasing the architecture size considerably. Concretely, this suggests that for counting tasks it may be important to delegate to tools [] such as code execution that do not have the same limitations.

## 8 Limitations

While we provide the first results on upper and lower bounds for counting, these are not yet tight, which would have been the ideal result. Specifically, we show impossibility for $d<m$ for MFE with one layer, but do not show that with more layers (e.g., two), though we conjecture this is true. Proving it would require novel technical tools, as it is not clear that the communication complexity argument is extendible to this case. For QC, we show that the inversion based architecture has inherent limitations for one layer, and here too it would be interesting to prove lower bounds for additional layers. In terms of empirical evaluation, we restricted our training experiments to small architectures, but it would be interesting to explore these tasks for models closer to those used in practice. Additionally, it would be interesting to see how pretrained models perform on these tasks after fine-tuning. Finally, it would be interesting to use a mechanistic interpretability approach to check which computation is actually being implemented in trained transformers (either pretrained on language or from scratch on counting).

### Acknowledgements

This work is partially supported by ISF grant no 1156/23 and the Blavatnik Research Foundation.

## Appendix A Proofs from Sec. 4.4

###### Proof of Proposition 4.3 .

Here we provide additional information about the CountAttend solution. Recall that the idea is for the last token $x_{n}$ to attend to earlier tokens such that tokens identical to $x_{n}$ will
eventually receive a weight close to $1/{c_{x_{n}}}$, the count of $x_{n}$ in the sequence. In what follows, we consider the scale of the logits that will provide this result at sufficient precision.

The attention weight of a unit-norm token embedding $\boldsymbol{v}_{i}$ with itself is $e^{T\boldsymbol{v}_{i}\cdot\boldsymbol{v}_{i}}=e^{T}$, and the attention weight of $v_{i}$ with $v_{j}$ is $e^{T\boldsymbol{v}_{i}\cdot\boldsymbol{v}_{j}}\leq e^{TJ}$ where $J$ is an upper bound on the dot product between any two vectors in $\mathbb{R}^{d}$ among a set of $m$ vectors (e.g., as obtained from analysis of random vectors as in Sec. C).

Now consider the sum of the attention weights (i.e. the denominator of the softmax in the attention module).
Let $c_{x_{n}}$ denote the number of occurrences of $x_{n}$ in the context, and
let $c^{\prime}$ denote the number of tokens
$x_{i}$ such that $x_{i}\not=x_{n}$.
We get that the sum of the attention weights is
$c_{x_{n}}e^{T}$ plus a quantity bounded by
$c^{\prime}e^{TJ}.$
If we divide this by $e^{T}$ then we get
that the normalization factor equals to $c_{x_{n}}$
plus “noise” bounded by
$c^{\prime}e^{T(J-1)}.$
From this we can recover $n_{0}$ if
$c^{\prime}e^{T(J-1)}<\frac{1}{2}.$
We clearly satisfy this inequality if

$$ $T\geq\frac{\log(2n)}{1-J}.$ $$

Substituting the bound we have for $J$ for a random embedding (see Sec. C) we get that we need $T$ such that:

$$ $T=\Omega\left(\frac{\log(2n)}{1-\sqrt{\frac{\log m}{d}}}\right).$ $$

Using the above, we obtain that the output of the attention is $1/c_{x_{n}}$ to within $0.5$ accuracy in the inverse.
To get $c_{x_{n}}$ we need an MLP that inverts $1/x$.
This can be done as follows.

It is well known that we can
implement a “delta function” using four ReLU neurons.
For example we can approximate a delta function of height $h$ between $a$ and $b$, by
$\frac{h}{\epsilon}(\max(0,x-a)-\max(0,x-(a+\epsilon)-\max(0,x-b)+\max(0,x-(b+
\epsilon)))$
for some sufficiently small $\epsilon$.
We use $4$
ReLU neurons to implement a
“delta function”
between $1/(k-1/2)$ and
$1/(k+1/2)$ of height $k$ for each $k=1,\ldots,n$.
∎

###### Proof of Proposition 4.3 .

## Appendix B Proof from Sec. 5

###### Proof of Thm. 5.1 .

Our proof relies on the following set disjointness lower bound [].
(It is similar to a lower bound argument in Sanford et al. [], but simpler since we assume that all arithmetic in the transformer is performed exactly by registers of $p$ bits.)
Alice and Bob are given inputs $a,b\in\{0,1\}^{n}$, respectively. Their goal is to compute $\max_{i}a_{i}b_{i}$ by sending single bit messages to each other in a sequence of communication rounds.
The lower bound says that any deterministic protocol for computing $\max_{i}a_{i}b_{i}$ must have at least $n$ rounds of communication.

We construct a reduction from the set disjointness problem to the MFE task. We assume for ease of notation that the length of the context is $2n$, and also assume that $m>3n$. If $m<3n$ then we set the context size to be $n^{\prime}=m/6$ and continue the proof as is with $n^{\prime}$ instead of $n$. Note that since the lower bound is given by $\Omega(\min\{m,n\})$, using $n^{\prime}$ instead of $n$ will provide a lower bound that depends on $m$. In fact, $\min\{m,n\}$ can be viewed as the “effective ” dictionary size, which is the maximal number of different tokens that a transformer sees given an input sequence of length $n$.

Assume that Alice and Bob received inputs $a,b\in\{0,1\}^{n}$. Suppose we have the following distinct tokens in our dictionary (which is possible by our assumption on $m$): $s_{1},\dots,s_{n},y_{1},\dots,y_{n},z_{1},\dots,z_{n}$.
We consider the following input context $x_{1},\dots,x_{2n}$ to the transformer. For
$j\in\{1,\dots,n\}$, if $a_{j}=1$ we set $x_{j}=s_{j}$, and
otherwise we set $x_{j}=y_{j}$.
Similarly,
for every bit in $b$ in place $j\in\{1,\dots,n\}$, if $b_{j}=1$ we set $x_{n+j}=s_{j}$, otherwise we set $x_{n+j}=z_{j}$. We also assume there is some query token $x_{0}$ known to both Alice and Bob and different from the rest of the tokens. This is the last token and we assume that the desired output token corresponds to it.

Note that if the most frequent element in the context $x_{1},\dots,x_{2n}$ appears twice, then this token must be $s_{\ell}$ for some $\ell\in\{1,\dots,n\}$ which means that
$a_{\ell}b_{\ell}=1$. Otherwise if the most frequent element appears only once and then $\max_{i}a_{i}b_{i}=0$.

Suppose there exists a $1$-layer transformer with $h$ heads followed by an MLP of arbitrary size that solves the MFE task for all inputs $x_{1},\dots,x_{2n}$. Assume the embedding dimension of each token is $d$, namely $s_{i},y_{i},z_{i}\in\mathbb{R}^{d}$ for every $i\in\{1,\dots,d\}$. Also, denote the weights of the heads by $Q_{j},K_{j},V_{j}$ for each $j\in[h]$, and assume w.l.o.g. that they are of full rank (i.e. rank $d$), otherwise our lower bound would include the rank of these matrices instead of the embedding dimension (which can only strengthen the lower bound). We design a communication protocol (following the construction in []) for Alice and Bob to solve the set disjointness problem:

- 1.
Given input sequences $a,b\in\{0,1\}^{n}$ to Alice and Bob respectively, they calculate $x_{1},\dots,x_{n}$ and $x_{n+1},\dots,x_{2n}$, respectively.
- 2.
Alice computes the
$p$ bit
representation of
$s_{j,a}=\sum_{i=1}^{n}\exp(x_{i}^{\top}K_{j}^{\top}Q_{j}x_{0})~{},$
for each head $j$ and transmits them to Bob.
The number of transmitted bits is $O(ph)$.
- 3.
Bob finishes the computation of the softmax normalization term for each head $j\in[h]$ and sends it to Alice, namely he computes:
$s_{j}=s_{j,a}+\sum_{i=n+1}^{2n}\exp(x_{i}^{\top}K_{j}^{\top}Q_{j}x_{0})~{}.$
The number of transmitted bits is again $O(ph)$.
- 4.
For each head $j\in[h]$ Alice computes the first part of the attention matrix which depends on her input tokens and transmits it to Bob. Namely, she computes:
$t_{j,a}=\frac{\sum_{i=1}^{n}\exp(x_{i}^{\top}K_{j}^{\top}Q_{j}x_{0})V_{j}x_{i}
}{s_{j}}~{}.$
The number of transmitted bits is $O(dph)$, since $x_{i}\in\mathbb{R}^{d}$, and the assumption that $V_{j}$ is full rank.
- 5.
Bob can now finish the computation of the attention layer. Namely, he computes:
$t_{j}=t_{j,a}+\frac{\sum_{i=n+1}^{2n}\exp(x_{i}^{\top}K_{j}^{\top}Q_{j}x_{0})V
_{j}x_{i}}{s_{j}}~{}.$
Finally, Bob passes the concatenation of the vectors $t_{j}$ for $j=1,\dots,h$
through the MLP. This step does not require any additional communication rounds.

By the equivalence between the set disjointness and the most frequent element problem that was described before, Bob returns $1$ iff the inputs $\max_{i}a_{i}b_{i}=1$, and 0 0 otherwise. The total number of bits transmitted in this protocol is $O(dph)$, hence by the lower bound on the communication complexity of set disjointness we must have that $dph\geq\Omega(n)$.
∎

###### Proof of Thm. 5.2 .

For the case of $d=m$, we consider the embedding vectors to be equal to $e_{i}$ for each token $i\in[m]$, namely the standard unit vectors. We use a single attention head with query matrix $Q=0$, and value matrix $V=I$. In this case, it is easy to see that for any input sequence $x_{1},\dots,x_{n}$, the output of the attention layer is a vector $\boldsymbol{v}\in\mathbb{R}^{d}$ which is the histogram over the different tokens. Namely, if token which is mapped to $e_{i}$ appeared $c_{i}$ times, then $(v)_{i}=c_{i}$.

To find the most frequent token, we only need an MLP that outputs the maximum over a vector of numbers. To do so we can use the construction from []. Namely, a one hidden layer MLP with width $O(d^{2})$ (Thm 3.3 therein).

∎

###### Proof of Thm. 5.1 .

###### Proof of Thm. 5.2 .

## Appendix C Inner Product of Random Vectors

Let $\boldsymbol{v}_{1},\ldots,\boldsymbol{v}_{m}$ be random unit vectors in $\mathbb{R}^{d}$
where each coordinate in any vector is $\pm 1/\sqrt{d}$ with probability $1/2$, independently.
Hoeffding’s inequality (C.1) implies that with probability
$1-1/poly(m)$,
$|\boldsymbol{v}_{i}\cdot\boldsymbol{v}_{j}|=O(\sqrt{\frac{\log m}{d}})$ for all pairs $i,j\in[m]$, $i\not=j$.

Indeed,
$\boldsymbol{v}_{i}\cdot\boldsymbol{v}_{j}$
is a sum of $d$ random variables of values $\pm 1/\sqrt{d}$ so by (4) we have

$$ $Pr\left(\boldsymbol{v}_{i}\cdot\boldsymbol{v}_{j}\geq t\right)\leq 2e^{-\frac{ dt^{2}}{2}}$ $$

and therefore for $t=O(\sqrt{\frac{\log m}{d}})$ we get that the dot product is large than $t$ with polynomialy small probability for all pairs $\boldsymbol{v}_{i},\boldsymbol{v}_{j}$.

###### Proposition C.1 .

Hoeffding’s inequality states that if $X_{1},\ldots,X_{n}$ are independent random variables such that
$a_{i}\leq X_{i}\leq b_{i}$ then

$$ $Pr\left(|S_{n}-E[S_{n}]|\geq t\right)\leq 2e^{-\frac{2t^{2}}{\sum_{i=1}^{n}(b_ {i}-a_{i})^{2}}}$ (4) $$

where $S_{n}=X_{1}+\ldots+X_{n}$.

Figure: Figure 3: Evaluation of Gemini 1.5 on the MFC task. See Sec. D.
Refer to caption: extracted/5907230/figures/mfc_morning.png

###### Proposition C.1 .

## Appendix D Empirical Results on the MFC Task

In Sec. LABEL:sec:gemini we provided results for running Gemini 1.5 on the QC task. Here we provide results for the MFC task. Experimental setting is similar to QC, with the samples drawn in the same way. The prompt has the following format: “You are given an array of length 5. Find the count of the most frequent element in the array. The array is [1,1,2,2,3]. How many times does the most frequent word in the array appear? Respond in just one number. No additional text.”. Another difference from QC was that the sequence length was set such that the expected MFC would be 40, for the given vocabulary size $m$. This was done via simulation (since no closed form expression is available for expected MFC). For the binary baseline in this case we cannot choose an MFC task, because in the binary case the MFC will be larger than our desired expected count (i.e., 40). Instead, we take the binary baseline to be the minimum frequency count, with expected minimum equal to $40$. Results are provided in Fig. 3, and show that results deteriorate with vocabulary size (also with respect to the binary baseline) again suggesting that vocabulary size affects the complexity of this task, as our theoretical results suggest.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. In your examples, you are counting the number of appearance of certain letters instead of tokens. Therefore, tokenization plays an important part in this task. However, tokenization is not discussed in the paper.
2. The proposed architecture only considers one single layer with one head without normalization layers. This is not the standard Transformer architecture with MLP and skip connections.
3. Lines 181-183 argue that replication of the input results in the same output. Is there any theoretical proof? Is this true in practice?
4. Please elaborate on why the assumption in Eq(1) holds in reality.
5. In the entire paper, including the experiment section, the training of the model and the training data are not discussed. Do pre-training, training data, and fine-tuning affect the ability of Transformers? What if a dataset is constructed with (sequence, count) pairs?
soundness: 2
strengths: This paper studies an important question on the counting ability of Transformers. An interesting construction is proposed to address the query counting problem using Transformer architecture.
summary: This paper studies an important question on the counting ability of Transformers. An interesting construction is proposed to address the query counting problem using Transformer architecture.
weaknesses: Although some theoretical discussion are provided for the proposed construction, the construction itself is only a toy model and may be too simple to reflect the ability of realistic Transformers. Also, the fact that this particular construction cannot achieve certain tasks does not indicate that there does not exist a construction that can. Plus, there are too many loose ends in the proofs (see Questions below).

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. If LayerNorm is used in the architecture, could this generate more parameter efficient architecture?

2. Is there a way to investigate what solution Transformers really converge to in training other than probing? For example, will the two construction differs meaningfully on some out of distribution test?
soundness: 4
strengths: 1. The presented theory is very clear. The authors explain their theoretical contribution with very intuitive argument.

2. The theoretical argument that vocabulary size and context length jointly blocks the learning of counting is well supported by empirical experiments.
summary: This paper presents a theoretical study on how a Transformer can learn the solution of counting. They consider two typical solutions. The first one is a histogram solution that keeps a histogram of the number of different types of tokens, and this requires a dimension linearly scaling with the vocabulary size. The second solution is by first calculating the inverse of the number of times the query tokens appears and then use the MLP to calculate the inverse function. Intriguingly, they show if the feedforward layer is of depth 1, the required width to represent the inverse function scales linearly with the context length.They empirically verified their theory on pretrained LLMs.
weaknesses: 1. The work is mostly constructive so it remains unclear whether Transformers will converge to either of the solution. A mechanistic investigation as mentioned in the conclusion will be a great supplement for the paper.

2. The width bottleneck in the second construction seems to hold only for 1-layer MLP.

3. Technically, the argument that position encoding is necessary only holds for encoder-based model or causal model with 1-layer, a point that should be made clear in the paper.

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. How can the conclusions of this paper generalize to real-world tasks, such as math, code, and so on?
2. Do the limitations observed with small dimensions still apply to multi-layer transformers?
soundness: 2
strengths: 1. This work focuses on the counting task for the language models. The authors provide both theoretical and empirical results to demonstrate the limitations of LLMs when the dimension is small.
2. This paper is well-presented and easy for the readers to follow.
summary: This work investigates the ability of LLMs in counting tasks. This paper focuses on two basic tasks, QC and MFE, and demonstrates that the transformers can solve these tasks if the models' size is large enough, but face limitations when the dimension is small. Empirical and theoretical results highlight the importance of understanding these limitations to improve transformer architectures.
weaknesses: 1. **Lack of Generality:** While the paper focuses on the counting task, its impact on real-world applications is unclear. The conclusions are specific to counting and may not generalize well to broader contexts.
2. The study primarily analyzes one-layer transformers, leaving the capabilities of multi-layer transformers unexplored. Further **theoretical** investigation is needed to understand how additional layers might influence performance on counting tasks.

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: To clarify the unique contributions of your work compared to "Overcoming a Theoretical Limitation of Self-Attention," it would be beneficial to highlight additional perspectives or distinct aspects that differentiate your approach.  As they propose nearly identical Transformer design in counting but with more in-depth experimental analysis.
soundness: 3
strengths: This paper is well-written, presenting a complex topic in a clear and structured manner. the use of well-defined examples, such as the "Query Counting" task and the "Most Frequent Element" problem, aids in illustrating the limitations and possibilities of Transformer architectures in a concrete way.
summary: Paper investigates the ability of Transformer models to perform simple counting tasks, specifically focusing on counting the occurrences of tokens in a sequence. It explores whether Transformers can count effectively and, if so, under what conditions.
weaknesses: Overall, many of the assumptions in this paper are overly simplistic and do not align well with real-world Transformer design or training outcomes. Based on my experience and prior research on Transformer expressiveness, the findings presented here—particularly those concerning cases where d>m have been hinted at in earlier works and are not surprising, given the simplified single-layer Transformer model used in this study. However, real-world Transformer training is far more complex and has been shown to perform poorly on counting tasks, regardless of dimensionality.
Moreover, the experimental design in this paper is weak, with insufficient results to substantiate the authors' claims. The lack of detailed experimental methodology and the vagueness of the reported results undermine the argument. While I agree with the proposed "ideal case" of a histogram approach to counting, this notion is too trivial to warrant significant mention. Actual Transformer training involves many additional considerations, and practical optimization rarely converges to such an ideal case, as evidenced by many experiments I have done and those of other researchers.
I suggest the authors conduct a deeper exploration into the challenges of counting with Transformers. As it stands, the current version provides limited insights into the expressiveness of Transformers, on top of the existing papers.

Detailed problem: 
1. Counting Problem in Transformers: Previous research has explored the counting problem in Transformers, but this work lacks a comprehensive review of such studies. A notable contribution to this area is the paper "Language Models Need Inductive Biases to Count Inductively", beyond transformer they also study linear attention such as Mamba and RWKV. Additionally, counting on out-of-distribution (OOD) data often requires incremental dictionary size (training to count to n but testing on count to n+k). Several studies, therefore, focus on the "parity" problem in Transformers, where models count the even or odd occurrences of a token, representing an early form of counting research. For instance, "Overcoming a Theoretical Limitation of Self-Attention" proposes a similar approach to counting as your paper, which is not acknowledged here. A recent paper, "Counting Ability of Large Language Models and Impact of Tokenization," provides an extensive overview of counting in Transformers, which could guide a more thorough literature review on counting with Transformer. There are lots of Parity-related theoretical work and experimental work with Transformers, which study counting, they all need to be discussed. 

2. Expressiveness Limitations in Transformers: Prior theoretical and empirical findings suggest that Transformers lack the expressive capacity to act as counter machines. For example, DeepMind's work "Neural Networks and the Chomsky Hierarchy" shows that Transformers struggle even with parity tasks, which are simpler than general counting. Moreover, "Language Models Need Inductive Biases to Count Inductively" and "Overcoming a Theoretical Limitation of Self-Attention" argue that complete Transformer architectures cannot perform counting tasks without specific inductive biases. Common positional encodings (like sinusoidal or absolute encoding) fail in this regard. My own experiments corroborate this—Transformers do not converge to a counting model without cetain biases. This limitation is examined thoroughly in "Overcoming a Theoretical Limitation of Self-Attention," which offers a similar solution to the one proposed in your paper.

3. Limited Theoretical Scope: The theoretical analysis in this paper overlooks several critical components in modern Transformer design that impact expressiveness, such as layer normalization, positional encoding, floating-point precision, residual connections, and advanced activation functions. Many recent studies cover these in greater depth. By simplifying the Transformer to a single attention module, this paper overlooks architectural nuances, which limits the validity of its conclusions. I recommend engaging with recent theoretical work on Transformer expressiveness in light of these elements (matter of fact, many research show that some of the modules here greatly change the theoretical limits of Transformer). The simplified version of your work is similar to MLP, where each node simply connect to every other previous node, and therefore results are trivial. 

4. Precision Considerations: Precision is not adequately addressed here. While the paper mentions the need for 
d>m it omits the scenario where infinite precision enables Turing completeness with finite d. This has been shown in works like "Attention is Turing Complete" and recent chain-of-thought (CoT) research by Tengyu Ma and others. With higher precision, more information can be compressed within floating-point numbers, contrary to this paper's assumption of clear-cut feature vectors.

5.  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. Both "Language Models Need Inductive Biases to Count Inductively" and "Counting Ability of Large Language Models and Impact of Tokenization" discuss this in detail. This is critical for two reasons: (a) natural language counting often requires sequential comprehension, necessitating a step-by-step approach; and (b) in comparison, recurrent models like RNNs handle counting tasks more naturally than Transformers (almost easily achieve 100\% OOD accuracy according to above papers), highlighting the role of recurrence in inductive counting. Recurrence signals a model’s inductive capacity for counting, suggesting that neural networks generally adopt an inductive counting approach without explicit biases.

6.Prior research shows that positional embedding is crucial for OOD counting. This paper assumes that any positional encoding will equip the model with positional information for counting. However, specific positional encoding designs can distort theoretical assumptions. For example, transformations of values like 1 and 0.5 through positional encoding (either additive in Cosine or geometrical shifting as in ROPE) can cause information loss and lead to unexpected behaviors. These aspects, covered in previous studies, are not discussed here.

6. Experiment lacking specifications. Your experiment does not say how in-distribution training and out of distribution testing is done. Within DIstritbuion, the counting behaves different than OOD, as shown by above mentioned papers. 


7. The experimental section lacks clarity regarding in-distribution (ID) training and OOD testing. Counting behaves differently in ID versus OOD contexts, as shown in the referenced papers. Providing explicit details on the data split methodology would strengthen the validity of the results.

8. LLM-Based Counting: Counting in LLMs presents unique challenges. Not only does it involve inductive methods, but factors like tokenization and chain-of-thought (CoT) prompting also affect the theoretical counting capacity of the Transformer architecture. Studies by Tengyu Ma and others in "Counting Ability of Large Language Models and Impact of Tokenization" delve into these factors. The absence of these considerations in your experiments reduces the credibility of the findings. CoT prompting can unexpectedly enhance the expressiveness of Transformers by mimicking recurrent behaviors similar to RNNs. Although the paper discusses LLMs, it does not mention use of CoT and its effect, which is worth addressing given its relevance to theoretical expressiveness.

9 Gap Between Theory and Practice: The theoretical proofs claim the existence of a Transformer capable of counting, albeit simplified. However, existence does not guarantee practical trainability. As previously mentioned, Transformers rarely converge to this type of behavior in practice without specific biases. Floating-point representations of features can also complicate interpretation, indicating a substantial gap between theoretical assertions and empirical outcomes. More experimental evidence is needed to substantiate the theoretical claims made in this paper.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. Do the constructions require infinite precision or only log precision? 
2. For QC, which are the mechanisms of the transformer to perform the task in practice?
3. The lower bounds in Lemma 4.4 and Theorem 5.1 are proved with constant-depth models. 
Do the lower bounds hold for deep models (up to poly(n) depth)? If not, do the lower bounds hold for constant-depth models with CoT?
soundness: 3
strengths: 1. The paper analyzes the capability of transformers by considering simple yet representative counting problems, which is an interesting perspective.
2. In counting tasks, the paper highlights the relations between the embedding dimension, the model size, the dictionary size, and the sequence length theoretically. These results present an architectural limitation of the transformer and provide insights on how to overcome the issues.
summary: This paper studies transformers' counting capability, focusing on two counting tasks, i.e., QC and MFE. 
For analysis, the paper considers transformers' expressive power. The paper shows two possible
solutions for QC, i.e., Histogram and CountAttend, and highlights the relations between the embedding dimension, the model size, the dictionary size, and the sequence length in the two solutions, respectively. For MFE, the paper further 
proves bounds on the size of the embedding compared to the size of the dictionary. Experiments evaluate the dependence between the transformer model size and its ability to perform counting tasks both models trained from scratch and a pretrain LLM.
weaknesses: 1. The theoretical analysis is restricted to the expressive power of the transformer. It is unclear whether the proposed solutions for the counting tasks are learnable.
2. While the theoretical analysis shows two possible solutions for QC, the experiments do not demonstrate the mechanisms of the transformer to perform the task.

Meta-review: As the title suggests, this paper investigates the counting capabilities of Transformer models, focusing on simple counting tasks such as Query Counting (QC) and Most Frequent Element (MFE). The authors claim that the ability of transformers to perform these tasks is significantly influenced by the vocabulary size and the embedding dimension. They provide theoretical results showing that the embedding dimension must scale with the vocabulary size for effective counting. Empirical results support these theoretical findings, demonstrating a phase transition in performance as predicted.

Strengths:
- **Novel Theoretical Insights**: The paper provides new theoretical insights into the limitations of Transformers for counting tasks, particularly highlighting the impact of vocabulary size.
- **Empirical Validation**: The theoretical claims are supported by some empirical results, which show a clear phase transition in performance.
- **Clarity and Presentation**: The paper is well-written and presents complex theoretical concepts in a clear and structured manner.

Weaknesses:
- **Simplistic Assumptions**: The theoretical analysis is based on overly simplistic assumptions (eg one-layer transformer) that may not hold in real-world scenarios.
- **Limited Experimental Scope**: The experiments are not comprehensive enough to fully support the theoretical claims. Specifically, the lack of OOD evaluation and detailed experimental methodology weakens the empirical evidence.
- **Lack of Mechanistic Insights**: The paper does not explore the mechanisms by which Transformers perform such counting tasks, which could provide deeper insights into the model's capabilities and limitations, and leaves many questions open. 
- **Concerns about Validity of Experimental Setting**: 

Most important reason to reject: overly simplistic model with strong assumptions and limited generality, and therefore likely to have modest impact in the community. 

Overall, this seems to be a borderline paper, with interesting somewhat novel results and insights that I believe might not raise to the level of acceptance to ICLR.

Decision: Reject