Paper: Carrying over Algorithm in Transformers (ICLR-2024)

Reviewer comment (Experiments):
While I agree that digit arithmetics constitutes a valid testbed to build mechanistic accounts of the inner workings of deep learning models, the current task setting seems oversimplified. To establish the emergence of a putative carrying over algorithm, the authors should demonstrate that such algorithm allows to systematically solve the target problem, for example by testing the model over problem instances involving longer digits and more operands than those observed during training.

Author response:
It is true that exploring the task for larger integer addition is a very important one and that the carrying over algorithm works for integer addition of any length. As we mentioned in appendix G, the model trained on 3 digit addition does badly on 5 digit addition. However, we believe this does not mean the carrying over algorithm is not implemented, it just means that the carrying over algorithm is implemented for 3 digit addition and partly for 5 digits. To be more specific, the model can still do the summations at each digit and carry a one, but gets 'confused' when more digits are present. We will add a discussion about longer lengths to the paper.

Reviewer comment (Experiments):
The authors frame their research as an investigation of mathematical reasoning abilities of LLMs, arguing that a fruitful approach is to study simplified architectures and then extrapolate the results to larger models. While I agree that this might be a reasonable approach, the authors provide only “suggestive evidence” that their findings extend to one LLM (Alpaca). I think it would be more effective to directly frame this contribution as an investigation of the computational capabilities of simple Transformers architectures, and mention as a side point their putative application to LLMs.

Author response:
We agree that the mentioning of LLMs is mostly motivational and study of Alpaca 7B provides suggestive evidence. Clearly more study is necessary to make a more concrete link between the small and large models. See also our response to your last question.

Reviewer comment (Writing):
English phrasing and grammar could be improved.

Author response:
Regarding the English phrasing/grammer, we will go over the text again to fix that.

Reviewer comment (Novelty):
The “Related works” section only mentions a few studies related to the emergence of grokking in Transformer architectures. This section should be extended by including relevant papers that have investigated the capability of Transformer architectures to learn arithmetic problems. For example, others have proposed how to improve the standard Transformer architecture to improve generalization on simple arithmetic tasks [1], possibly by exploiting grid-like problem representations that promote the emergence of carrying over algorithms [2]. More recent work has also studied how systematic length generalization in basic integer arithmetic might be improved using relative position embeddings [3].

Author response:
We thank the referee for mentioning these other references, they look indeed very interesting for studying ways to improve the standard Transformer architecture. However, we would like to mention that in our work, we were not necessarily focussing on improving the standard Transformer (we only included RoPE) as to the best of our knowledge many much larger models still use the standard Transformer (and often with a rotary embedding). This might be the reason why Refs. [1, 2] have not caught our attention, but certainly we should have included Ref. [3].

Reviewer comment (Experiments):
To demonstrate that Transformers are learning an algorithm, their behavior should be tested also with out-of-distribution problem instances (as usually done in the related literature, see papers mentioned above). This would include longer operands than those used during training (> 3-digits) but also more operands than those used during training (> 2, in this case).

Author response:
We agree with this comment and will add some additional discussion to the paper.

Reviewer comment (Evaluation):
A related point is the issue of memorization. The authors should better explain why they believe that one-layer models are memorizing the solutions, while two-layer models are actually learning a systematic behavior.

Author response:
We thank the referee for brining this up, the memorization discussion in the appendix B just addresses the one-layer model, which does not learn the addition problem perfectly. It does however, do the summation of the digits. What we meant with memorization here is just the fact that the residual stream seems to have large spread and not many localized pockets as we saw in the two-layer model. This indicates that there is not a special subspace in the residual stream dedicated to a particular feature of the data (like carried one necessary or not), but rather it is just treating each example differently. However, we do notice some interesting ordering, so it is also not just memorizing either. Perhaps memorization is not the correct terminology here, we will rewrite this part accordingly. The 2-layer model learned something systematic, as we illustrated in Fig. 1, because i) the residual stream is much more divided up into pockets related to specific feature of the data. ii) 2-layer models 'understand' where to add a carried one (from the attention patterns and the residual stream in the second layer) and iii) also have a particular component (the final MLP) that does that actual carrying of the one. We will try to clarify this further in the paper.

