Position Coupling: Improving Length Generalization of Arithmetic Transformers Using Task Structure

We are grateful for the reviewer’s effort and constructive feedback. Below we summarize your feedback/questions and address these one by one.

W1. Scope of Application: the generalizability of position coupling to broader and more complex tasks is not convincingly shown.

• We agree that the applicability of position coupling to broader tasks has not yet been explored in this paper. However, we want to note that enhancing the length generalization ability of models, even for the addition tasks, is considered an important problem as addressing and improving the arithmetic abilities of models can lead to a better understanding of the model’s capabilities.
• Definitely, the next step of our work is to extend the application of position coupling to complex, real-world tasks. To do so, we plan to develop a method called “automatic coupling”. Unlike position coupling which requires manual assignment of position ID for each token, it will automatically learn the task structure and assign appropriate position IDs without human intervention. This advancement will enable broader applications beyond simple arithmetic/algorithmic tasks. We believe our findings can serve as a stepping stone towards this future development.

W2. Complexity and Practicality: Position coupling introduces additional complexities in PE, which may limit its practical applicability.

• The introduction of position coupling does add complexity, particularly in designing task-specific couplings. However, once the design is determined, the additional overhead due to its implementation is negligible, as explained in lines 659-662.
• Furthermore, in contrast to approaches such as [1] and [2] that employ index hinting which requires doubling the input sequence length, position coupling does not increase the sequence length. As a result, we believe that there are no significant additional complexities (e.g., time, memory) introduced during the training phase.
• Additionally, we believe that the development of “automatic coupling”, as mentioned earlier, has the potential to fully address the complexity issue.
• If there are remaining concerns regarding additional complexity introduced by our method, please let us know without hesitation.

W3. Proper Citation of Related Work: “Randomized Position Encodings Boost Length Generalization of Transformers” (ACL 2023) is not properly cited.

• We agree that there are some similarities in the underlying concepts, and therefore we will add the citation and provide a more detailed comparison.
• We note that our method (described in lines 113-115) differs from that of Ruoss et al., 2023. Our method randomly selects the start position ID and assigns a sequence of consecutive numbers. In contrast, Ruoss et al., 2023 assigns a sequence of increasing integers, which are generally not consecutive.

Q1. Can position coupling be generalized to other PE schemes than APE?

• Thank you for an insightful question. During the rebuttal period, we realized that our proposed coupling method could be extended to Relative PEs such as RoPE and T5’s relative bias. The relative position between the query and key is determined by the difference in position IDs that were assigned by our position coupling method. Specifically for RoPE, we conducted experiments and observed that position coupling enhances the length generalization capability of RoPE (See Fig. 4 of the PDF file in our Global Response). We believe that this approach has significant potential for the research of length generalization and we will add experiments in our next revision.
• Instead of adapting our method to other PE methods than APE, it is also possible to integrate position coupling with existing RPE methods such as standard RPEs, RoPE, or FIRE [3]. In this approach, RPE methods are used independently alongside position coupling, which may provide hope for application to general LLM models. Thus, we think this is another promising direction to further improve the applicability of our proposed methodology, so we will conduct some experiments and consider adding the results to our final manuscript.

We hope our response has adequately addressed the reviewer's concerns, and we would appreciate it if you could reconsider your assessment.

---

References

[1] Zhou et al., What algorithms can transformers learn? a study in length generalization. ICLR 2024.
[2] Zhou et al., Transformers can achieve length generalization but not robustly. arXiv preprint, 2024.
[3] Li et al., Functional interpolation for relative positions improves long context transformers. arXiv preprint, 2023.

We deeply appreciate the reviewer's valuable questions and rich comments, and we hope our response relevantly addresses all points raised in the review.

W1. Some of the contributions are limited to 1-layer Transformers.

• We first highlight that our theoretical construction (Thm 5.1) is not limited to 1-layer models because it naturally extends to multi-layer cases.
• As you pointed out, our impossibility result on NoPE may not hold for a multi-layer model. However, the main purpose of this result is to show a provable gap in the expressive power between position coupling and NoPE. Therefore, we believe that the inapplicability of the impossibility result in multi-layer models does not weaken the motivation for using position coupling.

W2. With enough knowledge of the problem, we don’t need ML.

