Counting in Small Transformers: The Delicate Interplay between Attention and Feed-Forward Layers
Authors: Freya Behrens
Luca Biggio
Lenka Zdeborova
Venue: ICML 2025 poster (2025)
Keywords: learning theory
representation learning
algorithmic tasks
attention
associative memories
implicit bias
learning regimes
Primary area: deep_learning->theory
TL;DR: We show that transformers adopt distinct counting strategies, relation-based and inventory-based, which shape learning regimes and are influenced by architectural choices, impacting performance and robustness in a simple histogram task.

Abstract: Next to scaling considerations, architectural design choices profoundly shape the solution space of transformers. In this work, we analyze the solutions simple transformer blocks implement when tackling the histogram task: counting items in sequences. Despite its simplicity, this task reveals a complex interplay between predictive performance, vocabulary and embedding sizes, token-mixing mechanisms, and feed-forward layer capacity. We identify two theoretical counting strategies transformers adopt, relation-based and inventory-based counting, each defining distinct learning regimes for the task. These strategies dictate how functionality is distributed between attention and feed-forward layers. We further show that adding softmax and beginning-of-sequence tokens allow for more robustness when embedding dimensions are comparatively small. Empirical introspection of trained models closely confirms both the learning regimes of the various architectures and the formation of these strategies during training. We demonstrate how a basic task that requires only aggregation and selection is significantly impacted by minor design changes.

Counting in Small Transformers:
The Delicate Interplay between Attention and Feed-Forward Layers
Freya Behrens
Luca Biggio?
Lenka Zdeborová
Abstract
been 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. (2017). At its core, a single trans- former block primarily alternates two basic components: the coken-mixing attention mechanism and a standard fully con- nected multi-layer perceptron. At a high level, the attention mechanism mixes the tokens, while the multi-layer percep- tron applies 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 al., 2021; Dordevic et al., 2024; Gu & Dao, 2024). 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., 2021; Delétang et al., 2023; Abbe et al., 2023; Ouellette et al., 2023).
Next to scaling considerations, architectural de- sign choices profoundly shape the solution space of transformers. In this work, we analyze the solu- tions simple transformer blocks implement when tackling the histogram task: counting items in se- quences. Despite its simplicity, this task reveals a complex interplay between predictive perfor- mance, vocabulary and embedding sizes, token- mixing mechanisms, and pacity. We identify two theoretical counting strategies transformers adopt, relation-based and inventory-based counting, each defining distinct learning regimes for the task. These strategies dic- tate how functionality is distributed between at- tention and seed-forward layers. We further show that adding softmax and beginning-of-sequees tokens allow for more robustness when embed- ding dimensions are comparatively small. Empir- ical introspection of trained models closely con- firms both the learning regimes of the various architectures and the formation of these strategies during training. We demonstrate how a basic task that requires only aggregation and selection significantly impacted by minor design changes.
layer ca-
In this work, we investigate this question for algorithms that compare and aggregate information from mechanistic per- spective (Cammarata et al., 2020: Olah et al., 2020; Elhage et al., 2021; Michaud et al., 2023; Quellette et al., 2023) and focus on the histogram task (Weiss et al., 2021) as a prototypical problem. It consists of predicting the number of appearances of each token in the input sequence pro- cessed by the model - counting and solving it requires both comparison and aggregation. Even modern language models with up to 8B parameters currently fail to solve this task robustly and efficiently in-weights, rather than via chain-of tf--uhhthtt (Appendix A). This failure, despite the task's apparent simplicity, motivates a study of the relative role of different architectural components and their impact on the final solutions implemented by the model in con- trolled setting. To this end, we focus on models following the architectural template of primitive transformer blocks, i.e. alternating token-mixing attention mechanism and a mulli-layer perceptron.
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., 2021; Girdhar et al., 2019; Brown et al., 2020). While scale has
'Statistical Physics of Computation Laboratory, icole poly- technique fédérale de Lausanne (EPFL). Lausanne, Switzer- land 'Department of Computing Sciences. Università Boc- coni, Milan, Italy, Correspondence to: Freya Behrens <freya.behrens @ epfl.ch>.
In our analysis, we provide explicit theoretical constructions (parameter configurations) for a range such architectures reaching perfect accuracy in a model-dependent hyperpa- rameter regime. In a subsequent step, we compare these
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).

---

