Autonomous Knowledge Graph Exploration with Adaptive Breadth-Depth Retrieval

AI summary

Overview

Research area: Retrieval-augmented generation and knowledge graph (KG) retrieval using large language model (LLM) agents.

Technical level: Advanced. The paper assumes familiarity with knowledge graphs, retrieval metrics, BM25, embedding retrievers, tool-use agents, and model distillation (though the core idea is intuitive).

Scope: The paper introduces ARK (Adaptive Retriever of Knowledge), a training-free, tool-using retriever that lets an LLM switch between broad lexical search and multi-hop neighborhood traversal over text-rich knowledge graphs, and shows that its behavior can be distilled into a small model.

What This Paper Is About

Retrieving evidence from a knowledge graph requires two different search behaviors: broad search when a query mentions many loosely related entities, and deep multi-hop traversal when the answer only appears after following relational links. Existing methods tend to do one or the other: similarity-based retrievers are broad but shallow, while traversal-based agents depend on picking the right seed entities and can get stuck exploring the wrong neighborhood. ARK's goal is to let a language model decide, step by step, how much breadth versus depth each query needs, without any task-specific training or preset hop depth.

Key Contributions

1. A training-free retrieval framework. ARK equips a language model with a minimal, graph-native tool interface consisting of exactly two operations: global lexical search over node descriptors and neighborhood exploration of a node's one-hop neighbors. Multi-hop traversal emerges from composing these two operations, and no task- or graph-specific training is required.

2. Adaptive breadth-depth control without fragile seeding. Rather than selecting seed entities in advance or fixing a maximum hop depth, the agent alternates between global anchoring and targeted expansion based on the query and what it has already retrieved. Global search remains available throughout the trajectory, which the authors argue avoids the anchoring errors of purely traversal-based agents.

3. Strong results on STaRK. Across the three heterogeneous STaRK graphs, ARK reaches 59.1% average Hit@1 and 67.4 average MRR, improving average Hit@1 by up to 31.4% and average MRR by up to 28.0% over retrieval-based and agent-based training-free methods.

4. Label-free trajectory distillation. The authors distill ARK's tool-use trajectories from GPT-4.1 into a Qwen3-8B student via supervised fine-tuning on teacher trajectories only, with no ground-truth relevance labels, improving Hit@1 by +7.0, +26.6, and +13.5 absolute points over the base 8B model on AMAZON, MAG, and PRIME respectively, while retaining up to 98.5% of the teacher's Hit@1 rate.

Main Findings

• Best average performance on STaRK. Per-graph, ARK scores Hit@1 / R@20 / MRR of 55.82 / 60.61 / 64.77 on AMAZON, 73.40 / 84.47 / 79.87 on MAG, and 48.20 / 69.46 / 57.68 on PRIME, for averages of 59.14 Hit@1, 71.51 R@20, and 67.44 MRR. It is the strongest method assessed overall; on PRIME it is surpassed only by the RL-trained GraphFlow (51.39 Hit@1, 79.71 R@20, 61.37 MRR).

• Training-free baselines trail by wide margins. Averages for training-free retrieval-based methods are BM25 (27.85 Hit@1, 43.57 R@20, 36.68 MRR), ada-002 (26.96 / 45.88 / 36.79), GritLM-7B (31.85 / 47.34 / 41.61), and KAR (45.01 / 56.11 / 52.67). Training-free agent baselines are Think-on-Graph + GPT-4o (20.22 / 42.73 / 31.43) and Think-on-Graph + LLaMA3 (12.71 / 14.41 / 14.84).

• Trained methods gain less on this benchmark. Averages for methods requiring training on the target graph are mFAR (49.63 / 71.00 / 60.20), MoR (48.93 / 66.14 / 58.77), GraphFlow (46.11 / 57.68 / 54.89), and AvaTaR (37.55 / 50.18 / 45.53).

• Tool use adapts to query type. STaRK annotates queries as primarily textual or primarily relational, and these annotations are not given to ARK. ARK's tool-call allocation nonetheless tracks them: on AMAZON it relies almost entirely on global search (87.7% of calls), while on MAG and PRIME it shifts toward neighborhood exploration (65.3% and 52.3% of calls respectively).

• Neighborhood exploration drives the relational gains. In an ablation on a random 10% subset of the test split, removing neighborhood exploration entirely drops MAG from 79.2 to 30.5 Hit@1 and PRIME from 49.2 to 23.1 Hit@1, while AMAZON falls more moderately from 58.5 to 54.5. Disabling query-based ranking within the neighborhood drops results more mildly (AMAZON 56.0, MAG 72.1, PRIME 44.7 Hit@1). Disabling type-based filtering hurts most on the heterogeneous PRIME graph (42.2 Hit@1 versus 49.2 for the full setup). Global search is not ablated because the authors state it is required to enter the graph at all.

• More compute buys better retrieval, with parallelization being cheap. Quality improves monotonically as inference budget grows from shallow single-agent trajectories to deeper multi-agent configurations. Going from one to two agents gives substantial gains with only modest latency increases, because independent agents run in parallel and end-to-end latency is set by the slowest agent rather than the sum.

• Voting aggregation beats simpler fusion. With three parallel agents, ranking nodes by frequency across trajectories and breaking ties by earliest appearance gives Hit@1 / R@20 of 55.82 / 60.61 (AMAZON), 73.40 / 84.47 (MAG), and 48.20 / 69.46 (PRIME), outperforming ordering-based merging (55.39 / 60.06, 71.24 / 84.15, 44.70 / 68.78) and random merging (17.65 / 57.55, 38.04 / 83.03, 26.30 / 67.50).

• BM25 beats dense retrieval inside the loop. Replacing the lexical relevance function with a dense retriever (text-embedding-3-large) lowers results on AMAZON (47.13 versus 55.82 Hit@1) and PRIME (43.23 versus 48.20), and on MAG the dense variant is slightly higher on both metrics (75.88 versus 73.40 Hit@1; 85.11 versus 84.47 R@20). The authors interpret this as the agent compensating for weaker lexical matching by issuing refined subqueries and expanding through the graph.

• Distillation is data-efficient. Using 10% of teacher trajectories recovers roughly half of the total improvement from the full training set, and distilling the 600-query setting takes about 30 minutes on a single H100 GPU. Student size matters most on PRIME, where base models perform poorly and the teacher-student gap is largest; the 8B student absorbs long-horizon exploration patterns better than the 4B.

• Success correlates with selective expansion. On MAG and PRIME, failed trajectories either make no neighborhood calls at all (the agent never recognizes that relational evidence is needed) or show a long tail of many calls (the agent keeps expanding without converging). Successful trajectories rarely exceed ten neighborhood calls.

• Notable limitations. Agentic retrieval has higher latency than single-pass retrievers; the best configuration uses a large proprietary LLM; lexical global search can under-retrieve when vocabulary mismatches (paraphrases, domain aliases) occur; and evaluation is centered on text-rich KG benchmarks, so gains may not transfer to graphs with sparse text descriptions.

Methodology in Plain English

The paper treats retrieval from a knowledge graph as a conversation between a language model and a graph, rather than as a single scoring pass. The model is given two tools and nothing else.

The first tool, global search, runs a lexical query (BM25 over an inverted index of node text such as titles and descriptions) and returns the top-k matching nodes from the entire graph — default k = 5, though the agent may override it. This is the breadth tool: it finds entry points and handles queries that are mostly about matching language.

The second tool, neighborhood exploration, takes a node the agent already found and returns its adjacent nodes, filtered by optional node-type and edge-type constraints the agent specifies, and optionally ranked by a subquery the agent writes. It returns a fixed budget of k = 20 neighbors per call. This is the depth tool: the agent performs multi-hop traversal by calling it repeatedly, without any preset hop count, deciding when to stop based on what it has seen.

The agent runs for at most T_max = 20 steps and ends when it calls a dedicated finish action or hits the limit. Throughout, it maintains an ordered list of selected nodes; earlier selections rank higher. Nothing about the process is trained — it is a training-free tool-use loop.

To make the system more robust, the authors run n = 3 independent agents in parallel on the same query and fuse their rankings: nodes are ordered by how many agents retrieved them, with ties broken by the earliest appearance anywhere. Rankings for evaluation are truncated to the top 20 nodes for Recall@20.

For efficiency, they then distill the teacher's behavior. They run GPT-4.1 as the teacher three times per training query with temperature 0.7, cap the distillation set at 6,000 training queries per graph (at most 18,000 trajectories per graph, 94.4 million tokens total), and fine-tune a Qwen3-8B student with LoRA and next-token prediction for one epoch at a 16,384-token context length, learning rate 1×10⁻⁵, using AdamW and early stopping on the official validation split. Crucially, the loss is computed only on assistant-generated tokens — tool outputs and user messages are masked — and no relevance labels are used at all, so the approach works even when ground-truth evidence is unavailable. Training runs on a single NVIDIA H100 in about five hours.

Evaluation uses STaRK, which contains an AMAZON e-commerce graph (~1M entities, 9.4M relations; 1,035,542 entities and 9,443,802 relations in the table, 592,067,882 tokens), a scholarly MAG graph (1.9M entities, 39.8M relations; 1,872,968 entities and 39,802,116 relations, 212,602,571 tokens), and a biomedical PRIME graph built from PrimeKG (129K entities, 8.1M relations; 129,375 entities and 8,100,498 relations, 31,844,769 tokens). Metrics are Hit@1, Hit@5, Recall@20, and MRR (the paper notes Hit@5 appears in the appendix). The primary backbone is GPT-4.1, with GPT-4o and GPT-5.2 results reported in the appendix.

Why This Matters

Impact on research. The paper reframes graph retrieval as an interactive decision process rather than static ranking, and shows that a two-operation interface is enough to cover both breadth-heavy and depth-heavy queries. It also provides evidence that distillation from interaction trajectories alone — without relevance labels — can transfer an agentic retrieval policy into a small open model, which addresses the cost and reproducibility problems of proprietary-LLM agents. The compute-accuracy heatmaps give the community an interpretable way to report latency-quality trade-offs rather than single operating points.

Real-world applications:

• Scientific literature search: on MAG-style citation graphs, an agent can anchor on a paper and then expand to co-authors and their papers that match query keywords, supporting literature discovery that keyword search alone misses.
• Biomedical evidence gathering: on PRIME/PrimeKG-style graphs, typed neighborhood expansion can follow drug-gene-disease relations, where type constraints proved important for avoiding drift into unrelated graph regions.
• E-commerce and product question answering: on AMAZON-style catalogs, mostly descriptive queries are served almost entirely by global lexical search, suggesting a cheap, low-latency deployment mode for text-heavy catalogs.
• Enterprise and domain knowledge bases: any organization with a text-rich entity-relation store could expose two tools to an LLM and allow adaptive retrieval without building a labeled training set.

Industry relevance. The distillation result is the practical hook: a model trained from GPT-4.1 trajectories running on a single H100 in about five hours, and a 600-query distillation completing in about 30 minutes, offers a path from a costly proprietary agent to a deployable 8B model that retains up to 98.5% of the teacher's Hit@1. The paper's compute landscapes also give practitioners an explicit dial — agent count and trajectory depth — for meeting a fixed latency or cost budget.

Future Directions

• Closing the latency gap. Agentic retrieval requires multiple LLM calls per query, which the authors list as a primary limitation. Reducing the number of interaction steps without losing the adaptive behavior is an open problem.
• Robustness beyond lexical entry points. Because global search is lexical, paraphrases and domain-specific aliases can cause under-retrieval and prevent the agent from finding the right starting node. Hybrid or semantically grounded anchoring is a natural extension.
• Generalizing past text-rich graphs. The evaluation is centered on text-rich KG benchmarks; whether the same two-tool interface helps on graphs with sparse or templated node descriptions is untested.
• Better handling of the two observed failure modes. Failed trajectories either never use neighborhood exploration or overuse it badly. Detecting when relational evidence is needed, and when to stop expanding, could be learned or prompted more explicitly.
• Stronger safety practices. The ethical section calls for redaction of sensitive fields and bias audits, since lexical search and neighborhood ranking may under-retrieve for languages or communities with sparse descriptions or different naming conventions.

Target Audience

Researchers and engineers working on retrieval-augmented generation, knowledge graph question answering, or agentic tool use with LLMs will get the most from this paper, particularly those evaluating on STaRK or building graph-grounded systems. It is also relevant to practitioners who want an agentic retriever but cannot afford a large proprietary model at inference time, since the distillation section addresses that directly. Readers looking for a beginner-level introduction to knowledge graph retrieval should expect to look up BM25, rank fusion, LoRA fine-tuning, and the STaRK benchmark first.

Authors’ abstract

Retrieving evidence for language model queries from knowledge graphs requires balancing broad search across the graph with multi-hop traversal to follow relational links. Similarity-based retrievers provide coverage but remain shallow, whereas traversal-based methods rely on selecting seed nodes to start exploration, which can fail when queries span multiple entities and relations. We introduce ARK: Adaptive Retriever of Knowledge, a tool-using KG retriever that gives a language model control over this breadth-depth tradeoff using a two-operation toolset: global lexical search over node descriptors and one-hop neighborhood exploration that composes into multi-hop traversal. ARK alternates between breadth-oriented discovery and depth-oriented expansion without depending on a fragile seed selection, a pre-set hop depth, or requiring retrieval training. ARK adapts tool use to queries, using global search for language-heavy queries and neighborhood exploration for relation-heavy queries. On STaRK, ARK reaches 59.1% average Hit@1 and 67.4 average MRR, improving average Hit@1 by up to 31.4% and average MRR by up to 28.0% over retrieval-based and agent-based training-free methods. Finally, we distill ARK's tool-use trajectories from a large teacher into an 8B model via label-free imitation, improving Hit@1 by +7.0, +26.6, and +13.5 absolute points over the base 8B model on AMAZON, MAG, and PRIME datasets, respectively, while retaining up to 98.5% of the teacher's Hit@1 rate.