Positional Attention: Out-of-Distribution Generalization and Expressivity for Neural Algorithmic Reasoning

(6) While that can be due to appropriate inductive biases, the classic transformer should be able to learn these tasks too. If the positional embeddings of the classic transformer are learnable and concatenated to the inputs, then isn’t it possible to show that for every positional Transformer, one can construct a classic transformer that is functionally equivalent to it?

We kindly refer the reviewer to answer (3).

(7) Furthermore, as the authors also discuss in their related work section, we know that the classic transformer is a universal approximator for continuous and discrete functions, as well as Turing-complete. Hence, the classic transformer should also be able to solve the tasks considered in the paper.

We kindly refer the reviewer to answers (5) and (6).

(8) In particular, the weights can be very sensitive to the scale of input values. We observe a dramatic change in attention weights as soon as we give the model the same input but scaled so that the values lie outside the domain of the training data. Perhaps I am missing something, but shouldn’t scaling the input by $\lambda$ scale all elements of the pre-softmax attention matrix by $\lambda^2$. And as scaling the pre-softmax activations is equivalent to adjusting the temperature of the softmax, isn't this difference in attention predictable and expected?All that makes me think that the reported results might be caused by technical details in the implementation of the experiments, rather than due to positional attention being fundamentally better for some reason.

Due to the bias term in the encoding layer, key and query transformations, scaling the input by $\lambda$ does not scale the elements of the pre-softmax attention matrix by $\lambda^2$. The actual scaling depends on the learned bias terms and thus is not predictable. On the other hand, the reviewer is right in that the elements of the pre-softmax attention matrix will inevitably be affected by the scaling of the input. This is precisely the issue with standard transformers and it partially explains why they suffer from poor OOD performance. An important motivation for us to propose positional transformers is to solve this issue.

Let us clarify that we are not trying to claim that positional attention is fundamentally better. It is never our goal and we, too, do not believe that positional attention is fundamentally better in every aspect. In this paper our goal is to enable good OOD performance of the transformer architecture for algorithmic tasks, and we propose positional attention as a solution to effectively address the instability issue of self-attention described above. Such results are particularly relevant to the community focused on Neural Algorithmic Reasoning [3, 4].

(9) Improving the ability of a Transformer to solve algorithms in isolation seems unnecessary: the example tasks used in the paper can be solved exactly with just a couple of lines of code. There would be value if the model uses these tasks as a subroutine of more complex tasks, e.g., given a natural text description of purchases, group them by type and compute commulative sums for each group. However, the paper does not study whether the proposed positional encodings work for more general data distributions, e.g., NLP tasks, different formatting of the same tasks, non-PCOC tasks, or mixtures of tasks. My intuition is that the ridgidity of the attention mechanism that improves the performance on the tasks considered in the paper, would be detrimental to the performance on such tasks that require flexibility, nuance, and robustness. And as LLMs are useful precisely for these less-structured tasks —for the structured ones, one is better off writing a Python script— I feel that the paper has failed to propose a modification of the transformer architecture that would result in improved performance in any realistic case. As such, the paper would likely be of limited value to the community.

We appreciate the reviewer’s suggestion. Below, we describe a realistic application which requires complex and flexible reasoning. Briefly, we present the model with an input consisting of a mixture of alphanumeric and numeric elements corresponding to categories (alphanumeric) and prices (numeric). We then ask the model to output either the minimum/maximum/sum of a random subset of categories. We measure value generalization with respect to the category names (using unseen categories when testing) as well as the category prices (testing on larger prices). We present 3 experiments: one where the prompt is exclusively to calculate the minimum, one where the prompt is exclusively to calculate the sum and finally, one where the prompt either asks for minimum or maximum (which we call multitask). This task requires pattern matching, since the categories between the train and test data are different. It requires conditional reasoning, since the queries are about a sub-set of the categories in the samples. Finally, it requires algorithmic reasoning to compute the output.

For details on our experimental setting, please refer to Appendix C.6 of the revised manuscript. The alphanumeric part of the input is tokenized and passed through an embedding layer, while the numeric part is passed through a linear layer. Below, we present the median MSE and MAPE losses for both standard and positional Transformer on all three experiments.

Median MSE for standard transformer:

Median MSE for positional transformer:

Median MAPE for standard transformer:

Median MAPE for positional transformer

We note that, even in this complex relational task with mixed-type inputs, the positional Transformer still significantly outperforms the standard Transformer, demonstrating the potential utility of our application for certain real-world applications.

We thank the reviewer for taking the time to read our manuscript and provide insightful feedback. We hope that, if satisfied, the reviewer might consider raising their score. Below, we respond to the reviewer’s concerns in detail, starting with a key point highlighted in the review.

We appreciate the reviewer's interest in having a clear hypothesis in our paper. While we did not explicitly state this as a hypothesis, we did observe and discuss that transformers struggle with value generalization. Additionally, in our explanation of the architecture (e.g., as shown in Figure 1), we highlight the specific components that may contribute to overfitting.

For clarity, we have made this hypothesis more explicit in our manuscript (lines 185-192 in orange) and we restate it below.

We hypothesize that two factors contribute to poor scale generalization in standard Transformers:

1. The inclusion of input values in the computation of the attention weights;
2. The use of positional encodings in the input matrix

are potential sources of overfitting during training. To address this, we decouple input values from attention weight computation and remove positional encodings from the input matrix, leading to positional attention. In practice, we validate our empirically validate our hypothesis, demonstrating that positional Transformers achieve better scale generalization within specific algorithmic reasoning tasks.

As mentioned in our paper at lines 185-186, this hypothesis is inspired by the observation that many parallel algorithms do not rely on values within their communication frameworks. We further discuss how the presence of these components could explain the poor OOD performance of standard transformer models in Appendix E. We also provide an ablation study in Appendix G demonstrating how the removal of these two factors results in better performance.

We really appreciate the reviewer's request for stating a clear hypothesis, as this helps to establish a consistent structure throughout our paper.

(1) Section 5 of the paper focuses offers a theoretical explanation of why the proposed positional transformers can simulate any algorithm falling under the Parallel Computation with Oracle Communication (PCOC) model. While it is indeed quite interesting that there is a rich class of computation that can be approximated with fixed attention, it is not clear how this theory ties in with the rest of the paper.

Since the reviewer mentions "fixed attention", to clarify and ensure alignment, we reemphasize that the attention weights in our model are not fixed, but are a function of the positional encodings. The attention weights change at each layer. However, it is possible that we have misunderstood the reviewer’s point, and the reviewer meant that the attention is independent of the input values.

The introduction of PCOC is motivated by the need to demonstrate the expressivity of our positional transformers. We do this because our architecture removes important components from the standard transformer architecture, namely:

1. Removing input values from the computation of attention weights;
2. Removing positional encodings from the input matrix.

Section 5 investigates whether these restrictions do not compromise positional transformers’ ability to simulate parallel algorithms. To provide such expressivity guarantees, we examine how the proposed architecture can simulate some specific computational model. The fundamental differences between our proposed architecture and existing computational models require the development of a parallel model that better reflects the positional transformer architecture, hence the introduction of PCOC. Due to these differences between the architectures, we find this theoretical assessment necessary. We kindly encourage the reviewer to consider how the proposed architecture would lack substantiation without a formalized expressivity result.

(10) Overall, despite the clear writing and good presentation, the paper has major conceptual and methodological flaws. The restricted architecture that is proposed could work well for the small task set the paper considers, but is unlikely to work well in practical settings where flexibility, nuance, and robustness matter. Moreover, the paper does not hypothesise or explain why the prosed restricted architecture would perform better than the more general classical transformer that subsumes it. This question is also not addressed by the theoretical analysis of the paper. The leaves the central claim that “positional attention enhances empirical OOD performance” not well-substantiated.