Counting in Small Transformers: Attention and Feed-Forward Layers
algorithms with the performance and mechanistic behav- ior 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 de- pendence on the scale of the model's hyperparameters and the type of token-mixing mixing mechanism utilized. Our main contributions are as follows:
the inherent permutation equivariance of the histogram task. We refer to the vector Fyy for each position é = 1,... I,, as the mixed token. Note that we assume that all operations in the network are executed with infinite precision. We comment on this assumption when it becomes problematic.
Token Mixing. We consider two types of mixing mech- anisms A with different activation functions and a single head. We refer to the case where the function A is con- stant in X as linear mixing (lin), e.g. Ala (x) = A and A titm (x) softmax (A), where A E LR is learn- able matrix and the softmax operator is applied row-wise. The number of learnable parameters is therefore L
We identify two algorithmic strategies for solving the his- togram task: relation counting, which uses local pairwise token comparisons, and inventory-based count- ing, which counts all tokens in the alphabet and extracts the desired count at given position.
As an alternative mixing structure, which we refer to as dot- product mixing (dot), we consider the popular attention mechanism which constructs the matrix A to be explicitly dependent on the inputs, i.c.
We clarify how the emergence of these strategies depends on the architecture specific inductive biases. Relation- based counting is memory and compute-rfficient leverag- ing attention for comparisons. Inventory-basee counting relies on token-mixing for aggregation, but uses feed- forward layers to implement a comparison to the com- plete alphabet with higher computational demands.
1 Adot = vd
XWO
(2)
and Adot+aftm (x) = softmax (Aduc where We and k are learnable d x d matrices. Note that, without loss of generality, we assume the value matrix to be the identity. Then the number of parameters for dot-product mixing is 2d In line with previous work (Weiss et al., 2021), for architectures employing the dot-producd mixing, we also an- alyze models utilizing the so-called beginning- ooating-o--uqquune (BOS) token. This special token, indicated with the symbol $ is appended to the original input x resulting in a new sequence x (S, x1, I2, ... .I) of length I. 1. We will refer to the architecture that includes the BOS token as bos.
We show how the emhedding size required for perfect solutions relates to alphabet size and the maximal se- quence length. Due to the discrete nature of the task, near-orthogonality is sufficient in many cases. Gains in prediction accuracy compared to models of the same capacity can arise from the softmax operator and dot- product attention reducing noise from linear dependence.
Our code is publicly available at https//giitub.coo/SPPCC group/counting-attenn attentioo
2. Background and Notation
Feature Transformation. The feature transformation is a single hidden layer perceptron with ReLU activations. The hidden layer is of dimension p. The function ʃ is applied identically to every mixed token Ty for L' = 1,... L. as:
Architecture. As inputs, we consider sequences of to-
kens x = (0i, I2, ,T)) E T- from the alphabet T = {1....,") ,T). A sequence of outputs is y = (y1.... 99)) has the same length as the input sequence, where each out- put token y E (1,...,C,, with C S L. We analyze several 1-layer model architectures where token-mixing mech- anism is followed by per-token feature transformation. Formally, we consider a models F : T + C³ defined for the positions& = 1, 1.... ,L as
f( = ( + by W + bo
(3)
where f(E)) Ed + RC and where the weights have
the appropriate dimensions to accommodate a hidden layer of size p. i.c. W1 E Rddp,br E R", W2 E R2xC and by E R..
F(x)= = arg max f(z()o: f/= 3 (A(xxll (L) oE (E(1,..)))
3. Experimental Setup
with the token mixing matrix A : Rixd + R and the token-wise feature transformation f RRd + C The embedding is x E R where āe denotes its E-th row, is obtained by passing the input sequence into a standard embedding layer (lcarnable lookup-table) of dimension d. We refer to the embedding associated with token E T as e, E Rdd or exy E Rd for the embedding of the token ze at position& l. We do not include positional embeddings due to
Task and Dataset. We consider simple algorithmic task that is referred to as histogram: given sequence of to- kens, the goal is to return a sequence of the same length where each entry represents the number of times the corre- sponding input token appears in the entire sequence. For example, given x = [A, B, D, D, B, B), the output will be y = [1, 3, 2,2,3,3, 3]. We define the count of a token t in the sequence x at position as histx((). (f). In our experiments,
2

---

Counting in Small Transformers: Attention and Feed-Forward Layers
consider i..d. distributions of sequences of length L from input alphabet of size T, where I. I T. Our pling strategy relies first sampling set of partitions, and then assigning token to each partition (see App. D for d This allows for a close to uniform distribution over the values y.
such mutually orthogonal embeddings e, IR with norm of the overlap (ess ex) for distinct tokens & 8 and it is when = In such scenario, linear combination of token embeddings preserves magnitude (count information) about elements from the alphabet. A weighted sum of to- kens, denoted as = Σter aeet, can be broken down into the original tokens using projections on original token embeddings, where at = (er,e') [exl//
Models and Training. We measure the performance the histogram task of the four different variants of the token mixing models described Sec. 2, i.e. 1 iin and dot, with or without the softmax (+sftm). The token embeddings are jointly learned with the model parameters. We consider the dimension of the embedded tokens d, and the hidden layer size p of the feature transformation. We also consider the model bos(( sftm) where every imput sequence is prefixed with the BOS token prior to entering dot-product mixing layer (with softmax).
n10
10
Embeending
Models are trained with Adam with learning rate 10 -3 with cross-entropy loss for 500 epochs with batch size of
32. We consider the online learning setting where for each new epoch generate dataset fresh 10, 000 samples.
[dot] 10a
The accuracy is computed from 3, 000 independent samples.
embedding
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 of the feed-forward module.
[Lin] 102-
Tambeeddingg d
Fig. 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 the architecture. Fig. 2 further clarifies that the parameter efficiency under different architectures varies substantially. To investigate the underlying mechanisms we devise thcoretical 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 T). For hidden layer size p we distinguish the regime where models can sense only constant number of directions/features (p = 1) or one scal- ing as the alphabet size (p = T). In App. E.8 we show that a similar phenomenology persists for 2-layer transformers.
50
100
strong dependence
mean accuracy [%)
10'
10²
10'
10
hidden layer
hidden layer size
Figure 1. Accuracy on the histogram task for different I-layer ar- chitectures. Mean accuracy for varying embedding size d, hidden and
layer size p, for fixed ken mixing mechanisms dot, bos and 1in. (Left) Models with softmax; (Right) Models without softmax. Average over runs for every d,p 2,, 4, 6, 8, 12, 16, 23, 32, 45, 64, 91, 128). Vertical and horizontal white lines indicate p = and = respectively. White stars (dots) mark 100% (> 99%) accuracy configuration was found for at least one of the five runs.
the different to-
4.1. d T: Orthogonal token embeddings are separable
In the following, we use this property to theoretically con- struct the weights for all models that solves the task when d 2 Remarkably, the constructions require different numbers of hidden neurons p depending on the mixing
When the model dimension d equal or greater than the alphabet size T. tokens can be represented by embeddings that are mutually orthogonal. Assuming for there are
3

---

Counting in Small Transformers: Attention and Feed-Forward Layers
mechanism. This demonstrates the interplay of the mixing layer and the feature transform: for some mixing mecha- nisms, the latter needs to implement inventory-based count- ing (IC) (requiring p 2 T), and for others, relation-basee based counting (RC) (where p 2 1 is sufficient).
+sftn
no S sft
100
80
accuracy [%6
60
4.1.1. RELATION- BASED COUNTING: LEVERAGING DOT-PRODUCT MIXING
best
Lin
dot
When an extra gegiiningg-f-sequeeeen token (BOS is avail- able in bos, it can be used to extract information about token's ss count hi) in the attention layer of the net- work through its attention score (Kazemnejad et al., 2023). In the literature, the beginning (or end) of sequence to-
BO5
10a
10a
10
number of parameters
number of parametens
Figure 2. Test accuracy vs. total number of learned parameters. The data the same as generated for Fig. I, every data point the single experiment and we show the convex hull in solid lines.
kens have been linked to model-internal computations, such as counting. In (Weiss et al., 2021), 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 + s fm reach (close to) 100% accuracy whenever d > T, and we verify that relation-based counting algo- rithm can be theoretically implemented in these two archi- tectures by construction.
its activation can be interpreted as a proxy for the number of occurrences of zg, as 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. how the presence of the łвоS determines the final prediction through the application of f. Surprisingly, the dot model (without the softmax) reaches an empirical performance comparable bos in the regime d T and p 1, even though it does not have an extra token available.
Proposition 4.1 (RC with BOS token). For bos and bos+s ftm and a given L Z 2, there each exists a con- figuration of weights that solves the histogram lask at 100% accuracy, given that d 2 T 2 and p
We prove this by construction in App. B.2.3-B.2.2 and we provide the intuition of the proof in the following. For bos we set the /nos embedding to éBos = Σier es and take the mutually orthogonal token embeddings e, to have norm Assuming that too 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 Tx) will lead to attention matrix with entries:
Proposition 4.2 (RC with tagged embeddings). For dot and a given I,, > 2, there exists configuration of weights that solves the histogram task at 100% accuracy, given that d 2 T > 2 and 1.
T if m =1
We prove this in App. B.2.1. Intuitively, the construction uses single common direction Euit that is added to the otherwise mutually orthogonal token embeddings. A dot- product mixing then leads to aêm = az > O when æy is different from 2im. and agm = a= > when tokens are the same. Then, the number of counts can be easily extracted from the dot-product (Pem; .y)) of the counting token with the mixed token 3/ i.e. (eens, H/)) x 1+ histx (f)ax (L- histz(())aa4 We can, therefore, obtain perfect accuracy implementation in the regime where d > 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.
apm
1 if (f 1) 1) or (f,m m. > 1, = Tm)
0 if f. > 1, + Im
Projecting the mixed token Ze onto the lBos we obtain (T', <Bos) = T histx(() + 1, i.e. ébos is the single relevant direction for the prediction. Its magnitude relates linearly to histx((). A single hidden neuron p = 1 suffices and the output layer can transfer the count into a categor- ical representation. For bosts fsmm one needs to further account for the non-linearity of the softmax as described in App. B.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 boosssftm,, that tgos indeed plays a special role in the learned model: in the attention matrix
Dot-product attention with softmax fails to implement relation- -ased counting. Since the dot-product mecha- nism can naturally be used in relation- based counting, one might expect the dot +s ftm model to implement the same mechanism. However, and maybe surprisingly so, we em-
4

---

Counting in Small Transformers: Attention and Feed-Forward Layers
token embeddings
attention matrix
feedforward prediction
1.0
0.25
overlap for C Twith
Go
10
10
10
10
B
11
11
flaesos + (1 a)ep + es)
same (ey, et)
46
10
10
t6
10
10
0.20
different v- (e)
B
66
20
1
20
10
10
0.8
BOS (h, fsass
C
dansty
0.15
(Ovos, Gecel
10
10
grediction
6
10
0.10
20
10
0.4
3
-02
1
60
D E E
0.2
0.4
1.0
overtap
Figure 3. Relation-baut based counting with bos+sftm (T' = 32, L = 10, = 2, =
= 45). This model achieves 99.9% accuracy. It was selected as the best model from all our experiments with p = (Left) The tokens overlap (cosine similarity) with the same tokens, different tokens and the BOS all concentrate around different values. (Middle) This is reflected in the attention matrix after the application the row-wise softmax. The [BOS ('S') in the first column ar,0 becomes proxy for the count of zz. (Right) To demonstrate that the feedforward network is only sensitive this direction, we show its count predictions for a mix of tokens and fix = aéBos + (1 aD)e + eg. The contribution a of the BOS token to the intermediate n' is varied and D, B are two specific elements of the alphabet T. represented by their embeddings eD and ep. The y-axis shows the predicted class of the feedforward layer f(x') for given a. We measure ag,o, the actual contribution of the BOS token's for the sequence in the middle plot for the letters that occur 1, 3 and 5 times. The same experiment is repeated for different elements of the alphabet in App. E.4, showing independence of the count prediction on the other tokens present in the sequence.
pirically observe a marked difference between dot and dst+ in Fig. 1. dot+s only starts performing close to 100% accuracy when both the model dimension and the number of hidden neurons p are larger than the number of tokens T. To understand why it fails to learn for p show the attention matrix of ts f.mm in Fig. 4. Notably, it is based on the semantics, as +sftm) Im is higher when xy = Zmm. than otherwise. However, the normal- ization effect of the softmax activation prevents the develop- ment of a meaningful counter subspace that is needed in the relation-based algorithm. As a result of normalization, the attention scores are Σm dmm = 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 bosts tasssftn ffm circum- vented 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 fmm, the architecture fails to perform well for p 1 - it now needs to measure more than one direction in the foed-forward module.
4. 4.1. INVENTORY-BASED COUNTING: MEMORIZATION IN THE FRED-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 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 lintsftm and a given L,T > 2 there exists a configuration of weights which solves the histogram task for p 2 ' and d T.
We describe examples of such constructions in App. B.3.1 and B.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 Ajin we set a constant value a = 1/L so that the result of the mixing is simply position-independen linear combination of the input. The count (f) be extracted after the residual connec- tion where we add T' = 3, + exç- By setting the columns of the matrix (Wi), = es we can extract the count information up to the factor 1/a
In the following, we show that a solution of the histogram task can still be achieved through an inventory-basse count- ing algorithm with 2 T. We detail this in the follow- ing section, for the example of 1in. The statement for dot+s fmm is given in App. B.3.
hist(() x
Σ ReLU (67, -1)
a
EE
1
ReLU ((T), - -1) (āz, ezc)
tET
Note that, due to the 1 bias term, only the hidden neu-
5

---

Counting in Small Transformers: Attention and Feed-Forward Layers
resid. = 6g
resid. = 6c
resid. = @p
1.00
[dot sftm]
0.75
3
3
B
RB 20
0.50
D
23
D
20
0.25
2
D
29
0.00
EBCBBBDEEE
CBC
A
Fncer resid.
T= (8,C.D)
argmax: fe
[Lin+sftal
0.75
ap= 1p=1-aa-00
0.6
0.50
0.4
M
0.2
0.25
0
0.0
0.00
8 E
E D
0.0
0.5
1.0 0.0
0.5
1.0 0.0
0.5
1.0
da
as
Figure 4. haventory-hased counting with dot fot++ffm (p= 32, d 32) and Tin+s (p= 64, d = 64)). We have T = 32, = 10. The models achieves 99.4 99.47% and 100% accuracy respectively. (Leff columns) The attention matrix for two different sequences. It differentiates between same (red squares) and different tokens for dot + fmm but is invariant to the semantics for 1in. For dot do ftm. any counting direction that could emerge in token space is not informative due to the softmax normalization, so p 2 1' is required (see Fig. 1). (Right columns) We test the feed-forward layer f in isolation, by feeding it with an artificial mix of learned embeddings for the three tokens B, C and D. We plot the class predicted by f for inputs constructed as ø' aggB + agec + aped R. where = as + ac o and we change the residual R E {en, ec, ec,en) } from left to right, for cach of the models dot + ftm and 1 inn+sft The prediction strongly depends on the coefficient ar associated with the token 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.
ron for token (W1), ez = Iz that occurs in the residual connection has non-zero activation. The output layer W2 can then be designed to activate the correct output vector corresponding to the count (f) (see App. B.4). Since E [0. 1] and Σn Abm = 1 the same procedure can be implemented by a matrix which is passed through the softmax operator for 1 lin+sfe f.m. In practice, in this con- struction the feed-forward module correlated with the complete alphabet. acting as an inventory, or look-up table.
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 capac- ity being increased to p = T when inventory-based counting can in principle be implemented. In App. E.5 we investigate the singular value decomposition of W1, for learned models with 99% accuracy and p, d 2 T. We find that the largest T 32 singular values are larger than the surplus singular values when > T for models that can implement only IC. This behavior is less pronounced for models that can imple- ment 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, if it is superposition of the two, thus requires further investigation.
In Fig. 4, we inspect the attention matrix Alin and the feature transformation f which is learned for lints ftm in the regime where p ~ T' ~ d. The mixing has an off-diagonal of ~ 0.09 and a diagonal of ~ 0.16. 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 a corresponding to the token embedding fed through the residual connection. This behavior is reflected in 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.
4.2. d < 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 dif- ferent 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.
Superpositioned and selective implementations. Some of the models capabilities include one another. For example, the models that can implement clation-based counting for p = can also implement the solutions for sreetory-bbsedd counting for p 2 T. It is unclear, whether the memory- intensive solution is preferred when the memory is available,
This phenomenon is illustrated for dot in Fig. 5, where learned models with smaller d tend to overcount items in
6

---

Counting in Small Transformers: Attention and Feed-Forward Layers
the input, and observe a less spread distribution of overlaps. Nevertheless in Fig. I 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 predic- tion inherently more robust to the effect of noise stemming from entangled embeddings. This concept is illustrated in Fig. 8 in App. B.4 for the dot fm 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.
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 [log2(T + 1)] 2 d
In the following, we link concepts on optimally placing decision boundaries for noise robustness to characteriza- tion 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, 1974). This gives a means to understand the size of a given task with T, L requires at least.
Put simply, this construction requires that there are token embeddings for 6,8 = 1,..., T and 8 + with e >> 0 such that (ez, ex) = 1 and (ec, < 1 - c. This is fulfilled when every token is the binary encoding of its value, modulo mi- nor modifications due to the RC mechanism for bos s fmm. 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 re- quire 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.
Proposition 4.4 (Robustness via bounded mutual coher- ence). Given L 2 5, 2 2 and assuming that the Welch bound is attained for a given T, d, there exists a construction that solves the histogram task with
(1in, lintsftm p= T):
T()2-3)2
S d,
7-1+(2L-37) 37
(dot, bos; 1):
3))
+ +15a S d.
7-1+(21--3) 3)
(dot, bos: p T):
T()-1) (I--
S
In App. C.3.1 we show that this bound can be even further improved for bos a fm to a constant d = 4, but at the cost of increasing the temperature further as function of T, in addition to I,, requiring yet more accurate computa- tions. While in practice we are able to hand construct and code solution for the L,, T under consideration (available in the supplementary code), these models are not learned. However the very large numerical precision required might lead to learning difficult learning dynamies and be the rea- son why we do not observe any learned solutions of the histogram task in this regime.
We provide additional background and the proofs in App. C.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 given I. and T this provides an upper bound on the maxi- mal mutual coherence that is tolerated for 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 2 29, 30, 7. Gen- erally 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 con- struction for dot and p- = T which attains d = 12, provided in the supplementary code and in correspondence with Fig. 1. While this bound does not reach the d as in- dicated by the Welch bound, the mutual coherence of the embedding matrix we use is close to the maximally allowed value of M- 0.299 < 1/3.
4.3. Extensive T, L ablations and two layer architectures
Appendix E provides a number of ablations for T 15, 64 and L 5, 15, 30, for different architectures and hidden layer size p and embedding size d. Our predictions from the theory are largely consistent with the performance for different alphabet sizes T and small L, similarly to T = 32 and L = 10. However, for very large L = 30 it seems that the keeping the training budget constant lead to a weaker performance in the feasible regimes.
7

---

Counting in Small Transformers: Attention and Feed-Forward Layers
3=3
8=
d= 23
perfect explicit constructions and link them to introspection learned models. In concurrent but independent work to ours, Yehudai et al. (2024) explore very similar version of the histogram task, and find, as we do, that the numerical precision limits the counting size. While they do no examine different architectures, they find rigorous upper bounds.
8
57.71% acc.
87.67%
8- 98.28%
presictio
6-
4-
10-2
10-2
10
ground truth
ground truth
ground truth
2.0
9)
0.06
Memorization and Feed-forward Layers. The role of layers as memorization modules is investi- gated in the context of factual recall for language mod- els (Geva et al., 2021; Meng et al., 2022; Chughtai et al., 2024). Henighan et al. (2023) study a double decent phe- nomenon where the purpose of the feed-forward layer tran- sitions from storing data points to discovering generaliz- ing features as function of increasing training data di- versity (Raventos et al., 2023). In the histogram task, we observe a similar phenomenon as a function of the architec- ture: the forwaforward layer acts either as a look-up table or a feature detector for for a counting subspace.
1.5
to.
0.03
0.04
0.02
0.5
0.02
0.01
0.0
0.00
0.00
-5
50
100
200
overtap
overtop
overlap
Figure 5. Introspecting the regime with Entangled Embeddings with dot (T = 32, L= p0 32). We show trained instances of dot with 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.
Aligning Algorithm and Architecture. While theoretical work has defined the computational capacity of several (au- toregressive) neural networks (Weiss et al., 2021; Yun et al., 2019; Delétang et al., 2023; Liu et al., 2023), hallucinations and failure modes on seemingly trivial tasks in real-world transformers are the rule rather than an exception. Dziri et al. (2023) postulate that this may be due to misalignment be- tween 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 found evidence of superimposed computational graphs a single model (Elhage et al., 2022); our models shed light on how such entanglement functions in detail for orthornho embeddings.
In the same section, we show that for learned 2-layer atten- tion blocks the general phenomenology for the single layer models from Fig. is very coherent. Only in the critical regimes along the transition from not possible to possible transitions, there are some larger differences in average performance. This indicates that 2-layer networks do not necessarily operate differently in what they learned.
5. Related Work
Counting and Interpretability. The emergence of algo- rithmic capabilities in transformers (Olsson al., 2022: Power al., 2022) has led to numerous investigations aimed at reverse-enginen -egiie-enginnnriig trained models into human- understandable mechanisms (Zhong et al., 2023; Nanda et al., 2023; Quirke & Barez, 2024), including a variety of counting scenarios (Gould et al., 2023; Chollet et al., 2020; Ouellette et al., 2023; Cui et al., 2024; Golkar et al., 2024). In our work, we consider the histogram task introduced via the RASP(-L.) programming language (Weiss et al., 2021; Abbe et al., 2023). While RASP predicts that the histogram single layer transformers with one head require a BOS to- ken (Nye et al.. 2021), we find that this is not necessary by providing constructions. While many works in interpetabil- ity focus on causal interventions (Vig et al., 2020; Meng et al., 2022) to understand the computational mechanisms of models or assign relevance scores to their components (nostalgebraist, 2020; Elhage et al., 2021), we focus the hyperparameter scaling of several distinct models in relation to their performance and explicit constructions of different algorithms, similar to e.g. (Zhong et al., 2023; Quirke & Barez, 2024; Nichani et al., 2025). We give precise theo- retical conditions on the model configurations that lead to
6. Discussion & Conclusion
Limitations. Similar to other works in mechanistic inter- pretability (Zhong et al., 2023), we focus on 1-layer trans- formers simplified model for modern transformers. Our models are not autoregressive and do not account for the im- pact of causal masks or positional encodings. While more complex models could lead to more intricate interdepen- dencies between the components, potentially limiting the applicability of our findings to such architectures, it seems plausible that similar vector arithmetic could emerge in sub- spaces of large transformers (Gould al., 2023; Engels et al., 2025). Given its specificity, it is unclear if and how similar memory-architecture phenomena would emerge for different simple tasks (e.g. sorting or lookup). However, the learning regimes seem to transfer to two layer models, and even large auto regressege language models fail to reliably perform well on this task, hence motivating the study of
8

---

Counting in Small Transformers: Attention and Feed-Forward Layers
algorithmic emergence in smaller models.
Askell, A., Agarwal, S., Herbert-Voss A., Krueger, G.,
Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I,, and Amodei, D. Language models are few-shot learners. CoRR, abs/2005. 2020. URL https: peattps:/aaxxivorrt abs/2005.1415
Summary. We study how different components of simple transformer models contribute to the emergence of different solutions to the histogram task. The feasibility of solving the histogram task including the comparison and aggre- gation operations depends on the mixing mechanism, its interaction with the feed-forward transformation, and the softmax function in the attention. Relation-based count- ing, uses a dot product mixing mechanism and low-capacity feed-forward transformation, and inventory-based count- ing, which memorizes token embeddings in feed-forward weights, requires more parameters. We characterize feasi- bility regimes the phase space of embedding dimension d and feed-forward dimension p. confirming that models converge these mechanisms. Some regimes allow both strategies, and experiments show features of superimposed mechanisms. When d < T, tokens cannot form an orthog- onal basis for linear projection. Models exhibit varying robustness to noise from von-orthogonaliiy. The softmax activation minimizes token similarity in attention, reducing non-orthogonality effects effects-relevannt relevant in real- world models where 7 often exceeds model dimensions.
Cammarata, N., Carter, S., Goh, G., Olah, C., Petrov, M., Schubert, L.,, Voss, C., Egan, B., and Lim, S. K. Thread: circuits. Distill, 5(3)e24, 2020.
Chollet, F., Tong, K., Reade, W., and Elliott, J. Ab- straction and reasoning challenge, 2020. URL https: / com wwtp////oo/ti/io//m/oocc abstraction -nntttiii--nnn------eee wewd---rsnnnnccalleeeee.
Chughtai, B., Cooney, A., and Nanda, N. Summing up the facts: Additive mechanisms behind factual recall in LLMs, 2024. URL https: http//ppeeeeeeeeeeeeeeeeeeeeeeee net/ forum? waiumidddPggnDDEGG
Cui, H., Behrens, F., Krzakala, F, and Zdeborova, L. A phase transition between positional and semantic learn- ing in a solvable model of dot-product attention. In The Thirty-eighth Annual Conference on Neural Infor- mation Processing Systems, 2024. URL https:// openreview net / forum? Berum?id=BEWer WEIIEEBWIWPP.......
Future Directions. Examples of hallucinations and fail- ures in LLMs are as numerous as their successes. Though limited to one task, our analysis highlights how subtle ar- chitectural changes can drastically impact predictive power. close
We expect that similar mechanistic investigations to the regimes where models start failing will be extremely useful to understand how and why models fail in sometimes puzzling manners.
Delétang, G., Ruoss, A., Grau-Moya, J., Genewein, T., Wen- liang, L. K., Catt, E., Cundy, C., Hutter, M., Legg, S., Veness, J., and Ortega, A. Neural networks and the chomsky hierarchy. /1th International Conference on Learning Representations, 2023.
Acknowledgements
Donoho, D. L. and Elad, M. Optimally sparse represen- tation in general (nonorthogonal) dictionaries via sup minimization. Proceedings of the National Academy of Sciences, 10065:21177-220,, 2003. doi: 10.1073/pnas. 0437847100. URL https:/ /ww..www pna.. wnw.sorg//aee.o..o.....a..aaa org//or/ abs 1..a/10.100 1071 10730..../3/0m.... (am.aas.00387777 01/ 807
Luca Biggio acknowledges support from the AI for Science postdoctoral fellowship EPFL.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Dordevic, D., Bozic, V., Thommes, J., Coppola, D., and Pal Singh, S. Rethinking attention: Exploring shal- low feed-forward neural networks
alternative to attention layers in transformers (student abstract). Pro- ceedings of the AAAI Conference on Artificial Intelli- gence,
References
Mar. 2024. doi: 10.1609/ aaai.v38i21. 0338221.30436. URL hitps: / orr/ i index, index.ppp/ httpa//pppp/AAA//atttlleee/eeeee
Abbe, E., Bengio, S., Lotfi, A., and Rizk, K. Generalization the unseen, logic reasoning and degree curriculum. In ICML, 2023. URL https:/ http//axivvv.//////ee////a org/ abs/ 2301. 13105.
Dziri, N., Lu, X., Sclar, M., Li, X. I.., Jiang, I.., Lin, B. Y. Welleck, S., West, P., Bhagavatula, C., Le Bras, R., Hwang, J., Sanyal, S., Ren, X., Ettinger, A., Harchaoui, Z., and Choi, Y. Faith and fate: Limits transformers on compositionality. In Oh, A., Naumann, T., Globerson, A.,
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G.,

---

Counting in Small Transformers: Attention and Feed-Forward Layers
Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 70293-70332. Curran Associates, Inc., 2023.
Grattafiori, A., Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Vaughan, A., Yang, A., Fan, A., and et al. The llama 3 herd of models, 2024. URL https: arwii..r........................ abs/ /bs/2407.. 2838
Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y.. Chen, A., Conerly, T., DasSarma, N., Drain, D., Ganguli, D., Hatfield- Dodds, Z., Hernandez, D., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. A math- ematical framework for transformer circuits. Trans- former Circuits Thread, 2021. https:/transformer- circuits. htwp:p////22//aaammoo/oode..... html.
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference Lan- guage Modeling, 2024. URL https::////// /pppreveeeeee.ee.eee net/ httpmmmdd=Etsskeeee
Henighan, T., Carter, S., Hume, T., Elhage, N., Lasenby, R., Fort, S., Schiefer, N., and Olah, C. Superposition, memorization, and double descent. Transformer Circuits Thread, 2023.
Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodds, Z., Lasenby, R., Drain, D., Chen, C., Grosse, R., McCandlish, S., Kaplan, J., Amodei, D.. Wattenberg, M., and Olah, C. Toy models of superposition, 2022.
Jiang, A. Q. Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., de las Casas, D., Bressand, F., Lengyel, G., Lample, G., Saulnier, L., Lavaud, L. R., Lachaux, M.-
A., Stock, P., Scao, T. L., Lavril, T., Wang, T., Lacroix, T., and Sayed, W. E. Mistral 7b, 2023. URL https: //arxiv. org abs/bb//// /2300.2231..mmm 06825.
Engels. J., Michaud, E. J., Liao, I., Gurnee, W., and Tegmark, M. Not all language model features are one-dimensionally linear. In The Thirteenth In- ternational Conference on Learning Representations, 2025. URL http:::://pen.eeiee..eem. net/forum? (d=d63a4AM4hS
Jiang, Q.. Li, S., Bai, H., de Lamare, R. C., and Hc, X. Gradient-based algorithm for designing sensing matrix considering real mutual coherence for compressed sens- ing systems. IET Signal Processing, 11(4):35-363, 2017.
Fickus, M. and Mixon, D. G. Tables of the existence of equiangular tight frames, 2016. URL WtLhhtps://arri org/abs/Dl 1504. 00253.
Jyothi, R. and Babu, P. Telet: A monotonic algorithm to design large dimensional equiangular tight frames for applications in compressed sensing. Signal Processing. 195:108503, 2022.
Geva, M., Schuster, R., Berant, J., and Levy, O. Transformer feed-forward layers are key-value memories. In Moens, M.-F, Huang, X., Specia, L., and Yih, S. W.-t (eds.), Pro- ceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 5484-5495, On- line and Punta Cana, Dominican Republic, November 2021, Association for Computational Linguistics. doi: 10. 1116653/1120211.mmmpmmmm ac aclanthology org. wttp:///p2..mm11mmmmme/mm mam.mpmmmiii.mmmmmmmmmmmm 446.
Kazemnejad, A., Padhi, I., Natesan, K., Das, P., and Reddy, S. The impact of positional encoding on length generalization in transformers. Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https: https/////p/eeeeeeeeeee/ new forum? id-Drr12gejal.
URL https://
Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, C. Transformers learn shortcuts to automata. In The Eleventh International Conference on Learning Represen- tations, 2023. URL https: oppprevieneeeeeeeeeeeeeeeeeeeeee net/ forum? d
Girdhar, R., Carreira, J., Doersch, C., and Zisserman, A.
Video action transformer network. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 244-253, 2019.
Golkar, S., Bietti, A., Pettee, M., Eickenberg, M., Cranmer, M.. Hirashima, K., Krawezik, G., Lourie, N., McCabe, M., Morel, R., Ohana, R., Parker, L. H., Blancard, B. R.-S., Cho, K., and Ho, S. Contextual counting: A mech- anistic study of transformers on quantitative task, 2024. URL https::////.... /aaaiaiiivv///a///aaaaa/aaaaevva org/ abs/ 2406. 02585.
Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei. Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchi- cal vision transformer using shifted windows. CoRR, hts/2103.140300 2021. URL https: ptttss//arxxi..rr //aaxxxxvvvor///... abs
Meng, K., Bau, D., Andonian, A. J., and Belinkov, Y. Lo- cating and editing factual associations in GPT. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems,
Gould, R., Ong, E., Ogden, G., and Conmy, A. Successor heads: Recurring, interpretable attention heads the wild, 2023.
10

---

Counting in Small Transformers: Attention and Feed-Forward Layers
2022. URL httttp://operreviee... id=-h6AA56eR4,
Quirke, P. and Barez, F. Understanding addition in trans-
formers. In The Twelfth International Conference on Learning Representations, 2024. URL https:// openreview net. -orrumii--IIxxxxWW-
Michaud, J., Liu, Z., Girit, U., and Tegmark, M. The quantization model of neural scaling. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL id=3tbIw2qa88k
Raventos, A., Paul, M., Chen, F., and Ganguli, S. Pretrain- ing task diversity and the emergence of non-bayesian in-context learning for regression. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https: http:///pppeeeeeeeeeeeeeee net/ forum? 1d=BtAZ4a5XDg
wwtperewwwwwww.eeeeeeeeeeeeeee ner/forum?
Nanda, N., Chan, L., Lieberum, T., Smith, J., and Stein- hardt, J. Progress measures for grokking via mechanistic interpretability. In The Eleventh International Confer- ence on Learning Representations, 2023, URL ht https: //openreview. net. forummormmmmmmm 1dddddddd deddd=99F5-0 PbDDddW.
Strohmer, T. and Heath, R. W. Grassmannian frames with applications to coding and communication. Applied and Computational Harmonic Analysis, 14(3):257-27,, 2003. ISSN 1063-5203. doi:
Nichani, E., Lee, J. D., and Bietti, A. Understanding fac- tual recall in transformers via associative memories. In The Thirteenth International Conference on Learning Representations, 2025, URL https: ppppe/eeeerreeeeeeeeeeeee net/ htte//orummidddhSStPP0m
URL attps::/wwww.cclenceediee cciieee/Arricaee/////00323200000
nostalgebraist.
interpreting GPT: the LessWrong, January 2020.
Tolstikhin, I., Houlsby, N., Kolesnikov, A., Beyer, L., Zhai, X., Unterthiner, T., Yung. J., Steiner, A. P., Keysers, D., Uszkoreit, J., Lucic, M., and Dosovitskiy, A. MLP-mixer: An all-MLP architecture for vision. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, 2021. URL https: hettps://openreviiii wwwppeeeepeeeeeeepoepee net forum? i id=EI2KOXKdnP.
logit lens
URL
https://www www.leesswong.
com//posts/ActttBBDaaaaaat interppetinng-----the-----
Nye, M. I., Andreassen, A. J., Gur-Ari, G., Michalewski, H., Austin, J., Bieber, D., Dohan, D., Lewkowyez, A., Bosma, M., Luan, D., Sutton, C., and Odena, A. Show your work: Scratchpads for intermediate computation with language models. CoRR, 12.s01112.0014 14, 2021. URL htttps://arxiv www.wiiwww....r..aao.a..aooooeee abs 2112 00114.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. In Proceedings of the I International Conference on Neural Information Processing Systems, 17, pp. 6000-6010, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964.
Olah, C., Cammarata, N., Schubert, L., Goh, G., Petrov, M., and Carter, S. Zoom in: An introduction to cir- cuits. Distill, 2020. doi: //2291///iiill..00022.00
Vig, J., Gehrmann, S., Belinkov, Y., Qian, S., Nevo, D., Singer, Y., and Shieber, S. M. Causal mediation anal- ysis for interpreting neural NLP: the case of gender bias. CoRR, abs/2004. 12265, 2020. URL https: wwwxiv..ogg/aee abs 2004.0004.00 12265.
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A. Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds. Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. In-context learning and induction heads, 2022.
Weiss, G., Goldberg, Y., and Yahav, E. Thinking like transformers. In Meila, M. and Zhang, T. (eds.), Pro- ceedings of the 38th International Conference on Ma- chine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 11080 11090. PMLR, 18-24 Jul 2021. URL hww.s: httptttpss::pppeeeeeeeeeee http:///p//ode//eeeeeeeeceoce ngs ml press/ v139/ htt//wiis222 html.
Ouellette, S., Pfister, R., and Jud, H. Counting and algo- rithmic generalization with transformers. arXiv preprint 100.022310.08661. 2023.
Petzka, H., Trimmel, M., and Sminchisescu, C. Notes the symmetries of 2-layer relu-networks. In Proceedings of the northern lights deep learning workshop, volume 1, pp. 6-6, 2020.
Welch, I.. R. Lower bounds on the maximum cross correlation of signals (corresp.). IEEE Trans. Inf. Theory, 20:397-399, 1974. URL https://api.. semanticscholas org htt.aaaa.r..rrrrruussa000188811
Power, A., Burda, Y., Edwards, H., Babuschkin, I., and Misra, V. Grokking: Generalization beyond overfitting on small algorithmic datasets, 2022.
Ychudai, G., Kaplan, H., Ghandeharioun, A., Geva, M., and Globerson, When can transformers count to n?, 2024.

---

Counting in Small Transformers: Attention and Feed-Forward Layers
Yun, C., Bhojanapalli, S., Rawat, A. S., Reddi, S. J., and Kumar, S. Are transformers universal approxi- mators of sequenc--oo-sqqeence functions? CoRR, Abs/19 2019. URL http://arxiv.orrt abs/1912.1000
Zhong. Z., Liu, Z.,, Tegmark, M., and Andreas, J. The clock and the pizza: Two stories in mechanistic explanation of neural networks. In Thirty-eevntth Conference on Neural Information Processing Systems, 2023. URL ht tps: //apenreeeew net / forum? i id- S S5mmb@c1We. We.
12

---

Counting in Small Transformers: Attention and Feed-Forward Layers
Appendices
A Counting with large language models
B Explicit Constructions for Orthogonal Embeddings =
B.1 Overview
B.2 Relation-based counting
B.3 Inventory-based counting
B.4 Mapping scalar to categorical one-hot encoding
C Explicit Constructions for Linearly Dependent Embeddings d < T
C.I Overview
C.2 Explicit construction for bounded mutual coherence
C.3 Explicit Construction with binary representations and softmax
D Data Generation
E Additional Experiments
E.I Best Accuracy
E.2 Variability
E.3 Model with Random but Fixed Embeddings
E.4 BOS mixing token
E.5 Singular Value Decomposition of Wi
E.6 Varying the number of tokens in the alphabet
E.7 Varying the length of the sequence
E.8 Models with two layers
13

---

Counting in Small Transformers: Attention and Feed-Forward Layers
A. Counting with large language models
prompt language models with 7 and billion parameters to count the inputs of a list. We use the torch implementations of mi stral-7B-vo. 1 (Jiang et al., 2023) and Llama-3. 1-8B (Grattafiori al., 2024) with flcatl6 respectively and prompt in the english language. We prompt the model count items different lists of items which stem from groups of 12 items (letters, months, animals). The groups of items
letters: P' 'B', 'c', 'D', 'E', 'F', ' 'H, 'I, 'J, ',, (L']
animals: d'dog', 'at', r donkey', horse'' 'elephant', 'giraffe', 'lion', 'tiger', bear', 'wolf' zebra', ' monneyy
months: [' ['january', 'february', ' 'march', 'april', 'may', 1 june", 'july', august', 'september', 'october', 'november', ' 'december']
And the four prompt types are
1. (compare, position) [list] By comparing all ten items in the previous list to one another, we find that the [i'th] item of the list appears [prompt].
2. (compare, semantic) [lst]] By comparing all ten items in the previous list to one another, we find that [item] appears [prompt],
3. (summarize, position)
[list] By summarizing the occurrences of [category.name] in the previous 1ist of ten items, we find that the [i'th] item of the list appears [prompt],
4. (summarize, semantic)
[list] By summarizing the occurrences of [category.name] in the previous list of ten items, we find that [item] appears [prompt].
When we generate data we sample the ten items from the list either uniformly at random with replacement, or as we for the toy model experiments as specified in appendix
We then sample 100 lists for each model and generate all possible queries for different items and positions. We also change the category type and sampling strategy and report the confusion matrices of the ground truth over the predictions in Fig. 6.
We observe that the performance is not very good. Querying for given item semantically rather than its position consistently performs better. For the positional queries the models mainly predict the count 3. Depending on the type of object that counted, counting more than is correct only rarely for both models. The type item that counted also impacts performance. Generally 1lama-3.1 seems to perform better than Mistral-7B.
14

---

Counting in Small Transformers: Attention and Feed-Forward Layers
icompare, positional)
(compare, semantic) (summarize, posibionllisummariie semantic)
morthr recuurences
prodiction
2.5
0.0
prreenntionngruuunn.oo bubl
10.0
0.6
Tocourencer
prosiation
75
5.0
255
0..
0.0
10.0
arimals the accurences
prodictin
7.5
10
10
10
10
#occurnnces ground truth
Asccurences ground truth
#cscurences
#occurmices ground truth
ground truth
Figure 6. Mistral 7B-v0.1 predictive performance on counting tasks with temperature zero. Samples of 100 lists 10 items with cach sampling stralegy, where generate all possible queries for different items and positions for each prompt type.
(compare, positional)
(compare; semantic) (summarize, positionalksuummriie. semantic)
1.0
10.0
montrs recurences prediction
0.0
Pipredicion ground mruth
leterss the alshabet
0.0
reccurences pradiction
23
0.0
10.0
the recorrences prodition
7.5
3.0
2.5
0.0
*occurences ground truth
Aoccurencess ground truth
#ctcurentes ground truch
#occurences ground truth
Figure Llama 3.1 predictive performance on counting with temperature zero. Samples of 100 lists of items with each sampling strategy, where we generate all possible queries for different items and positions for each prompt type.
15

---

Counting in Small Transformers: Attention and Feed-Forward Layers
B. Explicit Constructions for Orthogonal Embeddings d = TT
B.1. Overview
In the parameter regime where d 2 T there is always an orthonormal basis size T Rd, 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
Yt ReLU = ReLU bi);
(4)
In many cases, a simple linear regression can map the scalar y/ to the correct count of tokens 2g, and we describe how to achieve this mapping to the classification problem in Section B.4.
In the following, we characterize which parameters W1,b, b1 in equation 4 allow for correct mapping each mechanism. Importantly, the architecture exhibits numerous symmetries due the ferd-forward ReLU network (Petzka al., 2020). To demonstrate feasibility, we select one specific implementation. In the main text observe that there is 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 RdxT is an orthonormal basis of IRd which we will use to create different forms of token embeddings.
The supplementary code at that have the weight configurations that are used to prove Propositions 4.2-4.3 and B.1, which allows one to test the devised weight configurations for fixed T", L, in practice.
attention contains executable pytorch models
B.2. Relation-based counting
B.2.1. [poT:p=1)
Proof of Proposition 4.2. We set d> > 2 with and p 1. We choose the embeddings of the tokens of the dot model
El = Eg + Êent Vt I....T
(5)
where the set E = [ē,]7=1 is an orthonormal basis of an arbitrary but fixed T-dimensional subspace of R4, and ēent Ei= ét. The key and query matrix are set the scaled identity WK Wq d2/41, and hence the mixing layer Adot can be viewed carrying the unmodified dot- dot-product operation between all pairs of tokens. The first layer weights W1, by E Rd can be fixed as
Wi é
+1); (Tn//(++)))) b, (1 + I(T 2)),
(6)
and the second layer weights W2, b2 E R follow the recursion
(W2), -1 T++1
(b2) 0;
(7)
(W2); -1 L+1
(b2)) = ((W2)--1 (W2)e) (E 0.5) + b
Ve 2...., L.
(8)
Given these parameters, it holds that for tokens
S 1, < their dot-product is
ea) = 22+T i 3 T T =f
(9)
Because our choice of the query and key matrices, directly follows that for tokens Ze, Em at positions and m from a given sequence x, their attention score is
22+T
if 21 Im
atm
(10)
T if σs Im
16

