Autonomous Knowledge Graph Exploration with Adaptive Breadth-Depth Retrieval

TL;DR

ARK tackles the dual challenge of breadth and depth in knowledge-graph retrieval for language models by providing a minimal, training-free tool interface with global search and one-hop neighborhood exploration. It enables adaptive, query-driven control over the breadth-depth tradeoff and demonstrates strong, transfer-friendly performance on the STaRK benchmark, including a label-free distillation of the policy to an 8B model with minimal loss in accuracy. The approach combines parallel agent execution and rank-fusion to stabilize retrieval, and its trajectories can be distilled into compact models to decrease inference cost while preserving most teacher performance. This work advances practical, generalizable KG grounding for LLM-based reasoning and retrieval across heterogeneous graphs and domains.

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, an agentic 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 agentic 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.

Figures (5)

• Figure 1: Overview of Adaptive Retriever of Knowledge. ARK interacts with a KG through a minimal two-tool interface: (a) For text-dominant queries, ARK emphasizes breadth by issuing Global Search to retrieve a broad set of candidates. (b) For relation-focused queries, ARK applies Neighborhood Exploration starting from a previously retrieved node (in this case, a drug) and expanding to related entities, enabling targeted relational retrieval. (c) For relation-dominant queries, ARK performs multi-hop retrieval by alternating Global Search and Neighborhood Exploration: it retrieves an initial node (e.g., an article), expands to related entities (e.g., co-authors), and continues expanding and filtering (e.g., papers connected to each author that match query keywords) to recover an ordered set of relevant evidence.
• Figure 2: Thin bars show the share of text- vs. relation-centric queries in STaRK; thick bars show ARK's tool-call use. These STaRK annotations are not provided to ARK; instead, ARK autonomously shifts tool use to match the dominant query type.
• Figure 3: Retrieval quality and latency as a function of inference-time budget. Heatmaps report Hit@1, Recall@20, and end-to-end latency (seconds) on each STaRK graph. Moving from the top left (shallow trajectories, single agent) to the bottom right (deeper trajectories, multi-agent) allocates more compute and improves retrieval performance at the cost of higher latency. Color scales are normalized within each graph and metric for readability.
• Figure 4: Evaluation of the same ARK pipeline on the STaRK test sets while varying only the LLM backbone (Qwen3-4B/8B base, Qwen3-4B/8B distilled, or GPT-4.1). "Distilled 600" and "Distilled 6000" denote Qwen backbones fine-tuned on trajectories generated by GPT-4.1 from 600 or 6000 training queries per graph, respectively (three trajectories per query; tool calls and observations only; no label supervision).
• Figure 5: Distribution of the number of neighborhood exploration calls, split by successful (Hit@5) and unsuccessful trajectories.