Reviewer comment (Experiments):
The authors argue that “Usually the carrying over algorithm is employed in a right to left fashion”. This is the “human way” to implement the carry over algorithm, which might be due to processing constraints (we prefer to process things sequentially when the problem requires deliberate reasoning). This might not apply to Transformers, which should learn to parallelize the computation whenever possible (also see [2]) In line with this observation, it would be interesting to better investigate model behavior on the “Non-carry” (NC) scenario (subset #1). In such case, I would expect the Transformer to execute the entire operations in parallel.

Author response:
The referee is absolutely correct here, we indeed meant here the human way and did not try to make any claim the Transformer should or should not do this. The parallelization is something we also see from the attention patterns. They provide the addition of the digits all in parallel. We thank the referee for mentioning this parallelism, which we have incorporated in sections 3 and 4. Any carried one is added later in final MLP of the two-layer network. This is indeed particularly clear in the non-carry task.

Reviewer comment (Writing):
What is the rationale for using three = at the end of the arithmetic operations, rather than just one?

Author response:
For the encoder-only models, the rational for the three equal signs was to read off the answer above these positions. One could also have done different things, for instance not use any equal signs at all, but we have not experimented with other choices.

Reviewer comment (Evaluation):
“The attention patterns are averaged over the test dataset”: shouldn’t the authors average only similar cases, to avoid mixing carry vs. non-carry cases (or even according to subsets identified by the authors at pg. 3)?

Author response:
Here we were interested in the overall structure of the attention patterns, i.e. whether the staircase pattern would develop or not and present the reader with the transitions. This did not require worrying about the different tasks, which was relevant for two or more layers, i.e. figure 3.

Reviewer comment (Evaluation):
Phase transitions: from Fig. 2 it is not evident that phase transitions are actually occurring. The figure only shows that at some point the attention pattern changes, but in order to be a phase transition such change should happen suddenly (e.g., between two consecutive epochs). Furthermore, following the authors’ narrative this should happen not just during two consecutive epochs, but during the epochs corresponding to the abrupt changes in the loss function. From the figure, this does not seem to be the case. Can the authors clarify this?

Author response:
Here we plotted the attention patterns within each phase and checked that indeed the pattern suddenly changes whenever the loss drops (or accuracy increases suddenly). We will clarify this a bit more in the main text.

Reviewer comment (Presentation):
In the figures, the curves representing trends of just one individual run should be replaced with curves representing mean and variance (or min-max) of the six multiple runs.

Author response:
We thank the referee for this comment, but the bold curves are the mean and so the variance can be roughly deduced from the plots since there are not that many. We can adjust this to the referee's suggestion.

Reviewer comment (Experiments):
Extending the present findings to LLMs is not easy, though the anecdotical evidence provided by the authors suggests that there might be some commonalities. Since a major limitation of Alpaca under the setup investigated by the authors was the overall low accuracy, it would be interesting to study whether prompting techniques could improve model performance, and maybe promote the emergence of a more interpretable algorithmic strategy. Also testing other LLMs might allow to establish whether more accurate models would use more systematic strategies at inference time.

Author response:
We thank the referee for mentioning this. We agree this is indeed a limitation. As we have discussed in replies to other reviews, we have tested our hypothesis on llemma-7b (https://arxiv.org/abs/2310.10631) as well. We will discuss this further in the paper and hope this addresses the question of the referee.

Reviewer comment (Experiments):
Except for Alpaca 7B, the authors should provide more discussion for other transformer models in LLM and other modalities tasks.

Author response:
We thank the referee for bringing this up. We agree that other modalities would be interesting to study. However, this would be particularly interesting for studying the training dynamics or if the different tasks have something in common, how they affect the implementation of any algorithm. The focus of this work was more on the final trained model instead for a single task. We believe the task considered here is already somewhat rich enough. However, we do believe that testing on other LLMs is important, so we have added also tested our hypothesis on llemma-7B (https://arxiv.org/abs/2310.10631). See reply to the other reviews. We will incorporate this in a revision.

Reviewer comment (Novelty):
The innovation is limited as the main ideas have been circulated outside carrying over algorithm in Transformers.

Author response:
We thank the referee for the vote of confidence here! We believe the novelty is there (see also replies to other reviews), but indeed it is a specialized setup, but we believe still relevant for the ICLR community.

Reviewer comment (Experiments):
It would be better for authors to provide a more detailed discussion on large-digit addition (e.g. 5-digit, 6-digit, etc), since it can be a more complex task and may have different situations with 3-digit addition.

Author response:
In appendix G, we studied 5 digit addition when the model is trained on just 3 digit addition. This discussion is rather short and we agree with the reviewer it would be good to extend the discussion there and in the main text, which we will do in a revision. Notice also that the 4 digit case has already a larger number of different situations and is discussed in appendix F.

Reviewer comment (Experiments):
In this work, the authors analyze the attention pattern for 2 heads in each layer for one-layer and two-layer encoder/decoder-only models. It would be interesting to consider more heads for each layer in these small models.

Author response:
Thank you for the suggestion, that is indeed be interesting, but already in the two headed case, we saw there was only one head that was doing the heavy lifting and so we think increasing the number of heads would just isolate again a single head (in each layer) that would have the same job as in the two headed case.

Reviewer comment (Reproducibility):
Results of the figure 13 need to be cross verified.

Author response:
Thank you for your review and comments! What do you mean with the results of figure 13 need to be cross verified?