A Context-Enhanced Framework for Sequential Graph Reasoning

TL;DR

This work tackles sequential graph reasoning by addressing the strong internal coupling between step outcomes. It introduces the Context-Enhanced Framework (CEF), which inserts a Preprocessor that maintains per-node and per-edge context states and uses context-enhancement and context-update functions to infuse historical information into the processor. The framework is designed to work with both GNN-based and Transformer-based processors and is validated on the CLRS Algorithmic Reasoning Benchmark, where it delivers state-of-the-art results and broad performance gains. The findings indicate that leveraging historical context improves reasoning across diverse tasks, with architecture-dependent dynamics in how the enhancements manifest, and suggest future directions toward modeling entire solution graphs and deeper CLRS-centric optimizations.

Abstract

The paper studies sequential reasoning over graph-structured data, which stands as a fundamental task in various trending fields like automated math problem solving and neural graph algorithm learning, attracting a lot of research interest. Simultaneously managing both sequential and graph-structured information in such tasks presents a notable challenge. Over recent years, many neural architectures in the literature have emerged to tackle the issue. In this work, we generalize the existing architectures and propose a context-enhanced framework. The crucial innovation is that the reasoning of each step does not only rely on the outcome of the preceding step but also leverages the aggregation of information from more historical outcomes. The idea stems from our observation that in sequential graph reasoning, each step's outcome has a much stronger inner connection with each other compared to traditional seq-to-seq tasks. We show that the framework can effectively integrate with the existing methods, enhancing their reasoning abilities. Empirical evaluations are conducted on the challenging CLRS Reasoning Benchmark, and the results demonstrate that the proposed framework significantly improves the performance of existing architectures, yielding state-of-the-art results across the majority of the datasets within the benchmark.

Figures (5)

• Figure 1: An illustration of the CEF framework. The solid and dashed arrows in the figure indicate the direction of data flow for the graph structure during the inference process. The blue arrows represent the historical information aggregation operation performed within the framework.
• Figure 2: Comparisons between our CEF-GMPNN and Triplet-GMPNN. The 30 tasks are arranged in descending order of improvement magnitude.
• Figure 3: Comparisons between our CEF-RT and RT. The 30 tasks are arranged in descending order of improvement magnitude. Note that on the Knuth-Morris-Pratt reasoning task, both algorithms exhibit very poor performance. Refer to \ref{['tab:results']} for detailed scores.
• Figure 4: Experimental results of training the HeapSort algorithm on the CEF-GMPNN model using different forget factors. The horizontal axis represents the forgetting rate of storing the $s_v$ context, while the vertical axis represents the forgetting rate of storing the $h_v$ context.
• Figure 5: Experimental results of training the Finding Maximum Subarray algorithm on the CEF-RT model using different forget factors. The horizontal axis represents the forgetting rate of storing the context about node features, while the vertical axis represents the forgetting rate of storing the context about edge features.