Transformers meet Neural Algorithmic Reasoners

## Abstract

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

## 1 Introduction

Figure: Figure 1: Our TransNAR architecture, with its direct synergy of Transformers and Neural Algorithmic Reasoners, yields clear improvements in out-of-distribution reasoning across wide categories of algorithmic tasks in CLRS-Text [], a textual version of the CLRS-30 benchmark []. Here, the $x$-axis indicates one of the eight algorithmic families of CLRS-30, and the $y$-axis spans the average execution accuracy across a dataset of out-of-distribution examples. TransNAR enables *emerging capabilities* in the particular out-of-distribution regime depicted here, with over 20% absolute improvement in several of the algorithm classes.
Refer to caption: extracted/5665435/icml2024/imgs/block-plot.png

Recent work motivated [] and showcased [, ] the effectiveness of graph neural networks [, GNNs] at robustly solving algorithmic tasks of various input sizes, both in and out of distribution—such systems are often referred to as *neural algorithmic reasoners* [, NARs]. Provided appropriate inductive biases are used, NARs are capable of holding perfect generalisation even on $6\times$ larger inputs than ones seen in the training set, for highly complex algorithmic tasks with long rollouts []. NARs are, however, still relatively *narrow* forms of AI, as they require rigidly structured formatting of inputs, and they hence cannot be directly applied to problems posed in more noisy forms—such as in *natural language*—even when the underlying problem is still algorithmic in nature.

Figure: Figure 2: Augmenting LLMs with algorithmic reasoning: a bird’s eye view of TransNAR. A large language model (LLM) consumes input tokens and produces output tokens, as common for a unimodal Transformer. The neural algorithmic reasoner (NAR) module is a graph neural network (GNN) pre-trained to execute various algorithmic computation on a collection of graph-based inputs []—the pre-training pipeline is denoted by faded arrows. Throughout its forward pass, the Transformer may access the embeddings computed by the NAR, by leveraging cross-attention (trained by learnable “glue” weights).
Refer to caption: extracted/5665435/icml2024/imgs/idea.png

Conversely, the current undisputed state-of-the-art approach for modelling noisy text data are Transformer-based [] language models [, ]. In spite of their unrivalled natural language understanding properties, they are also notoriously brittle when faced with even the simplest algorithmic tasks []—especially if out-of-distribution generalisation is required [].

It appears that *uniting Transformers with NARs* can lead to fruitful returns on both sides. In this paper, we explore this interface for the first time, building the TransNAR model.

#### Contributions

Our exploration proved fruitful. The key takeaways we present in this work are as follows:

- 1.
We propose a hybrid architecture combining the language understanding of a Transformer with the robustness of reasoning of a pre-trained GNN-based NAR. The Transformer uses the NAR as a *high-dimensional tool* that will modulate its tokens embeddings.
- 2.
We show, through an evaluation on CLRS-Text [], the text-based version of the CLRS-30 benchmark, that such a NAR-augmented large language model (LLM) exhibits improved and more robust reasoning capabilities out-of-distribution (Figure 1).

Our work presents one of the most comprehensive size generalisation challenges given to Transformers to date, and the introduction of NARs moves the needle significantly.

## 2 Related work

Our work sits at the intersection of several areas: neural algorithmic reasoning, length generalisation in language models, tool use, and multimodality. Here, we briefly survey various relevant works in each area. Due to the diversity of perspectives, to preserve brevity, we do not offer a comprehensive review of related work, but rather aim to provide an indication of specific works that inspired ours the most.

#### Neural algorithmic reasoning

NAR is, in general terms, the art of building neural networks that are capable of capturing algorithmic computation. Such capabilities can be amplified by careful choices in algorithmic alignment [], step-wise training [] or contrastive objectives [].

