Transformers meet Neural Algorithmic Reasoners
Graph Neural Network;Transformer;Neural Algorithmic Reasoning;Length Generalization
Yes
No ethics review needed.
Please refer to Point 3 and Point 4 in the weaknesses section.
1. The TransNAR approach is well-motivated, and the core concept is clearly articulated.  

2. The method achieves competitive performance compared to a vanilla Transformer.  

3. The study explores a distilled version of the TransNAR transformer, demonstrating improved out-of-distribution robustness over a transformer trained from scratch without GNN support.
This paper proposes combining transformers with graph neural networks (GNNs) to build enhanced neural algorithmic reasoners. The core approach involves a cross-attention module, enabling language transformers to leverage information from a pretrained GNN. Experiments on the CLRS-Text benchmark indicate that the proposed method improves both in- and out-of-distribution performance.
1. The novelty is somewhat limited, as combining text-based transformers with GNNs via cross-attention is well-explored in previous works, such as [1] and [2].  
[1] https://aclanthology.org/Q19-1002.pdf. Semantic Neural Machine Translation Using AMR.   
 [2] https://aclanthology.org/2020.tacl-1.2.pdf. AMR-To-Text Generation with Graph Transformer.   

2.Both Figures 4, 5 and 6 are not clear enough. I cannot clearly read them in a print version. Please consider to use at least a table to show the results. 

3. Only the Chinchilla base transformer is used in experiments. Testing with open-source pretrained LLMs like LLaMa and Mixtral could provide broader insights. 

4. Experiments are limited to a single dataset with pre-constructed input graphs, leaving questions about generalizability. Testing on automatically constructed graph datasets, which may be less accurate but useful, could strengthen the study’s applicability.
Thank you for your encouraging comments and for supporting our work!

> The proposed method is domain-specific. It's hard to imagine how this method would benefit broader domains, as the NAR requires a specific type of input.

It is indeed the case that TransNAR as presented is not applicable to problems where graphs are not provided or derived. The motivation of our work was to show that:

* The knowledge stored in NARs _can_ be meaningfully leveraged by Transformers (via TransNAR).
* The knowledge captured within TransNAR can be _distilled_ back into a general-purpose Transformer, even though they are purely autoregressive.

We believe our distillation result to be the proof-of-concept that can be used as a blueprint for empowering broader domain learning in the future!

> How can we combine this method with a pretrained language model? So that we can leverage the math reasoning ability in state-of-the-art language models.

The most direct application of our work here would include fine-tuning the language model using our distillation loss, potentially leveraging LoRA-style techniques for efficiency. We could also use TransNAR-style models to generate additional training data for fine-tuning.

> Have you compared with using GNN-based NAR alone?

We are very upfront about NARs’ performance in the early stages of our paper:

> _Provided appropriate inductive biases are used, NARs are capable of holding perfect generalisation even on 6× larger inputs than ones seen in the training set, for highly complex algorithmic tasks with long rollouts (Jürß et al., 2023)._

It is well-known that GNN-based NARs can generalise much better than LLMs on these tasks. However, their lack of auto-regression and causal masking would make them much harder to train at scale over general purpose text data. The aim of our paper is to see whether some of these benefits can be upstreamed into LLMs, which _can_ be trained at scale more easily and efficiently.

> How does a pretrained language model perform on these tasks?

We refer the Reviewer to the CLRS-Text paper by Markeeva et al., where two-shot performance for Gemini 1.5 Flash is reported. Generally speaking, pre-trained LLMs are quite poor at executing these algorithms even across small problem sizes, with sorting perhaps being the only semi-notable exception (up to length ~14).
Reply to Reviewer 79uu
We would like to thank you for commending the clarity of the paper and the simplicity of the approach, as well as proposing multiple interesting avenues for improvement! 

We would like to initiate the discussion early, as we are not certain we completely comprehended all your suggestions – keeping in mind that our experiments may take up to 4--5 days to complete, we want to make sure that we have understood well while there’s still plenty of time to gather additional results!

Note that this is not a complete response; we will be preparing updates to the paper in parallel while we are awaiting your thoughts.

> Of course a hybrid of a general-purpose model (transformer) and a task-specific model (NAR) will perform better on the specific task. Especially when the model requires special graph-structured data, which is the case for TransNAR.

We are in agreement here, though we are afraid you may not be considering our message in its entirety. Indeed, we agree it should be obvious that a hybrid model should have a better performance on this task. 