• We highlight that length generalization is a crucial issue in LLM. Plenty of existing research [1–4] (including us) regards arithmetic/algorithmic tasks as manageable but interesting test beds for studying length generalization because LLMs fail to length-generalize even on such simple tasks.
• The main message from our work is that incorporating task structure when designing positional encodings can effectively improve the length generalization capabilities. We believe our findings can serve as a stepping stone for future research in this area.

W3. Zero-padding is a problem-aware data processing.

• Note that zero-padding to match the operand lengths is prevalent in this research area ([3–5]).
• One can implement our method without relying on zero-padding. We present experimental results using a no-padding scheme in the attached PDF file (Fig. 1) in our General Response. While the no-padding position coupling approach fails in in-distribution generalization when trained on a 1-layer model (due to the increased complexity of the algorithm that the model should learn), combined with proper reversing of the number(s), it functionally extrapolates the performance of deeper models.

Q1. Should the operands in every testing example have the same length?

• You’re correct: we sampled the operands to be the same length while testing. To test the model (trained on zero-padding format) with the sample “20345+34=”, we could input it as “20345+00034=”.
• To address the concern, we tested on operands sampled with different lengths. See Fig. 3 of our PDF attached to the General Response. Each axis corresponds to the lengths of each operand. The results show that the model is capable of solving tasks even when the two operands have different lengths, although zero-padding is applied to ensure consistency in the input format. We will add this result to our revision.

Q2. Reversing the answer solely is a problem-aware data processing.

• Note that solely reversing the answer is also a common practice in this research area [3–5]. From now on, let us compare two different formattings: (a) solely reversing the answer and (b) reversing all the numbers.
• We first empirically compare (a) and (b) while applying position coupling (see Fig. 1 of the PDF in our General Response). For 1-layer models, both formats exhibit near-perfect length extrapolation and show little difference. Conversely, for 6-layer models, a noticeable performance gap emerges. With zero-paddings, (b) performs better than (a); but if there’s no padding, (a) performs much better than (b). There seems to be no clear winner between the two.
• The similarity between (a) and (b) for 1-layer models is expected. If we assign the position IDs based on the significance of the digits as usual, there is NO effective difference between (a) and (b) for a 1-layer model in terms of its prediction, which can be deduced from Prop 5.2. Accordingly, our Theorem 5.1 based on (a) can also be applied to (b) without any modification.
• The difference between (a) and (b) for deeper models is also expected. In multi-layer models, the causal attention mask causes each token embedding to depend on the embeddings of preceding tokens after the first layer. Thus, unlike in the previous case, the predictions of the model may differ depending on the input format.

Q3. The optimization might not be the reason why deeper models perform worse.

• There are two key aspects of optimization: convergence and implicit bias. Our explanation primarily concerns implicit bias. To answer your first question, deeper models do fit the training samples just as well as the shallower models. Thus, convergence is not the issue. We also think that overfitting is not the case, as the trained models only struggle with longer sequences while achieving perfect accuracy on in-distribution samples.
• Therefore, we believe that the issue lies in the implicit bias of the models. Among the infinite number of solutions that achieve zero training loss, shallower models seem to possess a better implicit bias, which allows them to find solutions that generalize better for longer sequences.
• Specifically, we conjecture that the outstanding performance of the 1-layer model is due to its relatively restricted expressivity and that the model has no way to fit the training samples other than learning the true algorithm. However, for deeper models, the model can still fit the training samples without necessarily learning the true algorithm due to greater expressivity.
• For a detailed discussion on this, please refer to our General Response.

---

References

[1] Jelassi et al. Length generalization in arithmetic transformers. arXiv preprint, 2023.
[2] Kazemnejad et al.. The impact of positional encoding on length generalization in transformers. NeurIPS 2023.
[3] Zhou et al., What algorithms can transformers learn? a study in length generalization. ICLR 2024.
[4] Zhou et al. Transformers can achieve length generalization but not robustly. arXiv preprint, 2024.
[5] Lee et al. Teaching arithmetic to small transformers. ICLR 2024.

We thank the reviewer for the insightful comments. We provide our response to the reviewer's concerns.

W1. Theoretically, why do deeper models perform worse despite their expressivity?

• For a broader answer to the question, see our General Response.
• We hypothesize that the performance degradation is due to the bad implicit bias of deep models (learning shortcuts to only achieve in-distribution generalization) when learning a simple algorithm to solve the task. We believe exploring a theoretical explanation for the bad implicit bias of large models on low-complexity tasks is a promising research direction.

W2+Q1. How robust is position coupling to input formats? Can it be applied to tasks when specific formatting is undesirable?