Recently, it was demonstrated that: (1) it is possible to learn an NAR capable of executing *multiple* algorithms simultaneously in its latent space []—with the Triplet-GMPNN [] skillfully doing so for a collection of thirty algorithms across the CLRS benchmark []; (2) Once trained, such NARs can be usefully deployed in various downstream tasks: reinforcement learning [, ], self-supervised learning [], combinatorial optimisation [, ], computational biology [] and neuroscience [].

Our work’s use of NAR is mostly motivated by two of the works listed before: we use a relatively small, pre-trained, multi-task NAR [], and deploy it in a far more scaled environment: as shown by Numeroso et al. [], NAR should in principle be scalable to systems that are orders-of-magnitude greater than the NAR’s training distribution ($180,000\times$ in that particular case).

Figure: Figure 3: TransNAR hybrid architecture. Similar to Alayrac et al. [], we interleave existing Transformer layers with gated cross-attention layers which enable information to flow from the NAR to the Transformer. We generate queries from tokens while we obtain keys and values from nodes and edges of the graph. The node and edge embeddings are obtained by running the NAR on the graph version of the reasoning task to be solved. When experimenting with pre-trained Transformers, we initially close the cross-attention gate, in order to fully preserve the language model’s internal knowledge at the beginning of training.
Refer to caption: x1.png

#### Length generalisation in LLMs

While NARs can often strongly generalise to far greater test inputs [], LLMs have seen significantly less success in such scenarios. We attribute this to their autoregressive, causally-masked objective, which may not always correspond to the most logical order in which outputs of algorithms should be predicted. Just as a simple example, performance of various LLMs on multiplication can be significantly improved by predicting the result in reverse order []. Of course, on more complicated algorithms, it may be much harder to determine the best way to permute the input, and it may not be the most human-readable.

Knowledge of the above issues has led to a significant amount of effort being invested in building Transformers that can generalise in length. While length generalisation is not the only kind of distribution shift of interest to OOD reasoning, it is among the most easy such shifts to simulate. Accordingly, various works have attempted to induce length generalisation in LLMs, through the use of careful prompting [, ], randomised positional encoding [], curricula [] or scratchpads []. We firmly believe that an important trait of reasoning is robustness with respect to prompt quality—so long as the prompt unambiguously specifies the problem—and hence deliberately do not explore prompt modification approaches here; only randomised positions [] are leveraged out of the works above in our model.

#### Tool use and multimodality

Another way to obtain robust generalisation performance is to leverage a hard-coded algorithm (also known as a *tool*) by teaching an LLM to invoke its API []. Arguably, most of the major successes of reasoning with LLMs [, , ] can primarily be attributed to an LLM’s clever usage of a tool rather than the LLM itself, as a tool will by definition not have issues in generalising to diverse inputs.

Since our aim is to directly evaluate reasoning capabilities of LLMs, we explicitly do not permit tool use in our baselines. That being said, we envision the pre-trained NAR as a *modulator* for the Transformer’s embeddings which is more robust to OOD noise. Hence, we may observe the NAR as an *“internal tool”*: rather than using raw tokens, the Transformer and NAR can communicate using their embeddings, breaking the associated algorithmic bottlenecks [, ].

How to actually realise this communication and embedding exchange? For this, we turn to *multimodal LLMs* [] for inspiration, since we need to integrate signals coming from two different representations of algorithmic problems (text and graph). Specifically, our exchange operator is directly inspired by vision language models (VLMs) and the cross-attention operator used in Flamingo [], which offered a principled way of fusing information from text and image modalities.

## 3 TransNAR: Augmenting Transformers with a pre-trained GNN-based NAR

This section describes our hybrid TransNAR architecture (refer to Figure 3). TransNAR accepts a dual input consisting of a textual algorithmic problem specification (of $T$ tokens) and its corresponding CLRS-30-specific graph representation (of $N$ nodes) and outputs a textual response to the problem. We can assume that, once encoded, the textual input is stored in ${\bf T}\in\mathbb{R}^{T\times k}$, and the graph input is stored in ${\bf G}\in\mathbb{R}^{N\times l}$. Note that, for simplifying the equations to follow, we make an assumption that all of the information relevant to the graph version of the problem is stored in the nodes—which is often not true in CLRS-30 (there may be edge- and graph-level inputs as well) but it doesn’t change the underlying dataflow presented below.

