Counting in small transformers: The delicate interplay between attention and feed-forward layers

## Abstract

Abstract How do different architectural design choices influence the space of solutions that a transformer can implement and learn? How do different components interact with each other to shape the model’s hypothesis space? We investigate these questions by characterizing the solutions simple transformer blocks can implement when challenged to solve the histogram task – counting the occurrences of each item in an input sequence from a fixed vocabulary. Despite its apparent simplicity, this task exhibits a rich phenomenology: our analysis reveals a strong inter-dependence between the model’s predictive performance and the vocabulary and embedding sizes, the token-mixing mechanism and the capacity of the feed-forward block. In this work, we characterize two different counting strategies that small transformers can implement theoretically: relation-based and inventory-based counting, the latter being less efficient in computation and memory. The emergence of either strategy is heavily influenced by subtle synergies among hyperparameters and components, and depends on seemingly minor architectural tweaks like the inclusion of softmax in the attention mechanism. By introspecting models trained on the histogram task, we verify the formation of both mechanisms in practice. Our findings highlight that even in simple settings, slight variations in model design can cause significant changes to the solutions a transformer learns.

## 1 Introduction

Transformers are the key neural network behind many recent deep learning advances, most notably large language models (LLMs). Their success is partly due to their versatility in processing diverse data types, including text, images, and video, represented as sequences of tokens (Liu et al.,, ; Girdhar et al.,, ; Brown et al.,, ). While scale has been a key factor in unleashing the potential of these models, it is remarkable that their architecture still largely follows the same simple template of the original transformer model proposed by Vaswani et al., (). At its core, a single transformer block primarily alternates two basic components: the token-mixing attention mechanism and a standard fully connected multi-layer perceptron. At a high level, the attention mechanism mixes the tokens, while the multi-layer perceptron applies a nonlinear feature transformation identically to each token. Despite the widespread use of transformers, there is no clear consensus on the distinct roles of their components, how they interact, or if they can be substituted with alternative modules (Tolstikhin et al.,, ; Bozic et al.,, ; Gu and Dao,, ). In particular, the specific contribution of each architectural element to the model’s hypothesis space –the range of algorithms it can learn and implement in practice– remains opaque (Weiss et al.,, ; Delétang et al.,, ; Abbe et al.,, ; Ouellette et al.,, ).

In this work, we investigate this question from a mechanistic interpretability perspective (Cammarata et al.,, ; Olah et al.,, ; Elhage et al.,, ; Michaud et al.,, ; Ouellette et al.,, ) by considering the histogram task as a prototypical problem (Weiss et al.,, ). This task consists of predicting the number of appearances of each token in the input sequences processed by the model – counting. It encompasses two distinct fundamental algorithmic operations: comparison and aggregation. Despite its apparent simplicity, this task exhibits a rich phenomenology, allowing us to study the relative role of different architectural components and their impact on the final solutions implemented by the model in a controlled setting. To this end, we focus on models following the architectural template of primitive transformer blocks, i.e. alternating a token-mixing attention mechanism and a multi-layer perceptron.

In our analysis, we provide explicit constructions (parameter configurations) for a range of such architectures reaching perfect accuracy in a model-dependent hyperparameter regime. In a subsequent step, we compare these algorithms with the performance and mechanistic behavior of models trained from data. Our findings reveal that this class of models is capable of implementing
strikingly different solutions for the histogram task, with a strong dependence on the scale of the model’s hyperparameters and the type of token-mixing mechanism utilized. Our main contributions are:

- •
We identify two main algorithmic strategies that can be used to solve the histogram task perfectly: relation- and inventory-based counting. Relation-based counting uses local pair-wise comparisons between tokens in a given sequence to obtain the number of occurrences conditioned on a given position.
Inventory-based counting relies on the knowledge of the complete alphabet and counts the occurrences of all possible tokens to then extract the correct count for a given position.
- •
We show that the emergence of either mechanism during learning depends on the specifics of the architecture and the inductive bias it possesses in relation to the task. Relation-based counting is memory and compute-efficient as it can leverage an attention-like dot-product mixing mechanism for comparison operations. Inventory-based counting, instead, can be implemented based on an input-independent token-mixing mechanism. This weak inductive bias can be compensated via a feed-forward module with a large enough hidden layer that can memorize a lookup table to implement a comparison operation (inventory): the model-task misalignment can be closed at the cost of increased memory and compute requirements.
- •
When the embedding dimension is comparatively smaller than the size of the alphabet, we show that non-orthogonal embeddings can still result in some models attaining perfect accuracy.
Due to the discrete nature of the counting task, near-orthogonal embeddings may not have a detrimental effect on prediction performance. Additionally, major gains are possible for the softmax operator and dot-product attention which together can remove noise stemming from linear dependence in a semantic, token-dependent manner.
In this context, we also identify a curious regime where very small embedding dimensions, independent of the alphabet size, are in theory possible, but are never learned.