• See our General Response for details on the robustness to input formats.
• As we emphasized there, proper use of the input format is crucial for solving the task even with position coupling. Also, the model’s performance depends on the choice of input format. Thus, if we cannot apply any formatting, we should not expect significant success in solving the given task, not just in terms of length generalization.

W3. Compare with recent techniques for length generalization.

• Thank you for your suggestion. One notable length generalization result for the addition task before us appears in [1], combining FIRE (a relative PE method) and index hinting. They achieve near-perfect length generalization up to operand length 100 by training on up to 40-digit additions. (Recall that we achieve 30-to-200 generalization with a 1-layer model.) Despite their great performance, they require doubling the input length because of index hinting. In contrast, our method doesn't require doubling the input sequence, so we believe our method is more efficient to run.
• We also conducted experiments combining RoPE and our method and achieved an improved length generalization compared to vanilla RoPE: see Fig. 4 of PDF in our General Response.
• We will add these comparisons in our final manuscript.
• [1] Zhou et al., Transformers can achieve length generalization but not robustly. arXiv, 2024.

W4. Can position coupling be applied to a wider range of 2D or higher-dim tasks?

• As demonstrated in the paper, our method significantly helps Transformers to length-generalize on the tasks with a clear structure between token positions. Extending this to various multi-dim tasks is an interesting future work.
• A challenge in multi-dim tasks (not only for length generalization and for applying our method) is the exponential growth (in task dimension) of the number of tokens, which makes it difficult for the model to analyze queries and generate responses. Overcoming this dimensionality problem is an interesting future direction.

Q2. Are there specific task structures that render position coupling less effective or inapplicable?

• Let us give you some examples. Our preliminary experiments weren’t very successful for some tasks including (1) addition with a varying number of operands and (2) multiplication with varying lengths of both operands. We tried couplings similar to the ones applied to simple addition and Nx2 multiplication, respectively, but they weren’t effective (although we did not invest much effort in making it work).
  • The algorithm for solving (1) needs to attend to a varying (over problem instances) number of positions for generating tokens.
  • The algorithm for solving (2) needs to attend to positions in varying relative distances in terms of the “coupled” position IDs. On the contrary, our theoretical construction for simple addition and Nx2 tasks does not suffer from these difficulties, which is the key reason for successful length generalization; thus, we do not expect length generalization on tasks (1) and (2) without any advance in input formats.
• Some tasks don’t have any structure between specific positions (e.g., sorting and mode), thereby we cannot directly apply position coupling. (See our response to Reviewer ow6J.)

Q3. Can you formalize the notion of task structure and provide guidelines for designing proper coupling for different tasks?

• Defining task structure involves the relationship between the query and response, focusing on which tokens in the query influence the determination of each token in the response. Designing position coupling relies on human intuition, making it challenging to provide concrete guidelines. However, for tasks with unclear coupling structures, assigning position IDs (piece-wise) consecutively may be worth trying. It is shown to be empirically effective in the Nx2 multiplication task (which is later theoretically backed by Thm 6.1), although it's not intuitive how the coupled positions capture the task structure.

Q4. Why does using the same embedding layer for both position coupling modules perform better than separate layers?

• As mentioned in lines 655–657, we do not have a clear explanation for why sharing the embedding layer performs better. We conjecture it is due to the row-column symmetry of the Minesweeper generator task, meaning that transposing the board still results in a valid problem. We believe that tasks where rows and columns have different semantic structures might not benefit from using the same embedding layer.

Q5. Can we extend the evaluation of position coupling to more complex 2D tasks?

• Although our study focuses on simple arithmetic/algorithmic tasks, we are eager to explore more complex multi-dimensional tasks.
• For image-related tasks (e.g., involving ViT), we could apply position coupling similar to what we did for the Minesweeper generator task. For graphs and other non-sequential tasks, the position coupling scheme may need refinement, but extending it to these complex tasks is promising future work.

We hope our response has adequately addressed the reviewer's concerns, and we would appreciate it if you could reconsider your assessment.

We thank the reviewer for the positive review and valuable comments. Below, we address the reviewer’s concern.

W1. Position coupling is geared towards specific tasks in length generalization.

First, as mentioned in the conclusion of our paper, we focus on certain tasks with a handy structure between specific positions of tokens. Since vanilla Transformers often fail to learn the true structure of the tasks—even for simple algorithmic tasks—as observed by plenty of previous works [1–10], we aim to guide the model to properly learn the structure by coupling the positions, thereby improving length generalization. The structure between positions is important for employing our method, so we admit that it is not easy to apply our method to some tasks. However, we address the three tasks raised by the reviewer.