The forward pass of TransNAR unfolds as follows. First, we properly initialise the inputs by setting ${\bf T}^{(0)}={\bf T}$ and ${\bf G}^{(0)}={\bf G}$. Next, to compute the representation of a step $(t+1)$, the text (token) representations are fed to the current layer of the Transformer []:

$$ ${\mathbf{\Theta}}^{(t+1)}=\mathrm{FFN}\left(\mathrm{softmax}\left(\frac{({\bf T }^{(t)}{\bf Q}_{t})^{\top}{\bf T}^{(t)}{\bf K}_{t}}{\sqrt{d_{k}}}\right){\bf T }^{(t)}{\bf V}_{t}\right)$ (1) $$

where ${\bf Q}_{t},{\bf K}_{t}\in\mathbb{R}^{k\times d_{k}},{\bf V}_{t}\in\mathbb{R}^
{k\times k}$ are the key, query and value transformations, respectively, and $\mathrm{FFN}$ is a feedforward network. In a similar manner, the graph representations are fed to the NAR layer, implementing e.g. a standard max-MPNN []:

$$ ${\bf g}^{(t+1)}_{u}=\phi\left({\bf g}^{(t)}_{u},\max_{1\leq v\leq N}\psi\left( {\bf g}^{(t)}_{u},{\bf g}^{(t)}_{v}\right)\right)$ (2) $$

where $\psi,\phi:\mathbb{R}^{k}\times\mathbb{R}^{k}\rightarrow\mathbb{R}^{k}$ are learnable *message* and *update* functions, respectively, and $\max$ is the elementwise-max aggregation. Note that Equation 2 only provides pairwise interactions between nodes for brevity—in reality, our NAR is a Triplet-GMPNN [], which also contains triplet interactions and a gating mechanism. Further, note that there is no timestep index on the learnable parts of the NAR—at each step, a *shared* function is applied. This aligns well with the iterative, repeated nature of algorithmic computation on graphs.

Once both streams have prepared their representations, $\mathbf{\Theta}^{(t+1)}$ and $\mathbf{G}^{(t+1)}$, the node embeddings in the graph condition the Transformer’s token embeddings to produce the final outcome of the TransNAR block in the Transformer stream, inspired by Flamingo []:

$$ ${\bf T}^{(t+1)}=\mathrm{FFN}\left(\mathrm{softmax}\left(\frac{({\mathbf{\Theta }}^{(t)}{\bf Q}^{\times}_{t})^{\top}{\bf G}^{(t)}{\bf K}^{\times}_{t}}{\sqrt{d _{k}}}\right){\bf G}^{(t)}{\bf V}^{\times}_{t}\right)$ (3) $$

where ${\bf Q}_{t}^{\times},{\bf K}_{t}^{\times}\in\mathbb{R}^{k\times d_{k}},{\bf V}
_{t}^{\times}\in\mathbb{R}^{k\times k}$ are the key, query and value transformations of the cross-attention, respectively. No additional transformations are performed on ${\bf G}^{(t+1)}$ before concluding this layer.

This process repeats until the final, $N_{l}$-th layer, when the final text output is read out from ${\bf T}^{(N_{l})}$. The final output is converted into token logits by a prediction head produced by the final layer, which we supervise by means of a standard next-token prediction objective.

Prior to the start of TransNAR fine-tuning, we pre-train the NAR to robustly execute the thirty algorithms spanned by CLRS-30 [], in a manner similar to Ibarz et al. []. Such procedures are known to yield out-of-distribution generalisation at up-to-$4\times$ larger inputs in graph space. The parameters of the NAR are generally kept *frozen* during fine-tuning, as additional gradients would eliminate the model’s original robustness properties. This is also, similarly, the reason why no cross-attention is performed by the graph embeddings. The LLM itself may be pre-trained over large-scale datasets [], to establish its general language priors, though we recover the same experimental findings even if the LM is randomly initialised at the onset.