We appreciate the reviewer’s feedback, but we find this particular comment to be disproportionate, as it does not present evidence of major conceptual or methodological flaws in our work. The authors thoroughly examined the reviewer's comments, and could not identify any specific issues that substantiate the claims of major flaws in the proposed architecture or methodology. Below, we address the specific points raised:

(10.1) [...] but is unlikely to work well in practical settings where flexibility, nuance, and robustness matter

This statement appears unfair, as we demonstrate a realistic application in (9), which indeed requires flexibility, nuance, and robustness. We hope the reviewer appreciates the effort that was put in addressing their concerns, especially regarding the more realistic settings, which seems to be one of the reviewer's main point of concern.

(10.2) Moreover, the paper does not hypothesise or explain why the prosed restricted architecture would perform better than the more general classical transformer that subsumes it. This question is also not addressed by the theoretical analysis of the paper.

While we acknowledge that an explicit hypothesis is not stated in a dedicated section, as mentioned, we did observe and discuss that transformers struggle with value generalization. These factors provide a rationale for why the restricted architecture offers advantages in specific contexts.

Furthermore, as stated in (2), the goal of the theoretical analysis was to address the expressability of the architecture, as some components are removed from the standard transformer's architecture.

(11) Why are you parametrising the fixed attention matrices A by the Wq and Wk matrices instead of directly parameterising A (in Eq. 2)? Is it to reduce the model size? How would your results change if you parametrise A directly?

This structure complies with the traditional attention structure and expresses some factorization of an attention matrix A. For fixed-length experiments with the dimension of positional encodings being equal to the sequence length, the two parameterizations (via $W_q$ and $W_k$ or directly using A) are equivalent since P is square. However, in variable-length experiments (Section 6.1), the attention matrix depends on a different P for each length, which breaks this equivalence. For fixed lengths, using positional encodings with number of columns less than n also restricts the attention patterns that can be represented, effectively imposing a low-rank condition on A. Appendix C.3 includes experiments with positional encodings where P is not square.

Once again, we would like to thank the reviewer and invite them to respond with further comments.

References:
[1] Sanford et al. Transformers, parallel computation, and logarithmic depth. ICML 2024
[2] Hoffmann et al. Training Compute-Optimal Large Language Models. NeurIPS 2022
[3] Bevilacqua et al. Neural Algorithmic Reasoning with Causal Regularisation. ICML 2023
[4] Rodionov et al. Neural Algorithmic Reasoning Without Intermediate Supervision. NeurIPS 2023

(2) Most critically, how does this theory explain the observation that fixed attention result in better empirical performance?

The expressivity results in our paper ensures that positional Transformers are at least theoretically capable of performing the algorithmic tasks of interest in this paper. It is not designed to address the learnability potential of the architecture, which is a separate question discussed in Appendix F. In short, we do not provide formal proofs on this matter, as such results would require precise control over the training process, which, to our knowledge, no existing theoretical framework adequately captures for models such as transformers. Additional gaps in this area are further discussed in Appendix F. We hope the reviewer recognizes that establishing theoretical OOD generalization bounds is beyond the scope of this work and is left as a direction for future research.

That said, as mentioned previously, we propose a hypothesis for why positional Transformers demonstrate better OOD performance for the tasks studied here. This hypothesis is supported by extensive numerical experiments designed to validate it.

(3) Doesn't the referenced work of Sanford et al. imply that transformers with dynamic attention should also be able to approximate algorithms from the same class?

While the transformer architecture can approximate algorithms within the same class, this does not guarantee that transformers can, through training, discover solutions that perform effectively in the out-of-distribution settings described in our paper. Although two methods may have the same expressivity (as we further elaborate in (5)), this does not imply that they will perform equally well after training. As noted, the components of standa

(4) The paper’s main thesis is that "positional attention enhances empirical OOD performance", however it does not propose a hypothesis explaining as to why that might be the case

We kindly refer the reviewer to the initial discussion in the author response.

(5) I find this conclusion of the paper to be surprising and difficult to believe without further details. The authors construct a strictly less expressive architecture that has better performance.

The reviewer's statement that positional transformers are strictly less expressive for simulating parallel algorithms is incorrect. As outlined in our paper, the expressivity of the proposed architecture is demonstrated through its ability to simulate the PCOC model. Furthermore, [1] establishes that computations performed by transformers are achievable by the MPC computational model. In Appendix A, we note that while PCOC uses data-agnostic communication, it can simulate MPC provided sufficient memory is available. This suggests that positional transformers are as expressive as MPC and standard transformers. However, achieving this equivalence for certain tasks may require additional depth or width in the architecture.

In addition, as the reviewer may be aware, we point out that in practice it is typically not the case that the most expressive models have the best performance. For example, in traditional statistical learning, regularization terms are often added to a linear model to actually reduce model complexity, avoid overfitting and consequently improve performance. Even in the context of Large Language Models, depending on the size of the training data, it has been shown that larger, and hence more expressive, models often have worse performance than smaller (and hence strictly less expressive) models with appropriate model sizes (e.g. see Figure 3 in [2]).

I would like to thank the authors for the extremely detailed and comprehensive response. I will try to respond to as much as I can.

On the hypothesis: Thank you for clarifying your hypothesis. Just to be on the same page, I do agree that you might be indeed seeing improved OOD performance on the tasks you are evaluating on. However, my concerns whether that applies to the Transformer architecture in general remain. One can come up with an architecture that performs really good on a specific task. However, the challenge is to do that while maintaining the performance on all other tasks. Nevertheless, I do see the theoretical value of your work.

On the attention weights: Perhaps I am missing a key detail, but I still fail to see how the attention weights are not fixed for a given input sequence length. Moreover, assuming causal masking is applied, that means that (up to the softmax normalisation) the attention weights of any sequence of length $n$ are a $n\times n$ sub-matrix of the attention weights of any sequence of length $m>n$. The parameterization you use might reduce the number of parameters that need to be stored (and the rank) but it does not change the fact that the mixing of information throughout the model is fixed (given a sequence length).

While the transformer architecture can approximate algorithms within the same class, this does not guarantee that transformers can, through training, discover solutions that perform effectively in the out-of-distribution settings described in our paper. Although two methods may have the same expressivity (as we further elaborate in (5)), this does not imply that they will perform equally well after training.

While this is correct, the current paper also does not investigate the trainability of the Positional Transformer. Possibly, you can frame this as that you are providing inductive biases that are more aligned with the tasks you are considering.

In Appendix A, we note that while PCOC uses data-agnostic communication, it can simulate MPC provided sufficient memory is available. This suggests that positional transformers are as expressive as MPC and standard transformers. However, achieving this equivalence for certain tasks may require additional depth or width in the architecture.

While I haven't read Appendix A in detail, I am concerned by the fact that your proof appears to rely on an external oracle to provide the routing destinations (which feels like the difficult part) and that you are routing all the messages to all the nodes, which seem to mean that 1. your model size needs to grow with the input length, and 2. you are effectively showing that a single processor can simulate a multi-process task. That would further imply that, beyond the input layer, you might not even need the transformer, and possibly could achieve the same with an MLP.

On the additional experiments: While I greatly appreciate the effort that you've invested in designing and evaluating the additional experiments in Appendix C.6, I am afraid they still appear to miss the point. Looking at the examples on lines 1410-1420, they appear to be strictly formatted in a very precise and standardised way. As such, the model need not know the context, as it has every piece of information exactly where it expects it exactly the way it expects it. Which is probably why you constructed the task this way. My assumption is that if you add an extra nuisance token in the beginning, say ["Cat-1", "Cat 1", 1.2...] your performance will go close to random guess. I understand that you are not claiming that this experiment shows universal practical value of your proposed architecture but my concerns that it is unlikely to work well in practical settings where flexibility, nuance, and robustness matter.

