A Context-Enhanced Framework for Sequential Graph Reasoning

1 Introduction

Recent years have witnessed a steady growth in exploring the reasoning capacity of deep learning [He et al., 2023; Jiang et al., 2022; de Rezende et al., 2021]. This research domain holds profound significance in the field of artificial intelligence - it enables AI not only to recognize patterns but also to deduce new knowledge from existing information, thereby broadening the capacity for complex problemsolving. As a variety of reasoning tasks can be formulated into graph problems, there is an increasing focus and attention on deep graph reasoning [Wang et al., 2023; Ling et al., 2022;

∗Corresponding authors.

Zhu, 2022]. This trend delves into the study of graphstructured data and aims to build specialized network structures that can learn to reason over them.

In this paper, we consider sequential graph reasoning, one of the most challenging tasks in deep graph reasoning. Diverging from conventional reasoning tasks like graph (node) classification [Wu et al., 2019], this type of task goes beyond the one-shot model and targets step-by-step reasoning over graphs. The applications of seq-graph reasoning span across various domains in artificial intelligence. For instance, in AIpowered math problem solving [Lample and Charton, 2020; Huang et al., 2022; Drori et al., 2020], neural networks are expected to provide the entire process of solving a mathematical problem, rather than a mere final answer. Another emerging application is neural algorithm learning [Diao and Loynd, 2023; Ibarz et al., 2022; Jain et al., 2023], where networks are anticipated to sequentially imitate individual steps of classical graph algorithms.

There have been many recent advances in seq-graph reasoning [Li et al., 2023; Zhao et al., 2024; Gerasimova et al., 2023]. While the proposed neural architectures exhibit diversity, they generally adhere to a common sequential paradigm. In the framework, graph-structured hidden states are maintained in memory. At each step, the outcome of the preceding step serves as input, integrating with the hidden states for reasoning. Variations among existing works arise in the specifics of their integration processes. Numerous graph networks (e.g. GAT [Velickovic et al., 2018] and MPNN [Gilmer et al., 2017]) that excel at one-shot reasoning have been investigated.

The current paradigm is essentially a direct combination of the techniques for processing serialized data and those for handling graph-structured data. We notice that such a straightforward combination may fall short of yielding competitive reasoning capabilities due to the distinctive nature of seq-graph reasoning. Compared to traditional seq-to-seq tasks [Yu, 2017], each step's outcome in seq-graph reasoning has a much stronger intrinsic connection with the others. For example, in automatic geometric problem solving [Schreck and Mathis, 2016], the outcome sequence forms an entire problem-solving process, which can even be represented by an extra reasoning graph. However, most existing works fail to capture these intrinsic connections.

1.1 Our Contributions

Motivated by the observation above, this paper proposes a framework to capture the contextual patterns of step outcomes in seq-graph reasoning. To evaluate the reasoning ability of this framework, we conduct empirical evaluations on the challenging CLRS Reasoning Benchmark [Velickovic et al., 2022]. This benchmark consists of 30 different algorithmic tasks for probing the reasoning capabilities of graphstructured networks. The contributions of the paper are summarized as follows:

• We introduce a Context-Enhanced Framework (CEF) for seq-graph reasoning. The framework extends the existing paradigm and exhibits notable flexibility, allowing for adjustments based on available memory space during implementation.
• We show that the framework can effectively integrate with various existing architectures. Specifically, the paper categorizes existing methods into two types: GNNbased and Transformer-based, and provides detailed CEF implementations for both.
• We perform a diverse range of experiments on the CLRS benchmark. The results demonstrate that, for both types of existing architectures, our framework significantly enhances their performances across the majority of the datasets within the benchmark. When integrated with the latest architectures, we achieve new state-of-the-art results.

1.2 RelatedWork