Section 2 provides the necessary background and notation. In Section 3 we describe our experimental setup, followed by our theoretical and experimental results(^1^11All results and code to reproduce them is available at [https://github.com/SPOC-group/counting-attention](https://github.com/SPOC-group/counting-attention)) in Section 4. Section 5 discusses the related literature. Section 6 presents the limitations, conclusion and open questions of this work.

## 2 Background and Notation

##### Architecture.

As inputs, we consider sequences of tokens $\mathbf{x}=(x_{1},x_{2},\cdots,x_{L})\in\mathcal{T}^{L}$. Each token stems from the set $\mathcal{T}=\{1,\cdots,T\}$ of size $T$. The corresponding sequence of outputs $\mathbf{y}=(y_{1},\cdots,y_{L})$ has the same length as the input sequence, where each output token belongs to the output alphabet $\mathcal{C}$ of size $C$, i.e. $y_{\ell}\in\{1,...,C\}$, with $C\leq{L}$.
In this work, we analyze several 1-layer model architectures where a token-mixing mechanism is followed by a per-token feature transformation. This setup includes the case of a single transformer block where the dot-product attention mechanism is followed by a token-wise feed-forward network.
Formally, we consider a model $F:\mathcal{T}^{L}\to\mathcal{C}^{L}$ defined for the positions $\ell=1,\cdots,L$ as

$$ $\displaystyle F(\bar{\mathbf{x}})_{\ell}=\operatorname*{arg\,max}_{c\in\{1, \cdots,C\}}f(\bar{x}^{\prime}_{\ell})_{c}\,;\,\,\bar{x}^{\prime}_{\ell}=\bar{x }_{\ell}+[\mathbf{A}(\bar{\mathbf{x}})\bar{\mathbf{x}}]_{\ell}$ (1) $$

with the token mixing matrix $\mathbf{A}:\mathbb{R}^{L\times d}\to\mathbb{R}^{L\times L}$ and the token-wise feature transformation $f:\mathbb{R}^{d}\to\mathbb{R}^{C}$. The embedding $\bar{\mathbf{x}}\in\mathbb{R}^{L\times{d}}$, where $\bar{x}_{\ell}$ denotes its $\ell$-th row, is obtained by passing the input sequence $\mathbf{x}$ into a standard embedding layer (learnable lookup-table) of dimension $d$.
We refer to the embedding associated with token $t\in\mathcal{T}$ as $e_{t}\in\mathbb{R}^{d}$ or $e_{x_{\ell}}\in\mathbb{R}^{d}$ for the embedding of the token $x_{\ell}$ at position $\ell$. We do not include positional embeddings due to the inherent permutation equivariance of the histogram task. We refer to the vector $\bar{x}_{\ell}^{\prime}$, for each position $\ell=1,\cdots,L$, as the mixed token.
Note that we assume that all operations in the network are executed with infinite precision. We comment when this becomes problematic.

##### Token Mixing.

We consider two types of mixing mechanisms $\mathbf{A}$ with different activation functions. We refer to the case where the function $\mathbf{A}$ is constant in $\bar{\mathbf{x}}$ as linear mixing (lin), e.g.

$$ $\displaystyle\mathbf{A}_{\mathrm{lin}}(\bar{\mathbf{x}})$ $\displaystyle=A\,,\,\,\,\,\,\,\,\,\,\mathbf{A}_{\mathrm{lin+sftm}}(\bar{ \mathbf{x}})=\text{softmax}(A)\,,$ (2) $$

where $A\in\mathbb{R}^{L\times L}$ is a learnable matrix and the softmax operator is applied row-wise. The number of learnable parameters is therefore $L^{2}$.
As an alternative mixing structure, which we refer to as dot-product mixing (dot), we consider the popular attention mechanism which constructs the matrix $\mathbf{A}$ to be explicitly dependent on the inputs, i.e.

$$ $\displaystyle\mathbf{A}_{\mathrm{dot}}(\bar{\mathbf{x}})=\frac{1}{\sqrt{d}} \bar{\mathbf{x}}W_{Q}W_{K}^{T}\bar{\mathbf{x}}^{T}\,,\hskip 15.00002pt\mathbf{ A}_{\mathrm{dot+sftm}}(\bar{\mathbf{x}})=\text{softmax}\left(\frac{1}{\sqrt{d} }\bar{\mathbf{x}}W_{Q}W_{K}^{T}\bar{\mathbf{x}}^{T}\right)\,,$ (3) $$

where $W_{Q}$ and $W_{K}$ are learnable $d\times{d}$ matrices, and the softmax function is applied row-wise. Note that, without loss of generality, we assume the value matrix to be the identity. The number of parameters for dot-product mixing is $2d^{2}$. In line with previous work (Weiss et al.,, ), for architectures employing the dot-product mixing, we also analyze models utilizing the so-called beginning-of-sequence (BOS) token. This special token, indicated with the symbol $\$$, is appended to the original input $\mathbf{x}$ resulting in a new sequence $\tilde{\mathbf{x}}=(\$,x_{1},x_{2},\cdots,x_{L})$ of length $L+1$. We will refer to the architecture that includes the BOS token as bos.

##### Feature Transformation.

The feature transformation is a single hidden layer perceptron with ReLU activations. The hidden layer is of dimension $p$. The function $f$ is applied identically to every mixed token $\bar{x}^{\prime}_{\ell}$ for $\ell=1,\cdots,L$,
as:

$$ $\displaystyle f(\bar{x}^{\prime}_{\ell})=\mathrm{ReLU}(\bar{x}^{\prime}_{\ell} W_{1}+b_{1})W_{2}+b_{2}$ (4) $$

where $f(\bar{x}^{\prime}_{\ell}):\mathbb{R}^{d}\to\mathbb{R}^{{C}}$
and where the weights have the appropriate dimensions to accommodate a hidden layer of size $p$, i.e. $W_{1}\in\mathbb{R}^{d\times{p}},b_{1}\in\mathbb{R}^{{p}},W_{2}\in\mathbb{R}^{p
\times{C}}$ and $b_{1}\in\mathbb{R}^{{C}}$.

## 3 Experimental Setup

##### Task and Dataset.

We consider a simple algorithmic task that is referred to as histogram: given a sequence of tokens, the goal is to return a sequence of the same length where each entry represents the number of times the corresponding input token appears in the entire sequence. For example, given $\mathbf{x}=[A,B,D,D,B,B]$, the output will be $\mathbf{y}=[1,3,2,2,3,3]$.
We define the count of a token $t$ in the sequence $\mathbf{x}$ at position $\ell$ as $\mathop{hist}_{\mathbf{x}}(\ell)$.
In our experiments, we consider i.i.d. distributions of sequences of length $L$ from an input alphabet of size $T$, where $L\leq T$.
Our sampling strategy relies on first sampling a set of partitions, and then assigning a token to each partition (see App. C for details). This allows for a close to uniform distribution over the values of $\mathbf{y}$.

##### Models and Training.

We investigate the performance on the histogram task of the four different variants of the token mixing models described in Sec. 2, i.e. lin and dot, with or without the softmax (+sftm), where the token embeddings are jointly learned with the model parameters. Their relevant hyperparameters
are the dimension of the embedded tokens $d$, and the hidden layer size $p$ of the feature transformation. Additionally, we consider the model bos(+sftm) where every input sequence is prefixed with the BOS token prior to entering a dot-product mixing layer (with softmax). Previous studies (Weiss et al.,, ; Kazemnejad et al.,, ) have demonstrated that transformer networks consistently attend to BOS tokens, despite their lack of semantic content, and we explore this point in our experiments.
All models are trained with Adam with a learning rate of $\nu=10^{-3}$ on the cross-entropy loss for $500$ epochs with a batch size of $32$. We consider the online learning setting where for each new epoch we generate a dataset of $10,000$ data samples.
We compute the accuracy attained by each model based on a set of $3,000$ independent data samples, which covers a large range of all possible input sequences.

## 4 Learning Regimes in Counting

In order to understand the contributions of the different architectural components, we analyze the performance of the above-stated models with varying mixing mechanisms in different learning regimes characterized by the embedding dimension $d$ and the number of hidden neurons $p$ of the feed-forward module.

Figure: Figure 1: Performance on the histogram task for different 1-layer transformer architectures. Mean accuracy for varying embedding dimension $d$, hidden layer dimension $p$, for fixed $T=32$ and $L=10$ for the different token mixing mechanisms dot, bos and lin. (Top) Models with softmax; (Bottom) Models without softmax. Average over $5$ runs for every $d,p\in\{1,2,3,4,6,8,12,16,23,32,45,64,91,128\}$. Vertical and horizontal white lines indicate $p=T$ and $d=T$ respectively. White stars mark the parameter configurations, where a $100\%$ accuracy configuration was found during training in at least one of the five runs. White dots mark the same for $\geq 99\%$ accuracy configurations.
Refer to caption: extracted/5909847/figures/memorizing_counts-phase-diagram-T32-mean-best-marked.png

Figure: Figure 2: Accuracy vs. Parameter count. The data is the same as generated for Fig. 1, every data point is a single experiment and we show the convex hull in solid lines.
Refer to caption: extracted/5909847/figures/memorizing_counts-params-T32.png

Fig. 1 shows the accuracy attained by learned models for sequences of length $L=10$ with $T=32$ different input tokens. We observe that the models exhibit both high and low
accuracy across various parameter regimes, with a strong dependence on the architecture.
Fig. 2 further clarifies that the parameter efficiency under different architectures varies substantially.
To investigate the underlying mechanisms we devise theoretical constructions and mechanistic interpretations of the learned solutions. We delineate two regimes in each of the parameters:
for the embedding dimension $d$ we distinguish the regime of non-orthogonal embeddings ($d<T$) and of possibly orthogonal embeddings ($d\geq T$).
For hidden layer size $p$ we distinguish the regime where models can sense only a constant number of directions/features ($p=1$) or one scaling as the alphabet size ($p=T$).

### 4.1 d ≥ T 𝑑 𝑇 d\geq T italic_d ≥ italic_T : Orthogonal token embeddings are separable

When the model dimension $d$ is at least as big as the number of tokens $T$, tokens can be represented by embeddings that are mutually orthogonal to one another. Assuming all tokens $t\in\mathcal{T}$ have such mutually orthogonal embeddings $e_{t}\in\mathbb{R}^{d}$ with a norm of 1, the overlap is $\langle e_{s},e_{t}\rangle=0$ for distinct tokens $t\neq s$ and it is 1 when $t=s$. In such a scenario, a linear combination of token embeddings preserves magnitudal - count - information about single tokens. By leveraging knowledge about the embeddings of the alphabet, a weighted sum of tokens, denoted as $e^{\prime}=\sum_{t\in\mathcal{T}}\alpha_{t}e_{t}$, can be broken down into the original tokens using projections on the original token embeddings, where $\alpha_{t}=\langle e_{t},e^{\prime}\rangle/\|e_{t}\|^{2}_{2}$.

In the following, we use this property to theoretically construct the weights for all models that solves the task when $d\geq T$.
Remarkably, the constructions require different number of hidden neurons $p$ depending on the mixing mechanism. This demonstrates the interplay of the mixing layer and the feature transform: for some mixing mechanisms, the latter needs to implement inventory-based counting (IC) (requiring $p\geq T$), and for others, relation-based counting (RC) (where $p\geq 1$ is sufficient).

#### 4.1.1 Relation-based counting: Leveraging dot-product mixing

When an extra beginning-of-sequence token $t_{\mathrm{BOS}}$ is available in bos, it can be used as a to extract information about a token’s count $\mathop{hist}_{\mathbf{x}}(\ell)$ in the attention layer of the network through its attention score Kazemnejad et al., ().
In the literature, the beginning (or end) of sequence tokens have been linked to model-internal computations, such as counting.
In Weiss et al., (), it is shown that the RASP language can solve the histogram task with one layer and one attention head. We confirm empirically that bos and bos+sftm reach (close to) 100% accuracy whenever $d>T$, and we verify that a relation-based counting algorithm can be theoretically implemented in these two architectures by construction.

###### Proposition 4.1 (RC with BOS token) .

For bos and bos+sftm and a given $L\geq 2$, there each exists a configuration of weights that solves the histogram task at 100% accuracy, given that $d\geq T>2$ and $p=1$.

We prove this by construction in App. A.2.3-A.2.2 and we provide the intuition of the proof in the following.
For bos we set the $t_{\mathrm{BOS}}$ embedding to $e_{\mathrm{BOS}}=\sum_{t\in\mathcal{T}}e_{t}$ and take the mutually orthogonal token embeddings $e_{t}$ to have norm 1.
Assuming that $t_{\mathrm{BOS}}$ is at the first position of the sequence of now length $L+1$, a simple dot-product operation in the attention mechanism (with $Q,K=d^{\frac{1}{4}}\mathbb{I}_{d}$) will lead to an attention matrix with entries:

$$ $\displaystyle a_{\ell m}=\begin{cases}T&\mathrm{if}\hskip 2.84544pt\ell=m=1\\ 1&\mathrm{if}\hskip 2.84544pt(\ell>1,m=1)\,\,\mathrm{or}\hskip 2.84544pt(\ell, m>1,x_{\ell}=x_{m})\\ 0&\mathrm{if}\hskip 2.84544pt\ell,m>1,x_{\ell}\neq x_{m}\\ \end{cases}.$ $$

Projecting the mixed token $\bar{x}^{\prime}_{\ell}$ onto the $t_{\mathrm{BOS}}$ we obtain $\langle\bar{x}^{\prime}_{\ell},e_{\mathrm{BOS}}\rangle=T+\mathop{hist}_{
\mathbf{x}}(\ell)+1$, i.e. $e_{\mathrm{BOS}}$ is the single relevant direction for the prediction. Its magnitude relates linearly to $\mathop{hist}_{\mathbf{x}}(\ell)$. A single hidden neuron $p=1$ suffices and the output layer can transfer the count into a categorical representation.
For bos+sftm one needs to further account for the non-linearity of the softmax as described in App. A.2.2.

In the learned models, some instances in the given regime indeed achieve 100% accuracy. While their weights do not correspond exactly to the relation-based counting algorithm described previously, they exhibit similar properties. In Fig. 3, we show for bos+sftm, that $t_{\mathrm{BOS}}$ indeed plays a special role in the learned model: in the attention matrix its activation can be interpreted as a proxy for the number of occurrences of $x_{\ell}$, as it has different values for tokens that occur a different amount of times.
Other entries of the attention matrix are comparatively low when the compared tokens are the same and high when they are different.
The comparison operation naturally provided by the dot-product allows the model to extract the count of the same tokens, for each token in the sequence.
We also show in Fig. 3 how the presence of the $t_{\mathrm{BOS}}$ determines the final prediction through the application of $f$.

Figure: Figure 3: Relation-based counting with bos+sftm ($T=32,L=10,p=2,d=45$). This model achieves $99.9\%$ accuracy. It was selected as the best model from all our experiments with $p=2$. (Left) The tokens overlap (cosine similarity) with the same tokens (red), different tokens (grey) and the BOS (light blue) all concentrate on different values. (Middle) This is reflected in the attention matrix after the application of the row-wise softmax. The $t_{\mathrm{BOS}}$ (‘$’) in the first column $a_{\ell,0}$ becomes a proxy for the count of $x_{\ell}$. (Right) To demonstrate that the feedforward network is only sensitive to this direction, we show its count predictions for a mix of tokens $\alpha e_{\mathrm{BOS}}+(1-\alpha)e_{D}+e_{B}$, where the contribution $\alpha$ of the BOS token is varied and $D,B$ are two specific elements of the alphabet $\mathcal{T}$. The same experiment is repeated for different elements of the alphabet in App. D.4. We mark the $a_{\ell,0}$ obtained from the left as vertical lines, the prediction is correct for all counts independent of the precise token.
Refer to caption: x1.png

Surprisingly, the dot model (without the softmax) reaches a an empirical performance comparable to bos in the regime $d\geq T$ and $p=1$, even though it does not have an extra token available.

###### Proposition 4.2 (RC with tagged embeddings) .

For dot and a given $L,T>2$, there exists a configuration of weights that solves the histogram task at 100% accuracy, given that $d\geq T>2$ and $p=1$.

We prove this in App. A.2.1.
Intuitively, the construction uses a single common direction $e_{\text{cnt}}$ that is added to the otherwise mutually orthogonal token embeddings.
A dot-product mixing then leads to $a_{\ell m}=a_{\neq}>0$ when $x_{\ell}$ is different from $x_{m}$, and $a_{\ell m}=a_{=}>0$ when tokens are the same.
Then, the number of counts can be easily extracted from the dot-product $\langle e_{\text{cnt}},\bar{x}_{\ell}^{\prime}\rangle$ of the counting token with the mixed token $\bar{x}_{\ell}^{\prime}$, i.e. $\langle e_{\text{cnt}},\bar{x}_{\ell}^{\prime}\rangle\propto 1+\mathop{hist}_{
\mathbf{x}}(\ell)a_{=}+(L-\mathop{hist}_{\mathbf{x}}(\ell))a_{\neq}$. We can, therefore, obtain a perfect accuracy implementation in the regime where $d\geq T$ with only a single hidden neuron. This is in line with the observed empirical performance by dot even without access to a BOS token.

###### Proposition 4.1 (RC with BOS token) .

###### Proposition 4.2 (RC with tagged embeddings) .

##### Dot-product attention with softmax fails to implement relation-based counting.

Since the dot-product mechanism can naturally be used in relation-based counting, one might expect the dot+sftm model to implement the same mechanism. However, and maybe surprisingly so, we empirically observe a marked difference between dot and dot+sftm in Fig. 1. dot only starts performing close to 100% accuracy when both the model dimension $d$ and the number of hidden neurons $p$ are larger than the number of tokens $T$.
To understand why it fails to learn for $p=1$, we show the attention matrix of dot+sftm in Fig. 4. Notably, it is based on the semantics, as $(\mathbf{A}_{\mathrm{dot+sftm}})_{\ell m}$ is higher when $x_{\ell}=x_{m}$ than otherwise. However, the normalization effect of the softmax activation prevents the development of a meaningful counter subspace that is needed in the relation-based algorithm.
As a result of normalization, the attention scores are $\sum_{m}a_{\ell m}=1$, so any direction present in all tokens (and by the symmetry of the task, it would need to be present in all tokens) would be uninformative after the token mixing – its weight would be one regardless of the input sequence and would therefore not carry information about the count.
Before, the model bos+sftm circumvented this problem by adding the extra token with a special functionality that does not need to be counted.
Because this is not possible for dot+sftm, the architecture fails to perform well for $p=1$ – it now needs to measure more than one direction in the feed-forward module.
In the following, we show that a solution of the histogram task can still be achieved through an inventory-based counting algorithm with $p\geq T$. We detail this in the following section, for the example of lin. The statement for dot+sftm is given in App. A.3.

Figure: Figure 4: Inventory-based counting with dot+sftm $(T=32,L=10,p=32,d=32)$. This model achieves 99.47% accuracy. (Left) The attention matrix for a given sequence differentiates between similar and different tokens. However in this case, any counting direction that could emerge in token space is evidently not usable, as $p\geq T$ is required (see Fig. 1). (Right) This is reflected in the output from the feature transformation $f$, shown here for a linear combination of three different tokens from the alphabet, $B,C,D$. The prediction strongly depends on the coefficient $\alpha_{t}$ associated with the token $t$ present in the residual connection and only weakly on the others. The non-linear scaling of the decision boundaries is due to the softmax activation function.
Refer to caption: x2.png

#### 4.1.2 Inventory-based counting: Memorization in the feed-forward layer

When the feed-forward hidden layer has one neuron for each distinct token available in the alphabet, it can detect as many directions. This allows the feed-forward layer to extract the information of any token direction separately and thereby implement a custom comparison operation that works for all of the tokens in the alphabet. While this is less parameter efficient and requires memorizing the complete alphabet, it enables the model to solve the task.

###### Proposition 4.3 (IC with memorization in the feed-forward layer) .

For lin and lin+sftm and a given $L,T>2$ there exists a configuration of weights which solves the histogram task for $p\geq T$ and $d\geq T$.

We describe examples of such constructions in App. A.3.1 and A.3.2.
Again, several solutions exist due to symmetries, and in the following we give an intuition for one of them.
In the linear mixing layer $\mathbf{A}_{\mathrm{lin}}$ we set a constant value $a=1/L$ so that the result of the mixing is simply a position-independent linear combination of the input.
The count $hist_{{\mathbf{x}}}({\ell})$ can be extracted after the residual connection where we add $\bar{x}_{\ell}^{\prime}=\bar{x}_{\ell}+e_{x_{\ell}}$. By setting the columns of the matrix $(W_{1})_{t}=e_{t}$ we can extract the count information up to the factor $1/a$

$$ $\displaystyle\mathop{hist}_{\mathbf{x}}(\ell)=\frac{1}{a}\sum_{t\in\mathcal{T} }\text{ReLU}\left(\left\langle\bar{x}_{\ell}^{\prime},(W_{1})_{t}\right\rangle -1\right)=\frac{1}{a}\sum_{t\in\mathcal{T}}\text{ReLU}\left(\langle\bar{x}_{ \ell}^{\prime},e_{t}\rangle-1\right)=\frac{1}{a}\langle\bar{x}_{\ell},e_{x_{ \ell}}\rangle$ $$

Note that, due to the $-1$ bias term, only the hidden neuron for token $(W_{1})_{t}=e_{t}=x_{\ell}$ that occurs in the residual connection has a non-zero activation.
The output layer $W_{2}$ can then be designed to activate the correct output vector corresponding to the count $\mathop{hist}_{\mathbf{x}}(\ell)$ (see App. A.4). Since $a\in[0,1]$ and $\sum_{m=1}^{L}a_{\ell m}=1$ the same procedure can be implemented by a matrix which is passed through the softmax operator for lin+sftm.
In practice, in this construction the feed-forward module is correlated with the complete alphabet, acting as an inventory, or look-up table.

Figure: Figure 5: Inventory-based counting with lin+sftm ($T=64,L=10,p=128,d=128$). The model achieves $99.97\%$ accuracy. (Left) Attention matrix learned by Adam, which is constant in the input sequence $\mathbf{x}$. The different score on the diagonal assigns a different weight to the token at the current position $\ell$ than to all other tokens. (Right) Predictions on an artificial mix of learned embeddings for the three tokens $B,C$ and $D$. The prediction depends on the token in the residual connection, but is largely independent of the presence of other tokens in the mixing. This indicates that $f$ projects the mixed token onto the alphabet $\mathcal{T}$ and is able to extract tokens due to orthogonality.
Refer to caption: x3.png

In Fig. 5, we inspect the attention matrix $\mathbf{A}_{\mathrm{lin}}$ and the feature transformation $f$ which is learned for lin+sftm in the regime where $p\sim T\sim d$. The mixing has an off-diagonal of $\sim 0.11$ and a diagonal of $\sim 0.08$. Feeding the feature transformation $f$ with a weighted combination of $3$ tokens, $B,C,D$, we observe that the final prediction of the network depends mainly on the coefficient $\alpha_{t}$ corresponding to the token embedding fed through the residual connection.
Notably, this behavior is close to Fig. 4 (right) and suggests that the feature transformation must have encoded the information of the token embedding in its weights, hence requiring at least $p=T$ hidden neurons.

###### Proposition 4.3 (IC with memorization in the feed-forward layer) .

##### Superpositioned and selective implementations.

Some of the models capabilities include one another. For example, the models that can implement relation-based counting for $p=1$ can also implement the solutions for inventory-based counting for $p\geq T$.
It is unclear, whether the memory-intensive solution is preferred when the memory is available, or if the efficient solution is learned nonetheless.
Curiously, in Fig. 1, we observe that the model dot (which is capable of RC) witnesses a very slight decrease in maximal learned performance from $100\%$ accuracy to $99\%$ despite its capacity being increased to $p=T$ when inventory-based counting can in principle be implemented.
In App. D.5 we investigate the singular value decomposition of $W_{1}$, for learned models with $\geq 99\%$ accuracy and $p,d\geq T$. We find that the largest $T=32$ singular values are larger than the surplus singular values when $p>T$ for models that can implement only IC.
This behavior is less pronounced for models that can implement RC, where the largest singular value is often relatively much larger than the following $T=32$, but still show a small dip after the $T=32$ singular values.
Understanding which algorithm is implemented in this regime, or if it is a superposition of the two, thus requires further investigation.

Figure: Figure 6: Introspecting the Regime with Entangled Embeddings with dot ($T=64,L=10,p=128$). We show examples of dot for $T=64,L=10,p=128$ for varying the model dimension $d$. (Top) The confusion matrix of ground truth and predicted counts. (Bottom) The overlap distribution between same and different token embeddings.
Refer to caption: x4.png

### 4.2 d < T 𝑑 𝑇 d<T italic_d < italic_T : Non-orthogonal embeddings and the discrete nature of counting

The scenario where $d<T$ fundamentally differs from the one explored in Section 4.1 because the embeddings for different tokens can no longer be mutually orthogonal. Some token pairs then have a non-zero overlap due to their linear dependence, causing the mixing of tokens to entangle count information across different directions in the embedding space. This phenomenon is illustrated for dot in Fig. 6, where learned models with smaller $d$ tend to overcount items in the input, and observe a less spread distribution of overlaps. Nevertheless in Fig. 1 we observe a number of results that empirically show almost perfect accuracy solutions with $d<T$ both for models with RC or IC.
Indeed, the discrete nature of the histogram task, i.e. the fact that every token can only be mapped to $L$ distinct counts, makes the prediction inherently more robust to the effect of noise stemming from entangled embeddings. This concept is illustrated in Fig. 7 in App. A.4 for the dot+sftm model. As long as the value of the logits in the final output layer falls within the margin between two counts the model still solves the task with perfect accuracy.
The relative size of this margin decreases when $L$ is increased, making the task harder when more classes need to be distinguished.

In the following, we link concepts on optimally placing decision boundaries for noise robustness to a characterization of this entanglement noise, measured by the mutual coherence of the token embedding set (i.e., the maximum absolute overlap between pairs of distinct embeddings). The mutual coherence of a set of $T$ vectors of dimension $d$ is lower bounded by the Welch bound (Welch,, ). This gives a means to understand the size of $d$ a given task with $T,L$ requires at least.

###### Proposition 4.4 (Robustness via bounded mutual coherence) .

Given $L\geq 5,T\geq 2$ and assuming that the Welch bound is attained for a given $T,d$, there exists a construction that solves the histogram task with

- (lin, lin+sftm; $p=T$): $\left\lceil\frac{T(2L-3)^{2}}{T-1+(2L-3)^{2}}\right\rceil\leq d$,
- (dot, bos; $p=1$): $\left\lceil\frac{T(2L-3)^{2}}{T-1+(2L-3)^{2}}\right\rceil\,+1\leq d$,
- (dot, bos; $p=T$): $\left\lceil\frac{T(L-1)}{T-1+(L-1)}\right\rceil\leq d$.

We provide additional background and the proofs in App. B.2. The idea is to use constructions analogous to the RC and IC with orthogonal embeddings, while keeping track on how the errors of non-zero overlaps between pairs of different embeddings propagate through the model.
For a given $L$ and $T$ this provides an upper bound on the maximal mutual coherence that is tolerated for a perfect solution. This can be connected to the dimensionality $d$ via the Welch bound. Evaluating the bounds for the setting in Fig. 1,
we obtain, in the order of the above list, $d\geq 29,30,7$.
Generally it is hard to generate matrices that attain the Welch bound and manually we did not succeed to find them for $d=29,30$. However we can indeed create an explicit construction a for dot and $p=T$ which attains $d=12$, as provided in the supplementary code and in correspondence with Fig. 1. While this bound does not reach the $d$ as indicated by the Welch bound, the mutual coherence of the embedding matrix we use is close to the maximally allowed value of $\mathcal{M}=0.299<1/3$.

The previous results apply specifically to models without the softmax operator in the token mixing step – models with this non-linearity can be more robust and attain even smaller $d$, as clearly visible in Fig. 1.
The idea is that a softmax function with a high enough inverse temperature can non-linearly scale down the attention scores for different token pairs relative to those of the same tokens. Thereby, the noise introduced in the dot-product layer through pairs of different embeddings becomes arbitrarily close to zero after applying the softmax.

###### Proposition 4.5 (Robustness via softmax error-reduction) .

Given $T,L>2$, there exist weight configurations that solve the histogram task for the parameter combinations (bos+sftm; $p=1$) and (dot+sftm; $p=T$) with $\lceil\log_{2}(T+1)\rceil+2\leq d$.

Put simply, this construction requires
that there are token embeddings for $t,s=1,\ldots,T$ and $s\neq t$ with $\epsilon>0$ such that

$$ $\displaystyle\langle e_{t},e_{t}\rangle=1\,\,\,\textrm{ and }\,\,\,\langle e_{ t},e_{s}\rangle<1-\epsilon\,.$ (5) $$

This is fulfilled when every token is the binary encoding of its value, modulo minor modifications due to the RC mechanism for bos+sftm.
Setting the softmax temperature high enough as a function of $L$ allows for the contributions from non-equal tokens to be decreased relative to the ones of same tokens.
Evaluating this function for Fig. 1, we obtain $d=7$, which closely corresponds to the most parameter efficient solutions of the histogram task that we observe.
As $L$ grows, we require stronger concentration from the softmax by adjusting its temperature. Since real-world networks execute finite computations, computational instabilities or collapses might occur.
It is therefore not clear that this correspondence will hold for all values of $L$.
In App. B.3.1 we show that this bound can be even further improved for bos+sftm to a constant $d=4$, but at the cost of increasing the temperature further as a function of $T$, in addition to $L$. This might be the reason why we do not observe any learned solutions of the histogram task in this regime.

###### Proposition 4.4 (Robustness via bounded mutual coherence) .

###### Proposition 4.5 (Robustness via softmax error-reduction) .

## 5 Related Work

##### Mechanistic Interpretability and Counting.

The emergence of algorithmic capabilities in transformers (Olsson et al.,, ; Power et al.,, ) has led to numerous investigations aimed at reverse-engineering trained models into human-understandable mechanisms (Zhong et al.,, ; Nanda et al.,, ; Quirke and Barez,, ). Previous studies have investigated a variety of histogram tasks and the mechanisms behind them (Gould et al.,, ; Chollet et al.,, ; Ouellette et al.,, ; Cui et al.,, ). In our work, we consider the histogram task introduced within the context of the RASP(-L) programming language (Weiss et al.,, ; Abbe et al.,, ). Weiss et al., () predict that single layer transformers with one head require an additional BOS token as a scratchpad (Nye et al.,, ) to be able to solve the task. However, we find that the task does not necessarily require the BOS token and we give explicit constructions for several of such one-layer architectures. Our main focus is the interpretation of the hyperparameter scaling of several distinct models in relation to their performance and explicit constructions of different algorithms, similar to the studies in Zhong et al., (); Quirke and Barez, (). We give precise theoretical conditions on the model configurations that lead to perfect explicit constructions.
While many works in this area focus on causal interventions (Vig et al.,, ; Meng et al.,, ) to understand the computational mechanisms of models or assign relevance scores to their components (nostalgebraist,, ; Elhage et al.,, ), our approach primarily involves gaining insights through direct introspection of the model’s components.

##### Memorization and Feed-forward Layers.

The role of feed-forward layers as memorization modules has been investigated in the context of factual recall for language models (Geva et al.,, ; Meng et al.,, ; Chughtai et al.,, ). Henighan et al., () study a double decent phenomenon where the purpose of the feed-forward layer transitions from storing data points to discovering generalizing features as a function of increasing training data diversity (Raventos et al.,, ). In the histogram task, we observe a similar phenomenon as a function of the architecture: the feed-forward layer acts either as a look-up table or a feature detector for a single direction in embedding space – the counting subspace.

##### Aligning Algorithm and Architecture.

While theoretical work has outlined the computational capacity of a range of (autoregressive) neural networks (Weiss et al.,, ; Yun et al.,, ; Delétang et al.,, ; Liu et al.,, ), hallucinations and failure modes on seemingly trivial tasks in real-world transformers are the rule rather than an exception. Dziri et al., () postulate that this may be due to a misalignment between the computational graph of a model and the task itself. In this work, we show that subtle differences in components such as the mixing type and layer width play a crucial role in terms of algorithmic alignment. Previous work discovered evidence for the superposition of different computational graphs in a single model (Elhage et al.,, ) – we complement this analysis with a toy model that is able to disentangle non-orthogonal, hence superimposed, embedding directions in some parameter regimes.

## 6 Discussion & Conclusion

##### Limitations.

Similar to other works in mechanistic interpretability (Zhong et al.,, ), we focus on 1-layer transformers as a simplified model for modern transformers. Our models are not autoregressive and do not account for the impact of causal masks or positional encodings. While more complex models could lead to more intricate interdependencies between the components, potentially limiting the applicability of our findings to such architectures, it seems plausible that similar vector arithmetic could emerge in subspaces of large transformers (Gould et al.,, ; Engels et al.,, ). Given its specificity, it is unclear if and how similar memory-architecture phenomena would emerge for different simple tasks (e.g. sorting or lookup).

##### Summary.

We study how different components of simple transformer models contribute to the emergence of different solutions to the histogram task. Our analysis shows that
the parameter regimes where solving the histogram task is feasible for these models is influenced by the choice of the mixing mechanism and its inter-dependency with the feed-forward transformation, as well as the softmax activation function in the attention mechanism. We identify two distinct algorithmic approaches that 1-layer transformers can utilize to solve the histogram task: relation-based counting and inventory-based counting. The relation-based method employs a dot product mixing mechanism combined with a low-capacity feed-forward transformation and relies on the presence of an appropriate counter direction within the token embedding space. In contrast, the inventory-based method involves memorizing the token embeddings within the feed-forward module’s weights, thus requiring more parameters. By characterizing the feasibility regimes of these mechanisms in the phase space defined by the embedding dimension $d$ and the hidden dimension $p$ of the feed-forward module, we confirm that learned models converge to solutions resembling these mechanisms. In certain regimes, both strategies can potentially be implemented, and our experiments indicate that some learned models exhibit features of superimposed algorithmic mechanisms. In the regime where the embedding dimension $d$ is smaller than the alphabet size $T$, tokens cannot form an orthogonal basis and solve the task directly via a linear projection. Despite this, we find that the considered models exhibit different levels of robustness to the noise stemming from non-orthogonality. Our analysis precisely characterizes how different models cope with this aspect and identifies less stringent feasibility regimes in terms of the embedding dimension. In particular, we find that the softmax activation can be very effective in minimizing the effective similarity between distinct tokens after a comparison opearation through the attention layer, hence reducing the impact of non-orthogonality.
This is particularly relevant to real world models, where the alphabet size is usually much larger than the model dimension.

##### Future Directions.

At this moment, examples for hallucinations and failures of LLM’s are as numerous as their success stories. Even though we only analyze the feasibility regime of a single task, this small example already exhibits a rich phenomenology. It shows that a number of subtle modifications to a models architecture can influence its predictive power drastically. The prime example is the softmax function which becomes a curse or a blessing depending on slight differences in the setup. We expect that similar mechanistic investigations at or close to the regimes where models start failing will be extremely useful to understand how and why models fail in sometimes puzzling manners.

## Appendices

[sections]
\printcontents[sections]l1

## Appendix A Explicit Constructions for Orthogonal Embeddings d = T 𝑑 𝑇 d=T italic_d = italic_T

### A.1 Overview

In the parameter regime where $d\geq T$ there is always an orthonormal basis of size $T$ in $\mathbb{R}^{d}$, these explicit constructions give the correct prediction for all input token sequences. For all the models we describe below, we define the sum of the hidden layer neurons as:

$$ $\gamma_{\ell}=\sum_{i=1}^{p}\text{ReLU}(z_{\ell,i})=\sum_{i=1}^{p}\text{ReLU}( W_{1}\bar{x}^{\prime}_{\ell}+b_{1})_{i}$ (6) $$

In many cases, a simple linear regression can map the scalar $\gamma_{\ell}$ to the correct count of tokens $x_{\ell}$, and we describe how to achieve this mapping to the classification problem in Section A.4.
In the following, we characterize which parameters $W_{1},b_{1}$ in equation 6 allow for a correct mapping in each mechanism. Importantly, the architecture exhibits numerous symmetries due to the feed-forward ReLU network (Petzka et al.,, ). To demonstrate feasibility, we select one specific implementation.
In the main text we observe that there is no one-to-one correspondence between our explicit constructions and the learned weights, even though both functions achieve the same perfect accuracy.
Throughout, unless otherwise specified, we assume that $E\in\mathbb{R}^{d\times T}$ is an orthonormal basis of $\mathbb{R}^{d}$, which we will use to create different forms of token embeddings.
The supplementary code at [https://github.com/SPOC-group/counting-attention](https://github.com/SPOC-group/counting-attention) contains executable pytorch models that have the weight configurations that are used to prove Propositions 4.2-4.3 and A.1, which allows one to test the devised weight configurations for fixed $T,L,d$ in practice.

### A.2 Relation-based counting

#### A.2.1 ( dot ; p = 1 𝑝 1 p=1 italic_p = 1 )

###### Proof of Proposition 4.2 .

We set $T=d>2$ with $L\geq 2$ and $p=1$. We choose the embeddings of the tokens of the dot model as

$$ $\displaystyle e_{t}=\tilde{e}_{t}+\tilde{e}_{\text{cnt}}\,\,\,\,\forall t=1,\dots T$ (7) $$

where the set $E=\{\tilde{e}_{t}\}_{t=1}^{T}$ is an orthonormal basis of an arbitrary but fixed $T$-dimensional subspace of $\mathbb{R}^{d}$, and $\tilde{e}_{\text{cnt}}=\sum_{t=1}^{T}\tilde{e}_{t}$.
The key and query matrix are set to the scaled identity $W_{K}=W_{Q}=d^{1/4}I_{d}$ and hence the mixing layer $\mathbf{A}_{\mathrm{dot}}$ can be viewed as carrying out the unmodified dot-product operation between all pairs of tokens.
The first layer weights $W_{1},b_{1}\in\mathbb{R}^{d}$ can be fixed as

$$ $\displaystyle W_{1}=\tilde{e}_{\text{cnt}}/(T+1)\,;\,\,\,b_{1}=-(1+L(T+2))\,,$ (8) $$

and the second layer weights $W_{2},b_{2}\in\mathbb{R}^{L}$ follow the recursion

$$ $\displaystyle(W_{2})_{1}$ $\displaystyle=-1+\frac{1}{L+1}\,,\,\,\,$ $\displaystyle(b_{2})_{1}$ $\displaystyle=0\,;$ (9) $\displaystyle(W_{2})_{\ell}$ $\displaystyle=-1+\frac{\ell}{L+1}\,,\,\,\,$ $\displaystyle(b_{2})_{\ell}$ $\displaystyle=\left((W_{2})_{\ell-1}-(W_{2})_{\ell}\right)(\ell-0.5)+b_{\ell-1 }\,,$ $\displaystyle\forall\ell=2,\dots,L.$ (10) $$

Given these parameters, it holds that for tokens $1\leq t,s\leq T$ their dot-product is

$$ $\displaystyle\langle e_{t},e_{s}\rangle=\begin{cases}2+T&\mathrm{if}\hskip 2.8 4544ptt\neq s\,,\\ 3+T&\mathrm{if}\hskip 2.84544ptt=s\,.\end{cases}$ (11) $$

Because of our choice of the query and key matrices, it directly follows that for tokens $x_{\ell},x_{m}$ at positions $\ell$ and $m$ from a given sequence $\mathbf{x}$, their attention score is

$$ $\displaystyle a_{\ell m}=\begin{cases}2+T&\mathrm{if}\hskip 2.84544ptx_{\ell} \neq x_{m}\,,\\ 3+T&\mathrm{if}\hskip 2.84544ptx_{\ell}=x_{m}\,.\end{cases}$ (12) $$

Hence, the mixed token after applying the residual connection is

$$ $\displaystyle\bar{x}_{\ell}^{\prime}=\bar{x}_{\ell}+\sum_{m:x_{\ell}=x_{m}}(T+ 3)\bar{x}_{m}+\sum_{m:x_{\ell}\neq x_{m}}(T+2)\bar{x}_{m}$ (13) $$

so that computing

$$ $\displaystyle\bar{x}_{\ell}^{\prime}W_{1}$ $\displaystyle=\left\langle\bar{x}_{\ell}^{\prime},\frac{\tilde{e}_{\text{cnt}} }{1+T}\right\rangle$ (14) $\displaystyle=\left\langle\bar{x}_{\ell},\frac{\tilde{e}_{\text{cnt}}}{1+T} \right\rangle+\sum_{m:x_{\ell}=x_{m}}(T+3)\left\langle\bar{x}_{m},\frac{\tilde {e}_{\text{cnt}}}{1+T}\right\rangle+\sum_{m:x_{\ell}\neq x_{m}}(T+2)\left \langle\bar{x}_{m},\frac{\tilde{e}_{\text{cnt}}}{1+T}\right\rangle$ (15) $\displaystyle=1+hist_{{\mathbf{x}}}(\ell)(T+3)+(L-hist_{{\mathbf{x}}}(\ell))(T +2)$ (16) $\displaystyle=hist_{{\mathbf{x}}}(\ell)+1+L(T+2).$ (17) $$

Then the single hidden unit has the value $\gamma_{\ell}=\mathrm{ReLU}(\bar{x}_{\ell}^{\prime}W_{1}+b_{1})=hist_{{\mathbf
{x}}}(\ell)$.
It is easy to show (analogous to Fig. 7) that the output logits $c=\gamma_{\ell}W_{2}+b_{2}$ with $c\in\mathbb{R}^{L}$, correctly identify the count for integer values $x\in[1,\dots,L]$. This is because we constructed our recursion such that at a given input $x=\ell$ we have that $(W_{2})_{\ell}(\ell-0.5)+(b_{2})_{\ell}=(W_{2})_{\ell-1}(\ell-0.5)+(b_{2})_{
\ell-1}$ and $(W_{2})_{\ell}>(W_{2})_{\ell-1}$, so it holds that

$$ $\displaystyle\operatorname*{arg\,max}_{i=1\dots L}c_{i}(y)=\begin{cases}1&y=1 \\ 2&y=2\\ \dots\\ L&y=L\,,\end{cases}$ (18) $$

which gives the correct classification output for all possible inputs, and hence solves the histogram task at 100% accuracy.
∎

Note, however, that this weight configuration is only one example, and some symmetries in the model can lead to different but also 100% correct algorithms. This is especially important as we compare the regime outlined in the Theorem with the weight configurations learned.

###### Proof of Proposition 4.2 .

#### A.2.2 ( bos+sftm ; p = 1 𝑝 1 p=1 italic_p = 1 )

###### Proof of Proposition 4.1 for bos+sftm .

We set $T=d>2$ with $L\geq 2$ and $p=1$ and consider the model dot+sftm.
Note that in this model every sequence $\mathbf{x}$ is prefixed with $t_{\mathrm{BOS}}$ before it is fed into the embedding and then the mixing layer. Again we use mutually orthogonal embeddings.
$E=\{\tilde{e}_{t}\}_{t=1}^{T}$ is an orthonormal basis of an arbitrary but fixed $T$-dimensional subspace of $\mathbb{R}^{d}$, and $\tilde{e}_{\text{cnt}}=\sum_{t=1}^{T}\tilde{e}_{t}$.
We set $e_{\mathrm{BOS}}=\sum_{t=1}^{T}E_{t}$, where $e_{t}=E_{t}$ and the latter is a column of $E$.
Analogous to the background token from Proposition 4.1 there is only one direction $p=1$ to detect in the feedforward model, so we set

$$ $\displaystyle W_{1}=e_{\mathrm{BOS}};\,\,\,b_{1}=-1.$ (19) $$

For a given token $x_{\ell}$ we have that in the dot-product mechanism $\langle e_{\mathrm{BOS}},e_{x_{\ell}}\rangle=1$, $\langle e_{x_{\ell}},e_{x_{m}}\rangle=1$ if $x_{m}=x_{\ell}$ and 0 0 otherwise. Due to the softmax, the mixing coefficient is $a=e/((k_{x_{\ell}}+1)e+(L-k_{x_{\ell}}))$ (where $e$ is Euler’s number) for comparing $x_{\ell}$ to $t_{\mathrm{BOS}}$ and to all the tokens where $x_{\ell}=x_{m}$, and $b=1/((k_{x_{\ell}}+1)e+(L-k_{x_{\ell}}))$ otherwise, where, $k_{x_{\ell}}=hist_{\mathbf{x}}{(\ell)}$. Hence, the mixed token is:

$$ $\bar{x}^{\prime}_{\ell}=ae_{\mathrm{BOS}}+ak_{x_{\ell}}\bar{x}_{\ell}+\sum_{x_ {m}\neq x_{\ell}}b\bar{x}_{m}+\bar{x}_{\ell}.$ (20) $$

Applying $W_{1}$ and $b_{1}$, we obtain:

$$ $\begin{split}\gamma_{\ell}&=aT+ak_{x_{\ell}}+b(L-k_{x_{\ell}})\\ &=aT+ak_{x_{\ell}}+1-a(k_{x_{\ell}}+1)\\ &=a(T-1)+1\end{split}$ (21) $$

since $(k_{x_{\ell}}+1)a+(L-k_{x_{\ell}})b=1$ by normalization via the softmax function.
The value of $\gamma_{\ell}$ has a dependence on $k_{\ell}$ through $a$ and can be readout into the correct classification as shown Fig. 7.
∎

###### Proof of Proposition 4.1 for bos+sftm .

#### A.2.3 ( bos ; p = 1 𝑝 1 p=1 italic_p = 1 )

###### Proof of Proposition 4.1 - bos .

We set $T=d>2$ with $L\geq 2$ and $p=1$. The construction of the embeddings and $e_{\mathrm{BOS}}$ is analogous to the construction from Section A.2.2 for bos+sftm in the same setting.
However, since no softmax is applied, the mixing coefficients as outputs of $\mathbf{A}_{\mathrm{dot+sftm}}$ for comparing ($x_{\ell}$, $t_{\mathrm{BOS}}$) or $(x_{\ell},x_{m})$ where $x_{\ell}=x_{m}$ is $a=1$.
For $x_{\ell}\neq x_{m}$ it is $b=0$. Then from inserting these values in equation 20 and applying $W_{1}=e_{\mathrm{BOS}}$ and $b_{1}=-T$ we obtain

$$ $\displaystyle\gamma_{\ell}$ $\displaystyle=k_{x_{\ell}}.$ (22) $$

This clearly allows again the single neuron to be read off to the correct result similar to the construction from equation 8.
∎

Note that there is a simple alternative construction that uses the tagged embeddings from the constructive proof of Prop. 4.2.

###### Alternative Proof of Proposition 4.1 - bos .

We set $T=d>2$ with $L\geq 2$ and $p=1$. We note that by setting $t_{\mathrm{BOS}}$ to zero we can achieve equivalence to the model dot. Since according to Prop. 4.2 there exists a weight configuration for dot which solves the histogram task, this configuration will also solve the histogram task for bos with $t_{\mathrm{BOS}}=0$.
∎

###### Proof of Proposition 4.1 - bos .

###### Alternative Proof of Proposition 4.1 - bos .

### A.3 Inventory-based counting

#### A.3.1 ( lin ; p = T 𝑝 𝑇 p=T italic_p = italic_T ).

###### Proof of Proposition 4.3 - lin .

Assume that $T=d>2$ with $L>2$ and $p=T$ and the goal is to find a weight configuration for the model lin.
As embeddings we directly use the orthonormal basis with $T$ vectors $e_{t}$ in $\mathbb{R}^{d}$, where vectors are the embeddings are for the $T$ tokens. We set

$$ $\displaystyle\mathbf{A}_{\mathrm{lin}}=\left[\begin{matrix}a&a&\cdots&a\\ a&a&&\\ \vdots&&\ddots&\\ a&&&a\end{matrix}\right];\,\,\,W_{1}=E;\,\,\,\,b_{1}=-1,$ (23) $$

where $a=1/L$. We start by writing $z_{\ell,t}$ for $t\in\{1,...,p=T\}$, the $t$-th activation of the first hidden layer of the feed-forward module

$$ $z_{\ell,t}=\sum_{m=1}^{L}a_{\ell m}\langle e_{x_{m}},e_{t}\rangle+\langle e_{x _{\ell}},e_{t}\rangle-1\,.$ (24) $$

If $e_{t}=e_{x_{\ell}}$, we have

$$ $z_{\ell,t}=k_{x_{\ell}}a+1-1=ak_{x_{\ell}}\,,$ (25) $$

where, $k_{x_{\ell}}=hist_{\mathbf{x}}{(\ell)}$, applying the ReLU to this scalar keeps its value unchanged. If $e_{t}\neq e_{x_{\ell}}$, we have

$$ $z_{\ell,t}=ak_{e_{t}}+0-1=ak_{e_{t}}-1\leq 0\,.$ (26) $$

The right hand side of the above equation is negative given our choice of $a$, hence applying the ReLU returns 0. This means that, for each token in the input sequence, the contributions of orthogonal tokens cancel, leaving us with a single hidden hidden neuron activated. Hence the count can be read off from $\gamma_{\ell}$. Since only one neuron is activated at a time, the readout from the same procedure as in bos+sftm can be applied to all hidden neurons $z_{\ell,t}$ simultaneously, instead of only one.
This allows the model to solve the histogram task.
∎

###### Proof of Proposition 4.3 - lin .

#### A.3.2 ( lin+sftm : p = T 𝑝 𝑇 p=T italic_p = italic_T )

###### Proof of Proposition 4.3 - lin+sftm .

Assume that $T=d>2$ with $L>2$ and $p=T$. With the statement already proven for lin, we note that we can construct $\mathbf{A}_{\mathrm{lin+sftm}}$ such that it is equivalent to $\mathbf{A}_{\mathrm{lin}}$ from equation 23 via

$$ $\displaystyle\mathbf{A}_{\mathrm{lin+sftm}}=\left[\begin{matrix}a&a&\cdots&a\\ a&a&&\\ \vdots&&\ddots&\\ a&&&a\end{matrix}\right]=\text{softmax}\left(\left[\begin{matrix}\alpha&\alpha &\cdots&\alpha\\ \alpha&\alpha&&\\ \vdots&&\ddots&\\ \alpha&&&\alpha\end{matrix}\right]\right),$ (27) $$

where $a=1/L$ which implicitly defines a choice of $\alpha$.
This means that the construction is equivalent to lin and it follows automatically that also lin+sftm can solve the histogram task.
∎

###### Proof of Proposition 4.3 - lin+sftm .

#### A.3.3 ( dot+sftm : p = d = T 𝑝 𝑑 𝑇 p=d=T italic_p = italic_d = italic_T )

###### Proposition A.1 (IC for dot+sftm ) .

For dot+sftm and
given $L,T>2$ there exists a configuration of weights which solves the histogram task for $p\geq T$
and $d\geq T$.

###### Proof for Proposition A.1 .

We assume $L,T>2$ and $p=T$
and $d=T$ and we consider dot+sftm.
As previously for dot in Prop. 4.2, we set the key and query matrix to the scaled identity $W_{K}=W_{Q}=d^{1/4}I_{d}$. We use an orthonormal basis of $\mathbb{R}^{d}$ to define the parameters $e_{t}\in\mathbb{R}^{d}$ for the the token embeddings. In $\mathbf{A}_{\mathrm{dot+sftm}}{}$ the pre-softmax mixing weights will be $1$ for equal and 0 0 for different tokens due to the unit-norm token embeddings. Defining $k_{x_{\ell}}=\mathop{hist}_{\mathbf{x}}(\ell)$ for brevity, after the softmax we have that

$$ $\displaystyle a_{lm}=\begin{cases}\frac{e}{(L-k_{x_{\ell}})+ek_{x_{\ell}}}&x_{ m}=x_{\ell},\\ \frac{1}{(L-k_{x_{\ell}})+ek_{x_{\ell}}}&\mathrm{else}.\end{cases}$ (28) $$

Hence, for $e_{t}\neq{e_{x_{\ell}}}$

$$ $\langle\bar{x}^{\prime}_{\ell},e_{t}\rangle=\frac{k_{e_{t}}}{(L-k_{x_{\ell}})+ ek_{x_{\ell}}}<1,$ (29) $$

while for $e_{t}={e_{x_{\ell}}}$

$$ $\langle\bar{x}^{\prime}_{\ell},e_{x_{\ell}}\rangle=\frac{k_{x_{\ell}}e}{(k_{x_ {\ell}}e+(L-k_{x_{\ell}}))}+1,$ (30) $$

where the extra summand comes from the residual connection. Hence, by setting

$$ $\displaystyle W_{1}=E;\,\,\,\,b_{1}=-1,$ (31) $$

and applying the ReLU activation, equation 29 will be 0, while equation 30 will implicitly give us the counts as:

$$ $k_{x_{\ell}}=(L\gamma_{\ell})/(-e\gamma_{\ell}+\gamma_{\ell}+e)$ (32) $$

While the final layer cannot immediately implement non-linear functions in $\gamma_{\ell}$, it can take advantage of the fact that $\gamma_{\ell}$ can take only $L$ different values, similar to how we constructed $W_{2}$ and $b_{2}$ in Section A.2.1. Since eventually we need to map the $L$ values of $\gamma_{\ell}$ to the counts $[1,\cdots,L]$ the linear output layer is sufficient to implement this non-linear discrete map. Fig. 7 shows an example for this map for a given example. This allows the model to solve the histogram task.

The statement for $p>T$
and $d>T$ follows as we can simply set the surplus of parameters in the hidden layer/embeddings to zero.
∎

###### Proposition A.1 (IC for dot+sftm ) .

###### Proof for Proposition A.1 .

### A.4 Mapping a scalar to a categorical one-hot encoding

Figure: Figure 7: Demonstration of a hidden neuron output $\gamma_{\ell}$ which is mapped to $L=10$ different neurons $c_{i}$ using a single output layer, according to the decision boundaries shown by the dotted lines. After applying the argmax function to the 10 output neurons, the highest value gives the discrete output. The solid lines mark the values a single hidden neuron would achieve for different counts in the explicit construction of the dot+sftm model.
Refer to caption: x5.png

It is straightforward to map a single scalar $\gamma_{\ell}$ to a series of neurons which activate one after another.
This is needed as the second part of the feed-forward parameters to transform the count measured by the sum of the hidden neurons $\gamma_{\ell}$ to the discrete categorical representation of the output vector. Every output logit is a linear function of the hidden neuron’s value. Since in our constructions we only map functions, where the ground truth output logit corresponds to an interval $[a,b]\in R$, the superposition of linear functions with increasing slope allows us to realize such a mapping. A visual sample is given in Fig. 7 for dot+sftm.
In Fig. 8 we show the outputs for the lin+sftm model with the best accuracy for $T=32$ for every $p,d$ ran in Fig. 1.
While it is possible to learn the count from one hidden neuron only using inventory-based counting for each neuron, for some examples the count information seems to be spread out over several hidden neurons: The output logits are non-linear in the count and can hence not rely on a single hidden neuron only.

Figure: Figure 8: The output neurons $c_{i}(\mathbf{x}_{\ell})$ visualized for examples of a learned version of lin+sftm for several model dimensions $d$ and hidden layer sizes $p$. Note that differently from Fig. 7, in this case the x-axis shows the number of occurrences $\mathop{hist}_{\mathbf{x}}(\ell=0)=1,\ldots,L-1$ of the token $t$ in an input sequence $\mathbf{x}=[t,\cdots,t,v,\cdots,v]$ that contains otherwise only a token $v\neq t$ (and not the activation of a hidden neuron). We show the activations $c_{i}$ of the final layer output neurons activations (logits) in terms of the number of occurrences of a given token in the input. The colors represent the different output predictions and are as in the explicit construction from Fig. 7. We show several activations for different tokens $t\in\mathcal{T}$, where $T=32$, and we highlight one of the example tokens $t$ with a wider line. While similar to the explicit construction from Fig. 7, the models with 100% accuracy are not necessarily linear in the count.
Refer to caption: extracted/5909847/figures/outputs_linear+sftm.png

## Appendix B Explicit Constructions for Linearly Dependent Embeddings d < T 𝑑 𝑇 d<T italic_d < italic_T

### B.1 Overview

In this section, we discuss the scenario when $d<T$, i.e. when the embeddings are necessarily linearly dependent. In that case, we can no longer assume that there exist embeddings with $\langle e_{t},e_{s}\rangle=0$ for all $t\neq s$.
Nonetheless, also in this regime for some models it is possible to provide explicit constructions of the weights that have 100% accuracy.
This relies on the fact that the prediction problem is inherently discrete, i.e. it chooses exactly one among $L$ classes.
When we examine $\gamma_{\ell}\in\mathbb{R}$ from equation 6 which is mapped to the discrete class through the readout layer (see for example Fig. 7), we notice that the class boundary (the gray dashed class borders) can be placed variably in the margin between the values that $\gamma_{\ell}$ assumes for different counts $k_{x_{\ell}}$ (solid lines).
In the following explicit constructions, our goal is to design embeddings with $d<T$ in such a way that we maximize the aforementioned margin: there will be pairs of token embeddings in the alphabet that have non-zero similarity $\langle e_{t},e_{s}\rangle$, and in equation 6 this will create non-zero terms that will alter the value of $\gamma_{\ell}$.
This means that for every possible sequence with $k$ occurrences of token $x_{\ell}$, the hidden activation $\gamma_{\ell}$ will assume values in a certain range $[\gamma_{\ell}^{\mathrm{lower}}(k),\gamma_{\ell}^{\mathrm{upper}}(k)]$.
If these ranges overlap for different $k$, the count cannot be identified.
However, we construct embeddings such that every for every $k=1,\ldots,L-1$ it holds that

$$ $\displaystyle\gamma_{\ell}^{\mathrm{upper}}(k)<\gamma_{\ell}^{\mathrm{lower}}( k+1)\,,$ (33) $$

so we can still use a construction as in Fig. 7 to correctly compute the final count. In the remainder of this section, we introduce explicit constructions with $d<T$ for a given $L$, both for the cases where we have relation-based counting and inventory-based counting (the same argument as above transfers to $z_{\ell,t}$ from equation 24).
Notably, for the explicit constructions we propose, the function of the lowest achievable $d(p,T,L)$ differs across different mixing types. To summarize:

- •
For models with $\mathbf{A}$ constant in the inputs or models without softmax activation, our explicit construction relies on an embedding matrix with a small mutual coherence. The mutual coherence is a concept from compressed sensing and coding theory that ensures that the maximal similarity between pairs of vectors is small (Donoho and Elad,, ).
We can upper bound the mutual coherence that the margins of the construction can tolerate to still achieve perfect accuracy in terms of a given $L$.
At the same time, the mutual coherence of a set of vectors is naturally lower bounded in terms of the number of vectors $T$ and their respective dimension $d$, known as the Welch bound (Welch,, ). When this bound can be attained and $T,L$ are given, this leads to the following bounds on $d$ for the different models, as outlined in Prop. 4.4, as
(lin, lin+sftm; $p=T$): $\left\lceil\frac{T(2L-3)^{2}}{T-1+(2L-3)^{2}}\right\rceil\leq d$,
(dot, bos; $p=1$): $\left\lceil\frac{T(2L-3)^{2}}{T-1+(2L-3)^{2}}\right\rceil\,+1\leq d$,
(dot, bos; $p=T$): $\left\lceil\frac{T(L-1)}{T-1+(L-1)}\right\rceil\leq d$.
- •
For bos+sftm we rely on the fact that the softmax function accentuates the largest value and thereby can drive attention scores for equal tokens $a_{ii}$ higher relative to attention scores of non-equal tokens $a_{ij}$.
This distinguishes it from the previous case, and allows us to state Prop. 4.5 for which we describe an explicit construction that solves the histogram task with
(bos+sftm; $p=1$): $d\geq\lceil\log_{2}(T+1)\rceil+2$.
(dot+sftm; $p=T$): $d\geq\lceil\log_{2}(T+1)\rceil+2$.
Notably there is no explicit dependence on $L$ for the dimension. However, the smaller the dimension $d$ the more accurate computations and softmax numerical stability are required, as the softmax temperature depends on $L$. With infinitely precise computations we show it is even possible to achieve perfect accuracy with $d=4$, but for finite computations this might pose a problem when $L$ becomes too large.

### B.2 Explicit construction for bounded mutual coherence

We define the mutual coherence $\mathcal{M}$ of a set of $T$ unit norm vectors $\{v_{1},\ldots,v_{T}\}\subset\mathbb{R}^{d}$ as

$$ $\displaystyle\mathcal{M}=\max_{i\neq j}|\langle v_{i},v_{j}\rangle|\,.$ (34) $$

This value is lower bounded for a given matrix by the Welch bound (Welch,, )

$$ $\displaystyle\mathcal{M}\geq\sqrt{\frac{T-d}{d(T-1)}}=\mathcal{W}(T,d)\,,$ (35) $$

and equality can only be attained if $T<d^{2}$ (Strohmer and Heath,, ).
There is a large body of work in coding theory and compressed sensing concerning the existence and construction of a set of vectors that attains $\mathcal{M}$ at or close to $\mathcal{W}(T,d)$. Explicit constructions exist but are not known for every combination of $T$ and $d$. A list with existing constructions for the real space for small $T,d$ can be found in Fickus and Mixon, (), but otherwise gradient-based optimization has been used to find good candidate matrices (Jiang et al.,, ; Jyothi and Babu,, ).

In order to prove Prop. 4.4, we use the following idea: For a given $T$, $L$ and $p$, we can derive an upper bound on the mutual information of the embeddings in terms of $L$, which is required to obtain perfect accuracy.
The form of this upper bound depends on the precise mixing strategy and the choice of $p$.
Through the Welch lower bound on $\mathcal{M}$ we can in turn obtain a lower bound on $d$ in terms of $L$ and $T$.
Note that the Welch bound cannot be attained for $T<d^{2}$ and in this case the bound on $d$ is strict.

#### B.2.1 ( lin , lin+sftm ; p = T 𝑝 𝑇 p=T italic_p = italic_T )

###### Proof of Proposition 4.4 - lin .

To show the bound on $d$, we analyze the inventory-based construction for lin in equation 23.
Given that $p=T$, and $L>2$ is given, let us assume that there exists set of $T$ unit norm vectors $\{e_{1},\ldots,e_{T}\}\subset\mathbb{R}^{d}$ with mutual coherence $\mathcal{M}$.
We use these vectors as our embeddings.

The value $z_{\ell,t}$ for $t=x_{\ell}$, with $W_{1}=[e_{1},\ldots,e_{T}]$ and $b_{1}=-1$ is

$$ $\displaystyle z_{\ell,t}=ak_{x_{\ell}}+a\sum_{m:x_{m}\neq t}\langle e_{x_{m}}, e_{t}\rangle\,,$ (36) $$

and using that fact that the mutual coherence bounds the absolute value of the inner product

$$ $\displaystyle ak_{x_{\ell}}-a\mathcal{M}(L-k_{x_{\ell}})\leq z_{\ell,t}\leq ak _{x_{\ell}}+a\mathcal{M}(L-k_{x_{\ell}})\,.$ (37) $$

Similarly, for $t\neq x_{\ell}$ and $a=1/L$ it still holds that

$$ $\displaystyle z_{\ell,t}\leq ak_{t}+a(L-k_{t})\mathcal{M}-1+\mathcal{M}\leq 0\,,$ (38) $$

provided that $\mathcal{M}<1/(L+1)$, for the worst case where $k_{t}=L-1$.
This means that the ReLU sets all hidden neurons $z_{\ell,t}$ to zero when $t\neq x_{\ell}$, and are therefore no contribution to the final result.
Then, defining

$$ $\displaystyle\gamma_{\ell}^{\mathrm{lower}}(k)$ $\displaystyle=ak-a\mathcal{M}(L-k)\,,$ (39) $\displaystyle\gamma_{\ell}^{\mathrm{upper}}(k)$ $\displaystyle=ak+a\mathcal{M}(L-k)\,,$ (40) $$

we have that indeed for a sequence where $x_{\ell}$ occurs $k=1,\ldots,L-1$ times it holds that

$$ $\displaystyle 0\leq\gamma_{\ell}^{\mathrm{lower}}(k)\leq\gamma_{\ell}(k)\leq \gamma_{\ell}^{\mathrm{upper}}(k).$ (41) $$

The first inequality is required due to the ReLU and holds when $\mathcal{M}<1/(L-1)$.
From equation 33 we have the condition that for all $k=1,\ldots,L$ it holds that

$$ $\displaystyle\gamma_{\ell}^{\mathrm{upper}}(k)$ $\displaystyle<\gamma_{\ell}^{\mathrm{lower}}(k+1)\,,$ (42) $\displaystyle k+(L-k)\mathcal{M}$ $\displaystyle<(k+1)+(L-k-1)(-\mathcal{M})\,,$ (43) $\displaystyle\mathcal{M}$ $\displaystyle<\frac{1}{2(L-k)-1}\,,$ (44) and since we assume that there exist at least two different tokens in the sequence, minimizing the bound over $k$ leaves for $k=1$ $\displaystyle\mathcal{M}$ $\displaystyle<\frac{1}{2L-3}\,.$ (45) $$

which is valid provided that $L\geq{2}$.
Collecting all previous bounds on $\mathcal{M}$, we conclude that when $L\geq 4$ the above construction achieves the correct counts with $\mathcal{M}<\frac{1}{2L-3}$.

The Welch bound equation 35 gives an upper bound on $\mathcal{M}$ in terms of $T,d$ and therefore yields the final condition

$$ $\displaystyle d\geq\left\lceil\frac{T(2L-3)^{2}}{T-1+(2L-3)^{2}}\right\rceil$ (46) $$

under which the given weight configuration is able to solve the histogram task with perfect accuracy.
∎

For lin+sftm the construction and conditions transfer directly, when the constant $\mathbf{A}_{\mathrm{lin+sftm}}$ is constructed to match $\mathbf{A}_{\mathrm{lin}}$ exactly.

###### Proof of Proposition 4.4 - lin .

#### B.2.2 ( dot , bos ; p = 1 𝑝 1 p=1 italic_p = 1 )

###### Proof of Proposition 4.4 - dot , $p=1$ .

We assume that $L>2$ and $T$ given and we use a similar idea as the relation-based weight configuration from the proof of Prop. 4.2 for dot with $p=1$. For the token embeddings, we assume that we have a set of $T$ unit norm vectors with $\{v_{1},\ldots,v_{T}\}\subset\mathbb{R}^{d-1}$ with mutual coherence $\mathcal{M}$, where $d>2$.
We set the entries of the $T$ embedding vectors $e_{t}\in\mathbb{R}^{d}$ to be

$$ $\displaystyle e_{t}=\begin{bmatrix}\\ v_{t}\\ \\ \alpha\end{bmatrix}\,.$ (47) $$

The shared counting subspace is defined on the last coordinate of the vectors via $e_{cnt}=[0,0,\ldots,1/\alpha]$.
Then

$$ $\displaystyle\langle e_{t},e_{t}\rangle$ $\displaystyle=1+\alpha^{2}$ (48) $\displaystyle|\langle e_{t},e_{s}\rangle|$ $\displaystyle\leq\mathcal{M}+\alpha^{2}$ (49) $$

The mixed token with the residual connection at position $\ell$ for a given input sequence is

$$ $\displaystyle\bar{x}^{\prime}_{\ell}=\sum_{m=1}^{L}\langle e_{x_{m}},e_{x_{ \ell}},\rangle e_{x_{m}}+e_{x_{\ell}}$ (50) $$

and the single hidden neuron $\gamma_{\ell}$ for a bias term $b_{1}=0$ and $W_{1}=e_{cnt}$

$$ $\displaystyle\gamma_{\ell}=\langle e_{cnt},\bar{x}^{\prime}_{\ell}\rangle$ $\displaystyle=\sum_{m=1}^{L}\langle e_{x_{m}},e_{x_{\ell}},\rangle\langle e_{x _{m}},e_{cnt}\rangle+\langle e_{x_{\ell}},e_{cnt}\rangle$ (51) $\displaystyle=k_{x_{\ell}}(1+\alpha^{2})+\sum_{m:x_{m}\neq x_{\ell}}\langle e_ {x_{m}},e_{x_{\ell}},\rangle+1$ (52) $$

So that we can achieve for a given count $k$ the $\gamma_{\ell}^{\mathrm{lower}}(k)\leq\gamma_{\ell}(k)\leq\gamma_{\ell}^{
\mathrm{upper}}(k)$ with

$$ $\displaystyle 0\leq\gamma_{\ell}^{\mathrm{lower}}(k)$ $\displaystyle=k(1+\alpha^{2})-(L-k)(\mathcal{M}+\alpha^{2})+1\,,$ (53) $\displaystyle\gamma_{\ell}^{\mathrm{upper}}(k)$ $\displaystyle=k(1+\alpha^{2})+(L-k)(\mathcal{M}+\alpha^{2})+1\,.$ (54) $$

We achieve the upper bound from zero, when $\mathcal{M}<2/(L-1)$, assuming that $\alpha$ is close enough to zero so that is is negligible.
Finally, the condition from equation 33 yields

$$ $\displaystyle\mathcal{M}$ $\displaystyle<\frac{1}{2(L-k)-1}-\frac{2(L-k)-2}{2(L-k)-1}\alpha^{2}$ (55) $$

under the condition that $0<\alpha<\sqrt{\frac{1}{2(L-k)-2}}$. Again, assuming there exist at least two different tokens in the sequence, the r.h.s. of the above expression is minimized for $k=1$ as

$$ $\displaystyle\mathcal{M}$ $\displaystyle<\frac{1}{2L-3}-\frac{2L-4}{2L-3}\alpha^{2}$ (56) $$

which is always positive assuming $L\geq{2}$. This is the relevant bound when we have a large enough $L\geq 5$ and again $\alpha$ is close enough to zero.
Again, combining this with the Welch bound equation 35 leads to

$$ $\displaystyle d-1$ $\displaystyle\geq\left\lceil\frac{T(\frac{2L-3}{1-(2L-4)\alpha^{2}})^{2}}{T-1+ (\frac{2L-3}{1-(2L-4)\alpha^{2}})^{2}}\right\rceil\,,$ (57) $$

and when we choose $\alpha>0$ close to zero, as for lin before

$$ $\displaystyle d\geq\left\lceil\frac{T(2L-3)^{2}}{T-1+(2L-3)^{2}}\right\rceil\, +1.$ (58) $$

∎

This proof holds equivalently for bos when we set the BOS token embedding to zero.

###### Proof of Proposition 4.4 - dot , p = 1 𝑝 1 p=1 italic_p = 1 .

#### B.2.3 ( dot , bos ; p = T 𝑝 𝑇 p=T italic_p = italic_T )

We can decrease the required dimension $d<T$ even further than previously, when we have $p=T$ and implement inventory-based counting in the dot model (and equivalently in the bos model). In that case, the lower bound on $d$ becomes more loose, because we combine the ideas we saw in lin and $p=T$ for inventory-based counting and the effects on the margin in dot and $p=1$.

###### Proof of Proposition 4.4 - dot , $p=T$ .

In our construction, for a given $T$ and $L$, we assume that there is a set of $T$ unit norm vectors $\{e_{1},\ldots,e_{T}\}\subset\mathbb{R}^{d}$ with mutual coherence $\mathcal{M}$ upon which we build our embeddings.
Note that the only difference to the previous relation-based case $p=1$ is that this time there is no extra counting direction.
Importantly, we set $K=d^{1/4}I_{d}$ as before, but $Q=\frac{1}{L}d^{1/4}I_{d}$.
This gives an extra factor in the attention scores.
Further, we set $b_{1}=-1$ and the columns of $W_{1}\in\mathbb{R}^{d\times T}$ to the embeddings $e_{t}$, as we did for lin.
This results in a mixed token $\bar{x}^{\prime}_{\ell}$ according to equation 50.
The hidden neuron is

$$ $\displaystyle z_{\ell,t}$ $\displaystyle=\frac{1}{L}\sum_{m=1}^{L}\langle e_{x_{m}},e_{x_{\ell}}\rangle \langle e_{t},e_{x_{m}}\rangle+\langle e_{t},e_{x_{\ell}}\rangle-1\,$ (59) then with $t=x_{\ell}$ we have $\displaystyle z_{\ell,t}$ $\displaystyle=\frac{1}{L}\left(k_{x_{\ell}}+\sum_{m:x_{m}\neq t}\langle e_{x_{ m}},e_{t}\rangle^{2}\right)\,.$ (60) Note that the square in equation 59 is what differs from the $z_{\ell,t}$ in equation 36 . This is because the term $\langle e_{x_{m}},e_{t}\rangle$ is once introduced through the dot-product attention and once through the dot-product via $W_{1}$ . Conversely, with $t\neq x_{\ell}$ it becomes $\displaystyle z_{\ell,t}$ $\displaystyle=\frac{1}{L}\left(k_{t}\langle e_{t},e_{x_{\ell}}\rangle+\sum_{m: x_{m}\neq t}\langle e_{x_{m}},e_{x_{\ell}}\rangle\langle e_{x_{m}},e_{t} \rangle\right)+\langle e_{t},e_{x_{\ell}}\rangle-1$ (61) $\displaystyle\leq\frac{1}{L}\left(k_{t}\mathcal{M}+(L-k_{t})\mathcal{M}\right) +\langle e_{t},e_{x_{\ell}}\rangle-1$ (62) $\displaystyle\leq 2\mathcal{M}-1$ (63) if we set $\mathcal{M}<0.5$ , which we need anyways for $L\geq 2$ by the stronger upper bound on $\mathcal{M}$ that we derive in the following, we finally have for $t\neq x_{\ell}$ $\displaystyle z_{\ell,t}$ $\displaystyle<0\,.$ (64) $$

Again, negative $z_{\ell,t}$ are set to zero via the ReLU, and the final outcome $\gamma_{\ell}=z_{\ell,t=x_{\ell}}$ depends only on a single hidden neuron equation 59.
This eventually leads to

$$ $\displaystyle 0\leq\gamma_{\ell}^{\mathrm{lower}}(k)$ $\displaystyle=\frac{k}{L}\,,$ (65) $\displaystyle\gamma_{\ell}^{\mathrm{upper}}(k)$ $\displaystyle=\frac{1}{L}(k+\mathcal{M}^{2}(L-k))\,,$ (66) $$

and using the same concept as before, while minimizing over $k$ and applying the Welch bound, to the upper bound

$$ $\displaystyle\mathcal{M}$ $\displaystyle<\sqrt{\frac{1}{L-1}}\,.$ (67) $$

The final bound is more loose than it was for $p=1$ as we only require

$$ $\displaystyle d\geq\left\lceil\frac{T(L-1)}{T-1+(L-1)}\right\rceil\,.$ (68) $$

∎

###### Proof of Proposition 4.4 - dot , p = T 𝑝 𝑇 p=T italic_p = italic_T .

### B.3 Explicit Construction with binary representations and softmax

In our final analysis we examine the key difference between the models bos+sftm and bos – the softmax activation.
In order to show Prop. 4.4 we needed to construct embeddings with a low mutual coherence, because the term $\langle e_{t},e_{s}\rangle$ introduced an error on the mixed token, when $t$ and $s$ were not equal.
Now, with the softmax activation applied to the mixing coefficients, the model can use the non-linearity of this transform to its advantage to separate the relative error.

Recall the softmax function is

$$ $\displaystyle\textrm{sftm}(\mathbf{z})_{i}=\frac{e^{z_{i}}}{\sum_{j=1}^{n}e^{z _{j}}}\quad\text{for}\quad i=1,2,\ldots,n\,,$ (69) $$

and when we compute $\textrm{sftm}(\kappa\mathbf{z})_{i}$ we say it is a softmax with a inverse temperature $\kappa>0$.
When $\mathbf{z}$ of length $L$ contains only two different values, one with $k$ and the other with $L-k$ occurrences, then as $\kappa\to\infty$ the mass concentrates only on the larger value of the two, and sets the other to zero.
We use this intuition to create token embeddings that fulfill for all $t,s=1,\ldots,T$ and $s\neq t$

$$ $\displaystyle\langle e_{t},e_{t}\rangle$ $\displaystyle=1\,,$ (70) $\displaystyle\langle e_{t},e_{s}\rangle$ $\displaystyle<1+\epsilon\,,$ (71) $$

where $\epsilon>0$.

The idea is that the softmax with a high enough inverse temperature sets the term for different tokens, $\langle e_{t},e_{s}\rangle$, close enough to zero, essentially eliminating the noise.
Note that equation 70 is a weaker condition on the set of token embeddings than for example the bound of the mutual coherence in terms of the sequence length $L$ bos with $p=1$ in Section B.2.2.
It allows us to obtain perfect accuracy with smaller $d$.
In the following, we describe the construction of the matrix explicitly.
The supplementary code at [https://github.com/SPOC-group/counting-attention](https://github.com/SPOC-group/counting-attention) contains executable pytorch models that have the weight configurations that are used to prove Propositions 4.5 and the Remark for $d=4$, which allows one to test the devised weight configurations for fixed $T,L,d$ in practice.

#### B.3.1 ( bos+sftm ; p = 1 𝑝 1 p=1 italic_p = 1 )

###### Proof of Proposition 4.5 - bos+sftm .

For a given $T,L>2$ we set the embeddings vectors to the binary representation of the token index $t=1,\ldots,T$ in $d^{\prime}=\lceil\log_{2}(T+1)\rceil$ dimensions

$$ $\displaystyle e_{t}=\begin{bmatrix}\\ \mathop{\mathrm{bin}}(t)\langle\mathop{\mathrm{bin}}(t),\mathop{\mathrm{bin}}( t)\rangle^{-1}\\ \\ \alpha\\ 0\end{bmatrix}\,\,\mathrm{;}\,\,\,\,\,\,\,e_{BOS}=\begin{bmatrix}\\ \mathop{\mathrm{bin}}(0)\langle\mathop{\mathrm{bin}}(0),\mathop{\mathrm{bin}}( 0)\rangle^{-1}\\ \\ 1/\alpha\\ 1\end{bmatrix}\,.$ (72) $$

where $\mathop{\mathrm{bin}}(t)=[v_{1},\ldots,v_{d^{\prime}}]\in\{0,1\}^{d^{\prime}}$ with $t=\sum_{i=1}^{d^{\prime}}v_{i}2^{i-1}$. We select $\alpha>0$.
Then we have that

$$ $\displaystyle\langle e_{t},e_{t}\rangle$ $\displaystyle=1+\alpha^{2}\,,$ (73) $\displaystyle\alpha^{2}\leq\langle e_{t},e_{s}\rangle$ $\displaystyle\leq\sqrt{1-\frac{1}{d^{\prime}}}+\alpha^{2}\leq 1+\alpha^{2}-{ \epsilon}\,,$ (74) $\displaystyle\langle e_{t},e_{BOS}\rangle$ $\displaystyle=1\,,$ (75) $$

where $\sqrt{\frac{d^{\prime}-1}{d^{\prime}}}=\langle e_{2^{d^{\prime}}-1},e_{2^{d^{
\prime}}-2}\rangle$, which has the largest overlap among all possible non-equal pairs of tokens, and the lower bound comes from all coordinates being positive.
Using a readout on the direction only present in the $e_{BOS}$ token, namely, $W_{1}=[e_{cnt}]=[0,\ldots,0,1]\in\mathbb{R}^{d}$ and $b_{1}=0$, we construct

$$ $\displaystyle\gamma_{\ell}=\langle e_{cnt},\bar{x}^{\prime}_{\ell}\rangle$ $\displaystyle=\mathrm{sftm}(EE^{T}_{\ell})_{0}\langle e_{BOS},e_{cnt}\rangle+ \sum_{m=1}^{L}\mathrm{sftm}(EE^{T}_{\ell})_{m+1}\langle e_{x_{m}},e_{cnt} \rangle+\langle e_{x_{\ell}},e_{cnt}\rangle$ (76) $\displaystyle=\mathrm{sftm}(EE^{T}_{\ell})_{0}$ (77) $\displaystyle=\mathrm{sftm}([\langle e_{\ell},e_{BOS}\rangle,\langle e_{\ell}, e_{1}\rangle,\ldots,\langle e_{\ell},e_{L}\rangle])_{0}$ (78) $$

The goal of applying the softmax function is to diminish the contributions of error equation 74, while having the final dimension of the $e_{BOS}$ token be representative of the count of $x_{\ell}$.
The maximum error is induced when the upper bound equation 74 is attained for all tokens in the sequence $\mathbf{x}$ that are not equal to $x_{\ell}$.
The minimum error is obtained when these different tokens attain the lower bound.
Without loss of generality on the ordering, this implies that for a given length $L$ and a softmax activation function with an inverse temperature $\kappa$(^2^22In order to introduce the inverse temperature $\kappa$ of the softmax in the model, we scale the query matrix. We set $K=d^{1/4}I_{d}$, but $Q=\kappa d^{1/4}I_{d}$.), we have that

$$ $\displaystyle\gamma_{\ell}^{\mathrm{lower}}(k)$ $\displaystyle=\frac{e^{\kappa 1}}{e^{\kappa 1}+ke^{\kappa(1+\alpha^{2})}+(L-k) e^{\kappa(1+\alpha^{2}-\epsilon)}}\,,$ (79) $\displaystyle\gamma_{\ell}^{\mathrm{upper}}(k)$ $\displaystyle=\frac{e^{\kappa 1}}{e^{\kappa 1}+ke^{\kappa(1+\alpha^{2})}+(L-k) e^{\kappa\alpha^{2}}}\,.$ (80) $$

We explicitly need $\epsilon$ strictly greater than zero, since otherwise there is no information about the count in $\gamma_{\ell}$ when it becomes independent of the count $k$.
Notice, that this time it holds that $\gamma_{\ell}$ that correspond to higher values correspond to smaller counts, since a larger count corresponds to a larger denominator, i.e. a smaller $\gamma_{\ell}$.
Due to this inverse relationship, for this model, we want that for all counts $k=1,\ldots,L-1$ that it holds that

$$ $\displaystyle\gamma_{\ell}^{\mathrm{upper}}(k+1)<\gamma_{\ell}^{\mathrm{lower} }(k)\,.$ (81) $$

This can be achieved by setting the inverse temperature $\kappa$ accordingly.
In the following we show that there exists a $\kappa$ which fulfills equation 81 for all $d^{\prime}\geq 2$ and $L>2$.
Observe that $\gamma_{\ell}^{\mathrm{upper}}(2)<\gamma_{\ell}^{\mathrm{lower}}(1)$ implies the bounds for all other $k$.
We define the distance or margin as

$$ $\displaystyle\mathrm{dist}(\kappa)=\gamma_{\ell}^{\mathrm{lower}}(1)-\gamma_{ \ell}^{\mathrm{upper}}(2)\,.$ (82) $$

Since at $\kappa=0$ both $\gamma_{\ell}^{\mathrm{upper}}(2)=\gamma_{\ell}^{\mathrm{lower}}(1)=1/(L+1)$, the distance is zero. However then it becomes impossible to distinguish $k=1$ and $k=2$, as they receive the same weight. We therefore need the additional condition that $\gamma_{\ell}^{\mathrm{upper}}(2)\neq\gamma_{\ell}^{\mathrm{lower}}(1)$.
At $\kappa=0$, we observe that this function has a negative derivative, as

$$ $\displaystyle\frac{\partial}{\partial\kappa}\mathrm{dist}(\kappa)|_{\kappa=0}$ $\displaystyle=\mathrm{sftm}(\kappa z_{\mathrm{lower}})_{0}\left((z_{\mathrm{ lower}})_{0}-\sum_{i=0}^{L+1}(z_{\mathrm{lower}})_{j}\mathrm{sftm}(\kappa z_{ \mathrm{lower}})_{i}\right)$ (83) $\displaystyle-\mathrm{sftm}(\kappa z_{\mathrm{upper}})_{0}\left((z_{\mathrm{ upper}})_{0}-\sum_{i=0}^{L+1}(z_{\mathrm{upper}})_{j}\mathrm{sftm}(\kappa z_{ \mathrm{upper}})_{i}\right)$ (84) $\displaystyle=-\left(\frac{1}{L+1}\right)^{2}\left(\left[(1+\alpha^{2})+(L-1)( 1+\alpha^{2}-\epsilon)\right]-\left[2(1+\alpha^{2})+(L-2)\alpha^{2}\right]\right)$ (85) $\displaystyle=-\left(\frac{1}{L+1}\right)^{2}\left[(L-2)-(L-1)\epsilon\right]$ (86) $\displaystyle<0$ (87) $$

where the last bound is met when $0<\epsilon<0.5$ which is fulfilled already for $d^{\prime}=2$ and when $L>2$.
As the distance function is continuous, there exists a $\kappa$ close to zero for which the $\mathrm{dist}(\kappa)<0$.
Simultaneously, as $\kappa\to\infty$, we have that due to the concentration of the softmax probabilities on the largest entry, which here is $1+\alpha^{2}$, it holds that as $\kappa\to\infty$ we have $\mathrm{dist}(\kappa)\to 0$.
At the same time, the function approaches infinity from the positive regime.
For large enough $\kappa$ we have $\gamma_{\ell}^{\mathrm{upper}}(2)<\gamma_{\ell}^{\mathrm{upper}}(1)$.
When we select the smallest possible $\kappa>0$, we avoid computing functions with large exponential terms.
To find the non-trivial root of $\mathrm{dist}(\kappa)$ numerically, we consider a simplification of equation 82.
We define $u=e^{\kappa}$. Then it holds that we can solve

$$ $\displaystyle{\mathrm{dist}(\kappa)=0=(L-1)u^{(1-\epsilon)}-u-(L-2)}$ (88) $$

numerically for $\kappa>0$.
This shows that we can find an explicit construction with 100% accuracy with $p=1$ and $d^{\prime}>2$ for the bos+sftm when we have

$$ $\displaystyle d=\lceil\log_{2}(T+1)\rceil+2\,.$ (89) $$

For example, for the case of $L=10$ and $T=32$ this allows for a dimension $d=7$ with $\alpha=0.01$ (and for $T=31$ with the same settings $d=6$ suffices).
∎

Remark ($d=4$). In principle, it is enough to have some $\epsilon>0$ that ensures that overlaps between different token embeddings are strictly less than one. In principle, we can find an arbitrary number of tokens $T$ that satisfy this condition for just $d^{\prime}=2$.
Take for example the following construction. For $t=1,\dots,T$ tokens with $T$ odd we can design the set of embeddings

$$ $\displaystyle v_{t}=\begin{bmatrix}\sqrt{\frac{t}{T}}\\ \sqrt{\frac{T-t}{T}}\end{bmatrix}\,.$ (90) $$

Each $\langle e_{t},e_{t}\rangle=1$ and for $t\neq s$ the overlap $\langle e_{t},e_{s}\rangle\leq\langle e_{(T+1)/2},e_{(T-1)/2}\rangle=\sqrt{T^{
2}-1}/T$.
This implies that $\epsilon\to 0$ as $T\to\infty$ at a rate $1/T$. Since smaller $\epsilon$ imply larger values of the temperature to solve equation 88, this might become problematic when this exceeds the accuracy of computations.
Previously, for the binary representation construction from equation 72, we had that $\epsilon$ shrinks at a rate $\sim 1/\log_{2}(T)$.
For the intermediate regime between $\log_{T}(T)+1$ and $\log_{2}(T)$ dimensions, one can generalize this principle to arbitrary bases, e.g. $\log_{3}(T)>2$, resulting in a smaller dimension but also less favorable (smaller) $\epsilon$ – this construction thus comes with a clear trade-off.

###### Proof of Proposition 4.5 - bos+sftm .

#### B.3.2 ( dot+sftm ; p = T 𝑝 𝑇 p=T italic_p = italic_T )

###### Proof of Proposition 4.5 - dot+sftm .

For this model, the explicit construction is analogous to the previous one. Instead of using $p=1$ we use $p=T$.
The selection of the embeddings is analogous, but instead of a counting direction we read off all the weight directions separately with $T=d$. Not having a counting direction also saves the additional two dimensions required for bos+sftm with $p=1$. In the feed-forward layer with $W_{1}$ the explicit construction considers again $z_{\ell,t}$ for every token $t\in\mathcal{T}$.
The selection of the temperature is also analogous, with the exception that one has $L$ terms in the softmax instead of $L+1$.
∎

###### Proof of Proposition 4.5 - dot+sftm .

## Appendix C Data Generation

Every sample $\mathbf{x}=(x_{1},\cdots,x_{L})$ is generated recursively as follows, starting from size $K=L$ and alphabet $\mathcal{T}^{\prime}=\mathcal{T}$:

- 1.
Sample an integer $k$ uniformly from $[1,\cdots,K]$.
- 2.
Sample a token $t$ uniformly from $\mathcal{T}^{\prime}$.
- 3.
Set $x_{i}=t$ for all $i=k,\cdots,K$.
- 4.
Set $\mathcal{T}^{\prime}=\mathcal{T}^{\prime}\setminus\{t\}$ and $K=k$.
- 5.
If $K\neq 0$, repeat from 1.
- 6.
Set $\mathbf{x}=\texttt{shuffle}(\mathbf{x})$.

In contrast to sampling the elements of each sequence uniformly at random from the alphabet, this simple strategy enables us to better control the distribution of counts in the training dataset.

## Appendix D Additional Experiments

### D.1 Best Accuracy

In Fig. 9, we show the best reached accuracy during training over the five sample runs.
This gives insights into the feasibility of implementing a counting solution for a given combination of parameters $T,d,p$ of a model.

Figure: Figure 9: Experiments from Fig. 1 ($T=32$), we show only the best accuracy during training reached from the 5 randomly initialized runs per model/hyperparameter configuration.
Refer to caption: extracted/5909847/figures/memorizing_counts-phase-diagram-T32-best.png

### D.2 Variability

In Fig. 10 we explore the influence of initialization on the performance via the variability of the final accuracy for several runs. Especially in the $p,d<T$ regime where bos+sftm is able to reach an accuracy relatively close to 100%, the variability of the accuracies resulting from different initializations is quite large.

Figure: Figure 10: Experiments from Fig. 1 with $T=32$, standard deviation of the accuracy reached after training from the 5 randomly initialized runs per model/hyperparameter configuration.
Refer to caption: extracted/5909847/figures/memorizing_counts-phase-diagram-T32-std.png

### D.3 Model with Random but Fixed Embeddings

In Fig. 11, we repeat the experiments of Fig. 1, but for embeddings that are frozen throughout training (also 5 runs).
In the regime $d<T$ where there is no mutual orthogonality possible, the random embeddings result in worse performance than the learned ones. Especially for bos+sftm, learning the embeddings increases the performance strongly in some regimes. This indicated that the models indeed learn adapted embeddings here.

Figure: Figure 11: The difference between learned and random embeddings for $T=32$. Orange indicates that the random embeddings perform better on average. Purple indicates that the learned embeddings perform better on average. Experimental settings as in Fig. 1.
Refer to caption: extracted/5909847/figures/memorizing_counts-phase-diagram-T32-compare.png

### D.4 BOS mixing token

In Fig. 3 in the main, we describe how the $t_{\mathrm{BOS}}$ is the main predictor for the count. Here, we provide more evidence by showing how the count predictions for mixed tokens $\bar{x}^{\prime}$ output by the feature transform $f$ are invariant to the type of other token present in the mixed token.
The results for four different tokens are shown in Fig. 12.

Figure: Figure 12: For the same model as in Fig. 3, we vary the inputs to the feature transformation $f$ to show it is independent on the precise input sequence, but only depends on the prevalence of $t_{\mathrm{BOS}}$. We vary the inputs between the learned tokens $[B,C,D,E]$.
Refer to caption: x6.png

### D.5 Singular Value Decomposition of W 1 subscript 𝑊 1 W_{1} italic_W start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT

In Fig. 13 we show the distribution of singular values of $W_{1}$ for several runs of the model to investigate whether models that are capable of both IC and RC are implementing the more memory heavy IC or the same solution that they can find for $p=1$ with RC.

Figure: Figure 13: Singular values of $W_{1}$. We show the results for all models from Fig. 1 with $T=32$, where $p,d\geq T$ and the accuracy is at least 99%. Some qualitative differences are visible for bos and dot.
Refer to caption: x7.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: See weakness.
soundness: 2
strengths: The author studied an interesting topic about the relationship between the attention and the feedforward layer.
summary: In this paper, the author studied the relationship between attention and the feedforward layer through the task of histogram task.
weaknesses: 1. The paper’s narrative could use some improvement for better clarity and flow.

2. This paper studies the relationships among components in transformers within the context of histogram tasks. However, there isn’t enough evidence to show that this empirical study on histograms alone can adequately represent the properties of transformers in general. The cited work, Thinking like Transformers, also uses histograms but includes reverse sorting and other tasks to provide a broader context.

3. Another concern is the significance of these findings—why are they important, and do they provide meaningful insights?

4. There’s little discussion on how these findings could guide the adjustment or application of transformer structures in practical contexts.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: Could you explain in detail what the "online learning setting" in Line 153 means?
soundness: 3
strengths: 1. This paper demonstrates how various components of the transformer affect the model's learning ability, which can help understand the internal mechanisms of  transformer

2. The authors provide a simple yet effective testing environment, making it possible to evaluate the impact of the model architecture on solutions more clearly.

3. The visualizations are well done, making it easier to understand.
summary: This paper explores the performance of small Transformer models in performing counting tasks (specifically histogram tasks), analyzing how the collaboration of different components in the model architecture affects the final solution.
The authors found that in the histogram task, the predictive performance of the model is closely related to the vocabulary size, embedding dimension, token mixing mechanism and the capacity of the feed-forward block. Small Transformers can implement two counting strategies, the choice of which is influenced by hyperparameters and synergies between components. Even small architectural adjustments can lead to significant changes in the solutions learned by the model.
weaknesses: 1. The paper analyzes only a simple single-layer transformer structure, while scaled-up models in real applications may exhibit different behaviors. Additionally, this paper overlooks the impact of the autoregressive mechanism and positional encoding.

2. Although the paper reveals the impact of architectural details on model performance, it lacks clear guidance for practical applications. For instance, it analyzes how relation-based counting, inventory-based counting, and softmax significantly influence the model's learning strategies and performance. However, I find it challenging to relate these methods to transformer models in real-world applications. I believe that adjustments in model design often need to be validated in practical settings, yet the paper does not provide specific application examples or optimization suggestions.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Have you done any analysis or insight you can share on: 
1. what behaviour can we expect in terms of the same hyperparameter choices (embedding dimension, Hidden Layer Size, and Vocab Size) whether as a  2 or 3 layer blackbox and/or an isolation of one of the transformer layers in a multi-layer setting?
2. Have you done any initial analysis to share about whether the same relationships for parameters and potential solving strategies hold for other tasks (such as the sorting or lookup problems you referred to or any other tasks)
soundness: 3
strengths: * The paper is well-written, easy to follow, and provides sufficient analysis for its claims.
* Discovers and Discusses two strategies which 1-layer transformers can use to perform the counting task (i.e., relation-based and inventory-based) and in doing so considers a range of possible hyperparameter selections that affect this strategic choice.
* During its analysis further studies the role of BoS token showing that together with softmax in the counting task, can help models with smaller p and d sizes to still reach perfect accuracy.
summary: This work investigates how different components of transformer architectures contribute to solving the counting task (creating a histogram of token occurrences in a sequence). The key research question is: how do architectural choices like embedding dimension, hidden layer size, and mixing mechanism (dot product attention vs linear) affect the solutions a transformer can learn? The authors discover and theoretically back up two distinct strategies transformers can use: (1) relation-based counting (which is more efficient but requires specific architectural features) and (2) inventory-based counting (which requires more parameters but is more flexible) - and show how seemingly minor architectural choices like including softmax in attention can dramatically affect which strategy emerges. They map out exactly what parameter regimes allow each strategy to work, including surprising findings about how transformers can still function even when the embedding dimension is smaller than the vocabulary size.
weaknesses: * Limited scope and applicability, specifically: (1) Only examines single-layer transformers, (2)Focuses on just one task (histogram counting)
* No clear path to extending insights to more complicated settings more specifically multi-layer transformers
* Lack of practical impact in the sense that it is hard to find a direct translation from this paper findings to real-world transformer design for real-world tasks

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: - In your opinion, what do these findings tell us about the hypothesis space of transformers?
soundness: 3
strengths: - The explanations, logic and experiments are sound
summary: In an attempt to improve the interpretability of different elements of transformers, the authors use transformers to solve a histogram task. They incorporate two different counting strategies (relation-based and inventory-based), and replace different components of transformers in a controlled setting to study the relative role of each component. They find that both the attention mechanism (without applying softmax) and a sufficiently large FFN are capable of solving the histogram task. They also find that this task is sensitive to hyperparameters.
weaknesses: - If I understood correctly, the main contributions are that for the histogram test: 1. using the attention mechanism but without the softmax step (relation-based counting) is capable of high accuracy, and 2. Using learnable-matrix multiplication instead of attention (inventory-based counting) is not capable of high accuracy, but using a FFN can mitigate this by acting as a lookup table. It would be helpful if the authors would elaborate on why these findings are important.
- It would also be helpful if the authors provide some intuition on why they chose the histogram test. To be specific, transformers are designed for prediction and regression tasks. It's not clear why the authors chose to study a task that diverge from conventional transformer use cases. Perhaps there are interesting theoretical or practical insights that motivated this direction?

Meta-review: The paper investigates how different components of the transformer architecture, e.g. feedforward and attention layers, affect performance in counting tasks, specifically a histogram task in this paper. The reviewers appreciated the analysis of the learning ability in simple to understand scenarios and appreciated that the paper contains well done visualizations and structure. However, there were several shared concerns on i) a potentially limited analysis given the focus on single-layer transformer architectures, ii) limited discussion on how to leverage insights and lack of clear guidance for practical utility, iii) a too narrow investigation of just the histogram task. 
The authors have mainly provided clarifications in the discussion period, many of which are textual in nature, and the updates made to the paper seem to be rather limited (there seems to only be a single sentence that changed). Respectively, some reviewers have raised their score slightly, whereas others remained slightly less convinced. Overall, the reviews all remain at borderline ratings after the discussion period. The AC believes that some of the aspects discussed and outlined are initial good steps to further improve the paper. However, the AC also agrees with the sentiment that the paper would benefit from thoroughly addressing above three points in a revised paper version. The AC thus recommends to reject the paper in its present form.

Decision: Reject