---

Counting in Small Transformers: Attention and Feed-Forward Layers
Hence, the mixed token after applying the residual connection is
5
Et
+
Σ
(T + 3) +
Σ
(T + 2)
(I1)
so that computing
Eent
W1
T/h 1 1 T
(12)
Be.
Eent
Σ
(T (T+3)
Zmm 1 T
Eent
+
Σ
(T
im 1+1 T
Eent
(13)
m:zefam
- + x (T) + + (L hisl (T) +
(14)
histx (4) + ++ +
(15)
Then the single hidden unit has the value ア ReLU (2 = histx (f). It easy to show (analogous Fig. 8) that the output logits с 7/2 W + by with R,, correctly identify the count for integer values z E [l..... L]. This is because we constructed our recursion such that at given input = we have that (W2)e(E 0.5) + (b2)s = ( 0.5)+ (b2)e-1 and (W2)e > (W2)e-1- so it holds that
1
y 1
2
y =
arg ma..x(x))) 1=
(16)
L y= L.
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 different but also 100% correct algorithms. This is especially important compare the regime outlined in the Theorem with the weight configurations learned.
B.2.2. (BOS+SFTM;
Proof of Proposition 4.1 for bostsftm. We set = d 2 with L 2 2 and and consider the model dot + fmm. Note that in this model every sequence x is prefixed with toos before it fed into the embedding and then the mixing layer. Again we use mutually orthogonal embeddings. E = {ē,}ī" an orthonormal basis of an arbitrary but fixed T-dimensional subspace of Rdd, and éent = Σ é We set éBos = Σi=l Fg, where e, F, and the latter a column of E. Analogous to the background token from Proposition 4.1 there only one direction = to detect in the feedforward model, so we set
W1 = CBOS; bj -1.
(17)
For a given token we have that in the dot-product mechanism (EBOS, ezz) = 1, (Exes Cxm) = if Im and 0 otherwise. Due to the softmax, the mixing coefficient is a = e/((kz, + 1)e + (L - ka,)) (where is Euler's number) for comparing We to /rO and to the tokens where #e = Im. and 1/((kiz + 1)e + (L -kx)) otherwise, where, k.c hf). (f). Hence, the mixed token is:
1 aéBos + aliz,
Σ Fm t00
ba Tt.
(18)
Applying Wi and b, we obtain:
2t- = aT + akiz, + b(L - knr))
al' + akiz + a(kee + 1)
(19)
a(T 1) +
since (kz, + 1)aa + kb, bb= 1 by normalization via the softmax function. The value of ye has dependence on k through a and can be readout into the correct classification as shown Fig. 8.
17

---

Counting in Small Transformers: Attention and Feed-Forward Layers
B.2.3. (BOS:p=1)
Proof of Proposition 4.1 bos. We s' d with L > 2 and p 1. The construction of the embeddings and ebos is analogous to the construction from Section B.2.2 for bos+sftm in the same setting. However, since no softmax is
applied, the mixing coefficients outputs of Adot-sftm for comparing (Iz, tBos) (æe, Zm) where z 2m. is 1.
For ze + Im it is b = 0. Then from inserting these values in equation 18 and applying W1 eBos and bj = - we obtain
Yt
kgg
(20)
This clearly allows again the single neuron to read off the correct result similar to the construction from equation
Note that there is simple alternative construction that uses the tagged embeddings from the constructive proof of Prop. 4.2.
Alternative Proof of Proposition 4. bos. We set T d > 2 with L and p We note that by setting (BOS to zero achieve equivalence to the model dot. Since according to Prop. 4.2 there exists weight configuration for dot which solves the histogram task, this configuration will also solve the histogram task for bos with tbos 0.
B.3. Inventory-baser counting
B.3.1. (LIN; p T).
Proof of Proposition 4.3 lin. Assume that T > 2 with L > 2 and p T and the goal is to find weight configuration for the model 1in. As embeddings we directly use the orthonormal basis with T vectors in Rd, where vectors are the embeddings are for the T tokens. We set
a
a
a a
A
Wi E; b1 1,
(21)
where = 1/T. We start by writing ze, for E (1,.... T). the 1-th activation of the first hidden layer of the feed-forward module
28,
Σ m==1
abm (exm e2) + (exss 1.
(22)
If eq Exer we have
21,1 = kiz, a + = akæz
(23)
where, kiz = histx((), applying the ReLU to this scalar keeps its value unchanged. e * exg, we have
zg,t alie, + 0
ak.,
(24)
The right hand side of the above equation is negative given our choice of hence applying the ReLU returns 0. This means that, for each token in the input sequence, the contributions of orthogonal tokens cancel, leaving with single hidden hidden neuron activated. Hence the count can be read off from Ye. Since only one neuron is activated at time, the readout from the same procedure as in bos+sfte can be applied to all hidden neurons zt,t. simultaneously. instead of only one. This allows the model to solve the histogram task.
B.3.2. (LIN+SFTM: P=
Proof of Proposition 4.3 lintsftm Assume that T' = d > 2 with L > for 1 Lin, we note that we can construct Ain+atum such that is equivalent to A from equation 21 via
2 and p = T. With the statement already proven
a
a
a
A Ali ffmmm
sofumax
(25)
where a = 1/L which implicitly defines a choice of a. This means that the construction is equivalent to lin and it follows automatically that also lints ftm can solve the histogram task.
18

---

Counting in Small Transformers: Attention and Feed-Forward Layers
B.3.3. p
Proposition B.1 (IC for dot fsftm. For dot+sftm and given L, T 2 there exists configuration of weights which solves the histogram task for p 2 T and d 2 T.
Proof for Proposition B.I. We assume L,T and p= = dot in Prop. 4.2, we set the key and query matrix to the scaled identity WK = WQ d We use orthonormal basis of Rd to define the parameters 6 IR for the the token embeddings. In A the pre-softmax mixing weights will be 1 for equal and 0 for different tokens due to the unit-norm token embeddings. Defining kæt = histx(E) for brevity, after the softmax we have that
and d T and consider dot+sftm. As previously for
3m It.
(26)
else.
(1-
Hence, for et + eze
ker
(Tee, Et)
1,
(27)
(L. kore) + exize
while for e Exe
kz,6
(4/2, exx) = (kzze + (L kat))
+ +1,
(28)
where the extra summand comes from the residual connection. Hence, by setting
W1 = E; bj --1
(29)
and applying the ReLU activation, equation 27 will be 0, while equation 28 will implicitly give us the counts
kec = (Lr))(--e7 + ve +
(30)
While the final layer cannot immediately implement non-linear functions in Ye. it can take advantage of the fact that Y can take only different values, similar to how we constructed W2 and b2 in Section B.2.1. Since eventually we need map the L values of to the counts [1.... L] the linear output layer is sufficient implement this non-linear discrete map. Fig. 8 shows an example for this map for given example. This allows the model to solve the histogram task.
The statement for > T and d > T follows as we can simply set the surplus of parameters in the hidden layer/embeddings to zero.
B.4. Mapping a scalar to a categorical onc-hot encoding
30
29
28
output neuron
27
8
26
10
25
0.0
0.2
0.4
0.6
0.8
1.0
Y1
Figure 8. Demonstration of a hidden neuron output which is mapped to L = 10 different neurons c, using single output layer, according the decision boundaries shown by the dotted lines. After applying argmax function output neurons, the highest value gives the discrete output. The solid lines mark the values single hidden neuron would achieve for different counts in the explicit construction of the =st tm model.
19

---

Counting in Small Transformers: Attention and Feed-Forward Layers
is straightforward map single scalar "e series of neurons which activate one after another. This is needed the second part of the feed-forward parameters transform the count measured by the sum of the hidden neurons "Y the discrete categorical representation of the output vector. Every output logit is linear function of the hidden neuron's 's value. Since in our constructions we only map functions, where the ground truth output logit corresponds to an interval (a. bj E R, the superposition of linear functions with increasing slope allows us to realize such a mapping. A visual sample is given in Fig. 8 for dot + ftm. In Fig. 9 we show the outputs for the lints ftm model with the best accuracy for = 32 for every p, ran in Fig. While is possible to learn the count from one hidden neuron only using inventory- counting for each neuron, for some examples the count information seems be spread out over several hidden neurons: The output logits non-linear in the count and can hence not rely single hidden neuron only.
0.5
44.01
100 0074
00%
100
10
100
-00
20-
-033
100
500
3000
30
1.30
2.00
27.04%
31.34%
44.11%
9..99%%
1.00 00%
00.000%
50
30
500
-003
200
000
1000
24.59%
35. 75.87%
47.47%
60%
80. 00..44%
59 67.677
100 00%
100
200
-000
200
600
46-
11.9
36.27%
50.38%
17.69%
20-
0.0
-50
-10
100
000
150
1.50
300
200
25.38%
25.73%
48 48.96%
52.21%
54.76%
55.77%
57.24%
20-
-50
-50
100
34.15%
42.98%
33.32%
47.97%
20.
0
10
20
-00
26.25%
33.79%
52. 30%
25.32%
37.550%
35 0440
30.55%
1.0
18
10
-10
7.5 5.
75
2.5
2.0
7. 50 7.5
2.5
75
Figure 9. The output neurons c, (xx) visualized for examples of learned version of win+iiiiiit m for several model dimensions d and hidden layer sizes p. Note that differently from Fig. 8, in this case the x-axis shows the number of occurrences histx (E 0) I,..., L of the token in an input sequence x = [t, ... neuron). We show the activations e; of the final layer output neurons activations (logits) in terms of the number of occurrences of given token in the input. The colors represent the different output predictions and are as in the explicit construction from Fig. 8. We show several activations for different tokens T. where T 32, and we highlight one of the example tokens with a wider line. While similar to the explicit construction from Fig. 8, the models with 100% accuracy are not necessarily lincar in the count.
u..... vi that contains otherwise only a token v 3 t (and not the activation of a hidden
20

---

Counting in Small Transformers: Attention and Feed-Forward Layers
C. Explicit Constructions for Linearly Dependent Embeddings d < T
C.1. Overview
In this section, we discuss the scenario when d < T, i.c. when the embeddings are necessarily linearly dependent. In that case, we can no longer assume that there exist embeddings with (et, c.) = for all * Nonetheless, also in this regime for some models is possible 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 I. classes. When we examine he IR from equation 4 which is mapped to the discrete class through the readout layer (see for example Fig. 8), we notice that the class boundary (the gray dashed class borders) can be placed variably in the margin between the values that Y assumes for different counts ka, (solid lines). In the following explicit constructions, our goal to design embeddings with < 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 (s), and in equation 4 this will create non-zero terms that will alter the value of ye. This means that for every possible sequence with k occurrences of token zt. the hidden activation Ye will assume values in certain range (ylower (k)), 7t hppeer (R:)]. If these ranges overlap for different k the count cannot be identified. However, construct embeddings such that every for every k- 1,..., L - 1 it holds that
upper (k) < lower (k + 1). 7
(31)
so we can still use construction as in Fig. 8 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 Zg, from equation 22). Notably, for the explicit constructions we propose, the function of the lowest achievable d(p, T, I.) differs across different mixing types. To summarize:
For models with A constant in the inputs or models without softmax activation, explicit construction relies on an embedding matrix with a small mutual coherence. The mutual coherence is concept from compressed sensing and coding theory that ensures that the maximal similarity between pairs of vectors is small (Donoho & Elad, 2003). We can upper bound the mutual coherence that the margins of the construction can tolerate to still achieve perfect accuracy in terms of given L. At the same time, the mutual coherence of a set vectors is naturally lower bounded in terms of the number of vectors T and their respective dimension d, known as the Welch bound (Welch, 1974). When this bound can be attained and T, I. are given, this leads to the following bounds on d for the different models, outlined in Prop. 4.4, as (lin, lin+sftm; p = T):
T(2L- 3)2
S d.
T T- T-1+ (2L-333
-2)
(dot, bos; p 1):
+1 d,
(21-33
(dot, bos: p = T):
T-1+(L-1] Sdd.
For bosts ftm we rely on the fact that the softmax function accentuates the largest value and thereby can drive attention scores for equal tokens ag higher relative to attention scores of non-equal tokens azj. This distinguishes 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
(pos+sftm: 1): d 2 [log (T + 1)] + 2. (dot+sftm;; p T): Togz (T + 1)] +
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 even possible to achieve perfect accuracy with 4, but for finite computations this might pose a problem when L becomes too large.
C.2. Explicit construction for bounded mutual coherence
We define the mutual coherence M of set of T unit norm vectors fui,..., uT} Rd
M max (1a, 0,)l.
(32)
ifj
21

---

Counting in Small Transformers: Attention and Feed-Forward Layers
This value is lower bounded for given matrix by the Welch bound (Welch, 1974)
T- d 1) W(T,d), M 2 V
(33)
and equality can only be attained if < d² (Strohmer & Heath, 2003). There is a large body of work in coding theory and compressed sensing concerning the existence and construction of a of vectors that attains M at or close to W(T, d). Explicit constructions exist but are not known for every combination of and d. A list with existing constructions for the real space for small T, d can be found in (Fickus & Mixon, 2016), but otherwise gradient-based optimization has been used to find good candidate matrices (Jiang al., 2017; Jyothi & Babu, 2022).
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 M we can in turn obtain a lower bound on d in terms of and T. Note that the Welch bound cannot be attained for T d2 and in this case the bound on is strict.
C.2.1. (LIN, LIN+SETN; T)
Proof of Proposition 4.4 lin. To show the bound on d, we analyze the inventory-basee construction for 1in in equa- T, T, and > 2 is given, let us assume that there exists set of T unit norm vectors {ei,-.., er. C Rd tion 21. Given that p p=
with mutual coherence .M. We use these vectors as our embeddings.
The value og for Zy. with W .e......e. eTl and bj -1 is
28,t = akæz +
Σ
(exm:Ce)
(34)
ma 7tt
and using that fact that the mutual coherence bounds the absolute value of the inner product
akæA a.M(L kur) S Ze, S akk, + +MM(L kKL)
(35)
Similarly, for / + and a 1/L it still holds that
21, S ak2 + a((L -
(36)
provided that M < 1/(L+ 1). for the worst case where k, L 1. This means that the ReLU sets all hidden neurons Στ,z to zero when * zg. and are therefore no contribution to the final result. Then, defining
Tt
Jower (k) = ak aM(I. - k),
(37)
upper (k) = ak + aM(L k),
(38)
we have that indeed for sequence where 2, occurs k L. times it holds that
0 < 7 lower (k) S (k) < 7/
upper (k)).
(39)
The first inequality required due the ReLU and holds when M 1/(L 1). From equation 31 we have the condition that for all 1, 1,..., holds that
lt
upper (k) < lower (k- + 1),
(40)
t + (L - K)M < (k: + +)+ (l--k- (I----
(41)
1
M < 2(L - -)= -1
(42)
and since we assume that there exist at least two different tokens in the sequence, minimizing the bound over k: leaves for
1
M <
(43)
2L -3 3
22

---

Counting in Small Transformers: Attention and Feed-Forward Layers
which is valid provided that L Collecting all previous bounds on M, conclude that when L > 4 the above construction achieves the correct counts with M 22.-3
The Welch bound equation 33 gives an upper bound on M in terms of T, d and therefore yields the final condition
T(2L -3)
d
(44)
T-1+ (2L
under which the given weight configuration is able to solve the histogram task with perfect accuracy.
For 1 lintsftm the construction and conditions transfer directly, when the constant tin constructed to match Ajin exactly.
C.2.2. ROS:p=
Proof of Proposition 4.4 dot, 1. We assume that L > 2 and T given and we use similar idea as the relation-based weight configuration from the proof of Prop. 4.2 for dot with = 1. For the token embeddings. we assume that we have a set of unit norm vectors with "U,,...,.y C Rd-1 with mutual coherence M, where d >> We the entries of the T embedding vectors ez Rd to be
Ex
(45)
The shared counting subspace is defined on the last coordinate of the vectors via Cent = [0, 0,....... ..... 1/a]. Then
(er, e.) = =1+ a² Ke,, e.)l S M a a²
(46)
(47)
The mixed token with the residual connection at position ( for given input sequence is
E'
(exm> Excs )amm + Eze
(48)
m=11
and the single hidden neuron 'e for bias term bj = and W1 = Cent
7t (Cent, x//)
(Crm Cxre: (exm Cent) + (exer Cent)
(49)
m==1
kig (1 + + +
Σ
(ezm Caer?) +
(50)
m:am Five
So that we can achieve for a given count k the lower (lower(k) S (e(): S 7/
upper (k) with
S 7 lower (k) k11 a²) (I. (4-1)(M + +3)
(51)
upper (k:) = ks1 + a² + - (M + a²) +1.
(52)
We achieve the upper bound from zero, when .M- < 2/(L 1), assuming that a is close enough to zero so that is is negligible. Finally, the condition from equation 31 yields
2(T k) M < 2( kk)
1
k)- - 2(L k) a
(53)
under the condition that 0 a / 2(L-R)-2 Again, assuming there exist at least two different tokens in the sequence, the
r...s. of the above expression is minimized for k = 1
1
2L 4 2
M < 2L-3
(54)
2L
3
23

---

Counting in Small Transformers: Attention and Feed-Forward Layers
which is always positive assuming L 2. This is the relevant bound when have large enough L 2 and again is close enough Again, combining this with the Welch bound equation 33 leads to
21-3 )² T( T- 4ja
d - d-12
(55)
T-1+
(2L 1j02 )
and when we choose a > 0 close to zero, as for 1in before
T(2L - -3
+ 1.
(56)
T T (2L-3))
This proof holds equivalently for bos when we set the BOS token embedding to zero.
C.2.3. (DOT, pos:
can decrease the required dimension d even further than previously, when we have p T and implement inventory-bran bassd counting the dot model (and equivalently in the bos model). In that case, the lower bound on d becomes more loose, because combine the ideas saw in 1 lin and p T for inventory- counting and the effects on the margin in dot and 1.
Proof of Proposition 4.4 dot, P- In construction, for given T and L, we assume that there is set of T unit norm vectors fei....,,rr } C Rd with mutual coherence M upon which we build our embeddings. Note that the only difference to the previous relation-based case = 1 is that this time there is no extra counting direction. Importantly, we set K = / la as before, but Q = / This gives an extra factor the attention scores. Further, we set b = - and the columns of Wi E Rx the embeddings ez2 as we did for 1in. This results mixed token zy according equation 48. The hidden neuron
L
28,1
Σ m==
(ezm end exm, + (Ee, Ezi)
(57)
L
then with / zz have
1 L
Zg,t
Raz
+
Σ 17::
(Cam ca\²
(58)
Note that the square equation 57 is what differs from the žτ, equation 34. This is because the term (czm. ex) once introduced through the dot- dot-product attention and once through the dot-product via W₁. Conversely, with t # xe it becomes
29,1
T.
1
Aeler,eme +
Σ mE
(eam Exee (emmer))
+ (ee, Care) -
(59)
1 (k,M+ + (L - kx) + (et,
< T.
(60)
< 2M 1
(61)
we set M < 0.5, which we need anyways for T, 2 2 by the stronger upper bound on M that we derive in the following, we finally have for + It
26, 0.
(62)
Again, negative ze, are to via the ReLU, and the final outcome 7/6 26,t= depends only single hidden neuron equation 57. This eventually leads to
lower (k)
(63)
L
7/
upper (k) (k: + M k)),
(64)
24

---

Counting in Small Transformers: Attention and Feed-Forward Layers
and using the same concept as before, while minimizing over k and applying the Welch bound, to the upper bound
(65)
L-
The final bound is more loose than it was for as we only require
T(L - T 1-1+(1-1)
d
(66)
C.3. Explicit Construction with binary representations and softmmax
In our final analysis we examine the key difference between the models bos+s ftm and bos the softmax activation. In order to show Prop. 4.4 needed construct embeddings with low mutual coherence, because the term (et, ex) introduced an error on the mixed token, when i and 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
e
sftm(z), 22;
Σj= e for 1.2......
(67)
and when we compute sktm) we say it is a softmax with inverse temperature x >> 0. When z of length L contains only two different values, one with & and the other with L - k occurrences, then as k + oo 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 8 ......? T and 8 *
(e),e) = 1, (e), <1+c,
(68)
(69)
where 0.
The idea is that the softmax with high enough inverse temperature sets the term for different tokens, (et, ez), close enough to zero, essentially eliminating the noise. Note that equation 68 is weaker condition on the set of token embeddings than for example the bound of the mutual coherence terms of the sequence length bos with p = 1 in Section C.2.2. allows us to obtain perfect accuracy with smaller In the following, we describe the construction of the matrix explicitly.
The supplementary code at supplementary material contains executable pytorch models that have the weight contigurations 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.
C.3.1. (BOS+SPTM;
Proof of Proposition 4.5 -bbbbossttttt For a given T, L > 2 we set the embeddings vectors to the binary representation of the token index = 1,..., T in d' = [Toga (T + 1)] dimensions
bin(t)(bin(t),bbnn bin(t))
bin (On bin(
Et
;
EOS
(70)
1/a
0
where bin(t) (......, vwl E {0, d with t= Ed' 11 We select a > 0. Then we have that
(e,,et) = 1 a²
(71)
a² < (er, 6x) < 1
+ + <1 -
(72)
d'
(et, BBOs) =,
(73)
25

---

Counting in Small Transformers: Attention and Feed-Forward Layers
where (Ezer -11 E2a' -2), 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 eos token, namely, W1 = [Cent] [0,....0,, 0, E IRd and b = 0, we construct
Y (Cent; x/y) = ( )o ( (eBos, Cent) + Σ Aftt(EFF) )m+ : Cent) + Cant)
(74)
= =ft sftmm[[ee, eBo)), (ee, ei),...,(ees) ez))o
(75)
(76)
The goal of applying the softmax function is to diminish the contributions of error equation 72, while having the final dimension of the BoS token be representative of the count of xy. The maximum error is induced when the upper bound equation 72 is attained for all tokens in the sequence that are not equal to zy. 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 I, and a softmax activation function with an inverse temperature x², we have that
e
Jower (1:) = exl + keM(L+a9) + (I. k)em(1+a2-o)
(77)
ex1
upper (pper(k) 7t"
(78)
ex1 + ke*()++²)
+ (L - k))mn??
We explicitly need strictly greater than zero, since otherwise there no information about the count " when it becomes independent of the count k. Notice, that this time it holds that "e that correspond higher values correspond to smaller counts, since larger count corresponds to larger denominator, i.e. a smaller zye Due to this inverse relationship, for this model, we want that for all counts k = L,...,, L 1 that it holds that
upper (k 1) 78
Jower (k))
(79)
This can be achieved by setting the inverse temperature x accordingly. In the following we show that there exists a K which fulfills equation 79 for all d' 2 2 and L > 2. Observe that upper (2) < 2/
Jower (1) implics the bounds for all other k. We define the distance or margin as
dist(k) = Jower (1) - 7t.
upper (2).
(80)
Since at 6 = 0 both 7/8 upper (2) = lowoerr (1) 1/(L + 1), the distance is zero. However then it becomes impossible "t
to distinguish &- = 1 and k = 2, as they receive the same weight. We therefore need the additional condition that upper
1f. (2) # ) Jower (1). At k = 0, we observe that this function has a negative derivative, as
a dist(k)) = (ftmm 22ower )o
(Zlower )o
Zlower ) jfttm
(81)
Ōk
2=
sftm sftmm (KZupper))
(Zupper)o )o
(Zupper)
(82)
1
2
([(1 + a²) + +L 1)(1 + az -[2(1 + a² (L 2)a²))
(83)
T++1
2
= L
L+1)
-)) - (L-1)cl
(84)
< 0
(85)
where the last bound is met when 0- < e < 0.5 which is fulfilled already for d' = 2 and when > 2. As the distance function is continuous, there exists a i close to zero for which the dist(n) < 0. Simultaneously, as x. -> oo, we have that
'n order to introduce the inverse temperature of the softmax in the model, we scale the query matrix. We set K = d Ia. but
Q = x/d/4Aa la.
26

---

Counting in Small Transformers: Attention and Feed-Forward Layers
due to the concentration of the softmax probabilities on the largest entry, which here is a², it holds that - have (x) - 0. At the same time, the function approaches infinity from the positive regime. For large enough we have "p h (2) < 7/2 h (1).
When we select the smallest possible k > O, we avoid computing functions with large exponential terms. To find the non-trivial root of (ks numerically, we consider a simplification of equation 80. We define = e* Then it holds that
we can solve
- u- dist(k) = 0 (L -u- (L-2)
(86)
numerically for k >> O. This shows that we can find explicit construction with 100% accuracy with p and d' > 2 for the s ftm when we have
d = [log2 (7 1)1 +
(87)
For example, for the case of 10 and T= 32 this allows for dimension d= with a = 0.01 (and for T =3 with the same settings suffices).
Remark (d 4). In principle, is enough to have some 0 that ensures that overlaps between different token embeddings strictly less than one. In principle, we can find an arbitrary number of tokens T that satisfy this condition for just d' 2. Take for example the following construction. For t = 1,..., T tokens with T odd we can design the set of embeddings
4 T-1
2t
(88)
Each (e), = 1 and for t * s the overlap (ec, e.) < (E(T+1)// -1T-1 VT2 This implies that e + O as T - oo at a rate 1/T. Since smaller e imply larger values of the temperature to solve equation 86, this might become problematic when this exceeds the accuracy computations. Previously, for the binary representation construction from equation 70, we had that e shrinks at a rate ~ /og2 (T). For the intermediate regime between (T) and log2(2) dimensions. one can generalize this principle to arbitrary bases, e.g. log;(T)) > resulting in smaller dimension but also less favorable (smaller) this construction thus comes with clear trade-off.
C.3.2. (DOT+SFIM; p
Proof of Proposition 4.5 of using = 1 we use = T. The selection of the embeddings 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+s with p 1. the feed-forward layer with W1 the explicit construction considers again ž for every token T. The selection of the temperature also analogous, with the exception that one has terms in the softmax instead of T. 1.
For this model, the explicit construction is analogous to the previous one. Instead
D. Data Generation
Every sample x = (01,,". (,)) is generated recursively as follows, starting from size K L and alphabet T =T:
1. Sample an integer & uniformly from [L, KK.
2. Sample a token / uniformly from T'
3. Set x, = for all i k, ...... K.
4. Set T' = T' \ (t) and K = &.
5. If K + 0, repeat from 1.
6. Set x = shuffle(x)..
In contrast to sampling the elements of each sequence uniformly at random from the alphabet, this simple strategy enables us better control the distribution of counts in the training dataset.
27

---

Counting in Small Transformers: Attention and Feed-Forward Layers
E. Additional Experiments
E.1. Best Accuracy
In Fig. 10, 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 a model.
(ros) deleproduct lan
[is]] despprdiit attection
linee mainn
embending dimension
[N] AARROO 2000 (Sa)
10+
2
hidden layer size
Figure 10. Best accuracy during training, Experiments from Fig. 1 (7' 32), show only the best accuracy during training reached
from the 5 randomly initialized runs per model/hhyperprrmeeer configuration.
E.2. Variability
In Fig. 11 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 + ffm is able to reach an accuracy relatively close to 100%, the variability of the accuracies resulting from different initializations is quite large.
[n]
derigmncdd BOG
deégecticl allention
Unwer
embedding dimension
30a
15
10
10a
hidden layer P
Figure Variability to randomness initial conditions and optimization. Experiments from Fig. 1 with T 32, standard deviation of the accuracy reached after training from the 5 randomly initialized runs per model/hyperparrameer configuration.
28

---

Counting in Small Transformers: Attention and Feed-Forward Layers
E.3. Model with Random but Fixed Embeddings
In Fig, 12, 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+s sf sttm, learning the embeddings increases the performance strongly in some regimes. This indicated that the models indeed learn adapted embeddings here.
unttalicing ollon
122
12'
307
307
121
hilden Iager 11
Figure 12. Random but fixed embeddings (not learned). The difference between learned and random embeddings for 32, Orange indicates that the random emheddings perform hetter on average. Purple indicates that the learned embeddings perform better on average. Experimental settings as Fig. 1.
E.4. BOS mixing token
In Fig. 3 in the main, we describe how the tBos is the main predictor for the count. Here, we provide more evidence by showing how the count predictions for mixed tokens 3' output by the feature transform are invariant to the type of other token present in the mixed token. The results for four different tokens are shown in Fig. 13.
Trag
Feruol dectel
Tamur 111- K00 kg0
Kang 5-rahy rect
Dampu E Aa kml
Figure 13. Visualization of feature transform. 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 (BOs. We vary the inputs between the learned tokens [B, C, D, E).
29

---

Counting in Small Transformers: Attention and Feed-Forward Layers
E.5. Singular Value Decomposition of W1
In Fig. 14 show the distribution of singular values of W for several runs of the model investigate whether models that capable of both IC and RC implementing the more memory heavy IC or the same solution that they find for p 1 with RC.
St-protud diertion latten
ozloorotuck sbarton
Inoer moing
10
20-
4511
10-
20
1.5
15
100
0
25
50
75
100
125
25
50
75
100
125
25
50
75
100
125
singular value index
Figure 14. Singular values of Wi. We show the results for all models from Fig. with T = 32, where p. 2 T and the accuracy is at
least 99%. Some qualitative diffcrences are visible for and dot.
E.6. Varying the number of tokens in the alphabet
In Fig. 15 and 16 we change the number of tokens that may occur in the alphahet from T= = 32 the experiments from the main text 15 and T 64 respectively. The sequence length kept fixed at 10.
[bos]
Idotl
Bin] mixing
product
product
10a
embedding dimension
-80
accuracy [%6
10a
60
10a
10a
hidden layer size
Figure 15. Alphabet size T' = 15. Experiments in Fig. 1, except that T' = 15 instead of T = 32 was used. The sequence length is
fixed to L 10.
30

---

Counting in Small Transformers: Attention and Feed-Forward Layers
[bos]
[864]
Bin]
dol-product attention BOS token
dot-product attertion
linear mixing
10a
100
embe dding dimension d
meean accuracy [%]
40
101
10a
10a
10a
10a
10
10a
10a
10a
10a
hidden layer size
Figure 16. Alphabet size T 64. Experiments in Fig. 1, except that T 64 instead of T used. The sequence length fixed to L = 10.
E.7. Varying the length of the sequence
In Fig. 17, 18 and 19 we change the length of the sequence to L. = 5, 15, 30 respectively, while keeping the number of tokens fixed to 7 = 32. The sequence length is kept fixed at L = 10. Fig. 17 shows that the learned model can make use of the discrete nature of the task especially for 1in (without softmax) and for the dot-product style models with the softmax. While 15 in Fig. is only slightly worse than the case of L 10 that was used as the example in the main text, the performance T. 30 Fig. 19 very much deteriorated. Even in the cases without the softmax performance is far from what would be possible theoretically. We hypothesise that the optimization with the increased number of possible inputs becomes much harder. Interestingly, Fig. 19 even seems show double descent behavior along the embedding dimension axis.
[bos]
Idotl product
Bin] mixing
embedding dimension
accuracy [%6
60
hidden layer size
Figure 17. Sequence length L 5. Experiments in Fig. 1, except that L = 5 instead of L 10 was used. The alphabet size is fixed
31

---

Counting in Small Transformers: Attention and Feed-Forward Layers
[bos]
[864] dot-product a attention
Bin]
dot-product attention BOS token
linear mixing
10a
embe embedding dding dimension d
with
10a
100
meean accuracy [%]
40
10a
10a
10a
10a
10
10a
10a
10a
hidden layer size
Figure Sequence length T. 15. Experiments in Fig. 1, except that 15 instead of 10 used. The alphahet size fixed to T' = 32.
[bos]
[dot] dot-product attention
[in]] linear
dot-product attention 8OS token
107
with softmax
101
embedding dimension
80
mean accuracy [1%]
60
10a
20
no softmax
101
10a
10a
10a
10a
101
10a
hidden layer size
Figure 19. Sequence length L 30. Experiments as in Fig. 1, except that L 30 instead of 10 was used. The alphabet size fixed to T' = 32.
E.8. Models with two layers
In this section, we look at the case where we have models that have an extra layer, i.c. instead of the logit output layer after the feed-forward part, we add another layer with the same dimensionality as the previous layer - the same mixing and the same hidden layer size to then lead into the classification. The parameters are not shared between the layers. Note that this model does not have an extra residual in the ML.P layers.
32

---

Counting in Small Transformers: Attention and Feed-Forward Layers
We train the model in the setting as in the main and report the results for the different architectures, this time with layers, in Fig. 20. To compare easily with the previous set-up, show the difference between the single and double layer case in Fig. 21. Remarkably, the general picture does not seem to change significantly. Indeed, the two layer model is generally better, extending the range where perfect models can be found slightly, but the general trend remains. Given this coarse grained experiment we hypothesize, that the extra layer aids the optimization process, and improves robustness the regions where and the softmax is used to disentangle non-orthogonal embeddings. More generally, these results are not as comprehensive as our previous results as they are not supported theoretically beyond a single layer. They warrant more detailed further work with more layers and realistic settings.
dot-product attention BOS token
dot-product attention
linear
107
100
embedding dimension d
80
mebi eccures (9)]
109
10a
20
10a
10a
10a
10
10a
10a
101
10a
10a
101
hidden layer size
Figure 20. Two-layer performance. Experiments as in Fig. with T 32, L. = 10 but for models where the attention block is repeated twice described Section E.8.
[bos]
Idot) product attention
Binl mixing
dot product 8OS
(%) (Koerrree Jarrz
embedding dimension
10
-10
-20
10a
10a hidden layer size p
10a
10a
Figure 21. Comparison between one and two-layer models. Difference between the accuracy of a single and two layer attention model, for different mixing layers and hyperparameter sctups. Experiments Fig. for single layer attention model, and Fig. 20 for the two layer model.
33


Reviews:
[{"review_id": "z4XUWS5rNL", "reviewer": "Reviewer_S9FV", "summary": "This paper explores the delicate interplay between the attention mechanism and the feed-forward layers, and further offers deep insights into how subtle architectural choices can drive algorithmic behavior in Transformer-base models. \n\nAs an example, the authors investigate how small transformer models tackle the histogram task - counting token occurrences in a sequence. The paper identifies two primary counting strategies: relation-based counting, which leverages attention for local token comparisons, and inventory-based counting, where the feed-forward layer memorizes token identities to aggregate counts. The authors provide explicit theoretical constructions for both strategies and back them up with extensive experiments that analyze how factors like embedding dimension, hidden layer size, and softmax application affect performance and robustness.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer S9FV,\n\nWe thank you for taking the time and effort to carefully evaluate our work, including the larger part of the supplementary material. We are glad you found it clear and rigorous as well as providing an inspiration for studying how transformers solve other algorithmic tasks; and that you appreciated the analysis of embedding orthogonality and mutual coherence.\nWe are also in agreement that investigating these mechanisms in a more practical context with larger models is exciting future work that hopefully allows us to connect our theoretical insights with real-world applications! \n\nIn the following we would like to answer your questions:\n\n> In the paper, the term Token Mixing is used to generalize across the self-attention mechanism. Could the authors clarify whether this abstraction is primarily conceptual, or if it implies a formal equivalence? Are there theoretical reasons to treat attention as a subclass of token mixing, especially in the context of algorithmic tasks?\n\nAs you have noticed, we use the term token mixing to refer to the mechanism in the network that applies a weighted sum of tokens along the token dimension (sentence length), and we use feature mixing which analogously applies on the feature dimension. On a conceptual level, this clarifies that the two different blocks act along different dimensions on the input - this framing has been used before, e.g. in https://arxiv.org/abs/2105.01601. For the token mixing specifically, the way in which we define it helps us to formally unify the different flavours of self-attention that we examine in our experiments, e.g. using an activation or not, as well as the linear attention. We can therefore say that attention is formally a form of token mixing.\n\n\n > Given that the current study focuses on single-layer architectures and a relatively simple task (histogram counting), do the authors anticipate that the identified mechanisms (RC and IC) would generalize to multi-layer models or more complex algorithmic tasks?\n\nFor the running example of the histogram task, in Appendix E.8., we briefly discuss that the phenomenology of the phase transitions we observed for single-layer models also transfers to the setting with two layers. This is reasonable, since it is formally possible to construct a single layer of token + feature mixing as the identity function, hence the construction for the one-layer case naturally generalizes to more than one. While it could very well be possible that an extra layer affords solving the task with e.g. fewer dimensions, we do not observe strong evidence of this. However, at the same time it is unclear which exact mechanism is at play in the learned models, and how the RC or IC mechanisms would distribute over the layers. \nAs for more complex algorithmic tasks, we do expect these mechanisms to generalize to tasks that involve a counting operation on a finite alphabet.\n\nWe hope that our response gives you some additional insights on our work as well as the surrounding literature. If there are further questions or concerns you would like to discuss, feel free to get back to us.\n\nBest,\nThe authors", "claims_and_evidence": "The paper’s claims about the feasibility of perfect histogram task solutions under different parameter regimes (e.g., d >= T, p = T) are supported with constructive proofs and empirical results. The observed phase transitions in model accuracy (Figure 1) closely match the theoretical predictions, enhancing credibility.", "methods_and_evaluation_criteria": "While the paper does not introduce a new method in the conventional algorithmic sense, its design of the histogram task as a diagnostic probe, and the use of a structured hyperparameter grid search, serve effectively as methodological tools to reveal the architectural properties of small transformers.", "theoretical_claims": "Probably.", "experimental_designs_or_analyses": "The main experiments (Fig. 1 and 2) demonstrate the impact of different hyperparameter combinations on model performance, and the analysis of the paper starts here. They clearly validate how the model performs under different settings (Sec. 4.1 and 4.2). Afterwards, the authors provide the attention matrix and feedforward prediction to help explain why the model works or struggles in each setting.", "relation_to_broader_scientific_literature": "The paper advances the field by providing explicit constructions and detailed phase-space analyses that clarify how subtle design choices lead to different computational strategies in transformers. The paper identifies two distinct counting strategies and deepens the understanding of how architectural choices affect the solution space. Moreover, the discussion on the impact of embedding orthogonality and mutual coherence is also inspiring.", "essential_references_not_discussed": "Not applicable.", "other_strengths_and_weaknesses": "Strengths:\n1.\tThe paper not only gives a clear theoretical construction of how small Transformers can achieve counting tasks (including relation-based and inventory-based counting), but also verifies these theoretical predictions through rigorous experiments. In addition, the authors provide sufficient theoretical derivation and clear experimental demonstration.\n2.\tBy carefully analyzing the role of attention mechanisms and feed-forward layers under different hyperparameters, the paper provides valuable insights into the inner workings, parameter efficiency, and algorithmic implementation of the Transformer. In addition, the authors' analysis provides inspiration for studying how to better apply Transformer-like models to other tasks.\n\nWeaknesses:\nThere are certain limitations in the scope of tasks and discussions on practical applications. This article mainly focuses on the relatively simple task of counting, and most of the experiments are based on a single-layer Transformer. Although this helps to gain a deeper understanding of the basic mechanism, the applicability of its conclusions in more complex practical tasks or multi-layer, large-scale models still needs further verification or remains questionable. In addition, the article mainly focuses on the discussion of mechanisms and theoretical levels, and lacks discussion on how to apply these findings to actual large-scale models or solve real-world problems.", "other_comments_or_suggestions": "No.", "questions_for_authors": "1.\tIn the paper, the term Token Mixing is used to generalize across the self-attention mechanism. Could the authors clarify whether this abstraction is primarily conceptual, or if it implies a formal equivalence? Are there theoretical reasons to treat attention as a subclass of token mixing, especially in the context of algorithmic tasks?\n\n2.\tGiven that the current study focuses on single-layer architectures and a relatively simple task (histogram counting), do the authors anticipate that the identified mechanisms (RC and IC) would generalize to multi-layer models or more complex algorithmic tasks?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "n/a"}}, {"review_id": "N7Sg4jZuAJ", "reviewer": "Reviewer_fzMo", "summary": "This paper investigates how small transformer architectures implement counting mechanisms in a controlled histogram task. The study identifies two distinct counting strategies: relation-based counting, which leverages local pairwise token comparisons, and inventory-based counting, which relies on memorization through a feed-forward network. The choice of counting strategy is influenced by hyperparameter configurations, such as embedding size, token-mixing mechanism, and the presence of softmax in attention.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer fzMo,\n\nWe thank you for taking the time and effort to evaluate our work. We are glad you found it clear, systematic and rigorous and that you also appreciate the controlled setting that the histogram task provides for our analysis. While we agree with you that these very same properties limit the generalizability of our analysis, we do think that there are some broader findings that we can take away from this -- we discuss them in the context of the questions you posed:\n\n1. > What is the potential for extending it to more complex structures, such as multi-layer architectures and multiple attention heads? \n\nAs mentioned in the response to reviewer S9FV, an empirical analysis of counting was conducted for 2-layer transformer architectures in appendix E.8, where we observed a similar phenomenology as for the case of a single layer. However, it remains unclear whether the mechanisms implemented by a two layer net correspond to the IC and RC mechanisms described for a single layer. We also believe that several layers help with the robustness of the network to entangled embeddings and allow for more self-corrections when there is noise from overlapping token embeddings, and extending our theoretical analysis to this setting would be interesting future work.\nIn the current version of our work we have not investigated multiple heads but preliminary experiments suggest a similar picture and we will provide specific results in a camera-ready version. Indeed, the question of what the specific functions of the heads are during counting has been analyzed in a similar context very recently in the literature, see https://arxiv.org/abs/2502.06923. \n\n\n2. > Can the proposed method be generalized to more complex tasks and larger datasets?\n\nWe interpret this question as asking how far the RC and IC mechanisms, as well as our analysis on entangled embeddings, generalize to more complex tasks and datasets. We anticipate that the algorithmic circuits that we discovered can be composed with other functions, which allows them to theoretically be applied to other datasets and tasks as well, when processing the input requires counting as a “subroutine”. In terms of the entangled embeddings and how they can be corrected via softmax and self-attention, our analysis is perhaps more general. It provides a framework to understand for which specific architectures (almost) orthogonality is needed for counting.\nMore generally, we hope that rigorous and systematic empirical studies on how different architectures and parameterizations influence the algorithms is more broadly recognized as a tool of choice in the regime where tasks are complex and large yet broader simulations of tasks are feasible.\n\n3. > What insights do the findings presented in the paper provide? How might these insights inspire future research in architecture design?\n\nOur findings reinforce the notion that small details in the architecture matter, which can have fundamental impacts on how given models achieve on given tasks, in agreement with other works that touch upon similar themes (e.g. https://arxiv.org/abs/2402.01032). In addition, our observation on the beginning of sequence token, which can drastically reduce the parameter size and input dimension needed for the counting task, adds to the growing evidence that “free tokens” help transformers to execute more complicated functions. Finally, it would be interesting to better understand how robustness and self-correction, as seen via the self-attention in the single layer transformer, play out in larger networks, and whether tweaks of the architecture can further reinforce this behaviour.\n\nWe hope that our response clarifies how we expect our analysis to transfer to more complex settings and how it affirms specific research directions in architecture design. If you have further questions or concerns you would like to discuss, feel free to get back to us.\n\nBest,\nThe authors", "claims_and_evidence": "The claims made in the submission are generally supported by clear and convincing evidence. The paper provides both theoretical constructions and empirical experiments to substantiate its findings.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make sense for the problem. The histogram task is a well-defined and controlled setting for studying counting mechanisms in transformers.", "theoretical_claims": "The arguments are logically structured, leveraging dot-product attention properties and embedding orthogonality.", "experimental_designs_or_analyses": "The task is simple and well-suited for isolating the effects of model components. The authors systematically vary key hyperparameters (embedding dimension, feed-forward width, and attention mechanisms) and analyze accuracy trends, supporting their claims. However, while their empirical results align with theoretical expectations, further testing on more complex tasks or different datasets could strengthen generalizability.", "relation_to_broader_scientific_literature": "The paper builds on mechanistic interpretability and transformer analysis, aligning with prior work on counting tasks, attention mechanisms, and neural network generalization.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "1. The architecture is confined to a single-layer transformer, which may limit its practical applications. \n2. The experiments were conducted on a simple task and dataset.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. What is the potential for extending it to more complex structures, such as multi-layer architectures and multiple attention heads?  \n2. Can the proposed method be generalized to more complex tasks and larger datasets?  \n3. What insights do the findings presented in the paper provide? How might these insights inspire future research in architecture design?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "226LdONtkR", "reviewer": "Reviewer_UA3n", "summary": "This paper investigates the counting mechanism behind transformer blocks using the histogram counting task as a case study. Two types of counting tasks were studied: relation-based counting leveraging attention for pairwise token comparisons, and inventory-based counting using feed-forward layers to memorize token counts via orthogonal embeddings. The study reveals that some minor architectural choices, such as embedding size, token-mixing mechanisms, and softmax application, significantly influence the model's ability to count accurately. The study highlights the delicate balance between attention and feed-forward layers in small transformers and provides theoretical insights into the architectural determinants of counting capabilities.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer UA3n,\n\nWe thank you for taking the time to evaluate our work and give us feedback. We are especially pleased  that you find the aspect of how softmax and the dot-product attention influence robustness interesting. We agree that this point deserves further discussion and for a camera-ready version we would like to use part of the extra page to give this aspect more space.\n\nThank you also for pointing out that Figure 4 (right) is difficult to understand from the caption. Let us elaborate here: \n\nRecall that the output of the token mixing mechanism (one of lin, dot, dot_BOS etc.) computes weights of the matrix $A(\\bar x)$ to mix the input embeddings $\\bar x$, where we denote a single weight as $a_{ij}$. So the token at position $\\ell$ after the token mixing is  $\\bar x^\\prime_\\ell = \\sum_{i = 0}^L a_{i\\ell} \\bar x_i + \\bar x_\\ell $ . \n\nIf we know which letters were in the token sequence (recall that $x_\\ell$ is the letter of the alphabet at position $\\ell$ and $\\bar x_\\ell$ is the embedding of that letter in the $d$-dimensional space), we can also write the mixed token directly as a sum over the alphabet $\\bar x^\\prime_\\ell = \\sum_{t \\in \\mathcal T} \\alpha_t e_t + e_{x_\\ell}$, where $\\alpha_t = \\sum_{i =0}^L \\delta(x_i = t) a_{i\\ell}$. \nIn this context, Figure 4 examines how the feature extractor depends on different compositions of tokens that are mixed with different magnitudes of $\\alpha_t$ directly. We select three letters, and combine weighted sums of them to see how the values $\\alpha$ relate to the final count that is predicted by the feed-forward layer. We look at the case where the softmax function is part of the token mixing so we have $\\sum_{t \\in \\mathcal T} \\alpha_t = 1$ guaranteed. This allows us to isolate the decision boundaries of the feature extractor along the different counts in terms of $\\alpha$ in the plots on the right-hand side.\n\nWe find this a useful perspective, because it allows us to qualitatively compare how the feature extractors differ that are learned by the lin+sftm and dot+sftm. We clearly observe that in one case (dot+sftm) the decision boundaries scale non-linearly in alpha, whereas in the case of lin+sftm they are almost linear. This verifies our hypothesis that the architectures lead to different solutions and shows the analogy with our manual constructions. We also discuss at a later point in the manuscript, that these different scalings influence the robustness of the feature extractor.\n\nWe hope that this response clarifies the introspection experiment from Figure 4 further and we will expand on this aspect in an improved version of the manuscript to make it more friendly to the reader. If you have further feedback, doubts or concerns you would like to discuss, don’t hesitate to get back to us.\n\nBest,\nThe authors", "claims_and_evidence": "**Claims are supported by evidence.** \n\nThis study claims that counting performance in transformers depends on architectural design choices. Empirical results first confirm that attention facilitates pairwise comparisons (relation-based counting), while feed-forward layers store token counts (inventory-based counting). Theoretical analyses then further demonstrate that softmax normalization and embedding size are two key variables.", "methods_and_evaluation_criteria": "No specific method is described by the paper.", "theoretical_claims": "The paper developed several theories to analyze the countability under ideal conditions. The reviewer does not check the theory very carefully, but assumptions on embedding orthogonality and softmax effects seem make sense.", "experimental_designs_or_analyses": "The paper only conducts toy-level experiments. The experimental design for the histogram task using variants of token mixing models, including bos, lin, and dot, and with/without softmax, was analyzed.", "relation_to_broader_scientific_literature": "NA.", "essential_references_not_discussed": "Sufficient.", "other_strengths_and_weaknesses": "# Strengths:\n\n- **An enlightened work**: \nAs far as I can tell, this is the first work analyzing the countability of transformer blocks, which advances our understanding behind transformer mechanisms. The discussion of d and T points out an optimization direction for subsequent work. \n\n- **Theory contribution**: \nThe paper derives theoretical bounds for minimal embedding dimensions and the role of softmax in error reduction.\n\n# Weaknesses:\n\n- **Unfriendly figure illustrations**:\nThe explanations of some figures are not immediately understandable to the reader. Perhaps they could be expressed in more friendly way (e.g., Figure 4).\n\n- **Practical applicability**:\nThe current experimental setup may be too simple. The observations may not generalize well to other real-world counting tasks, e.g., visual counting.", "other_comments_or_suggestions": "In real-world scenarios, d may be significantly smaller than T, and computations cannot converge to infinite precision. The counting robustness may be discussed further.", "questions_for_authors": "I cannot see the point of Figure 4 (Right). A more friendly explanation is expected.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DSU036JYe9", "reviewer": "Reviewer_Jx43", "summary": "This paper shows how a small Transformer can implement robust counting by arranging token embeddings with sufficiently low overlap and then leveraging architectural components (like softmax and BOS tokens) to preserve that separation under mixing. The authors demonstrate that different choices in hyperparameters control the network’s expressiveness to store and retrieve count information.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer Jx43,\n\nWe thank you for reading and evaluating our work, and providing us with feedback. We are glad you generally appreciate the soundness of our work and the approach to change hyperparameters to understand which parts of the architecture are impactful.\n\nYou are indeed correct, in that our analysis rests on perfectly trained models (or as perfect as it gets, when the model does not afford a perfect solution). While we did observe some peculiarities to the training dynamics of the different architectures, we have not formally investigated the training dynamics themselves and therefore not shared results in the manuscript. These informal observations showed that e.g. linear mixing seems to observe a staircase behaviour in the loss, learning one letter of the alphabet at a time. In contrast, the dot-product mixing did not show such a staircase. Intuition from other toy problems (such as induction heads) also lead us to believe that the way the single counting subspace that arises in the models with BOS emerges apruptly. While we find these very preliminary intuitions gained during our experiments very interesting, they warrant a more formal treatment which we hope to tackle in future work.\n\nWe hope that this answer somewhat satisfies your curiosity. However, since the above statements are very speculative and we would prefer not to include them in the manuscript. If you have further questions, feedback or concerns you would like to discuss, don’t hesitate to get back to us.\n\nBest,\nThe authors", "claims_and_evidence": "The authors claimed that the histogram task can be learned via relation-based approach, or an inventory-based approach.\n\nThe claims are supported by:\n1. Explicit weight constructions to realise the expressiveness for perfect outputs.\n2. Experiments on synthetic sequences.", "methods_and_evaluation_criteria": "The method is based on a specific simple Transformer architecture. It seems to be reasonable for understanding the effects of difference hyper-parameters.", "theoretical_claims": "The propositions in Section 4 seem to be sound.", "experimental_designs_or_analyses": "The experiment designs are reasonable.", "relation_to_broader_scientific_literature": "This work analysed the roles of BOS tokens and some other Transformer configurations, which is related to earlier work on RASP(-L) (Wess et al 2021).", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n\n- The analyses seem very rigorous.\n\nWeaknesses:\n\n- A toy-like model on a specific task\n- Little analyses on the training dynamics. The analyses seem to be based on perfectly trained end-models.", "other_comments_or_suggestions": "It would be interesting to have more analysis on how the geometric features change during the training.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: All of reviewers lean toward acceptance post-rebuttal. The AC checked all the materials and concurs that the paper has made meaningful and insightful analysis of how small Transformers implement counting in two mechanisms -- relation-based and inventory-based. While there are limitations (e.g., the focus on single-block), the work presents a valuable initial attempt to analyze the countability of Transformer blocks, and therefore should be accepted. Please incorporate necessary changes in the final version.
Decision: Accept (poster)