Graph Neural Networks. Graph neural networks have emerged as a powerful tool for effectively capturing and leveraging the structural relationships present in graphstructured data. Numerous neural architectures have been proposed in this domain, such as Graph Convolution Networks (GCN) [Kipf and Welling, 2017], Graph Attention Networks (GAT) [Velickovic et al., 2018], Message Passing Neural Networks (MPNNs) [Gilmer et al., 2017] and so on. These architectures excel not only in one-shot graph reasoning [Ma et al., 2023] but also demonstrate comparable performance in seq-graph reasoning tasks [Ramesh et al., 2023]. Recently, [Ibarz et al., 2022] introduced a novel GNN architecture called triplet-GMPNN. They augment the classical MPNN to perform message passing toward edges and have achieved remarkable results in neural algorithmic reasoning.

Graph Transformers. Transformers and its variants play a crucial role in domains such as natural language processing, computer vision, and time series analysis [Vaswani et al., 2017]. Consequently, leveraging Transformers for tasks related to graph reasoning has emerged as a recent research focus, leading to the development of several variants of Transformers tailored for handling graph-structured data [Wu et al., 2021; Hussain et al., 2022]. A recent advancement in this direction is proposed by [Diao and Loynd, 2023]. They introduced Relational Transformers, which enhance attention computation by aggregating edge and graph feature information with node features.

Beyond One-Shot Graph Reasoning. The architectures mentioned above are fundamentally tailored for one-shot reasoning. When addressing seq-graph reasoning tasks, they necessitate the paradigm stated in the introduction. There have been some recent advances in exploring competitive paradigms for seq-graph reasoning [Numeroso et al., 2023; Jain et al., 2023; J¨urß et al., 2023]. Notably, the studies most closely related to our research were recently introduced by [Jain et al., 2023] and [J¨urß et al., 2023], where the GNN architectures are augmented with a priority queue and a stack, respectively.

1.3 Paper Organization

The paper is organized as follows. Section 2 provides preliminaries, including a problem definition and a paradigm widely used in the literature. In Section 3, we introduce our framework, and subsequently, Section 4 details the concrete implementations of our framework for different types of existing approaches. Empirical evaluations are presented in Section 5. Finally, Section 6 concludes the paper.

2 Preliminaries

The section formalizes the problem setup and introduces a widely-used paradigm. In a seq-graph reasoning task, we are given graph-structured data and required to perform reasoning over it at most T times. The given data is represented by a graph G = (V,E) with n nodes and m edges, where each node v ∈ V has feature x(0) v and each edge e ∈ E has feature x(0) e . For the notational simplicity, define X(0) V := {x(0) v }v∈V , X(0) E := {x(0) e }e∈E, and X(0) = {X(0) V ,X(0) E }. Further, for each training data, we are provided with the ground truth Y(t) of each reasoning step t ∈ [T ].

A widely used paradigm in the literature is the encodeprocess-decode paradigm [Ibarz et al., 2022]. In this paradigm, there are three modules: encoder E, processor P, and decoder D. In each step t ∈ [T ], the input tensors X(t-1) undergo three successive transformations E,P,D and then yield X(t) , which serves as the input for the next step.

Encoder Module. There are usually two encoder networks in this module 1 , each tasked with encoding node features and edge features respectively. Thus, after feeding X(t-1) to the encoder module, we obtain latent features L(t) for nodes and edges in the graph:

Processor Module. This module is the central component of the entire paradigm, where most prior works primarily vary. It consists of a (standard) graph-structured network (e.g., graph neural networks or graph transformers), and maintains a hidden state for each node and edge. Denote by H(t-1) V ,H(t-1) E the hidden states of the node set and the edge set at the beginning of step t respectively. First, we concatenate the latent feature and the hidden state of each node and edge, obtaining their new representations:

1 Some prior works add one extra encoder network to handle the global feature of the graph.

And then, feed these representations into the graph-structured network, obtaining new hidden states:

Decoder Module. Without loss of generality, we can assume that this module contains two networks to decode the hidden states of nodes and edges respectively:

Given the ground truth Y(t) of each step, we can compute the total loss and thus, train the entire network.

3 CEF: AContext-Enhanced Framework

The section presents a framework that generalizes the standard encode-process-decode paradigm. Before stating the details, let us give some intuition first. Compared to the traditional seq-to-seq tasks, we observe that the ground truths in seq-graph reasoning have a much stronger inner connection with each other. For example, in neural graph algorithm learning, the goal is to imitate the execution of a classical graph algorithm like Prim's algorithm [Prim, 1957]. In this scenario, each Y(t) corresponds to an individual step of Prim's algorithm, and the entire ground truth sequence represents the process of solving the minimum spanning tree problem by Prim's algorithm, which is an optimization problemsolving procedure. Thus, similar to the consideration of momentum in historical gradients during gradient descent, in seq-graph reasoning, momentum concerning the outcome at each step should also be introduced and taken into account. This observation motivates our framework.

The framework builds upon the encode-process-decode paradigm and newly introduces a Preprocessor module between the encoder and processor. See Figure 1 for an illustration. This module is designed to capture dependencies and patterns on the latent feature sequence < L(1) , L(2) , ..., L(t) , ... >, which are also referred to as "context" information. Hence, we term the new framework a Context-Enhanced Framework (CEF). More precisely, the module maintains a context state for each node and edge, storing the information aggregated from historical latent features. For each node v (resp. edge e), denote by c(t) v (resp. c(t) e ) the context state at the beginning of step t. Note that c(t) v can be a direct concatenation of historical latent features or the outcome of applying some dimensionality reduction techniques (e.g. gating mechanisms) to them. For simplicity, C (t-1) V , C (t-1) E and C(t-1) are defined similarly.

The module involves two pivotal operations: context enhance function fenhance and context update function fupdate. The application of fenhance aims to augment latent features with contextual information, subsequently serving as input to the processor module:

Concurrently, fupdate is employed to update the context state of each node and edge:

The most critical aspect within the framework is how to set appropriate fupdate and fenhance so that they can integrate well with various types of processor modules. We notice that a popular approach to capturing patterns in historical latent features is through the use of the attention mechanism. Specifically, we can employ a node-wise (resp. edge-wise) concatenation function for fupdate and a QKV attention network for fenhance. However, this method proves to be quite resourceintensive. The demand for memory significantly escalates with an increase in the total number of reasoning steps. An experimental study on this implementation is provided in the paper's full version. In the main body, our focus shifts to leveraging gating mechanisms in the preprocessor module - a more space-efficient approach that demonstrates comparable empirical performance.

4 Applications

As outlined in Section 1.2, existing (processor) architectures fall within two categories: GNN-based and Transformerbased. This section considers both processor types and presents specific implementations of the framework for each.

4.1 Integration with GNN-based Processors

This subsection considers GNN-based processors. We follow the literature [Ibarz et al., 2022] and formulate this type of processor into the model below. Note that the existing GNNbased processors only maintain hidden states for nodes.

In each step t ∈ [T ],

Various processors may utilize distinct f1, f2, f3 and aggregation function ⊕. Since only node hidden states are preserved and computed in such processors, we also focus on the transformation of node latent features in the preprocessor module and let s(t) e = ℓ(t) e directly for each edge e ∈ E. We introduce a novel gating mechanism to preprocess each node latent feature. Further, the weight sharing technique [Garland and Gregg, 2018] is employed to make the module more memory-efficient and reduce the risk of overfitting. We implement the preprocessor as follows:

flinear is a linear layer that transforms c(t) v into a onedimensional scalar. Then, we apply function tanh and ReLU sequentially, yielding αv ∈ [0, 1). Interpreting αv as a forget factor, s(t) v and c(t+1) v are obtained by a linear combination of the context state and latent feature. We provide an experimental study and discussion (in the full version of the paper) to explain why we chose the sigmoid function. In summary, empirical findings suggest that when integrating with GNN processors, the latent features in seq-graph reasoning exhibit distinctive patterns and prefer a higher probability of the forget factor being 0.

4.2 Integration with Transformer-based Processors

This subsection considers Transformer-based processors. They maintain hidden states for both nodes and edges. We propose a model below to capture the common structures of such processors. Use N(v) := {u|(v, u) ∈ E} to denote the neighbors of v. In each step t ∈ [T],

