From Interpolation to Extrapolation: Complete Length Generalization for Arithmetic Transformers

Thank you for your insightful review and valuable feedback. We have revised the paper based on your suggestions. Before we address your specific questions, here is a summary of the improvements we made to the paper: 1. We simplified the extensive notations used in section 5 and added more visuals to explain the concept more clearly. 2. In section 8, we added a discussion on the scope and limitations of our work to clarify its applicability. 3. We conducted additional experiments with ListOps and discussed the results in section 8 and appendix C. 4. We added citations to the works in Systematic Generalization. Thanks for pointing them out! 5. We improved section 4 by showing the attention patterns in the main text. **Adding a task where the attention pattern should not be a straight line would be interesting** Thank you for such a great suggestion. ListOps is an excellent test case for learning more complex structures. We conducted additional experiments with ListOps and discussed the results in section 8 and appendix C. Unfortunately, with the model scales we experimented with and the amount of resources we had, we could only solve a simplified single-operator, single-depth ListOps task via ABS. Learning multi-depth ListOps may require scaling up the model dramatically, as the heuristic in the NDR paper [2] indicates. Our goal is certainly to study the learning of more complex tasks and we plan them as future research. **Does this mean that only the last decoder layer has the attention bias? Or do all of them have it, but set to the same as the last layer?** We apologize for the confusion. For the tasks presented in this work, the bias matrices are obtained from the last decoder layer and applied to all layers during re-train, on a per-head basis. More flexible configurations such as per-layer bias could work better for more complex tasks, which is what we currently are exploring. We added clarification on this issue at the end of section 5. **Given this, in the first eq. on P6, what are the additive p_i and p_j components? Do you do a custom method, where you inject the PE directly into the attention at each layer?** Again we apologize for the confusion. We apply PE at the start of the process as the original Transformer paper does [1], rather than injecting it at each layer. The notation we used is misleading and we have revised it in section 5. Thanks for pointing it out. **How do you handle the cross-attention? Does that have this bias as well?** We used ABC on cross-attention as well. It works the same way as self-attention, and Figures 4 and 5 are both cross-attention patterns. **Have you tried a less aggressive bias, where you don't set the elements to -inf, but to some finite negative number, such that the network has the chance to overcome the bias if it is good for the task? Do you think it would work in this case?** This is an excellent idea and we are in the process of trying it out on NLP tasks which we think is a more suitable situation. We believe that in NLP, ABC could be possibly used to capture some global or background dependency patterns. However, the patterns are likely to be less rigid and clear than those in arithmetic tasks. A "soft" ABC, i.e., something like $$attn = (1-\lambda) \cdot attn_0+\lambda \cdot bias$$ where bias acts as a global context while $attn_0$ captures the local dependency, might work better. This is our ongoing work. Thanks for your suggestion! **If the authors could run their method on any task that does not require strictly monotonic attention** We have run experiments on ListOps which we discussed earlier. Thanks again for the constructive feedback that helped make this paper better! We hope our responses addressed your concerns. **References:** [1] Attention is All You Need [2] The Neural Data Router: Adaptive Control Flow in Transformers Improves Systematic Generalization

I would like to thank the authors for their updates and clarifications. Most of my questions have been answered. I find the outcome of the ListOps experiments very surprising. The authors use a simplified version of ListOps, where there is only a single operation applied to a list of numbers. My original intent was the opposite: to run it on recursively embedded operations, to see if the model can allow non-monotonic attention patterns. The surprising part of the current results is that for this simplified version of ListOps, monotonic memory access patterns should be enough, so I think the method should work (and my prediction was that it would fail on the original, recursive version). Given these negative outcomes, I think the authors should emphasize the limitations of the paper more so that it can be clear to the reader that this is not a ready-to-go solution for length generalization. The other remaining weakness is not compared with relative positional encodings. Overall, I think the paper contributes an interesting new idea. I consider it worth publishing even if it is not clear yet if it is widely applicable to other tasks. It could serve as a base for interesting future research.

Thank you for your insightful review and valuable feedback. We have revised the paper based on your suggestions. Before we address your specific questions, the following is a summary of the improvements we made to the paper: 1. We simplified the extensive notations used in section 5 and added more visuals to explain the concept more clearly. 2. In section 8, we added a discussion on the scope and limitations of our work to clarify its applicability. 3. We conducted additional experiments with ListOps and discussed the results in Section 8 and Appendix C. 4. We added citations to the works in Systematic Generalization. 5. We improved section 4 by showing the attention patterns in the main text. **The important results from section 4 are hard to understand if one only reads the main paper.** We agree that section 4 was a bit crammed and some material was left in Appendix A. We have reformatted sections 4 and 5 to be more clear and easy to understand. **The claim that constrained attention, and ABC, offers a complete ...** We apologize for the confusion. By "complete length generalization" we didn't mean that we solved all length generalization problems. Instead, we meant that our solution solves some specific arithmetic algorithm learning tasks, allowing the model to generalize (completely) to long lengths on those tasks. Ours was the first solution to achieve such results. We clarified the notion in the paper and added a discussion on the scope and limitations of our work in section 8. **You use 1/6 layer transformers, what is the benefit of a shallow encoder, compared to a (smaller) 6/1 encoder, or 4/4 model?** We have conducted tests with 6/6, 3/6, 3/3, and 1/3 transformers. We see little to no difference regarding both interpolation and extrapolation performances. We suspect that this is due to the simplicity of our tasks. **You observe that it sometimes interferes with the sliding window attention, could you try learnable PE?** Your suggestion to examine learnable positional encodings (PE) is well-taken, but we did not try learnable PE because studies have shown that learnable PE does not achieve good results when doing length extrapolation. More specifically, as [1] mentioned, "We also experimented with using learned positional embeddings instead, and found that the two versions produced nearly identical results". The learnable PE referenced here cannot enable extrapolation because all the encodings are learned on the in-training length. Our decision not to include learnable PEs in our experiments was also influenced by hardware limitations. **Could you try RPE on the successor and parity experiments? Your conclusions suggest that it would help generalize.** The potential of relative positional encodings (RPE) to improve length generalization is a promising direction for future research. While we have not included RPE in our current experiments, we have contrasted ABS with the top-performing variant of RPE as per [2] in the appendix, showing the superiority of ABS in Parity. **In the experiments from Appendix A.5 (table 3), you try cyclic PE on top of a w=1 windowed attention, what happens if you use it without windowed attention?** We also think that cyclic PE has some potential to act as a sort of RPE. However, when we tested cyclic position indexing (CPI) without attention windows, it did not yield satisfactory results and the model did not even interpolate. CPI appears effective only on very short windows. **Have you tried $n×n$ multiplication (for smaller values of $n$)? of $n×k$ for $k$ larger than 3 (so that the model cannot memorize all cases as so many unary operations)?** We have attempted $n×n$ multiplication and $n×k$ multiplication with $k$ larger than 3. Unfortunately, these experiments did not yield successful interpolation, rendering ABC unusable in this context. We have addressed this limitation in the revised paper in section 8. **The Jelassi results in Table 3 and Figure 1 correspond to a different architecture (encoder-only, with shared layers), it would be better to rerun them, or at least to clarify this.** We apologize for the misleading caption, this has been revised. Thank you for your insightful response. **References:** [1] Attention is All You Need [2] Neural Networks and the Chomsky Hierarchy

Thank you for your insightful review and valuable feedback. We have revised the paper and the following is a summary of the main changes: 1. We simplified the extensive notations used in section 5 and added more visuals to explain the concept more clearly. 2. In section 8, we added a discussion on the scope and limitations of our work to clarify its applicability. 3. We conducted additional experiments with ListOps and discussed the results in Section 8 and Appendix C. 4. We added citations to the works in Systematic Generalization. 5. We improved section 4 by showing the attention patterns in the main text. Now addressing your questions: **The presentation of this paper looks good in the first two sections, but the presentation quality drops in the subsequent parts.** We have revised Section 5, reducing the complexity of notations and augmenting the textual explanations with additional visuals to facilitate a clearer understanding of the concepts presented. **While it can be helpful for simple tasks, I’m not convinced of its effectiveness in more general settings (e.g., for natural language)** We added a discussion on the scope and limitations of our work in section 8. The current sparse attention bias works for the arithmetic tasks we study. For more complex tasks, it is likely to require a scaling up of the model, as the heuristic in the NDR paper [3] indicates, and much more expressible ways to implement the bias. Although exploring ABC in more generalized settings is beyond the scope of this current paper, we intend to address this in our subsequent research. Sparse global attention has been shown to have a great effect on the transformer model's performance. More specifically, in [1], they demonstrated that the transformer model can still hold their abilities even when the entire attention mechanism is replaced with a static attention matrix generated from averages of the attention matrix across a corpus of data, only suffering from an 8% decrease in accuracy in some cases. These findings, alongside ABC’s conceptual similarity with Relative Positional Encodings (RPE), provide a promising foundation for future investigations into ABC's utility in complex NLP tasks. We are currently in the process of exploring more tasks in NLP. **Can you compare ABC with more advanced baselines?** The suggestion to benchmark ABC against more sophisticated baselines is well-taken. While ABC has shown remarkable performance in the experiments conducted, outperforming all relevant baselines in [2] and [4] and achieving exceptional accuracies of 99+ percent over extended sequence lengths (over 8 times the training length), we acknowledge the value of broader comparisons. The benchmarks that we compare against include results from [2] and [4], which are very recent works and include extensive experiments. For example, [4] considered five major positional encodings: none, classical sin/cos, RoPE, ALiBi, and the relative positional encodings of Transformer-XL, and report the best-performing configurations. We believe they are strong SOTA results. Yet the transformer did poorly on those tasks. This indicates that the current state of transformers on arithmetic learning is rather weak. And ABC enables the model to generalize on certain tasks which we believe is a big step forward. **References:** [1] How Much Does Attention Actually Attend? Questioning the Importance of Attention in Pretrained Transformers [2] Randomized Positional Encodings Boost Length Generalization of Transformers [3] The Neural Data Router: Adaptive ControlFlow in Transformers Improves Systematic Generalization [4] Neural Networks and the Chomsky Hierarchy

Thank you for your valuable feedback! We have uploaded a revised version of our paper and will address your questions below: **Writing Clarity:** We have made changes to improve readability and make the line of reasoning more clear. In particular, we simplified the extensive notations used in section 5 and added more visuals to explain the concept more clearly. **ABC's Generality and Use Cases:** In section 8, we added discussion on the scope and limitations of our work to clarify its applicability. Ours is an initial attempt to study the roles of attention in transformer’s learning. To this end, we need to find tasks on which transformers can generalize. Surprisingly, even for tasks as simple as addition, no previous transformer model ever succeeded. We take this as an indication that our understanding of the model’s learning mechanism is inadequate. Therefore, the successful cases obtained in this paper, either via ABS or ABC, even though only on simple recurrent patterns and task-specific models, solve a few long-standing difficult or even “impossible” tasks (e.g., Parity) and represent a significant step forward. **Additional Tasks:** Additional and more complex tasks appear to require composition of simple patterns. Learning such compositions may require scaling up the model dramatically, as the heuristic in the NDR paper (Csordás et al., 2021b) indicates. We are exploring them in our current research. **Factual Errors:** You are right on both points, and we have made appropriate modifications in the revision. Thank you very much for pointing them out! Those are very keen observations! Again, we thank you for your valuable feedback and we hope that our answers and revision addressed your questions.