Source-linked AI summary
GRAIN: Bridging Name and Narrative Shifts in Real-World Graph Reasoning through Invariance-Rewarded Agentic RL
Zike Yuan, Han Zhang, Jianzhi Yan, Le Liu, Cai Ke, Huozhi Zhou, Jian Xie, Jiran Yin, Yukun Cao, Yue Yu, Hui Wang, Ming Liu, Bing Qin
TL;DR
LLMs are fragile to identifier and formulation shifts when extracting graph structure from real-world text. GRAIN addresses this with a single-agent semantic-parsing and tool-execution pipeline trained using structural invariance rewards, achieving robust in-distribution and OOD performance while reducing the need for costly multi-agent coordination.
Problem
LLMs remain brittle to naming and task-form shifts because graph structure extraction from noisy text can overfit surface patterns.
Method
GRAIN uses single-agent reinforcement learning with explicit semantic parsing, tool execution, isomorphic narratives, and a reward validating intermediate graphs against ground-truth topology.
Results
GRAIN achieves state-of-the-art accuracy and efficiency, outperforming complex multi-agent systems while generalizing robustly to OOD scenarios and larger graphs.
Takeaways & Limitations
Invariant structural grounding supports robust graph reasoning across surface variations and beyond the training distribution.
Takeaways & Limitations
Embedding comprehensive graph-tool definitions in the system prompt consumes context and may distract the model from core reasoning.
Abstract
from arXiv · showhide
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.
1 Introduction
GRAIN addresses LLM brittleness to node-label and task-form shifts by combining explicit graph parsing with single-agent reinforcement learning and structural invariance rewards. The paper also introduces GRIT for controlled robustness evaluation.
- Motivation: LLMs remain brittle when naming schemes or task formulations change despite equivalent underlying graph structures.Open-source models degrade under non-canonical identifiers, while closed-source models show variance across naming schemes.
- Motivation: SFT can overfit surface patterns, limiting structural generalization under isomorphic and scale shifts.
- GRAIN: GRAIN decomposes reasoning into semantic parsing and algorithmic execution, imposing verifiable constraints on intermediate graph representations.
- GRAIN: GRAIN trains on isomorphic narratives with an invariance-oriented reward that validates recovered intermediate graphs against ground-truth topologies.
- GRIT: GRIT covers six graph problems across 31 real-world prototypes and supports controlled evaluation of naming and formulation shifts.
2 Related Work
Prior graph-reasoning approaches reveal sensitivity to serialization and scale, while tool-augmented multi-agent systems improve reliability at substantial latency and token costs.
- LLMs for Graph Reasoning: LLM graph-reasoning performance declines as graph size and task complexity increase, especially for NP-hard problems.
- LLMs for Graph Reasoning: Existing methods are optimized for standardized inputs and remain brittle to changes in graph serialization formats and prompting schemes.
- Serialization Sensitivity: Different graph encodings and edge orders can drastically alter reasoning accuracy, leaving naming robustness unresolved.
- Agents and Tool Learning: Multi-agent graph frameworks improve performance through coordinated roles but incur high latency and token costs from inter-agent communication.
3 GRIT: A Benchmark for Robust Graph Reasoning
GRIT evaluates graph reasoning robustness by varying identifiers and narrative formulations while preserving topology. Its controlled multi-view design supports diagnosis across tasks, scenarios, and graph scales.
- Benchmark Design: GRIT evaluates identifier shifts and task-form shifts using multi-view instances that preserve the underlying graph structures.
- Benchmark Design: The benchmark covers six graph tasks across 31 scenarios, including logistics and social-network narratives with syntactic embeddings and semantic distractors.
- Multi-view Construction: Each base instance yields 8 queries combining Standard or Realistic formulations with four naming schemes.The naming schemes are Canonical, Random IDs, Semantic, and Noisy/Mixed.
- Generation: A deterministic slot-filling engine injects topologies into scenarios, guaranteeing isomorphism between text and the underlying graph.
- Benchmark Splits: GRIT separates in-domain, length-generalization, and OOD evaluation through graph-size and held-out-scenario splits.Large Test uses N ∈ [40, 60], while OOD Test includes unseen scenarios and identifier realizations.
4 GRAIN
GRAIN is a single-agent reinforcement-learning framework that addresses identifier and task-formulation shifts by combining structured graph reasoning with tool execution. Its training rewards valid outputs, task correctness, and recovery of topology invariant to surface changes.
- Framework: GRAIN models the LLM as a language–graph–tool agent that parses noisy queries, selects algorithms, constructs graph structures, and executes tools.The tool library provides computational correctness while learning focuses on entity resolution, graph reconstruction, and tool or parameter selection.
- Framework: The pipeline separates planning, graph construction, deterministic tool execution, and continued reasoning or final answering through tagged segments.The graph segment specifies nodes, edges, weights, and task parameters; the runtime parses it for the graph-tool library and injects the result.
- Invariant training: GRAIN trains on isomorphic views that vary naming schemes and query styles while preserving each underlying graph and task structure.This environment generates Canonical, Random, Semantic, and Mixed naming schemes alongside diverse formulations.
- Invariant training: The gated reward combines task-answer correctness with a Jaccard-based structure similarity score while hard-penalizing malformed trajectories.For valid outputs, task correctness and structure similarity are equally weighted with λ_inv = 1.0.
- Optimization: SFT warm-starts the policy on automatically generated small-graph trajectories before RLVR further optimizes the gated objective with ARPO.The warm start teaches tag usage, tool syntax, and basic construction, while ARPO uses entropy-triggered partial rollouts and branch-aware advantages.
- Evaluation: Across six GRIT graph-reasoning tasks, GRAIN consistently achieves state-of-the-art accuracy across varying model scales.Table 2 reports accuracy percentages, with best results bolded and second-best results underlined.
5 Evaluation
GRAIN is evaluated across six graph tasks, robustness shifts, OOD settings, larger graphs, ablations, structure recovery, and deployment efficiency. It consistently outperforms supervised and multi-agent baselines while improving generalization and reducing inference overhead.
- Overall Performance: GRAIN achieves consistent state-of-the-art performance across six graph reasoning tasks and varying model scales.Evaluation uses GRIT, external benchmarks, and unified strict-match accuracy with macro-averages.
- Overall Performance: 95.76% accuracy on Qwen-3-Base surpasses SFT’s 89.10%, with gains of +4.97% on Vertex Cover and +15.83% on BFS.These results support the effectiveness of the invariance-oriented structural reward over standard supervised imitation learning.
- OOD Generalization and Scalability: GRAIN reduces the OOD accuracy gap from 15.77% for SFT-Tool to 7.80% and transfers zero-shot to GraphInstruct at 97.05% and GraphArena at 91.0%.The evaluation covers unseen naming–narrative combinations and larger graphs beyond the N ≤40 training range.
- OOD Generalization and Scalability: GRAIN maintains accuracy beyond |V| > 40, whereas baselines degrade rapidly as larger graphs create context overload.Its explicit pipeline limits the LLM to semantic parsing and information extraction while external algorithms handle computation.
- Ablation Study: Removing the Structure Invariance Reward increases the generalization gap from 1.14% to 11.14% and lowers accuracy to 84.73%.Removing Diverse Naming also raises variance from 4.64 to 341.83 and lowers accuracy to 79.21%; ARPO improves accuracy from 96.51% to 96.76%.
- Ablation Study: GRAIN recovers 98.67% weighted Edge-F1 and 97.28% of graphs exactly, while removing the invariance reward reduces exact recovery to 63.10%.Replacing ARPO with GRPO reduces exact recovery from 97.28% to 95.22%.
- Efficiency Analysis: GRAIN uses approximately 3.3k tokens and 44.4 seconds, cutting token usage by approximately 85% relative to GPT-5nano’s more than 17k tokens.This is reported as the lowest latency and a favorable accuracy–efficiency trade-off.
- Overall Performance: 16.45% higher accuracy and 24% lower latency distinguish GRAIN from the MA-GTS multi-agent baseline.The comparison is not an isolated single-agent versus multi-agent effect because training, initialization, and organization differ.
6 Conclusion
GRAIN addresses fragility in graph reasoning under identifier and task shifts with a single-agent reinforcement learning framework. On GRIT, it achieves strong accuracy and efficiency while generalizing to OOD scenarios and larger graphs.
- GRAIN uses a single-agent reinforcement learning framework with a Structure Invariance Reward to decouple graph topology from surface variations.The framework targets identifier and task shifts in graph reasoning.
- GRAIN achieves state-of-the-art accuracy and efficiency on the GRIT benchmark, outperforming complex multi-agent systems.
- GRAIN generalizes robustly to out-of-distribution scenarios and larger graphs beyond the training distribution.The result supports invariant structural grounding as the basis for this generalization.
Limitations
The framework is limited by how it manages graph-tool contexts. Static inclusion of comprehensive tool definitions lengthens prompts and may distract the model from core reasoning.
- Static inclusion of comprehensive graph-tool definitions consumes context-window space and may distract GRAIN from core reasoning.The authors propose dynamically managing tool definitions in future work, including through standards such as MCP.
A GRIT Benchmark Implementation Details
GRIT combines controlled graph generation, narrative injection, OOD split construction, and deterministic evaluation to study robust graph reasoning. The implementation also details GRAIN’s training pipeline and compares its parsing-and-tool-execution behavior with CoT and multi-agent approaches.
- Graph generation: Underlying graphs use Erdős-Rényi generation, bounded integer weights for weighted tasks, and rejection sampling for connectivity requirements.Undirected graphs are constrained to be connected, while TSP graphs are constrained to be strongly connected.
- Ground truth and tie-breaking: Deterministic evaluation fixes BFS neighbor order, computes exact Metric TSP solutions by dynamic programming, and obtains graph-coloring chromatic numbers with backtracking.
- Narrative injection: The narrative engine injects graph topologies into 31 domain templates under four naming schemes, including high-entropy random IDs and aliases in noisy or mixed settings.
- OOD evaluation: The OOD test set holds out 24 scenario templates and naming distributions, separating structural transfer from domain-text memorization.
- Training pipeline: GRAIN’s two-stage pipeline combines SFT for tool use and JSON graph representations with agentic reinforcement learning via ARPO.Training uses full-parameter fine-tuning and initializes ARPO from the tool-use SFT checkpoint.
- GRIT benchmark: GRIT covers six graph reasoning tasks across 31 real-world prototypes, pairing standard formulations with diverse narratives and namings for controlled OOD evaluation.
- Training pipeline: A two-phase curriculum stabilizes structured output and tool invocation on smaller graphs before optimizing across graphs ranging from 10 to 40 nodes.
- Qualitative comparison: In the qualitative case study, GRAIN parses noisy narratives into JSON adjacency lists and invokes Dijkstra rather than calculating paths internally.This separates semantic parsing from algorithmic computation.
D Empirical Study: Node-Naming Robustness on Synthetic and Real-World Graphs
The empirical study finds that LLM graph reasoning is highly sensitive to node naming and narrative surface form. Open-source models particularly degrade with random identifiers, while closed-source models can favor semantic names over canonical labels.
- Surface-level changes in node identifiers or narrative framing induce drastic performance fluctuations across models.Large error bars indicate substantial instability in the Shortest Path evaluation.
- Open-source models show a sharp performance drop on Random IDs relative to Semantic or Canonical names.The pattern suggests reliance on sequential text patterns or meaningful words rather than topology.
- Closed-source models perform slightly worse on Canonical names than Semantic names for real-world graphs.This pattern is consistent with overfitting to rich contextual wording.
- Qualitative analysis locates these failures primarily in the grounding phase that maps text to graph representations.This finding motivates GRAIN’s emphasis on invariant decision-making.
E Zero-Shot Evaluation on the G-REAL Dataset
GRAIN is evaluated zero-shot on the unseen G-REAL dataset alongside a multi-agent baseline on a simplified GRIT subset. The results support strong cross-dataset transfer for GRAIN and a fairer comparison across independent evaluation dimensions.
- Evaluation Purpose: The zero-shot evaluation is designed to test whether GRAIN’s performance extends beyond the GRIT training distribution and to verify baseline-comparison fairness.G-REAL is the original dataset introduced with the MA-GTS framework.
- Evaluation Setup: The evaluation samples 100 uniformly distributed problems per task across varied node sizes from G-REAL.The study also evaluates MA-GTS on a low-difficulty GRIT OOD subset with 10 ≤|V|≤20 because multi-agent runtime grows exponentially on large graphs.
- Cross-Dataset Results: GRAIN achieves 86.6% accuracy on the entirely unseen G-REAL dataset.Task-level accuracies include 98% for Graph Coloring, 80% for TSP, and 82% for Vertex Cover.
- Baseline Comparison: MA-GTS achieves 82.79% on the GRIT OOD small-graph subset.This result provides a comparison point on a separately selected subset rather than the unseen G-REAL evaluation.
F Comparative Failure Analysis: Single-Agent vs. Multi-Agent
The comparison examines how single-agent and multi-agent systems handle noisy graph narratives, explicit structure extraction, and tool execution. GRAIN uses an explicit parsing-and-tool pipeline, while multi-agent decomposition introduces substantial interaction overhead and can remain unreliable on compact-model transfer.
- Single-Agent Pipeline: GRAIN decomposes reasoning into semantic parsing and algorithmic execution instead of relying on freeform Chain-of-Thought generation.The explicit intermediate representation imposes verifiable structural constraints before computation.
- Case Study: In a shortest-path case, GRAIN constructs a structured graph, runs Dijkstra’s algorithm, and obtains a minimal walking distance of 16.The example uses a noisy narrative with renamed locations and an explicit graph representation before tool-style algorithm simulation.
- Efficiency Trade-off: Multi-agent graph reasoning decomposes tasks into coordinated roles, but multi-round interaction can create efficiency bottlenecks.A five-agent MA-GTS case study records 23.5s latency and 7487 tokens, whereas GRAIN is described as a single-agent pipeline.
- OOD Transfer: 81.25% versus 11.35% overall accuracy: GRAIN substantially outperforms MA-GTS on the same 960 OOD inputs with Qwen3-4B-Instruct.The result is a transfer diagnostic rather than an isolated single-agent versus multi-agent comparison because training, backbone initialization, and agent organization differ.
- Failure Analysis: Inference-only multi-agent decomposition does not automatically repair entity grounding, graph-state preservation, or native tool-call reliability in the compact-model setting.The diagnostic identifies end-to-end reliability as the main weakness, rather than complete inability to solve graph tasks.
G Intermediate Structure Error Decomposition
The error decomposition evaluates exact intermediate graph recovery on the 2,760-example GRIT Test set and relates recovery to graph complexity. Recovery declines with node count but remains comparatively stable across density groups, within a restricted connected-undirected setting.
- Evaluation Scope: Exact graph recovery is evaluated over all 2,760 GRIT Test examples.Table 19 decomposes exact recovery outcomes, while the test set contains connected undirected graphs.
- Exact Recovery Criterion: Exact recovery counts any parse or format failure, node-set mismatch, extra or missing edge, or edge-weight mismatch as non-exact.This stricter criterion differs from edge-set overlap and exposes errors that high Jaccard or Edge-F1 scores can mask.
- Node-Count Effects: 99.63% versus 93.72%: exact graph recovery decreases from graphs with 10–15 nodes to graphs with 31–40 nodes.Within GRIT, the reported degradation is associated more strongly with graph size than with density.
- Density and Scope: Exact graph recovery remains at least 95.07% across density tertiles.The dense Erdős–Rényi construction produces small diameters, so the findings do not cover the full range of graph topological complexity.