The model extends the QKV attention mechanism, taking into account the incorporation and subsequent updating of edge features. Different processors employ different f functions within the model. Certainly, we could adopt a context enhancement approach similar to the one in Section 4.1. However, we notice that, unlike existing GNN processors, the Transformer processors in the literature do not utilize gating mechanisms in the update of hidden states. Thus, to capture the contextual information of latent features and hidden states simultaneously, we propose a refined framework that seamlessly integrates the preprocessor and processor, allowing for mutual sharing of states. Initially, the preprocessor utilizes hidden states for context state updates.

Subsequently, the processor employs the resulting context states to update hidden states.

flinear-1 and flinear-2 are two linear layers that transform the context states of nodes and edges into one-dimensional scalars, respectively. Then we apply function sigmoid, yielding forget factors and updating context states. This process enables the context states to capture contextual information from both latent features and hidden states.

Next, we compute qv, kv, ϑv. While preserving the original computation for qv, we strategically substitute Z(t) with the context states C(t+1) during the computation of kv and ϑv.In essence, departing from conventional self-attention, we empower each node to aggregate attention information from the context states. Finally, we leverage qv, kv, ϑv to update the hidden state for each node and edge.

5 Experiments

This section validates the empirical performance of our framework and aims to investigate the following questions:

• Can the framework simultaneously enhance the reasoning capabilities of GNN-based and Transformer-based architectures? If so, is there a noticeable discrepancy in the performance improvements between these two architectural types?
• Consider an architecture faced with a reasoning task in which it inherently performs sub-optimally. If our framework is applied, can we expect an enhancement in performance? In simpler terms, does the framework push the original architecture towards a more "extreme" or improved state?

To this end, we apply the framework to two state-of-the-art architectures: Triplet-GMPNN [Ibarz et al., 2022], a GNNbased model, and RT [Diao and Loynd, 2023], a Transformerbased model. We refer to their implementations as CEFGMPNN and CEF-RT, respectively. We remark that in the main body, we focus on the exploration of the aforementioned two questions; the ablation experiments regarding implementation details are deferred to the full version of the paper. The code is available at https://github.com/Ghost-st/CEF.

5.1 Setup

The experiments are conducted on a machine equipped with an i7-13700K CPU and an RTX 4090 GPU. The results are averaged over 5 runs with random seeds 5, 18, 25, 30, and 42.

Datasets. We use the CLRSAlgorithmic ReasoningBenchmark [Velickovic et al., 2022], a proven benchmark that offers a unified evaluation (micro-F1 score) for assessing the (seq-graph) reasoning capabilities of neural networks. It is derived from a foundational algorithms textbook "Introduction to Algorithms" [Cormen et al., 2009] and consists of 30 algorithmic reasoning tasks, covering a wide range of algorithmic procedures including sorting, searching, dynamic programming, graph algorithms, string algorithms, and geometric algorithms. The inclusion of such a variety of tasks enables us to draw meaningful conclusions about the effectiveness of the proposed framework.

Baselines. The two most crucial baselines we compare against are Triplet-GMPNN [Ibarz et al., 2022] and RT [Diao and Loynd, 2023], which play a pivotal role in addressing the aforementioned questions. Moreover, we include Memnet [Sukhbaatar et al., 2015], MPNN [Gilmer et al., 2017], PGN [Velickovic et al., 2020], and NPQ [Jain et al., 2023] in experiments as important references for comparison.

Training Details. In our implementations, flinear, flinear-1 and flinear-2 are individual linear layers. To ensure fair comparisons, we employ the experimental hyperparameter settings used in triplet-GMPNN [Ibarz et al., 2022] and Relational Transformer [Diao and Loynd, 2023] for CEFGMPNN and CEF-RT, respectively. Specifically, for CEFGMPNN, we set the batch size to 32 and the network is trained for 10,000 steps by Adam optimizer with a learning rate of 0.001; while for CEF-RT, we set the batch size to 4 and the network is trained for 10,000 steps by Adam optimizer with a learning rate of 0.00025. We remark that with the introduction of a preprocessor module in our framework, the average training time increases by 3 to 5 minutes for both methods. This increment represents less than 10% of the total runtime, indicating the efficiency of our implementations. More details are deferred to the full version.

5.2 Comparison with GNN-based Processors