## 4 Experiments

**Table 1: Samples from CLRS-Text for various algorithmic tasks of problem size 4. Input and target parts of the examples are clearly specified. Note that variability is limited at size 4, meaning that some algorithms may have trivial answers for the given inputs. Such effects tend to quickly disappear with scaling the problem size.**
| Algorithm | Input & Target | #tokens |
| --- | --- | --- |
| Articulation points | articulation_points: | 63 |
|  | A: [[0 0 0 0], [0 1 0 0], [0 0 0 0], [0 0 0 1]] |  |
|  | is_cut: [0 0 0 0] |  |
| Binary search | binary_search: | 49 |
|  | key: [0.011 0.029 0.635 0.719], target: 0.122 |  |
|  | return: 2 |  |
| Insertion sort | insertion_sort: | 65 |
|  | key: [0.561 0.081 0.892 0.565] |  |
|  | pred: [0.081 0.561 0.565 0.892] |  |
| Jarvis’ march | jarvis_march: | 75 |
|  | x: [-1.22 -1.05 0.331 -1.55], y: [-1.48 1.39 0.899 0.1] |  |
|  | in_hull: [1 1 1 1] |  |
| KMP Matcher | kmp_matcher: | 37 |
|  | string: [0 0 0 1], key: [3 3 2 3] |  |
|  | match: 0 |  |
| Matrix Chain Order | matrix_chain_order: | 77 |
|  | p: [0.461 0.957 0.462 0.42] |  |
|  | s: [[0 0 0 0], [0 0 1 2], [0 0 0 2], [0 0 0 0]] |  |
| Task Scheduling | task_scheduling: | 61 |
|  | d: [2 3 3 4], w: [0.042 0.875 0.954 0.761] |  |
|  | selected: [0 1 1 1] |  |

Figure: Figure 4: TransNAR significantly outperforms the baseline Transformer. We compare TransNAR to its corresponding Transformer baseline on various algorithms and for various input sizes: $12$ is the largest size in-distribution. The other two sizes tested—$10$ and $14$—are out-of-distribution, with the former testing interpolation and the latter extrapolation. Note that in-distribution generalisation is much easier for Transformers, and as such, we have modified the $y$-axis for this setting only to the $[0.7,1.0]$ range. It is evident that, on most algorithmic tasks of interest, the TransNAR is capable of outperforming its baseline Transformer. Additionally, we see that this advantage is consistent across both training regimes: initial training and finetuning. The metric used is the CLRS score. Each model was trained with 4 random seeds. Error bars indicate $\pm 1$ standard deviation.
Refer to caption: extracted/5665435/icml2024/imgs/results.png

In our experimentation, we will demonstrate that the recipe offered by TransNAR admits significant benefits to out-of-distribution reasoning in language model architectures. In this section we provide details of our experimental setup.

Transformer architecture and initialisation. We use a decoder-only, 6 layers, transformer model from the Chinchilla family [] pretrained on MassiveText []. In particular we use a model of $70$ million parameters with a context size $2,048$. To showcase the suitability of our approach regardless of the starting point of training, we run two ablative variants. In the first, the Transformer weights are initialised with the outcome of the pre-training—emulating a *fine-tuning* scenario—and in the second, we use a fully random initialisation. In our figures and tables of results that follow, we will refer to these two setups as “Pretrained” and “Untrained”.

Randomized positional encoding. Previous work has emphasised the significant relevance of *randomised* positional embeddings in Transformers, especially for enabling more robust reasoning []. Corresponding to previous studies on the generalization capabilities of language models, randomised positional embeddings have indeed led to significant gains on both our baselines and TransNAR, allowing more interesting reasoning behaviour to emerge in both. As such, all our experiments in this paper will use randomised positional embeddings. We provide more details in Appendix 6.