I do agree that the theoretical aspect of the paper is interesting. The authors have clearly put a great deal of effort and work in this paper. In light of that, I will increase my score to 5. However, as currently written, and in light of my above comments, it makes too strong claims on practical utility and importance. As such, I would encourage the authors to consider repositioning the paper to focus and expand on its theoretical contributions and resubmit.

References:
[1] Yan et al. Neural Execution Engines: Learning to Execute Subroutines. NeurIPS 2020
[2] Numeroso et al. Dual Algorithmic Reasoning. ICLR 2023
[3] Bevilacqua et al. Neural Algorithmic Reasoning with Causal Regularisation. ICML 2023
Rodionov et al. Neural Algorithmic Reasoning Without Intermediate Supervision. NeurIPS 2023
[4] Diao and Loynd. Relational Attention: Generalizing Transformers for Graph-Structured Tasks. ICLR 2023
[5] Back de Luca and Fountoulakis. Simulation of Graph Algorithms with Looped Transformers. ICML 2024
[6] Zhang et al. Inductive Biases and Variable Creation in Self-Attention Mechanisms. ICML 2022
[7] Narasimha. The Batcher-Banyan self-routing network: universality and simplification. IEEE Transactions on Communications 1998
[8] Lee. A New Benes Network Control Algorithm. IEEE Transactions on Computers 1987

On the hypothesis: Thank you for clarifying your hypothesis. Just to be on the same page, I do agree that you might be indeed seeing improved OOD performance on the tasks you are evaluating on. However, my concerns whether that applies to the Transformer architecture in general remain. One can come up with an architecture that performs really good on a specific task. However, the challenge is to do that while maintaining the performance on all other tasks. Nevertheless, I do see the theoretical value of your work.

We thank the reviewer for acknowledging the theoretical value of our work. However, we find it hard to understand the reviewer’s concept of generality, and why our work is being entirely judged by “applicability in general”. We are interested in solving algorithmic reasoning tasks, which are of interest to a wide community [1,2,3,4,5].

Nevertheless, despite not being the focus of our paper, the reviewer seems to question that this applicability is not general enough. To satisfy the reviewer's concern, we provided additional experiments that, based on the reviewer's comment, match their conception of flexibility, nuance, and robustness. The reviewer then further argued with additional details, not part of the initial reply, for a more complex task, to which we also successfully address at the end of our response (under experiment details below).

The reviewer further states that our work “makes too strong claims on practical utility and importance”. We ask the reviewer where in the paper such claims have been made, since our work was purposely designed to tackle algorithmic tasks.

Considering the aspects mentioned above, as well as the additional experiments addressing their concerns, we respectfully ask the reviewer to reevaluate our work.

On the additional experiments: While I greatly appreciate the effort that you've invested in designing and evaluating the additional experiments in Appendix C.6, I am afraid they still appear to miss the point. Looking at the examples on lines 1410-1420, they appear to be strictly formatted in a very precise and standardised way. As such, the model need not know the context, as it has every piece of information exactly where it expects it exactly the way it expects it. Which is probably why you constructed the task this way. My assumption is that if you add an extra in the beginning, say ["Cat-1", "Cat 1", 1.2...] your performance will go close to random guess. I understand that you are not claiming that this experiment shows universal practical value of your proposed architecture but my concerns that it is unlikely to work well in practical settings where flexibility, nuance, and robustness matter.

We would like to point out that we structured the task this way because it aligns with the example proposed by the reviewer, and we closely followed their suggestion. Nevertheless, since the reviewer raised additional nuance concerns, we have implemented a modification of the new experiment where we inject nuisance strings in the input at random positions within the list. Taking it one step further, we also considered OOD generalization in the type of additional nuisance input. To be specific, we augment the training input list by injecting strings of the form Cat+{id} or Cat-{id} at random positions (id is one of the categories present in the true list chosen at random, + or - is also chosen randomly). For instance, this should generate the exact example proposed by the reviewer: ["Cat-1", "Cat 1", 1.2...], where “Cat-1” is the nuisance string that should be ignored in the computation. For testing, we inject strings of the form Cat_{id} or Cat*{id} with the same assumptions as before. Therefore, the train and test nuisance strings are different. Please refer to Appendix C.7 or the revised manuscript for more details. We highlight that even in this more nuanced setting, our model still outperforms the standard Transformer.

[...]

Median MSE for positional Transformer (with irrelevant structure):

Median MSE for standard Transformer (with irrelevant structure):

Median MAPE for positional Transformer (with irrelevant structure):

Median MAPE for standard Transformer (with irrelevant structure):

We hope that our detailed responses and the additional experiments address the reviewer’s concerns. We believe they demonstrate the flexibility, nuance, and robustness of our work, and trust that they will resolve any remaining questions.

On the attention weights: Perhaps I am missing a key detail, but I still fail to see how the attention weights are not fixed for a given input sequence length. Moreover, assuming causal masking is applied, that means that (up to the softmax normalisation) the attention weights of any sequence of length n are a n×n sub-matrix of the attention weights of any sequence of length m>n. The parameterization you use might reduce the number of parameters that need to be stored (and the rank) but it does not change the fact that the mixing of information throughout the model is fixed (given a sequence length).

Your thinking is correct. Our previous remark re-explaining the architecture was just there to ensure we were on the same page, since other reviews suggested a misunderstanding. The attention weights depend on the positional encodings, and because the positional encodings are fixed, the attention weights are also fixed for a given input length once trained. However, they can still vary across layers and attention heads.

“While the transformer architecture can approximate algorithms within the same class, this does not guarantee that transformers can, through training, discover solutions that perform effectively in the out-of-distribution settings [...]” While this is correct, the current paper also does not investigate the trainability of the Positional Transformer. Possibly, you can frame this as that you are providing inductive biases that are more aligned with the tasks you are considering.

We do investigate the trainability of positional Transformers empirically, and these experiments are also aligned with the motivation of our proposed architecture, which is to solve algorithmic reasoning tasks. Maybe the reviewer expected a more technical analysis of this inductive bias, but in our previous response, we stated the reasons why a general PAC-style proof for this out-of-distribution setting is difficult (also discussed in Appendix F with an example). This is not a specific limitation of our work, as existing theoretical frameworks do not capture this nuance. For in-distribution generalization, we could establish an improvement in the complexity bound using an analysis similar to [6] and by exploiting the fact that our architecture breaks the dependence between the attention weights and the input values. However, such analysis is only in-distribution and does not generalize to the OOD setting that we are interested in. Alternatively, we could have made limiting and less practical assumptions for the sake of a theoretical result (such as infinite depth/width for NTK-like results), but the impact of such a result would not cover the practical aspects of our experiments. Providing an analysis to thoroughly study out-of-distribution would require new techniques, which due to its complexity, merit a paper of its own.

In Appendix A, we note that while PCOC uses data-agnostic communication, it can simulate MPC provided sufficient memory is available. This suggests that positional transformers are as expressive as MPC and standard transformers. However, achieving this equivalence for certain tasks may require additional depth or width in the architecture. While I haven't read Appendix A in detail, I am concerned by the fact that your proof appears to rely on an external oracle to provide the routing destinations (which feels like the difficult part) and that you are routing all the messages to all the nodes, which seem to mean that 1. your model size needs to grow with the input length, and 2. you are effectively showing that a single processor can simulate a multi-process task. That would further imply that, beyond the input layer, you might not even need the transformer, and possibly could achieve the same with an MLP.

The focus of the paper is to establish the expressivity of Positional Transformers by simulating PCOC, which can implement various parallel algorithms. Our focus is not establishing efficient simulations across different parallel models. The discussion in Appendix A.1.1 is meant to be a complement of the limitations of PCOC in comparison to MPC, discussed in Section 5.2.

Note that the strategy described in Appendix A is just one way to establish equivalence between PCOC and MPC, in case anyone is interested in such equivalences. While this equivalence is established using a complete network (i.e. sending all data to all machines), there are other ways to establish this equivalence in a less trivial way. For example, one could use communication networks such as Batcher [7] or Benes [8] networks, which are known to be universal, and therefore can realize any desired connection pattern between machines. For example, Benes networks establish pairwise communications, instead of “n-to-n” as described in our discussion, and because of this, they utilize significantly less memory. As a consequence, such networks require more rounds to establish communication (O(log n) for Benes). In this example, no machine has access to all data, so you cannot simply use a single MLP to solve the task. Overall, we think that the value of establishing an equivalence given a different oracle is marginal. However, if the reviewer thinks that this is important, we would be happy to add this observation in the Appendix.

We thank the reviewer for reading our manuscript and for the thoughtful comments. Below, we address the reviewer’s concerns one by one. We also provide additional experiments on complex relational reasoning task that might be of the reviewer's interest later on.

We welcome further questions and hope our responses and experiments meet the reviewer’s expectations to increase their score.

(1) The model’s OOD generalization appears to primarily rely on two factors: the independence of the attention matrix from the input X, and the absence of one-hot encoding for the input.

In the factors described by the reviewer, we believe that the second factor might have two different interpretations. The reviewer might refer to “one-hot encoding” as some representation related to the input tokenization, extending beyond the numerical data used in our previous experiments. To this end, we refer the reviewer to the results in response (3) which further describe a more complex setting that deals with numerical and textual data. In this setting, our proposed architecture also shows favorable OOD results with respect to both data modalities.

Alternatively, the reviewer might refer to the one-hot positional encodings used in positional attention. To this end, we emphasize that our results do not require one-hot positional encodings. To substantiate this claim, in Appendix C.3, we present experiments utilizing other types of positional encodings (binary, sinusoidal), which also show good OOD generalization results.

(2) For example, in the cumulative sum task, positional attention only needs to equally weight the first i tokens onto the i-th token, which allows the output to approximately scale by the same factor as the input within a certain range.

Although the reviewer proposes a plausible conceptual solution for the cumulative sum in the fixed-length regime, we argue that this explanation undermines the applicability of our method. Specifically, this reasoning does not extend to the variable-length regime (i.e., training and testing on sequences of up to n values), where a single scale factor cannot be universally applied to sequences of different lengths. In Section 6.1, we also demonstrate favorable results in this more challenging setting. Moreover, the reviewer's concern seems to primarily focus on our simplest task. Our study includes other tasks that require complex reasoning, where our method also performs competitively. We address these aspects further in the next comment.

(3) However, this design introduces notable limitations: first, positional attention is unable to model relative relationships between different inputs X, limiting its applicability to tasks that require more complex relative positioning.

The proposed architecture cannot adaptively model relationships between rows of X only using attention. However, this does not mean it cannot generally model relative relationships. We appreciate the reviewer’s observation and take this opportunity to clarify the architecture’s capacity to address relative reasoning tasks.

Following [1], relational reasoning involves tasks where a model learns relationships between entities (e.g., numbers, text) independent of the specific entities observed. These tasks may include pattern matching, comparisons, and conditional selections. While our architecture performs well on relational reasoning tasks like sorting and finding minimums, which capture relative relationships within sequences, the reviewer seems unconvinced of its capability to handle more complex tasks. To address this, we present a more challenging setting, requiring relative reasoning over both numbers and text, which we describe below.

Briefly, we present the model with an input consisting of a mixture of alphanumeric and numeric elements corresponding to categories (alphanumeric) and prices (numeric). We then ask the model to output either the minimum/maximum/sum of a random subset of categories. We measure value generalization with respect to the category names (using unseen categories when testing) as well as the category prices (testing on larger prices). We present 3 experiments: one where the prompt is exclusively to calculate the minimum, one where the prompt is exclusively to calculate the sum and finally, one where the prompt either asks for minimum or maximum (which we call multitask). This task requires pattern matching, since the categories between the train and test data are different. It requires conditional reasoning, since the queries are about a subset of the categories in the samples. Finally, it requires algorithmic reasoning to compute the output.

For details on our experimental setting, please refer to Appendix C.6 of the revised manuscript. The alphanumeric part of the input is tokenized and passed through an embedding layer, while the numeric part is passed through a linear layer. Below, we present the median MSE and MAPE losses for both standard and positional Transformer on all three experiments.

Median MSE for standard transformer:

Median MSE for positional transformer:

Median MAPE for standard transformer:

Median MAPE for positional transformer

We note that, even in this complex relational task with mixed-type inputs, the positional Transformer still significantly outperforms the standard Transformer, demonstrating the potential utility of our application for certain real-world applications.

[...]

Furthermore, the capabilities of the architecture can also be analyzed through the lens of computational models. As outlined in our paper, the expressivity of the proposed architecture is demonstrated through its ability to simulate the PCOC model. Furthermore, [2] establishes that computations performed by transformers are achievable by the MPC computational model. In Appendix A, we note that while PCOC uses data-agnostic communication, it can simulate MPC provided sufficient memory is available. This suggests that positional transformers are as expressive as MPC and standard transformers. However, achieving this equivalence for certain tasks may require additional depth or width in the architecture.

Finally, it is important to note that our arguments are not intended to claim that positional attention is fundamentally superior. We do not believe it outperforms in every aspect. Instead, our goal is to enhance the out-of-distribution (OOD) performance of transformer architectures on algorithmic tasks. Such results are particularly relevant to the community focused on Neural Algorithmic Reasoning [3, 4]. We propose positional attention as a solution to address self-attention’s shortcomings in such tasks.

(4) Second, this model appears suitable primarily for tasks involving numeric values that do not require one-hot encoding, constraining its utility in a broader range of applications.

We believe the reviewer refers to “one-hot encoding” as a representation related to tokenization of the input, extending beyond the numerical data used in our previous experiments. If this interpretation is correct, we kindly refer the reviewer to response (3), where we demonstrate the effectiveness of our architecture in addressing a multimodal reasoning task.

(5) Several minor typos and formatting inconsistencies were observed, such as missing periods on lines 156, 356, and 364. Addressing these issues would improve the manuscript’s professionalism and readability.

We appreciate the reviewer’s attentiveness in identifying these typos and inconsistencies. The suggested changes have been incorporated into the revised manuscript.

(6) In my view, the proposed architecture resembles a learnable adjacency matrix version of a Graph Neural Network (GNN), even with the explicit incorporation of positional information. Could the authors compare the results with those achieved by a GNN under the same sentence length? It would be interesting to see if the proposed method significantly outperforms GNNs, as such a comparison could highlight the unique advantages of this approach

While the relationship between transformers and Graph Neural Networks (GNNs) has been well-established, it is important to clarify the differences between our architecture from GNNs. In short, positional Transformers “exchange messages” using a communication network that can change at every layer, whereas GNNs operate within a fixed, predefined graph structure at all layers.

To address the reviewer’s curiosity, we compared the performance of our architecture against Graph Convolutional Networks (GCNs) and Graph Attention Networks (GAT). Since the tasks tested have no underlying native graph, we tested these models on complete and star graphs. Notably, the original GAT architecture on a complete graph is similar to a standard transformer but differs in that the value, key, and query weights are shared in GAT. Below, we present results reporting the median MSE for solving the cumulative sum and sorting tasks. We added these additional results to Appendix H of the paper.

Task: Cumulative Sum

[...]

Task: Sorting

As the results indicate, neither GCN nor GAT works very well even for in-distribution data (OOD scaling factor = 1), let alone achieving OOD generalization. We believe this is because these tasks are inherently unsuitable for standard message-passing architectures. If the reviewer has an alternative GNN architecture in mind, please let us know.

Additional theoretical comparisons: From a computational perspective, the flexibility of attention allows positional Transformers to represent massively parallel algorithms, as there are no communication restrictions at each round. While standard and positional Transformers are aligned with parallel computational models due to their unrestricted communication, GNNs, on the other hand, are aligned with another type of computational model called distributed computing models [5], where communication is constrained by an input graph.

Comparisons between parallel and distributed computational models are a popular topic [6], but are outside the scope of our paper. In the realm of neural networks and algorithms, a comparison between transformers and Graph Neural Networks has been established in Sec 5 of [2]. In theory, transformers are more efficient in the number of layers needed to simulate parallel algorithms, as their communication is more flexible than message-passing architectures. Overall, each one has its advantages and disadvantages. Since we do not deal with graph-based tasks in our paper, we specifically focus on the transformer architecture.

We once again thank the reviewer and welcome any further questions.

References:
[1] Boix-Adsera et al. When can transformers reason with abstract symbols? ICLR 2024
[2] Sanford et al. Transformers, parallel computation, and logarithmic depth. ICML 2024
[3] Bevilacqua et al. Neural Algorithmic Reasoning with Causal Regularisation. ICML 2023
[4] Rodionov et al. Neural Algorithmic Reasoning Without Intermediate Supervision. NeurIPS 2023
[5] Loukas, A. What graph neural networks cannot learn: depth vs width. ICLR 2020
[6] Das Sarma et al. Distributed Verification and Hardness of Distributed Approximation. STOC 2011

For completeness, we restate the theoretical justification about why such tasks are possible for positional Transformers (repeated from the rebuttal above):

The capabilities of the architecture can also be analyzed through the lens of computational models. As outlined in our paper, the expressivity of the proposed architecture is demonstrated through its ability to simulate the PCOC model. Furthermore, [1] establishes that computations performed by transformers are achievable by the MPC computational model. In Appendix A, we note that while PCOC uses data-agnostic communication, it can simulate MPC provided sufficient memory is available. This suggests that positional transformers are as expressive as MPC and standard transformers. However, achieving this equivalence for certain tasks may require additional depth or width in the architecture.

Finally, it is important to note that our arguments are not intended to claim that positional attention is fundamentally superior. We do not believe it outperforms in every aspect. Instead, our goal is to enhance the out-of-distribution (OOD) performance of transformer architectures on algorithmic tasks. Such results are particularly relevant to the community focused on Neural Algorithmic Reasoning [2, 3]. We propose positional attention as a solution to address self-attention’s shortcomings in such tasks.

References:
[1] Sanford et al. Transformers, parallel computation, and logarithmic depth. ICML 2024
[2] Bevilacqua et al. Neural Algorithmic Reasoning with Causal Regularisation. ICML 2023
[3] Rodionov et al. Neural Algorithmic Reasoning Without Intermediate Supervision. NeurIPS 2023

We thank the reviewer for engaging in the discussion phase. We address the reviewer's additional concerns one by one below:

While I understand your results demonstrate strong OOD performance in algorithmic tasks, the underlying mechanism appears inherently tied to problems where OOD arises from straightforward transformations (e.g., scaling of numeric values).

We would like to clarify that the new experiments provided in the rebuttal and Appendix C.6, perform OOD in both categories and numerical values. This is mentioned explicitly in our rebuttal and Appendix C.6 in the revised version of the paper. We beg the reviewer to read it carefully since this will alleviate a lot of confusion.

The categories are alphanumeric, not numeric. Alphanumeric means that the input is a string containing both characters and numbers. For example, “Cat1” and “Cat2” correspond to Category 1 and Category 2, respectively. Therefore, one cannot simply perform straightforward transformations like scaling the input.

Such tasks, while valuable, do not necessarily generalize to more intricate reasoning challenges like mathematical theorem proving or tasks requiring abstract symbolic reasoning.

We would like to emphasize that, in this work, we are not interested in theorem proving tasks, and as such, those tasks are completely out-of-scope in our paper. Nowhere in our manuscript or rebuttal did we mention that we are interested in solving such tasks. We find it a bit disproportionate and unfair to be judged on tasks that we never mentioned or cared to solve in the first place.

These types of tasks often involve OOD scenarios that extend beyond simple numeric transformations

We would be grateful if the reviewer reads the rebuttal and Appendix C.6. We provided additional experiments to illustrate performance on tasks that combine alphanumeric (letters and numbers in string format) and numerical values.

such as logical derivations, multi-step reasoning, or structural pattern recognition.

As we mention in our rebuttal and Appendix C.6 the new experiments do indeed require all the above. Namely, the new tasks require pattern matching, since the categories between the train and test data are different, conditional reasoning, since the queries are about a subset of the categories in the samples as well as algorithmic reasoning to compute the output.

Without a clear extension to such domains, the claimed robustness of the proposed approach in OOD scenarios seems limited in scope.

We would be extremely grateful if the reviewer reads the rebuttal and Appendix C.6 for more details.

For instance, consider a task requiring the sum of all numbers preceding a special token (e.g., "0") in the output sequence.

The task provided in Appendix C.6 illustrates exactly this point. The categories, which are in alphanumeric format, and thus cannot simply be scaled, are completely different in the train and test distributions. This means that the model needs to (i) understand the appropriate symbols for categories unseen during training, (ii) isolate the unseen categories corresponding to the query, and (iii) perform the appropriate aggregation on the corresponding values.

Such tasks require the architecture to dynamically integrate positional and content-based cues, which the fixed positional attention mechanism may struggle to represent effectively.

This is exactly what the task in Appendix C.6 demonstrates. As demonstrated, positional attention still outperforms standard attention in the setting of our chosen task.

While your additional experiments address some relative reasoning scenarios, the inability to handle more nuanced tasks involving dynamic input-output positional dependencies remains a concern

We hope that we addressed this multiple times already. Please let us know if you have any further concerns.

We thank the reviewer for the comments. Below, we address their concerns point by point. If you have any additional questions or concerns, please let us know. We are more than happy to discuss them further.

This paper proposes decoupling the query and key vectors from the value vectors, and making them global for every layer.

The reviewer’s statement suggests the impression that we learn a single attention mask and apply it universally across all layers. We would like to clarify that this is not the case, as we will explain in detail below.

Upon revisiting certain text passages (lines 67 and 82), we recognize how the phrasing might have led to this interpretation. To prevent such misunderstandings, we have rectified these passages to ensure our intended meaning is clear and accurately represented.

To clarify, we would like to restate our proposal briefly. We propose to detach the computation of query and key vectors from the value vectors and the MLPs, but this is an independent process for each layer. Importantly, there is no global structure that is propagated throughout the entire model. Instead, each query and key matrix, corresponding to its respective attention head, independently derives its values.

The input to these query and key matrices consists of positional encodings, which are the same across all layers and attention heads. However, the outcomes of each attention head may differ because they are parameterized by distinct sets of weights.

With this clarification, we address the weaknesses and questions made by the reviewer.

The perspective of Transformers as a graph convolution network with a learned attention mask is not new.

Maybe we do not understand the implication the reviewer is trying to convey, but this widely-recognized perspective was never presented as a novelty in our work. Our paper does not involve graph neural networks or address graph-based problems.

Further, the decoupling of structure and layer-input is also not novel.

Our work significantly differs from [1] in several aspects, as outlined below:

• Global Mask [1] vs. Layer-Specific Attention (ours): In [1], a separate module computes a global mask represented as an undirected weighted graph derived from the input embeddings. In contrast, our approach does not involve a global mask or a dedicated module. Instead, attention weights are computed independently at every layer.
• Dependency on Input Embeddings: In [1], attention weights rely on input embeddings to compute the global mask. In our approach, attention weights are computed solely based on positional encodings, independent of the input embeddings.

Additionally, the setting of our work is fundamentally different, which could be another source of misunderstanding. Unlike [1], which focuses on dependency parsing for NLP tasks, our work investigates the ability of transformers to learn to execute algorithms, a capability with numerous potential applications [2].

The overall contribution of this paper is minimal, and I would suggest focusing on parameterisations of P, particularly task-specific, for this type of model.

We believe that the reviewer is judging the contribution of our work based on the perspective that our architecture is reminiscent of that in Graph Neural Networks. However, we re-emphasize that this is not the case. While we appreciate the reviewer’s suggestion, in general, the parametrization of P is not critical for scale generalization in learning algorithms, as long as they remain expressive. In practice, full expressivity may not always be required to represent all communication patterns. To further substantiate this point, as suggested, we also provide empirical results with different positional encodings and observe no difference in performance, as discussed further below.

The write-up on PCOC is also strongly reminiscent of message-passing in graphical models,

While both PCOC and message-passing models involve the “exchange of messages”, the mechanisms are fundamentally different. Specifically, PCOC exchanges messages using a network that can be changed at every round, while message-passing is restricted to a predefined graph structure. This flexibility allows PCOC to represent massively parallel algorithms since there are no communication restrictions at each round.

Standard and Positional Transformers are aligned with parallel computational models, since their communication is not restricted. In contrast, Graph Neural Networks - which are aligned with another type of computational model called distributed computing models [3] - have their communication restricted by the input graph.

Comparisons between parallel and distributed computational models are a popular topic [4], but are outside the scope of our paper. In the realm of neural networks and algorithms, a comparison between transformers and Graph Neural Networks has been established in Sec 5 of [5]. In theory, Transformers are more efficient in the number of layers needed to simulate parallel algorithms, as their communication is more flexible than the one of message-passing architectures. Overall, each one has its advantages and disadvantages. In our paper, since we do not deal with graph-based tasks, we specifically focus on the Transformer architecture.

I'm not sure what this perspective contributes, though I may be missing something here.

The introduction of PCOC is motivated by the need to demonstrate the expressivity of our positional transformers. We do this because our architecture removes important components from the standard transformer architecture, namely:

• Removing input values from the computation of attention weights;
• Removing positional encodings from the input matrix.

Section 5 investigates whether these restrictions do not compromise positional transformers’ ability to simulate parallel algorithms. To provide such expressivity guarantees, we examine how the proposed architecture can simulate some specific computational model. The fundamental differences between our proposed architecture and existing computational models require the development of a parallel model that better reflects the positional transformer architecture, hence the introduction of PCOC.

Could you have more information / variants of how you compute P? This seems to be the most crucial element of the structure, now that it’s been decoupled from each layer’s input. What tasks work with just an embedding table? What tasks would require some prior processing from another Transformer model? etc.

We emphasize that the matrix P representing positional encodings is not computed, but rather provided as part of the input. This does not imply that the attention weights are fixed. Instead, P is used by the attention layer to determine how information from the input X is combined.

​​Theoretically, representing parallel algorithms requires the given P to be expressive, meaning it must enable the representation of any communication pattern among the input rows. For instance, one-hot encodings fulfill this requirement. However, other variations can be used in practice.

To this end, we present results for the cumulative sum and sorting tasks, using two other alternatives: binary and sinusoidal positional encodings. Due to space constraints, we include results for only these two tasks; the full results for all tasks can be found in Appendix C.3, highlighted in purple.

Results for cumulative sum

Results for sorting

As observed, less expressive positional encodings, such as sinusoidal encodings with $n/2$ dimensions and binary encodings with $\log n$ dimensions, also demonstrate consistently strong performance in positional transformers. This suggests that, in practice, learning certain algorithms may not require expressive positional encodings.

Can you plot accuracy on the tasks you’ve benchmarked on? MSE doesn’t really seem informative here.

Before presenting the additional experiments on accuracy, we would like to clarify to the reviewer that our experimental setup for all tasks is inherently regressive, as we work with real numbers and not integers. Specifically, our model processes a fixed-size list of real numbers as input and outputs another list of the same size, also consisting of real numbers. Consequently, there is no definitive “accuracy” measure one could consider. However, we consider the following two metrics which we believe closely resemble accuracy measures:

1. We evaluate the model on lists containing integers (while training is still done using real numbers) and round the model’s output to the nearest integer (or nearest 0.5 for the median task). A prediction is considered “correct” if the rounded list exactly matches the ground truth list.
2. We evaluate the model on lists of real numbers, considering a prediction "correct" if each entry in the predicted list is within an absolute precision of 0.05 and a relative precision of 5% compared to the corresponding entry of the ground truth list.

Accuracies of standard Transformer using (1):

Accuracies of Positional Transformer using (1):

Accuracies of standard Transformer using (2):

Accuracies of Positional Transformer using (2):

The results are given in the table below and can also be located in Appendix C.8 of the revised manuscript (in purple). Given the unforgiving nature of those metrics (since one entry that fails to satisfy the criterion will cause the entire prediction to be labeled as “incorrect”) it is expected that as the OOD scale factor increases the accuracy will sharply decrease. However, we highlight that our architecture still grossly outperforms the standard transformer (which is consistent with the results presented when considering the MSE loss).

We thank the reviewer again and invite any further questions. We hope that, if satisfied, the reviewer will consider raising their score.

[1] Shen, Y., Tan, S., Sordoni, A., Li, P., Zhou, J., Courville, A. (2022). Unsupervised Dependency Graph Network

[2] Veličković, P., Blundell, C. (2021) Neural Algorithmic Reasoning

[3] Loukas, A. (2019). What graph neural networks cannot learn: depth vs width.

[4] Das Sarma, A., et al. (2011). Distributed Verification and Hardness of Distributed Approximation

[5] Sanford, C., Hsu, D., & Telgarsky, M. (2023). Representational Strengths and Limitations of Transformers

I understand that your paper does not involve graph neural networks or try to address those problems.

Great! Thank you for clarifying, and replying to our rebuttal as well.

GNNs can be viewed as a specific form of Transformers where the attention weights are fixed to 1 where the two nodes are connected. The connection between nodes are fixed given the graph structure, and therefore independent of the input.

Right! We agree with this.

By using a layer-independent and input-independent P, you are keeping the weights at each layer fully determined by P.

This is correct. The attention weights in our architecture are fully determined by P. This represents the matrix of positional encodings, which can be one-hot, sinusoidal or binary, for example. In particular, the attention weights are a function of P. For the reviewer's convenience, we posted the definition of positional attention below, but for more information please see Figure 1 and Section 4 in our paper.

This would seem to be analogous to a weighted graph that is determined by P.

If the reviewer implies that the attention weights are a parameterized function of P, then this statement is correct. If the reviewer means that the adjacency matrix is equal to P, then this is not what we do in this paper, since the algorithmic tasks we tackle (sorting, median, etc…) are not naturally related to graphs.

In that sense, the structure you are proposing is very similar to GNNs, with the added benefit of being able to learn the graph.

This is a very important observation of the reviewer. Indeed, from a GNN perspective, our architecture is able to learn a new weighted graph at every layer. This is not possible for GNNs, where the graph structure is fixed at each layer. This is crucial and really boosts the performance of our method. Although the tasks that we are interested in are not graph-based tasks, we took the experimental results a step further by comparing our method against GNNs as well in Appendix H of the revised manuscript. We provide a summary of the results below.

We compared the performance of our architecture against Graph Convolutional Networks (GCNs) and Graph Attention Networks (GAT). Since the tasks tested have no underlying native graph, we tested these models on complete and star graphs. Notably, the original GAT architecture on a complete graph is similar to a standard transformer but differs in that the value, key, and query weights are shared in GAT.

Below, we present results reporting the median MSE for solving the cumulative sum and sorting tasks.

Task: Cumulative Sum

Task: Sorting

As the results indicate, neither GCN nor GAT work very well even for in-distribution data (OOD scaling factor = 1), let alone achieving OOD generalization. We believe this is because these tasks are inherently unsuitable for standard message-passing architectures. If the reviewer has an alternative GNN architecture in mind, please let us know.

[…] While Lines 123-129 argue that "MPC requires input data to specify destinations for communication," these destinations are not truly "specified by input," as each destination is computed from the source with detail-designed mlp[1].

We appreciate the reviewer’s attention to this inconsistency in our text. To clarify, MPC does not require the destination to be explicitly specified in the input, as it can be inferred through local functions. A more precise way to convey our argument is that communication in MPC may not be agnostic to input values. We updated our manuscript to rectify the passage outlined by the reviewer. We point to passages 126-131, 231-251, and 730-740 in the revised manuscript encompassing these changes, highlighted in blue.

[...] Since (one-hot) positional information can easily implement any fixed attention pattern, this theoretical contribution seems limited, especially considering that [1] achieved similar results without utilizing positional information.

While we respect the reviewer’s perspective, we believe the comparison requires further elaboration. Although we acknowledge [1] as an excellent contribution, certain assumptions made in [1] may have been overlooked by the reviewer, as outlined below.

First, the reviewer states that no positional information is utilized in [1], but this is not entirely accurate. Their proof strategy requires row-specific MLPs (e.g., MLP_1, MLP_2, etc.) during the communication phase. While this approach is valid (see [2, 3]), it inherently encodes positional information through indices. In a more conventional setting, replacing these row-specific MLPs with a single, larger MLP capable of representing all local MLPs would necessitate distinct positional identifiers for each row input. Without this positional information, rows with identical inputs would produce identical outputs, contradicting the requirements of their proof.

The reviewer might further argue that relying on one-hot positional encodings could trivialize the communication aspect of our proof, as works such as [1] may not require such expressive positional representations. While the reviewer may view this as a limiting assumption, we argue that the proof technique of [1] also makes limiting assumptions that are subject to debate. For example, in [1] and [3], local MLPs effectively function as a look-up table that maps source and destinations. However, the complexity of this operation and its dependence on the number of inputs is not explicitly characterized in [1] nor [3]. Thus, despite avoiding positional encodings, these works depend on potentially complex operations within their MLPs.

Ultimately, the two different approaches are equally valid. In our proof, we leverage explicit positional information, which we consider more conventional in practice due to its empirical success and the avoidance of row-specific MLPs.

Furthermore, in practice, full expressivity may not always be required to represent all communication patterns. In our experiments, we do not observe changes in performance utilizing other types of positional encodings (see Appendix C.3 in the revised manuscript for additional results).

We sincerely thank the reviewer for their time and insights on our manuscript. Below, we address the reviewer’s questions and concerns one by one. Please let us know if you have more questions or concerns. We are happy to discuss them.

The argument that positional transformers can simulate PCOC seems somewhat not surprising, given that MPC can simulate PCOC (with more rounds), PCOC can be considered a special case of MPC. [...]

The fact that MPC could simulate PCOC only establishes a relation between two distinct parallel computational models. However, this connection does not provide insights into the expressive power of positional transformers, which remains our primary focus in our theoretical analysis. The authors find it challenging to discern the specific implications of the reviewer’s comment. We propose addressing the reviewer’s concern from two potential perspectives:

Why is our theoretical contribution necessary?

We would like to begin by clarifying the motivation behind our analysis. The positional transformer architecture differs from the standard transformer architecture by:

1. Removing input values from the computation of attention weights;

2. Removing positional encodings from the input matrix.

Section 5 investigates whether these restrictions do not compromise positional transformers’ ability to simulate parallel algorithms. Due to these differences between the architectures, we find this theoretical assessment necessary. We kindly encourage the reviewer to consider how the proposed architecture would lack substantiation without a formalized expressivity result.

Is our theoretical contribution a corollary of [1]?

The results in [1] do not imply our theoretical results. In this context, to provide expressivity guarantees, we examine if our proposed architecture can simulate some computational model. While MPC could be a candidate (which could simulate PCOC), its assumptions do not fully align with the proposed architecture, namely because:

• MPC allows input values to determine communication across machines.

• MPC imposes sublinear memory constraints w.r.t. the number of machines.

These fundamental differences require the development of a parallel model that better reflects the positional transformer architecture, hence the introduction of PCOC. Regarding the specific simulation proof technique utilized in our work, we emphasize that the results in [1] do not apply to our setting, and our theoretical results differ in several ways with important practical implications. We discuss this in detail below.

The primary difference between the two approaches lies in how communication is modeled. In our setting, communication is executed exclusively by the attention layer. The parameters in $W_Q$ and $W_K$ determine how information is propagated across machines, which aligns with the oracle assumption of our computational model. In contrast, as the reviewer mentions, the communication in MPC can be determined by local functions. Consequently, in the proof of [1], this is a two-stage process in which row-wise MLPs are first required to encode the destinations. In the second stage, $W_Q$ and $W_K$ transform this information into appropriate encodings (Proposition B.1 in [1]).

This approach would only directly translate to our setting with additional modifications to the architecture. In [1], MLPs generate destinations, while in our architecture, they do not influence the computation of attention weights. Adapting this strategy would involve transforming $W_Q$ and $W_K$ into MLPs and using them as look-up tables to map sources and destinations, as in [1]. This deviates from our intended design and the standard transformer architecture, which uses linear transformations inside attention and often relies on explicit positional information. Moreover, incorporating MLPs into the computation of attention weights unnecessarily increases model complexity.

Considering all these aspects, we believe that there are too many critical differences among the papers that justify a separate expressivity result for our work.

Regarding the experimental validation (Section 6): [...]

We agree with the reviewer that the observed results may stem from differences in input formats, but this aligns with our intended assumptions and architectural design. Specifically, our architecture separates input values and positional information into two distinct streams: (a) one dedicated to the value matrix and MLPs, and (b) another for the query and key matrices—an approach not feasible in standard transformer architectures.

Furthermore, we emphasize that the standard transformer architecture’s expressiveness depends on positional encodings [4], which are explicitly added to the input. Without these encodings, the standard transformer architecture becomes permutation-equivariant [5], which is unsuited for algorithmic tasks such as sorting (permutation invariant).

To investigate this further, we have conducted additional experiments exploring variations in the data fed into these inputs to examine whether performance improvements can be attributed solely to any specific factor. Before detailing these experiments, we wish to correct a crucial detail in the reviewer’s comment: our work does not exclusively use integer data and thus cannot adopt the tokenization scheme proposed. Nonetheless, we attempted to accommodate this feedback in variation (2). The variations include the following:

1. Standard Transformers: Numbers and positional encodings for the MLPs, value, query, and key matrices.
2. Standard Transformers with non-parallel input: Input numbers are placed in one-hot positions, that is, the input is Diag(X) where X is the input list to other models. No additional positional encodings are used.
3. Positional Transformers: Numbers for the MLPs and value matrix, and positional encodings for the query and key matrices.
4. Misaligned Positional Transformers: Numbers for the MLPs and value matrix; Positional encodings for the MLPs, value, query, and key matrices (adding positional encodings to the input).
5. Input-regularized Standard Transformers: Numbers for the MLPs, value, query, and key matrices; Positional encodings only in query and key matrices.
6. No Positional Encodings: Numbers for the MLPs, value, query, and key matrices.
7. Using RoPE Only: Numbers in MLPs, value, query, and key matrices, removing absolute positional encodings, and using only RoPE in standard transformers.

Below, we show the results for two tasks: cumulative sum and sorting. Their respective plots are also illustrated in Appendix G, highlighted in blue.

Results for cumulative sum:

Results for sorting:

[...]

As demonstrated, none of these variations can even achieve reasonably low test loss over OOD data. For example, the losses for all other variations are at least 100x higher than those of the Positional Transformer. This highlights the importance of the following 3 unique features of Positional Transformer, which enabled its good OOD generalization performance: (1) use positional encodings, (2) do not use positional encodings in the value matrix of the attention, (3) use only fixed positional encodings for computing the attention weights. Note that these align with algorithms that are typically used to solve algorithmic tasks.

Regarding the second suggestion (“predict target label y one by one”), we believe the reviewer may refer to an autoregressive approach, where each cumulative prediction is sequential (e.g., appending the i-th element at step i). While common in NLP, such approaches inject greater supervision than our algorithmic setting for cumulative tasks due to the incremental addition of elements. Also, this approach is inherently sequential, diverging from the parallel models explored here and in related works such as [1]. If we have misunderstood or if the reviewer has suggestions on how this setting could improve results, we welcome further discussion.

Once again, we sincerely thank the reviewer and welcome any further comments. We hope the clarifications and additional results provided are sufficient for them to consider improving their score.

[1] Sanford, C., Hsu, D., & Telgarsky, M. (2024). Transformers, parallel computation, and logarithmic depth

[2] Loukas, A. (2019). What graph neural networks cannot learn: depth vs width

[3] Sanford, C., Hsu, D., & Telgarsky, M. (2023). Representational Strengths and Limitations of Transformers

[4] Peréz, J., Barceló, P., & Marinkovic, J. (2022) Attention is Turing Complete

[5] Tsai, Y., Bai, S., Yamada, M., Morency, L., Salakhutdinov, R. (2019) Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel

Responding to the reviewer's concern about the encoding of the input we present complementary results on a more challenging setting, requiring relative reasoning over both numbers and text, which we describe below. Briefly, we present the model with an input consisting of a mixture of alphanumeric and numeric elements corresponding to categories (alphanumeric) and prices (numeric). We then ask the model to output either the minimum/maximum/sum of a random subset of categories. We measure value generalization with respect to the category names (using unseen categories when testing) as well as the category prices (testing on larger prices). We present 3 experiments: one where the prompt is exclusively to calculate the minimum, one where the prompt is exclusively to calculate the sum and finally, one where the prompt either asks for minimum or maximum (which we call multitask). This task requires pattern matching, since the categories between the train and test data are different. It requires conditional reasoning, since the queries are about a subset of the categories in the samples. Finally, it requires algorithmic reasoning to compute the output.

For details on our experimental setting, please refer to Appendix C.6 of the revised manuscript. The alphanumeric part of the input is tokenized and passed through an embedding layer, while the numeric part is passed through a linear layer. Below, we present the median MSE and MAPE losses for both standard and positional Transformer on all three experiments.

Median MSE for standard transformer:

Median MSE for positional transformer:

Median MAPE for standard transformer:

Median MAPE for positional transformer

We note that, even in this complex relational task with mixed-type inputs, the positional Transformer still significantly outperforms the standard Transformer, demonstrating the potential utility of our application for certain real-world applications.

We appreciate the reviewer’s reply, and respect their decision to maintain their score. In any case, since the rebuttal is public, we would like to provide more evidence that the reviewer’s assumption about representation of numbers is incorrect. All results are provided in Appendix C.7, and we also summarize them in the tables below. We would also like to thank the reviewer for providing a clearer description of the experiment that they would like to observe in their second reply.

We understand that the reviewer would like to see a comparison to the Transformer architecture, where a standard tokenization and digit representation is used in the Transformer architecture. Below we perform the exact tokenization and non-parallel digit representation which the reviewer requested. For example: a number like “1.32” is first tokenized to (‘1’, ‘.’, ‘3’, ‘2’), and then a learnable embedding layer is used to obtain non-parallel representations of each digit of the tokenized input. Then learnable positional encodings are created through another embedding layer, which has the same dimension as the non-parallel representation of the digits, and these positional encodings are added to the non-parallel representation of the digits in a standard fashion.

In the new experiment, we measure the performance of the same Transformer architecture setting as in previous experiments. Briefly, the Transformer architecture has very poor OOD performance, even for the simplest “min” task. We took the experiment a step further, and we fine-tuned the pre-trained large version of GPT2 (774 million parameters) from Hugging Face. This experiment also failed to beat our much smaller positional Transformer.

This is to be expected. Tokenization of numbers makes the OOD task even more difficult. That’s because, in order to achieve OOD generalization, the model needs to learn to combine digits that represent a number into an appropriate representation, and then perform the reasoning and algorithmic tasks. It is almost impossible for the training procedure to converge to such a solution just by observing the training dataset, and without seeing anything about the test dataset. The numbers between the train and test datasets are largely different. Therefore, this approach naturally fails in the OOD setting. It only works in-distribution.

On the contrary, representing numbers as actual numbers, and not tokens, removes the issue of having to re-combine tokens into an appropriate representation. This “task” is embedded into the simple fact of representing numbers as numbers. The learnable parallel representation of numbers is also not an issue for our architecture, since our architecture is clearly able to OOD generalize on all tasks that we experimented with. This is true even for the newer tasks, which include textual information, as we presented above and in Appendices C.6 and C.7.

We hope this clarifies the concern of the reviewer that our comparison with the Transformer was not sufficient. In fact, with the proposed tokenization approach, the Transformer performed worse than what we initially reported. The worse performance of the fine-tuned GPT2 model further reinforces the argument that tokenization of numbers for algorithmic reasoning tasks is not the best approach one can implement.

[...]

We present here the results of the experiments described earlier.

Results for the first setting of Appendix C.7, character-only tokenization

Median MSE for positional Transformer (first setting of Appendix C.7, character-only tokenization):

Median MSE for standard Transformer (first setting of Appendix C.7, character-only tokenization):

Median MAPE for positional Transformer (first setting of Appendix C.7, character-only tokenization):

Median MAPE for standard Transformer (first setting of Appendix C.7, character-only tokenization):

Results for the second setting of Appendix C.7, full tokenization

Median MSE for standard Transformer (second setting of Appendix C.7, full tokenization):

Median MSE for positional Transformer (second setting of Appendix C.7, full tokenization):

Median MAPE for standard Transformer (second setting of Appendix C.7, full tokenization):

Median MAPE for positional Transformer (second setting of Appendix C.7, full tokenization):

[...]

Results for the third setting of Appendix C.7, fined-tuned GPT2

Median MSE for fine-tuned GPT2 (third setting of Appendix C.7, fined-tuned GPT2):

Median MAPE for fine-tuned GPT2 (third setting of Appendix C.7, fined-tuned GPT2):