In this experiment, we compare our method CEF-GMPNN with Triplet-GMPNN across 30 reasoning tasks within the CLRS dataset. The results are shown in Figure 2. Bar charts illustrating average scores, with standard deviations denoted by black lines, are presented for each task. Additionally, the tasks are arranged in descending order of improvement magnitude to better showcase trends.

From Figure 2, we see that our framework demonstrates enhancement across all 30 reasoning tasks, implying that our framework pushes the Triplet-GMPNN architecture towards an improved state. Notably, in tasks related to sorting and strings, such as bubble sort, quick select, Knuth-Morris-Pratt, and heap sort, the framework significantly improves the performance of Triplet-GMPNN.

5.3 Comparison with Transformer-based Processors

In this experiment, we compare our method CEF-RT with RT across 30 reasoning tasks within the CLRS dataset. The results are shown in Figure 3.

We observe a trend different from the previous experiment. The framework enhances the performance of RT across the majority of CLRS datasets. Specifically, in tasks related to graph optimization, such as strongly connected components (SCC) and depth-first search (DFS), there is a noteworthy increase in test scores of over 30%. However, for sorting tasks where the original RT approach exhibits poor average scores with very large derivations, our framework leads to a decrease in performance. This phenomenon indicates that for Transformer-based Processors, our framework steers them toward a more "extreme" state.

5.4 Summarized Results and Analysis

We summarize the experiential results in Table 1. Due to space limitations, we use column "Prior Best" to represent the best results among Memnet [Velickovic et al., 2022], PGN [Velickovic et al., 2022], MPNN [Velickovic et al., 2022], and NPQ [Jain et al., 2023]. Note that [Jain et al., 2023] proposed four variants of the NPQ approach, and for each task, we have chosen the best scores among these variants as reference items.

In the table, we present the specific test scores for each task, along with summarizing the average scores. Additionally, in line with prior work [Ibarz et al., 2022], we provide the respective proportions of tasks achieving scores greater than 90%, 80%, and 60%. We see the following trends:

• Our framework yields state-of-the-art results across the majority of the CLRS benchmark, demonstrating higher average scores and lower deviations compared to prior work. The overall average score of CEF-GMPNN is

82.68%, which is a substantial improvement over previous results.

• The results show a broader range of tasks where we attain scores exceeding 60%, 80%, and 90%. Specifically, CEF-GMPNN yields a score of at least 90% in 18 tasks, 80% for 22 tasks, and 60% for 20 tasks.
• For both Triplet-GMPNN and RT, our framework enhances their performance in most tasks. However, the specific details of the enhancements differ. For TripletGMPNN, the improvements brought by the framework are comprehensive, showing enhancement across all tasks. In contrast, for RT, the framework exhibits more significant improvements in certain tasks compared to Triplet-GMPNN, but it also shows performance degradation in tasks where RT initially performs poorly. These results indicate that our framework pushes Triplet-GMPNN towards an improved state while making RT more "extreme" in its performance.

6 Conclusion

The paper studies seq-graph reasoning. We introduce a context-enhanced framework for such tasks and show that it can effectively integrate with various existing approaches. Our experiments reveal that the proposed framework consistently enhances the reasoning capabilities of both GNN-based and Transformer-based architectures, achieving new state-ofthe-art results on the CLRS benchmark.

There are several intriguing avenues for future research. One interesting direction involves refining our current framework. As discussed in Section 1, the outcome sequence may form an entire problem-solving or optimization process, potentially represented by an additional reasoning graph. Therefore, it is worthwhile to explore a framework capable of efficiently capturing such graph structures hidden in the data. Another promising direction could be concentrating on the CLRS benchmark, aiming to design novel architectures that could yield further performance improvements.

Acknowledgements

This work is supported by the National Key Research Project of China under Grant No. 2023YFA1009402, the National Natural Science Foundation of China (No. 62302166), the Scientific and Technological Innovation 2030 Major Projects under Grant 2018AAA0100902, the Shanghai Science and Technology Commission under Grant No.20511100200, the Dean's Fund of Shanghai Key Laboratory of Trustworthy Computing, ECNU, and the Key Laboratory of Interdisciplinary Research of Computation and Economics (SUFE), Ministry of Education.