Pre-training the NAR.
Following Ibarz et al. [], we pre-train a multi-task MPNN-based NAR on input problem sizes of up to 16, from the CLRS-30 benchmark []. Owing to its graph structure formulation, such NARs are capable of significant OOD generalisation—sometimes staying competitive on graphs that are $4\times$ the size. We will attempt to utilise such models through TransNAR, to convey this rich representational knowledge into text.

Combining cross-attention contributions from nodes and edges.
The NAR pre-trained by the method presented in Ibarz et al. [] produces both node and edge latent representations, and we cross-attend to both of them, as they may contain complementary useful information. To cross-attend over the edge features, ${\bf E}^{(t)}\in\mathbb{R}^{N\times N\times k}$, we apply Equation 3 one more time (with $\mathbf{\Theta}^{(t)}$ cross-attending over ${\bf E}^{(t)}$), with the caveat that we need to flatten the first and second axis of ${\bf E}$ into one, to make sure the dimensionalities match. We combine the cross-attention contribution from the node and edge embeddings provided by the pre-trained NAR by concatenation, followed by the application of a linear layer. We have attempted to use other reduction schemes such as summing the vectors, or applying a 2-layer MLP. We have also attempted different preprocessing schemes such as orthogonalising the contributions using the Gram‐Schmidt process to ensure their algebraic complementarity before combining them. However, none of these variations have brought improvements over our original approach.

Datasets.
We use the CLRS-Text benchmark [], the text version of the CLRS-30 benchmark [].
Table 1 showcases several samples from this dataset, along with their input size and number of tokens. Note that the textual representation is directly derived from the graph-based CLRS-30 in a deterministic manner, so the two datasets convey exactly the same information. However, due to the tokenised representation, there are stringent limitations on how large of a problem size we can evaluate on without running out of context length for the Chinchilla models.

Accordingly, we train our algorithms on smaller problem sizes—$[4,8]$ and $12$, and evaluate on problem sizes $10$ (*out-of-distribution—interpolation*), $12$ (*in-distribution*), $14$ (*out-of-distribution—extrapolation*).

It is worth noting that CLRS-Text is among the most challenging long-range reasoning tasks for language models, compared to the present evaluation landscape—a clear step-up in complexity from grade school math, mainly because it allows for explicitly controlling for out-of-distribution generalisation. Yet, there exists a clear polynomial-time-algorithmic description for each of them, meaning that they can be explained in relatively little parameters—certainly way less than a typical large language model of today!

The dataset comprises $10,000$ samples per algorithm per input size, making up a total of $2,400,000$ data points, split as per above into $70\%$ for training and $30\%$ for validation.

#### Training details.

We train all models over seven epochs of the training data with a batch size of $256$ and employ an Adam optimizer [] with a learning rate of $10^{-4}$. We apply randomized positional encoding with a maximal length of $8,192$ on top of Rotary Positional Encoding (RoPE) used in the base Chinchilla transformer []. As previously mentioned, for all TransNAR models, we keep the NAR frozen during training.

Figure: Figure 5: Shape Score: The TransNAR significantly outperforms its baseline in terms of producing correct shapes. This score sheds light on an obvious failure model of regular Transformers out-of-distribution: they fail to capture the seemingly trivial dependency between input size and output size, and so irrespective of the complexity of the algorithm itself. The TransNAR model manages to considerably alleviate this problem (with many emerging gains), albeit, these gains do not always lead to perfect scores, implying a fruitful direciton for future research.
Refer to caption: extracted/5665435/icml2024/imgs/shape_score.png

Evaluation metrics. We refrain from computing the accuracy of each model using exact string matching, on the grounds that this metric does not provide insights as to the causes of failure on a particular datapoint, and more critically, it fails to capture how close to correctness a given model output is (as observed by Veličković et al. []). Instead, we evaluate the performance of each model according to three metrics measuring capabilities of increasing complexity over the generated text:

