Transformers Can Do Arithmetic with the Right Embeddings

Summary

This paper introduces a simple yet effective encoding scheme that can be used to address the limitations of transformers at representing positional information, which is crucial in many algorithmic tasks such as those involving arithmetic operations. The authors propose an ad-hoc positional embedding, called “abacus embedding”, which encodes the location of each digit relative to the start of the current number and thus provides an explicit signal that the transformer can use to align digits. The effectiveness of the method is tested on addition, multiplication and sorting problems, with a particular focus on out-of-distribution test cases.

Strengths

I think that this work is interesting and relevant. Although addition, multiplication and sorting problems might be considered trivial test cases because they can easily be solved with symbolic algorithms, they constitute an important benchmark to evaluate the algorithmic reasoning skills of neural networks, as also attested by the increasing interest of the deep learning community on mathematical tasks. The paper is well-written, and the method is clearly presented. The generalization achieved in the addition task is quite impressive, showing that the abacus embeddings enable a generalization factor of 6x in the OOD regimen. Although simple and straightforward, the proposed method seems original.

Weaknesses

The abacus embeddings are defined according to the hyperparameter k, which is fixed a priori (e.g., k = 100). This limits the flexibility and generalizability of the proposed encoding scheme. The authors deploy different architectures / hyperparameters to learn different problems (addition vs. multiplication vs. sorting). Since they argue that their architecture modification “improves performance on multiple algorithmic reasoning tasks simultaneously” it would be important to show that different tasks can really be learned simultaneously, without the need to build ad-hoc models for each algorithmic task that needs to be solved. It is true that arithmetic operators are binary and thus “both addition and multiplication accept only two operands”. However, we can have a sequence of additions / multiplications, and it is well-known that also increasing the number of terms in arithmetic expressions causes troubles to transformers. Because of these key issues, I think that the impact and significance of this work are not strong enough for a top-tier venue like NeurIPS.

Questions

- How could we address the fact that the hyperparameter k needs to be fixed a priori? - Can we implement a unified model that can learn all these tasks simultaneously? - The proposed method achieves impressive OOD accuracy for addition, but only works “in distribution” for multiplication. It would be important to investigate this phenomenon more in depth. - At least for addition, it would be useful to test OOD generalization by also adding more operands besides increasing the length of each operand. - How does the present method compare to other recent proposals such as xVal (https://arxiv.org/abs/2310.02989)? - I agree that addition, multiplication and sorting are good benchmarks because they are simple yet challenging; however the authors could better stress that these tasks are part of a broader class of elementary numerical skills that transformers struggle to learn (for a comprehensive review, see https://www.mdpi.com/2076-3417/14/2/744).

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

The authors properly addressed the limitations of their study.

Summary

This paper studies a well-known problem, the length generalization issue of transformers in terms of doing arithmetic. This paper solves this problem via two natural strategies: (i) separate two operands via a newly proposed embeddings (Abacus Embeddings), and (ii) using looped Transformer architecture.

Strengths

- The problem is well-motivated. - The conjectures are very natural, and confirmed via extensive experiments. - Experiments are well-designed and complete. - The proposed solutions enjoy great performance. - Considers diverse downstream tasks, including addition, multiplication, and sorting.

Weaknesses

[Medium] The reason why looped transformer or recurrency helped with length generalization is still unclear, and in-depth analysis is needed. For instance, does the number of recurrence related to the length of the digits?

Questions

Why does recurrency in terms of the model architecture help with length generalization?

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

See weakness.

Summary

The paper studies the arithmetic capabilities of transformers and the problem of length generalization, specifically the ability to solve problems larger than the ones seen during training. It introduces Abacus Embeddings, a novel positional embedding that encodes the position of each digit relative to the start of the number. For multi-digit addition, Abacus Embeddings result in state-of-the-art generalization to sequences six times longer than the training sequences. Additionally, the paper explores the benefits of incorporating recurrent blocks, leading to further improvements. Finally, the paper demonstrates the effectiveness of Abacus beyond addition, showing success with in-distribution multiplication and array sorting tasks.

Strengths

- **Originality:** While previous studies have noted that positional encoding can negatively impact the arithmetic generalization capabilities of transformer architectures, to the best of my knowledge, the introduced embeddings, the analyses, and the results presented in this paper are original. - **Quality and clarity:** The work is technically sound. The experiments are well-designed and convincing, and the code for their implementation is provided in the supplementary material. The paper is clearly written. - **Significance:** The goal of improving the extrapolation and reasoning abilities of transformers is both timely and significant. The results are very good, achieving state-of-the-art performance for length extrapolation in multi-digit addition.

Weaknesses

1. The paper does not discuss the choice of the value used for the maximal offset randomization parameter $k$, which determines the distribution of the starting position of the first digit of the numbers during training. How was the value $k = 100$ chosen? Additionally, could higher values further improve extrapolation performance? 2. The sentence “our methods perform so well that we look beyond addition” (line 239) does not sound appropriate for a scientific paper. Please consider rephrasing it, for instance, “Given the strong performance of our method in the multi-digit addition task, we extend etc.”.

Questions

3. Differently from addition, when considering multiplication, Abacus Embeddings achieve high in-distribution accuracy but struggle out-of-distribution (OOD), even when one operand is of unitary/short length. Do you have any insights about what might cause this significant difference? Do you have any ideas or potential modifications to the method that could improve OOD generalization in this context? 4. Could you elaborate more on how your embeddings could be integrated into settings that involve mixing arithmetic with natural language data?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper adequately discusses its limitations. I do not foresee any potential negative societal impacts arising from this study.