We find it significantly less obvious that these improvements can be distilled back into a standard text-based Transformer, and this is also a key result of our work. This is because:

* Distillation is a procedure originally aimed at transferring knowledge from a larger model to a smaller one, with both assumed to share the same input and output spaces. Here we adapt it to upstreaming the OOD generalization capability of a hybrid model with multiple modalities, back into a unimodal LLM.
* We are distilling knowledge gained from a hybrid of a non-autoregressive (GNN) and autoregressive (LLM) model into a purely autoregressive (LLM) model, and it is well known that generalisation can be harmed by autoregression (see, e.g., Barbero et al. (NeurIPS’24)). The fact that distillation allows us to convey these benefits in spite of that is not obvious to us!

Would you be in agreement with these points? We are very happy to discuss them further!

> The distillation results were hard for me to process and verify. The error bars (constructed on just 3 samples!) are very large.

We appreciate this remark, and ask: would reducing the variance through increasing the number of seeds be sufficient to address your concern? 

We have already kicked off these experiments; the reason why we ask if this is sufficient _now_ is because these experiments are fairly memory intensive and slow to run. They require caching logits of the teacher model, which are quite expensive given the large vocabulary size of the underlying language model (~32k).

> An obvious baseline is missing: distilling a NAR model to Transformer.

Could you please clarify how this baseline could be constructed?

From our understanding of distillation, it is mainly applicable between two models when their output spaces match, so that the student can be trained to mimic either the outputs or the logits of the teacher.

The NAR and Transformer do not have compatible output spaces, however—the Transformer produces language tokens while the NAR produces output signals on a graph.

Are you suggesting that we manually convert the NAR outputs into text and then supervise the Transformer on this via next-token prediction?

> Give more context to the reader (CLRS-30, NAR, tabulation and analysis)

We appreciate the reviewer’s remark and commit to add a brief introduction to both benchmarks and illustrate with one sample, as well as adding a tabular version of the results and elaborating on the analysis – this will be added to a revision during the discussion period, and we will let you know once we’ve done so.

> Did you look into how the distillation data that TransNAR produces is different from the groundtruth?

The distillation data produced by the TransNAR are logits – not final predictions. They can only really be compared with ground-truth solutions after decoding into text. We can get a feel for how accurate the decoded data is (compared to the ground-truth) via the TransNAR performance metrics in our paper. 

Would you like us to do any particular qualitative study on top of this?
(First) Reply to Reviewer mvbM
Continuing our answers from Part II:

> Could you comment why you chose to call Q as Key and K as Query in equations (1) and (3)?

Thanks for the remark, this was an omission on our part – we will correct it and properly name these parameters in our next revision.

> How would you pronounce NAR? Please make sure it is consistent in the paper (for example in line 065 as "a NAR-" and as "an NAR" in many other places).

Thank you for pointing this out! We would pronounce NAR as [en-eh-ar] (spelled-out) and therefore “an NAR” is the version we would go for. We will harmonise this across the paper in the next revision.

This concludes our responses -- please let us know if there is anything you would like to discuss further! We thank you again for your supportive assessment of our work.
Reply to Reviewer gZZh (Part III)
Continuing our answers from Part I:

> Is every iteration the same as illustrated in Figure 3? It seems it is not as the final output is only from the last layer; for other layers, the outputs from both modules feed to the next layer as inputs. Could you revise Figure 3 to show the entire process accurately?

You are correct, and we will revise Figure 3 to showcase multiple layers instead of just one!

> Given that the proposed transNAR fails to improve on some of the problems, could you please comment on the limits of the NAR approach? While it can generalize over a larger range of lengths, NAR and transNAR would still lack the generalization and robustness of correct algorithms.

Of course, we will gladly comment on the limits of NARs in this regard in the paper. 

While NARs do offer superior generalisation on algorithmic tasks compared to most other neural network approaches, they are still not exact algorithms, and will make mistakes provided a sufficient distribution shift is encountered at test time.

Accordingly, we will emphasise that our aim here is to improve the algorithmic generalisation capabilities of LLMs, so that they can be applicable in the wider range of scenarios for algorithmic problems, while acknowledging that out-of-distribution generalisation is a very hard problem for neural networks in general, and we do not think of these models as a pure drop-in replacement for algorithmic computation.

> Reasoning is inherently sequantial but it seems the graphs are not. Could you comment on how the graphs emulate the sequential nature of reasoning steps?