- 1.
The shape score: a binary-valued metric capturing whether the output has the right shape. For example, if we consider a sorting task, the output should have exactly the same number of elements as the input. Similarly, if the output is a matrix, we ensure its shape is consistent with both the input and the task.
- 2.
The parse score: a binary-valued metric capturing whether the output is free from any illicit characters, for example, considering again a sorting task on a list of numbers, the output shouldn’t contain any letters of the alphabet.
- 3.
The CLRS score: The percentage of elements in the output that match the ground truth answer. This score is the one traditionally used in CLRS-30 [, ], hence its name. Note that we automatically assign a CLRS score of 0 if the shape score is 0, as there is no clear correspondence between output indices.

These multi-faceted scores are explicitly designed to capture the various failure modes of LLMs when learning to reason over text inputs: they may overly specialise to the training problem sizes (leading to incorrect shapes at test time), fail to cope with unseen number combinations (leading to incorrect parsing), and of course, produce incorrect or inconsistent outputs, captured by the CLRS score.

### 4.1 Results

We summarize our findings in Figure 4 (for CLRS score). Our results show that our TransNAR significantly outperforms the baseline Transformer overall, and on most individual algorithms, both in- and out-of-distribution. In particular, we see that our approach not only enhances existing out-of-distribution generalisation capabilities, but also causes the emergence of these capabilities when there was a complete lack thereof—reflected in the figure by zero or near-zero performance of the baseline [].

The analysis of shape score (Figure 5) provides an additional way to shed light on why TransNAR performed as well as it did. Recall, first, that CLRS score is necessarily zero if shapes do not match. Observing the shape scores achieved, it appears that grounding Transformer outputs in NAR embeddings significantly increases the proportion of inputs for which a Transformer will produce an output of the correct shape—indicating that this is one very specific failure mode that TransNAR helps alleviate.

We note, however, that there remain a few algorithms for which TransNAR is not able to outperform the baseline. A closer look at the results indicates that such tasks (Binary Search, Find Maximum Subarray, Minimum, and Quickselect) all involve an element of *searching* for a particular index in an input list. This hints at a unified failure mode: as these failures persist both when interpolating and extrapolating, the model as implemented is not able to generalise to novel *index boundaries* unseen in the training data. We therefore suspect that the use of *index hints*—as already demonstrated by Zhou et al. []—is a promising avenue for ameliorating this behaviour. Alternatively, it might be the case that the final NAR-computed hidden states are harder to decode by the cross-attention layers in a generalisable way, and therefore might require either giving an additional capacity to the cross-attention and/or performing a more progressive decoding in that: instead of having all cross-attention layers decoding from the final NAR-computed hidden states, s, we could have early cross-attention layers decode from hidden states coming from earlier message passing steps, and later cross-attention layers decode from the later message passing steps.

Lastly, we provide parse scores in Appendix 7—omitting them from the main text because, in most cases, parsing can be done at full accuracy.

### 4.2 Limitations

While our approach demonstrates favourable average performance under all out-of-distribution regimes we have evaluated, we highlight that TransNAR requires access to both textual and graph-representation inputs to be efficiently trainable and usable. While this limits TransNAR to cases where a particular ground-truth executor or simulator (or prior belief about one) is available, now that we know that TransNAR-like ideas are beneficial, future research can enable the deployment of such ideas into purely unimodal Transformers. For example, lifting the need for a second data stream can be done by *distilling* the knowledge acquired by the trained TransNAR model into a vanilla Transformer model.

## 5 Conclusions

We presented a Transformer-NAR hybrid architecture: a language model that combines the language understanding skills of a Transformer with the robust algorithmic reasoning capabilities of a pre-trained graph neural network-based neural algorithmic reasoner, to solve algorithmic tasks specified in natural language. We have demonstrated the superiority of our model over its Transformer-only counterpart on the CLRS-text benchmark, in the in-distribution, and more importantly, in two out-of-distribution regimes, with respect to the input problem size. We hope that future work will draw on our results and insights shared here, and further investigate expansions of interest, notably, datasets with more ambiguous problem specifications (as often encountered in the real world), and for which their corresponding equivalent solver-ready symbolic inputs are not given in advance.

