Transformers meet Neural Algorithmic Reasoners

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).

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?

Thank you for engaging with us on distillation experiments! We find them to be an important part of our paper and are willing to improve the presentation there.

On the first point, we believe there is a slight misunderstanding. We wrote:

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

Which is exactly offering to convert NAR outputs into text & train the Transformer directly on them -- we just weren't sure we understood correctly! We will try to kick off NAR teacher experiments in the time we have left.

On the other, we are delighted that our 10-seed distillation results have converged, and significantly reduced standard errors. Please see below for all soft- and hard-OOD results.

Key results:

• Distillation almost-always significantly improves performance compared to the baseline (irrespective of distillation loss coefficient).
• As we drift further out-of-distribution, distillation fully on logits (1.0) trumps partially combining distillation with next-token prediction (0.5).

Please let us know if these results (and their tabulation!) appropriately address your concerns!

Length: 10

Length: 14

Length: 6

Length: 16

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...

Thank you for the kind words about the motivation of our work and recognising the strength of our comparative evaluation!

As we have done with Reviewer mvbM previously, we would like to initiate the discussion with you early, as we are not certain we completely comprehended all your suggestions. Some of our experiments may take up to 4–5 days to complete, and 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.

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].

Thank you for pointing out these references to us, and we will certainly cite and explore them in our related work discussion! We were not aware of these specific works.

To be clear: we are not claiming to be the first to combine GNNs and Transformers in general, but we do claim to be the first to combine them in the context of (algorithmic) reasoning and out-of-distribution generalisation, which is a setting that is particularly harmful for decoder-only Transformers.

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.

We appreciate your concern about the figures’ readability, and commit to tabulating these results in our next revision! We will let you know once the paper has been updated.

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

We appreciate your remark on this!

While it would certainly be possible to include additional base Transformers to the mix, we are averse to it for two reasons:

• It would take a significant amount of time and resources, which we would rather dedicate to other experiments at this time.
• It is highly unlikely to lead to a different conclusion to what we already observed with Chinchilla.

To provide direct evidence towards the second point, we would like to point you to the results in the CLRS-Text paper (Markeeva et al., 2024), where results for fine-tuning a modern, open-source, Gemma 2B model on this data are provided.

Generally, it does not seem that scale or algorithmic improvements of the open 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, performance on most algorithms promptly collapses (also like Chinchilla).

As such, it is our expectation that, were we to e.g. repeat our experiments with Gemma 2B, we’d still observe significant collapse out-of-distribution (as Markeeva et al. did), and TransNAR would be able to supplement that.

We have no strong reason to expect any other open language model to behave differently.

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.

We appreciate the Reviewer’s concern!

Firstly, we’d like to note that CLRS-Text should not be considered a “single dataset” – realistically it encompasses thirty different algorithmic skills that are learned and evaluated simultaneously.

Second, the focus of our work is on out-of-distribution generalisation – specifically, length generalisation – of language models. It is absolutely standard to leverage carefully constructed synthetic tasks (where outputs can be computed at arbitrary input lengths) for this purpose; see, e.g., “Exploring Length Generalization in Large Language Models” from Anil et al. (NeurIPS’22), which uses two crafted tasks (Parity and Variable Assignment).

Taken together, the tasks in CLRS-Text are arguably more challenging than most tasks typically seen in length generalisation papers, encompassing more generic polynomial time procedures beyond simple arithmetic – spanning path-finding, geometric algorithms, string matching, sorting, minimum spanning tree finding, and more.

It is because of this that we believe our collection of datasets is sufficient for the purpose we aim to show. Please let us know what you think!

All that being said, we are always open to suggestions for strengthening our work’s experimental section – please, let us know if there are any specific datasets you would like us to explore, and we can attempt to pursue them in the remainder of the discussion period.