Dear reviewer Jx43,

We 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.

You 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.

We 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.

Best, The authors

Dear reviewer UA3n,

We 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.

Thank you also for pointing out that Figure 4 (right) is difficult to understand from the caption. Let us elaborate here:

Recall 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$ .

If 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}$. In 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.

We 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.

We 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.

Best, The authors

Dear reviewer S9FV,

We 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. We 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!

In the following we would like to answer your questions:

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?

As 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.

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?

For 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. As for more complex algorithmic tasks, we do expect these mechanisms to generalize to tasks that involve a counting operation on a finite alphabet.

We 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.

Best, The authors

Dear reviewer fzMo,

We 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:

1. What is the potential for extending it to more complex structures, such as multi-layer architectures and multiple attention heads?

As 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. In 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.

2. Can the proposed method be generalized to more complex tasks and larger datasets?

We 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. More 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.

3. What insights do the findings presented in the paper provide? How might these insights inspire future research in architecture design?

Our 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.

We 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.

Best, The authors