cs.AIAug 27, 2026
SaveGRAIN: Bridging Name and Narrative Shifts in Real-World Graph Reasoning through Invariance-Rewarded Agentic RL
Organizations: Harbin Institute of Technology, Shenzhen, China · Peng Cheng Laboratory, Shenzhen, China · Xidian University, Xi’an, China
Abstract
Despite their potential in standardized graph tasks, Large Language Models (LLMs) remain brittle to real-world shifts in node identifiers and task formulation. While deterministic graph tools are invariant to such shifts, extracting topological structures from noisy text is highly fragile for LLMs, which often overfit to surface patterns. Moreover, mitigating these parsing failures via multi-agent systems incurs prohibitive latency. To address this, we propose GRAIN, a single-agent framework optimized via reinforcement learning. GRAIN models reasoning as a semantic parsing and tool-execution pipeline, guided by a Structure Invariance Reward. By validating extracted intermediate graphs against ground-truth topologies, this reward forces the LLM to learn robust text-to-structure mappings rather than memorizing linguistic artifacts. We also introduce GRIT, a benchmark evaluating sensitivity to such linguistic shifts. GRAIN outperforms multi-agent baselines by 16.45% in accuracy with approximately 24% lower latency. Furthermore, it demonstrates superior structural generalization, halving the out-of-distribution (OOD) gap of SFT models (from 15.77% to 7.80%) and maintaining robustness on large-scale graphs beyond the training distribution.
Figures & tables
Figure 1: GRAIN mitigates sensitivity to naming and formulation shifts in real-world graph reasoning via a single-agent tool-use pipeline trained with isomorphism-aware structural rewards .
Figure 2: Node-label Sensitivity. Large error bars reveal high instability: open-source models degrade on Random IDs, while closed-source models struggle with canonical forms.
Figure 3: Overview of the GRAIN framework. (Top) Inference: The LLM builds a graph IR in <graph> , calls solvers, and outputs the answer. (Bottom) Training: ARPO trains under isomorphism-preserving name perturbations, with partial rollouts and a reward for structural invariance, correctness, and valid format.
| Split | Node Range | # Tasks | # Scen. | # Graphs | Total Qs. |
| Train | 4–40 | 6 | 31 | 2160 | 17,280 |
| Test | 10–40 | 6 | 31 | 360 | 2,760 |
| Large Test | 40–60 | 6 | 31 | 120 | 960 |
| OOD Test | 10–40 | 6 | 24 * | 180 | 1,440 |
| * Denotes distinct scenarios unseen during training. | |||||
Table 1: Statistics of the GRIT benchmark splits.
| Model Family | Method / Setting | S. Path | Coloring | TSP | V. Cover | BFS | Centr. | Avg. |
| Proprietary Models & Multi-Agent Systems | ||||||||
| GPT-5-nano | Zero-shot CoT | 72.18 | 43.54 | 3.33 | 29.17 | 16.25 | 15.21 | 29.95 |
| Tool-use CoT | 73.39 | 68.96 | 41.88 | 28.61 | 31.46 | 40.42 | 47.45 | |
| G1-3B | Zero-shot CoT | 11.46 | 8.96 | 0.00 | 5.28 | 1.04 | 8.33 | 5.85 |
| MA-GTS(GPT-4o-mini) | Multi-Agent | 85.63 | 87.50 | 91.25 | 68.33 | 55.62 | 87.50 | 79.31 |
| MA-GTS(Qwen3-4B-Ins.) † | Multi-Agent | 13.75 | 12.50 | 1.88 | 4.17 | 0.62 | 7.50 | 6.74 |
Table 2: Main Results on GRIT Tasks. We report the accuracy (%) across six graph reasoning tasks. The best results are bolded , and the second-best results are underlined . “Ins” denotes Instruct models. For Qwen-3-Base, we compare CoT and SFT under both text-only and tool-use settings. † denotes zero-shot transfer of the six-agent MA-GTS pipeline to Qwen3-4B-Instruct on the 920-example Test set; its micro-average is 6.85%. GRAIN consistently achieves SOTA performance.
Figure 4: Robustness Evaluation. GRAIN (Red) achieves consistent robustness across shifts, unlike SFT-Tool (Blue) which degrades significantly in noisy, real-world scenarios.
| Method | GRIT Distribution | Unseen Tasks | |||
| Test | OOD | Gap | G-Ins | Arena | |
| Baseline: Qwen3-4B-SFT | |||||
| Zero-shot CoT | 37.11 | 27.29 | 9.82 | 32.83 | 13.0 |
| Tool-use CoT | 89.10 | 73.33 | 15.77 | 92.06 | 88.0 |
| GRAIN-4B (Ours) | |||||
| SFT + ARPO | 95.76 | 87.96 | 7.80 | 97.05 | 91.0 |
Table 3: OOD Robustness & Generalization. We report the accuracy drop ( ) on GRIT and zero-shot performance on unseen benchmarks (GraphInstruct, GraphArena).
| Method | Rounds | Tokens | Latency (s) | Acc. | |
| (Avg) | (Avg) | Avg | p90 | (%) | |
| GPT-5-nano (Single-Agent) | |||||
| Zero-shot CoT | 1.00 | 17.1k | 117.3 | 200.2 | 29.95 |
| Tool-use CoT | 2.00 | 22.1k | 126.1 | 212.8 | 47.45 |
| Multi-Agent Framework | |||||
| MA-GTS (4o-mini) | 6.00 | 13.5k | 58.4 | 96.1 | 79.31 |
Table 4: GRAIN achieves peak accuracy (95.76%) with the lowest overhead (3.3k tokens, 44.4s), significantly outperforming proprietary and multi-agent baselines.
Figure 5: Performance as node count ( ) exceeds the training range ( ). GRAIN demonstrates robust length generalization, sustaining accuracy on larger graphs where baselines collapse.
| Method / Variant | Metrics | ||
| Gap | Var. ( ) | Acc. | |
| GRAIN (Full) | 1.14% | 4.64 | 96.76% |
| Component Analysis | |||
| w/o Struct. Reward | 11.14% | 14.65 | 84.73% |
| w/o Diverse Naming | 4.11% | 341.83 | 79.21% |
| Optimization Method | |||
Table 5: Ablation Study. Gap denotes the performance drop ( ), and Var. ( ) measures stability across naming permutations.
| Method / Variant | Node-F1 | Edge-F1 | Exact |
| w/o Struct. Reward | 96.03 | 94.95 | 63.10 |
| SFT-Tool | 98.45 | 97.96 | 90.21 |
| replace w/ GRPO | 97.23 | 97.20 | 95.22 |
| GRAIN (Full) | 98.76 | 98.67 | 97.28 |
Table 6: Intermediate Structure Recovery. Results on 2,760 GRIT examples. Edge-F1 denotes weighted edge recovery, while Exact requires the entire predicted weighted graph to match the ground-truth graph.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
| Parameter | Value |
| General Settings | |
| Base Model | Qwen-3-4B |
| Tuning Method | Full Fine-tuning |
| Precision | BF16 |
| Context Window | 15,000 tokens |
| DeepSpeed Stage | ZeRO-3 (Offload) |
Table 7: Implementation Details. Hyperparameters and settings used for full-parameter fine-tuning.
| Graph Task | Canonical Graph-Theoretic Formulation | Representative Real-World Problem Formulations (Noun Phrases) |
| Centrality | Compute centrality scores of one or several nodes in a graph and identify structurally “key” nodes. | City traffic hubs; emergency response centers; social network key persons; information relay stations |
| Shortest Path | In a weighted directed or undirected graph, find the minimum-cost path from a source node to a target node. | Data-center latency routing; logistics cost routes; robot energy-efficient paths; network routing paths |
| TSP | In a weighted complete graph, find a minimum-cost tour that visits every node once and returns to the start (decision version NP-complete, optimization version NP-hard). | Parcel delivery tour; sales representative tour; food delivery tour; garbage collection tour; maintenance inspection tour; school bus route |
| Min Graph Coloring | Color all nodes of a graph using as few colors as possible so that adjacent nodes have different colors. | Wi-Fi channel assignment; cellular frequency planning; radio channel allocation; event loudspeaker placement; exam timetabling; map region coloring; interference-aware equipment layout |
| Graph Traversal (BFS) | Perform breadth-first search from a given source node and output the visit order or layer structure. | Post-disaster inspection sweep; tourist exploration route; UAV area scanning; building safety inspection |
| Vertex Cover | Select a minimum set of nodes such that every edge has at least one endpoint in the set. | Server monitoring placement; road intersection cameras; railway security posts; social liaison selection; pipeline sensor placement; campus patrol posts |
Table 8: Detailed descriptions of the six graph reasoning tasks included in the GRIT benchmark. This table complements the main results by providing the canonical graph-theoretic formulation and representative real-world problem scenarios for each task.
| Parameter | Value |
| Optimization & Data | |
| Actor LR | |
| Train Batch Size | 16 |
| Mini-batch Size | 4 |
| Epochs | 2 |
| KL Coefficient | 0.0 |
Table 9: Hyperparameters for RL Training (ARPO). We employ the GRPO estimator with vLLM-based rollouts. Note the extended context window (36k total) to handle graph reasoning trajectories.
| Task | Model | Node Naming Schemes (Accuracy %) | Var ( ) | |||
| Canonical | Random ID | Semantic | Mixed | |||
| Closed Source Models | ||||||
| Synthetic Graphs | gpt-5-nano | 83.33 | 66.67 | 75.00 | 70.00 | 52.50 |
| deepseek-3.2-reason | 96.61 | 91.53 | 93.22 | 96.67 | 6.54 | |
| Qwen-plus | 80.39 | 82.00 | 82.35 | 84.31 | 2.59 | |
| gemini-2.5-flash-lite | 23.33 | 13.33 | 51.67 | 40.00 | 292.00 | |
Table 10: Performance comparison across different node naming schemes on Synthetic Graphs and Real-world Queries. Accuracy is reported in percentages, and variance (Var) is scaled by .
| Component | Content / Instruction |
| Role Definition | You are a helpful assistant that specializes in solving graph theory problems using a graph theory tool . Given a question, you need to first think about the reasoning process in your mind and then provide the answer. During thinking, you can invoke the graph theory tool to compute shortest paths, minimum spanning trees, topological sorting, or other graph algorithms if needed. |
| Protocol & Tags | The reasoning process and answer must be enclosed within specific XML tags. Follow these guidelines strictly: 1. Thinking: Start with <think> and explain your step-by-step reasoning process. End with </think> . 2. Tool Call: When you need to use the graph theory tool, output <graph> followed by a JSON-formatted query, then </graph> . 3. Tool Result: After the tool call, output <result> followed by the tool’s output, then </result> . 4. Final Answer: Finally, output <answer> followed by the final answer, then the final answer is \boxed{answer here}</answer> . |
| Tool Schema | The JSON query inside <graph> tags must follow this exact structure: ⬇ { "problem_type" : "problem_type_here", "graph" : { "type" : "adjacency_list", "data" : { "node1" : { "neighbor1" : weight1, "neighbor2" : weight2 }, "node2" : { "neighbor3" : weight3, "neighbor4" : weight4 } }, "directed" : true_or_false, "weighted" : true_or_false }, "parameters" : { "parameter1" : "parameter1_value", "parameter2" : "parameter2_value", "algorithm" : "algorithm_name" } } Ensure that problem_type , graph structure , and algorithm parameters are correctly populated based on the user query. |
Table 12: System Prompt for GRAIN (SFT & RL). This prompt instructs the model to act as a graph reasoning agent, defining the explicit "Think-Tool-Result" loop and the strict JSON schema required for interacting with the deterministic graph solver.
| Stage | Content |
| User Query | My younger sibling isn’t feeling well, and the nearest hospital is at Zaphrel Summit . Right now we’re at home, near an intersection I’ll call Xundral View . […] I checked a city map and noted down the walking distances: Knights Market to Xundral View is 7 km… [narrative omitted for brevity] … Dragons Gate to Trelvon Bay is 7 km. Starting from Xundral View, what route should we walk to reach Zaphrel Summit while minimizing the total walking distance? |
| Reasoning | <think> From an application perspective, this requires graph theory tool calculation. Graph parameters identified: • Nodes: 9 distinct locations (e.g., Knights Market, Xundral View). • Goal: Compute shortest path from Xundral View to Zaphrel Summit . • Edge Extraction: Parsing narrative to build adjacency list. Tool call explanation: I will input graph information in JSON format, including nodes, edges, start and end points. </think> |
| Action | <graph> ⬇ { "problem_type" : "shortest_path", "graph" : { "type" : "adjacency_list", "data" : { "Knights Market" : { "Xundral View" : 7, "Skyline Gardens" : 1, "Dragons Gate" : 3, "Trelvon Bay" : 6 }, "Xundral View" : { "Trelvon Bay": 1 }, "Draycon Arcade" : { "Copper Clock Square" : 2, "Vorlen Heights" : 8 }, "Copper Clock Square" : { "Trelvon Bay" : 7, "Vorlen Heights" : 5, "Zaphrel Summit" : 8 }, "Granite Plaza" : { "Trelvon Bay": 8 }, "Dragons Gate" : { "Trelvon Bay": 7 }, "Trelvon Bay" : {}, "Zaphrel Summit" : {} }, "directed" : false, "weighted" : true }, "parameters" : { "start" : "Xundral View", "end" : "Zaphrel Summit", "algorithm" : "dijkstra" } } </graph> |
| Observation | <result> Shortest path algorithm (Dijkstra) result: • Path: Xundral View Trelvon Bay Copper Clock Square Zaphrel Summit • Total Length: 16 </result> |
| Final Response | <think> Tool call successful. Calculation result: shortest path contains 3 nodes… total weight 16. </think> Answer: 16 |
Table 14: Case Study on Shortest Path Task. The model parses a noisy narrative, constructs a structured JSON representation, and invokes the tool to find the optimal path. The code block demonstrates the explicit structural constraints imposed by GRAIN.
| Stage | Content |
| User Query | My younger sibling isn’t feeling well, and the nearest hospital is at Zaphrel Summit . Right now we’re at home, near an intersection I’ll call Xundral View . […] (Narrative containing distances: Knights Market to Xundral View is 7 km, etc.) … Starting from Xundral View, what route should we walk to reach Zaphrel Summit while minimizing the total walking distance? |
| Model Response | <think> Task restatement: Starting from Xundral View (XV), find a walking route to Zaphrel Summit (ZS) minimizing total distance… The graph is undirected and unweighted in terms of edges with weights equal to distances. 1. Internal Graph Construction: • Vertex set V = {XV, KM, SG, DA, GP, DG, TB, CCS, VH, ZS}. • Edge set extraction: – (XV, TB) = 1, (KM, XV) = 7, (KM, SG) = 1, (KM, DA) = 2 – (TB, CCS) = 7, (TB, GP) = 8, (CCS, ZS) = 8 … [full list extracted] 2. Algorithm Simulation (Dijkstra): Method: Run Dijkstra’s algorithm from XV, tracking tentative distances. • Initialize: dist(XV) = 0; all others = . • Step 1: From XV, relax TB with dist . dist(TB) = 1 . • Step 2 (Next smallest TB): – Relax KM: dist(KM) = 7 – Relax CCS: dist(CCS) = 8 – … [relaxing other neighbors] • Step 3 (Next smallest KM): Relax SG via KM: dist(SG) = 8… • […omitting intermediate relaxation steps for brevity…] • Final Step: ZS reached with dist(ZS) = 16. No other path reduces this. Conclusion: The minimal walking distance is 16. </think> Answer: 16 |
Table 15: Baseline Case Study: GPT-5-nano (Chain-of-Thought). Unlike GRAIN which offloads computation to external tools, the baseline model attempts to internally simulate the execution of Dijkstra’s algorithm via text generation. While successful in this instance, this approach relies heavily on the model’s ability to maintain long-context state and perform precise arithmetic without external verification, which typically degrades on more complex graphs.
| Model / Framework | Evaluation Setup | Accuracy (%) |
| GRAIN-4B (Ours) | Zero-Shot on G-REAL (All sizes) | 86.60 |
| MA-GTS | Evaluated on GRIT OOD ( ) | 82.79 |
Table 17: Cross-Dataset Generalization Test. GRAIN demonstrates powerful zero-shot extraction capabilities on the completely unknown G-REAL dataset, outperforming the multi-agent baseline evaluated on a simplified subset of our benchmark.
| Structural Outcome | # Examples | Percent |
| Exact graph recovered | 2,685 | 97.28 |
| Graph parse / format failure | 33 | 1.20 |
| Node-set mismatch | 23 | 0.83 |
| Extra edges only | 9 | 0.33 |
| Missing edges only | 7 | 0.25 |
| Edge-weight mismatch only | 2 | 0.07 |
Table 19: Intermediate Structure Error Decomposition. Counts and percentages over all 2,760 GRIT Test examples.