The graphs themselves are not sequential (in fact they are assumed permutation-symmetric); the reasoning steps emerge through the temporal axis of repeatedly iterating a recurrent GNN layer (the NAR) over such a graph structure.

> One of the known issues with graphs is the oversmoothing issue. Could you comment how that affects the NAR and transNAR?

As per our reply above, we follow standard NAR practices and use the max aggregation function in our NAR—this significantly reduces vulnerability to oversmoothing due to sharp decisions over neighbourhoods.

> While the improvements over the baseline seem impressive, could ypu give the specific contributions beyond combining NAR and transformers?

We believe that the key additional contribution beyond combining the two models is distilling these capabilities back into simple Transformers (through our experiments in Section 4.2.).

The successful distillation allows us to transfer some of the knowledge gathered within TransNAR back into more standard Transformers, making this knowledge accessible in models that can operate in settings where graphs are not provided.

> Could you provide references for transformers' natural langaue UNDERSTANDING properties? The paper mentioned that transformers have unrivaled such properties (Section 1) and at the same time they are limited in generalization due to their autoregressive, causally-masked objective (Section 2). These two do not seem to go together well.

Thank you for pointing out this dichotomy; this is something we are well-aware and are happy to explicitly provide references.

The language understanding capabilities are unrivaled in decoder-only Transformers because they are able to leverage training at scale to make impressive inferences from relatively unstructured textual corpora—most other specialised models (including our NAR) require the data to be prepared in a very specific format before they will show meaningful returns. A standard reference for these abilities, and their emergence at scale, could be “Emergent Abilities of Large Language Models” (Wei et al., TMLR’22).

That being said, the impressive levels of natural language-based predictions tend to only appear _within the distribution_ the decoder-only Transformer has been trained on, and tend to suffer when pushed outside of this distribution. As a simple example, in our algorithmic tasks, the Transformers perform quite impressively on the problem sizes they’ve been fine-tuned on, yet their performance quickly collapses after we add even 1–2 nodes to the problems.

>  Would you open-source the code code for your transNAR and baseline models? As many of the papers are from one group due to the unfair advantages created, it is not acceptable to the research community.

Yes, we commit to open-sourcing a reference implementation of TransNAR as well as the relevant baseline models.

> More generally, would you be committed to the principle that scientific research should be for the common good and accessible to all?

Yes, we are strongly committed to that principle and intend to uphold it.

_Answers to be continued in the next message..._
Reply to Reviewer gZZh (Part II)
Thank you very much for the positive and supportive review, and for the high degree of detail in your remarks!

> The paper is sketchy in covering important technical details.

We will address the specific points you raised below, but if there are any other points you would like us to address that aren’t already covered by our rebuttal response, please let us know!

We intend to upstream any relevant changes into the paper later during the discussion period, and we will notify you when we’ve done so.

> For example, the abstract states a two-phase training procedure is used; yet the main paper does not mention the two-phase training procedure even once.

We apologise for not explicitly addressing the “two-phase” remark beyond the abstract, this was an oversight on our side and it will be fixed.

For completeness, we remark here that the two phases are referring to 
(1) pre-training the NAR (following Ibarz et al., LoG’22), and 
(2) fine-tuning the TransNAR using the NAR’s embeddings (following our paper).

> Related to this, the paper describes the system wide issues only from lines 213 to line 223 without mentioning some known issues.

Thank you for remarking this, we are very happy to supplement this discussion with additional issues that may arise in training such models.

> For example, the particular implementation of the transNAR uses 6 layers and it is well known beyond two to three layers oversmoothing becomes an issue in graph neural networks.

The GNN layers used in NARs (including ours!) prefer the max aggregation function – a standard choice since the NEGA paper (Veličković et al., ICLR’20) – and since max leverages a sharp decision across its neighbourhood, it is not vulnerable to over-smoothing, even over very long horizons.

> The paper freezes the weights for graph embeddings, but multimodal models rely on a shared embedding space; the differences between the two embedding spaces should be a factor, yet there is no mention about that.

Thank you for remarking this; our approach mimics that of Flamingo, which also freezes the weights for its image embedding model. We acknowledge that this is not the only way to perform multimodality successfully, and will mention this explicitly in our discussion, along with a note about the potential challenges in differing embedding spaces.

> One naive way to combine NAR and transformers for CLRS-text is to use start of the art transformers to convert the problems into the graphic form required by NAR and then apply NAR. 

This is indeed a very interesting approach and we will explicitly note it in our discussion! 