## 6 Effect of Randomized Positional Encoding

Using randomized positional encoding has benefitted both our model and the baseline. In particular, combining them with NAR hiddens led to improvements OOD, most prevalently in the interpoloation regime (at length 10), but also, to some extent, in the extrapoloation regime (at length 14). One result we found interesting, was that before instating randomized positional encoding, the OOD performance of our hybrid models was limited (in fact thresholded) by the performance of the base LLM. Concretely, if the base LLM achieved near-zero performance, the hybrid architecture would fatally share the same fate. We can see that this is no longer the case: if the base LLM uses randomized positional encoding, even if its performance is near-zero, that of the hybrid architecture can still be reasonably good. This is illustrated in the second column of the figure 4, for example on the Graham Scan, Jarvis March, MST Prim algorithms.

## 7 Parse Scores

Figure: Figure 6: Parse Score: We can see that for a few algorithms, the TransNAR architecture falls behind the baseline in the extrapolation regime likely due to an unsufficient capacity of the cross-attention in charge of decoding from the NAR’s outputs.
Refer to caption: extracted/5665435/icml2024/imgs/parse_score.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: Please refer to Point 3 and Point 4 in the weaknesses section.
soundness: 2
strengths: 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.
summary: 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.
weaknesses: 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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - 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).
soundness: 3
strengths: 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.
summary: 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.
weaknesses: 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.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Did you look into how the distillation data that TransNAR produces is different from the groundtruth?
soundness: 2
strengths: The paper was clear enough for me to understand the main idea. I appreciate that the paper contained a “Limitations” section.
summary: 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.
weaknesses: 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

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 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?
soundness: 3
strengths: The method is straightforward and well-motivated. Experiment results also show that it performs well on the targeted tasks compared to standard transformer architecture.
summary: 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.
weaknesses: 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.

Meta-review: The paper proposes a novel hybrid architecture called TransNAR, which combines the strengths of Transformers and Graph Neural Networks to enhance neural algorithmic reasoning. The authors introduce a two-phase training procedure where the Transformer cross-attends to node embeddings from a pre-trained GNN. The model is evaluated on the CLRS-Text benchmark, demonstrating significant improvements in algorithmic reasoning tasks, both in and out of distribution. Additionally, the paper shows that Transformers distilled from TransNAR models exhibit improved out-of-distribution generalization.


Strengths:
* Well-motivated approach, core-concept is clearly articulated (despite some missing technical details and presentation issues)
* Strong empirical results (although baseline appropriateness is questionable)
* Shows that the hybrid model can be distilled back into a transformer-only model that outperforms a text-only base transformer model

Weaknesses:
* As pointed out by multiple reviewers, the idea of combining Transformers with GNNs via cross-attention is not entirely new and has been explored in previous works. 
* Limited experimental scope: no evaluation with pretrained open-source LLMs, a single (meta-)dataset considered
* Missing baselines: 
* Poor presentation of results, hard-to-read bar charts with minimal difference, large error bars computed on few seeds

What's missing:
* As suggested by virtually all the reviewers, the presentation of results in this paper requires severe polishing. Font size, y-axis dynamic range, etc, large error bars due to small number of seeds, all make it hard to appreciate and assess the significance and soundness of the results. Additionally, some additional simple baselines, including those suggested by the reviewers, would have gone a long way towards contextualizing the performance of the proposed method. Finally, although the authors provided some additional results in the form of table during the rebuttal, those did not seem to have made it to the revised version of the paper. Additionally, the authors promised a "data generation baseline" in response to mvbM's review, which I could not find in the comments or revised pdf.

Overall, this paper is an interesting albeit incremental contribution that suffers from various issues that in my opinion need to be resolved before publication in a venue like ICLR.

Decision: Reject