Parity.

• Given a binary query sequence, the objective is to output 1 if there is an odd number of 1s in the query and to output 0 otherwise.
• Without any data processing, it is a difficult task for Transformers to achieve length generalization [1,4,5,7,9]. However, with the help of a scratchpad generated by leveraging the idea of unrolling the query, Transformers can achieve length generalization [9]. To utilize the power of position coupling, we opt for a different scratchpad from that used in [9].
• By unrolling the query, we can generate a sequence of partial results. Let’s put the partial result from the 1st to n-th token of the query into the n-th token of the response sequence. For example, if the query is given as 010011, the response with scratchpad would be 011101. Then, the last token of the response immediately becomes the answer for the original task. We train the model to generate the whole response (including scratchpad) for each query to solve the task step-by-step with next-token prediction. Now, we can naturally couple the n-th positions of the query and the response when we apply our method. For example, given an input sequence “010011=011101”, we can assign (3,4,5,6,7,8,2,3,4,5,6,7,8) as position IDs (when the starting ID is randomly chosen as 3).
• To showcase the efficacy of position coupling on the parity task with a proper input format, we compare 4 different settings: position coupling with/without scratchpad and NoPE with/without scratchpad. For the “position coupling without scratchpad” setting, we naively couple all the positions (except for the ‘=’ token) with the same position ID (e.g., “010011=1” can get (3,3,3,3,3,3,4,3)). We train the models on the queries of lengths 1–20 and test the lengths up to 100. We measure exact-match accuracy (including scratchpad if applicable) as well as the accuracy only for the single token at the position of the last token of the response except for the EOS token (called “parity accuracy”). The result of the experiments is shown in the PDF file attached in our General Response. Without scratchpads, both NoPE and position coupling cannot even achieve good in-distribution performances. Even with our scratchpads, without any position embeddings (NoPE), a 6-layer 8-head model showcases a very restricted length generalization capability up to the length of ~30. The model performs worse than random from the length of 40 because the model sometimes outputs tokens other than 0 or 1. Most importantly, our 1-layer 4-head model with position coupling and scratchpad achieves perfect length generalization up to length 100! We strongly believe that thanks to the combination of coupled position IDs and scratchpad enabling Transformers to learn a simple algorithm for solving the task with next-token prediction, we could achieve another remarkable length generalization result.
• We will add this result to our final manuscript with more ablations.

Sorting and Mode.

• The sorting task aims to generate a response equal to the sorted sequence of the given query; the mode task aims to find the most frequent token appearing in the given query.
• In both tasks, there is no exploitable structure between the positions of tokens. Thus, it is not straightforward to apply position coupling to solve these tasks, thereby we did not test our method on these tasks.
• Not only is it unnatural to couple the positions, but it is also unnecessary to do so. Vanilla Transformers already length-generalize well on these tasks [9].

In short, position coupling is an effective method if we can create a clear structure between positions with proper usage of input format; however, it is inapplicable if there is no such structure. Nonetheless, there are a lot of real-world tasks whose underlying structure between positions is vague or unavailable. We leave the research direction of extending our idea to such tasks by automatically discovering appropriate couplings of the positions as interesting future work.

Please let us know without hesitation if you have further questions or comments.

---

References

[1] Bhattamishra et al., On the ability and limitations of transformers to recognize formal languages. arXiv preprint, 2020.
[2] Kim et al., Have you seen that number? investigating extrapolation in question answering models. NeurIPS 2021.
[3] Nye et al., Show your work: Scratchpads for intermediate computation with language models. arXiv preprint, 2021.
[4] Chiang and Cholak., Overcoming a theoretical limitation of self-attention, arXiv preprint, 2022.
[5] Delétang et al., Neural networks and the Chomsky hierarchy, ICLR 2023.
[6] Kazemnejad et al.. The impact of positional encoding on length generalization in transformers. NeurIPS 2023.
[7] Ruoss et al., Randomized positional encodings boost length generalization of transformers, ACL 2023.
[8] Lee et al., Teaching arithmetic to small transformers. ICLR 2024.
[9] Zhou et al., What algorithms can transformers learn? a study in length generalization. ICLR 2024.
[10] Zhou et al., Transformers can achieve length generalization but not robustly. arXiv preprint, 2024.