In the meantime, we prompted a few state-of-the-art models to extract feature matrices that could be used as NAR inputs from the textual data. 

While they were able to extract some interesting graph-style featurisations of the text data, we note that our NARs require highly bespoke data, including the detailed internal states of the algorithms considered, which no LLM was able to extract at this time. As such, we will not be able to feed such outputs into NAR.

> Also the transformer model used in the paper is far from the state of the art; as a baseline, the authors should use the performance of a fine-tuned state-of-the art transformer model on the CLRS-text dataset. 

We acknowledge that our Chinchilla-style base model is not present state-of-the-art at such a model scale; however we do believe our comparison is fair, because the TransNAR model also uses that same Chinchilla-style model as the starting point.

To address your question, we would like to point you to the CLRS-Text paper (Markeeva et al., 2024), where results for fine-tuning a modern Gemma 2B model on this data are provided. 

Generally it does not seem that scale or algorithmic improvements of the Gemma 2B model have made a significant difference to the models’ OOD capability compared to Chinchilla: on most algorithms, we observe that Gemma (much like Chinchilla) can learn to fit the algorithms well in-distribution, and out-of-distribution most algorithms promptly collapse.

> In addition, it seems the input to the transformer model is only the text but the transNAR has dual inputs; the impact of the dual inputs should be factored in when stating the improvements.

We agree with you that the duality of its inputs are providing TransNAR with an advantage and are happy to directly address this. 

That being said, we also wish to point you to our distillation experiments (Section 4.2.) where these capabilities are transferred back into simple Transformers, leading to improved OOD performance while no longer requiring dual inputs.

_Answers to be continued in next message..._
Reply to Reviewer gZZh (Part I)
Yes
No ethics review needed.
- Is every iteration the same as illustrated in Figure 3? It seems it is not as the final output is only from the last layer; for other layers, the outputs from both modules feed to the next layer as inputs. Could you revise Figure 3 to show the entire process accurately?

- Given that the proposed transNAR fails to improve on some of the problems, could you please comment on the limits of the NAR approach? While it can generalize over a larger range of lengths, NAR and transNAR would still lack the generalization and robustness of correct algorithms.

- Reasoning is inherently sequantial but it seems the graphs are not. Could you comment on how the graphs emulate the sequential nature of reasoning steps?

- One of the known issues with graphs is the oversmoothing issue. Could you comment how that affects the NAR and transNAR?

- While the improvements over the baseline seem impressive, could ypu give the specific contributions beyond combining NAR and transformers? It seems to me a fair comparison is to use a straightforward combination as the baseline rather than the transformers alone, whose limitations are known in this area.

- Could you provide references for transformers' natural langaue UNDERSTANDING properties? The paper mentioned that transformers have unrivaled such properties (Section 1)  and at the same time they are limited in generalization due to their autoregressive, causally-masked objective (Section 2). These two do not seem to go together well.

- Would you opens-source the code code for your transNAR and baseline models? As many of the papers are from one group due to the unfair advantages created, it is not acceptable to the research community.

- More generally, would you be committed to the principle that scientific research should be for the common good and accessible to all?

- Could you comment why you chose to call Q as Key and K as Query in equations (1) and (3)?

- How would you pronounce NAR? Please make sure it is consistent in the paper (for example in line 065 as "a NAR-" and as "an NAR" in many other places).
The paper uses graphs and texts so that the new system can use the natural language inputs as well. The combination leads to significant performance improvements over the baseline systems. While the integration uses cross attention very similar to those in multimodal models, the benefits seem significant.
This paper introduces a transNAR architecture by having both texts and graphs as input and using cross-attention to combine the strengths of both. Systematic results on CLRS-text demonstrate significant improvements in a number of problems, even though no improvements are achieved in some of them.
The paper is sketchy in covering important technical details. For example, the abstract states a two-phase training procedure is used; yet the main paper does not mention the two-phase training procedure even once. Related to this, the paper describes the system wide issues only from lines 213 to line 223 without mentioning some known issues. For example, the particular implementation of the transNAR uses 6 layers and it is well known beyond two to three layers oversmoothing becomes an issue in graph neural networks. The paper freezes the weights for graph embeddings, but multimodal models rely on a shared embedding space; the differences between the two embedding spaces should be a factor, yet there is no mention about that.

A major strength of the paper is the improvements as highlighted in Figure 1. However, the baselines do not reflect the actual state of the art. One naive way to combine NAR and transformers for CLRS-text is to use start of the art transformers to convert the problems into the graphic form required by NAR and then apply NAR. Also the transformer model used in the paper is far from the state of the art; as a baseline, the authors should use the performance of a fine-tuned state-of-the art transformer model on the CLRS-text dataset. In addition, it seems the input to the transformer model is only the text but the transNAR has dual inputs; the impact of the dual inputs should be factored in when stating the improvements.
Yes
No ethics review needed.
Did you look into how the distillation data that TransNAR produces is different from the groundtruth?
The paper was clear enough for me to understand the main idea. I appreciate that the paper contained a “Limitations” section.
The paper presents a TransNAR neural architecture that is a hybrid of Transformer and Neural Algorithmic Reasoner (NAR) architectures. The Transformer part of the model is allowed to cross-attend to NAR states. The model requires a hybrid input: the Transformer part expects text and the NAR part consumes a graph. The paper tests the proposed architecture on CLRS-Text, which is a text version of an algorithmic benchmark called CLRS-30. The model is trained on both textual inputs and the corresponding graph. The paper’s main claims are that (1) TransNAR performs better than Transformer. (2) that TransNAR teacher can produce extra training data for the Transformer student, and that this data increases the generalization of the latter.
The paper proposes a very straight-forward idea and presents unsurprising results. Of course a hybrid of a general-purpose model (transformer) and a task-specific model (NAR) will perform better on the specific task. Especially when the model requires special graph-structured data, which is the case for TransNAR. 

The distillation results were hard for me to process and verify. The error bars (constructed on just 3 samples!) are very large. An obvious baseline is missing: distilling a NAR model to Transformer.

On the writing side, the paper could benefit from the following improvements: 
- Give more context to the reader:
   - better explain CLRS-30 benchmarks
   - better explain NAR architecture
- Include a rigorous textual analysis of performance numbers (e.g." in 3 cases out of 7 our model is significantly better than the baseline (p-value = 0.01)"), in addition to presenting all numbers on all tasks in large overwhelming figures
- Include a tabular version of the results
Yes
No ethics review needed.
1. How can we combine this method with a pretrained language model? So that we can leverage the math reasoning ability in state-of-the-art language models.
2. Have you compared with using GNN-based NAR alone?
3. How does a pretrained language model perform on these tasks?
The method is straightforward and well-motivated. Experiment results also show that it performs well on the targeted tasks compared to standard transformer architecture.
The paper proposes a new model architecture, TransNAR, that combines a Transformer model with a Neural Algorithmic Reasoner via a cross-attention mechanism. The motivation is to leverage the reasoning and generalization ability in NAR representation to improve the performance of the transformer model. The experiment shows promising results in the out-of-distribution regime and over 20% absolute improvement in several classes.
The proposed method is domain-specific. It's hard to imagine how this method would benefit broader domains, as the NAR requires a specific type of input. Combining this method with existing large language models would also require further experiments. Figure 5 and 6 are hard to read.
@inproceedings{
anonymous2024transformers,
title={Transformers meet Neural Algorithmic Reasoners},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=fk4czNKXPC},
note={under review}
}
Transformers have revolutionized machine learning with their simple yet effective architecture. Pre-training Transformers on massive text datasets from the Internet has led to unmatched generalization for natural language understanding (NLU) tasks. However, such language models remain fragile when tasked with algorithmic forms of reasoning, where computations must be precise and robust. To address this limitation, we propose a novel approach that combines the Transformer's language understanding with the robustness of graph neural network (GNN)-based neural algorithmic reasoners (NARs). Such NARs proved effective as generic solvers for algorithmic tasks, when specified in graph form. To make their embeddings accessible to a Transformer, we propose a hybrid architecture with a two-phase training procedure, allowing the tokens in the language model to cross-attend to the node embeddings from the NAR. We evaluate our resulting TransNAR model on CLRS-Text, the text-based version of the CLRS-30 benchmark, and demonstrate significant gains over Transformer-only models for algorithmic reasoning, both in and out of distribution. Finally, we empirically show that Transformer-only models distilled from TransNAR models also exhibit improved out-of-distribution generalization capabilities.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Graph Neural Network Transformer Neural Algorithmic Reasoning Length Generalization
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/dfb0ce862d46c080ba331b5eab32debf86be6245.pdf
learning on graphs and other geometries & topologies
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Transformers meet Neural Algorithmic Reasoners
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission