Co-adaptive Rule Routing and Constraining for Structured Knowledge Retrieval

AI summary

Overview

Research area: Information retrieval (cs.IR) — specifically structured/knowledge-graph retrieval, combined with agentic reasoning and reinforcement learning.

Technical level: Advanced. The paper assumes familiarity with knowledge graphs, contrastive learning, policy-gradient/GRPO-style reinforcement learning, and bi-level optimization.

Scope: A framework, CoRE, that learns to jointly route each query to a latent reasoning rule and train rule-constrained retrievers from downstream retrieval feedback alone, evaluated on three STaRK-QA datasets (AMAZON, MAG, PRIME).

What This Paper Is About

Structured knowledge retrieval works best when it follows an explicit reasoning path — for example, traversing Symptom → Diagnosis → Treatment in a clinical knowledge base, or Paper → Author → Paper in a citation graph. In practice, the correct rule for a given query is unknown and expensive to annotate, so existing methods either retrieve without structural constraints or rely on heuristically/LLM-constructed query–rule labels that may not match what the query actually needs. CoRE treats the rule as a latent variable and learns both which rule to use and how to retrieve under that rule from nothing but query–evidence pairs.

Key Contributions

1. Reasoning-constrained retrieval without query–rule supervision. The authors formulate rule selection as a latent routing problem and factorize structured retrieval into a rule-routing agent and a set of rule-constraining retrievers, removing the need for query–rule pair supervision or gold reasoning paths.

2. A co-adaptive training loop. CoRE alternates between (i) training each rule-constraining retriever with contrastive learning on the queries the router assigned to it, and (ii) updating the router with reinforcement learning using retrieval rank as the reward, so both components improve through a closed feedback loop.

3. Three router instantiations. The framework supports a BERT encoder with a classifier head (CoRE-BERT-C), an LLM encoder with a classifier head (CoRE-LLM-C), and an LLM encoder with a token-generation decoder (CoRE-LLM-G), all sharing the same retriever design.

4. Empirical gains plus new diagnostics. CoRE-BERT-C achieves the strongest average performance across AMAZON, MAG, and PRIME; ablations isolate the contributions of router and retriever adaptation; and marginalizing the routing policy over each dataset's queries exposes a distribution of rule usage that the authors argue can inform automatic rule refinement and discovery.

Main Findings

• Best average retrieval performance. CoRE-BERT-C reaches average H@1 54.47, H@5 80.12, R@20 75.24, and MRR 65.97 across the three datasets, the strongest average in Table 1. The best hybrid baseline, MoR, reaches average H@1 48.93, H@5 71.00, R@20 66.14, and MRR 58.77.

• Large gains on MAG and PRIME. CoRE-BERT-C scores MAG H@1 63.37 / MRR 74.17 and PRIME H@1 49.34 / MRR 61.27, versus MoR's 58.19 / 67.14 on MAG and 36.41 / 46.92 on PRIME. On AMAZON the margin is smaller (H@1 50.71 vs. MoR's 52.19).

• Not best on every individual cell. HYBGRAG reports a higher MAG H@1 (65.40) than CoRE-BERT-C (63.37), and MFAR reports a higher MAG R@20 (71.70) than CoRE-BERT-C (83.51 — actually CoRE is higher here). The paper's claim is about the strongest average, not uniform dominance; PRIME is where CoRE's lead is largest, which the authors attribute to its 65-rule space being hard for LLMs to navigate from general knowledge alone.

• Classification-based routers beat generation-based routing on average. CoRE-LLM-C averages H@1 52.35 / MRR 64.10 and CoRE-LLM-G averages 48.31 / 57.61, both below CoRE-BERT-C. The authors attribute the gap to classifier routers using reward feedback from all candidate rules, whereas generation-based routing relies on eight sampled rollouts per query, giving higher-variance gradients — most visible on PRIME, where eight rollouts cover little of the 65-rule space.

• Both adaptation directions matter. Removing router adaptation drops average H@1 from 54.47 to 16.00 with an unadapted BERT router and to 19.69 with a zero-shot GPT-5-mini router. Removing retriever adaptation (fixed BM25-traversal scoring) drops average H@1 to 48.08 and MRR to 58.87; the loss is largest on PRIME (H@1 38.86), which has more diverse multi-hop rules.

• Generation-based routers transfer across datasets. Routers trained independently on AMAZON, MAG, and PRIME were each evaluated on all three datasets with target-domain retrievers held fixed; every transferred router outperformed a zero-shot Qwen router, suggesting the training learns a transferable mapping from query semantics to reasoning patterns rather than memorizing dataset-specific assignments.

• Multi-domain pretraining plus target fine-tuning is best for a foundation router. Using Qwen2.5-32B-Instruct as the base router, the paper reports that multi-domain pretraining followed by target-specific fine-tuning gives the highest H@1 (Figure 4); the specific per-setting numbers are not included in the provided text.

• Routing distributions reveal rule usage. Marginalizing the routing policy over all queries in a dataset produces a rule-usage distribution that the authors use to motivate automatic rule refinement and discovery.

Methodology in Plain English

The system splits one hard problem into two simpler ones that are trained in alternation.

First, the researchers enumerate a set of candidate reasoning rules from the dataset schema — every ordered sequence of entity types reachable through permitted relations in the traversal direction. This gives 6 rules for AMAZON, 8 for MAG, and 65 for PRIME.

Second, a router looks at a query and picks a rule. It can be a BERT or LLM encoder feeding either a softmax over the candidate rules or a token-by-token generator that writes the rule out as text.

Third, given the chosen rule, a retriever performs a constrained depth-first search through the knowledge graph, keeping only entities reachable along that type sequence. For each candidate it builds a feature vector of BM25 similarities between the query and the text descriptions of every node on the traversal path, then scores it with a learnable function that also takes the rule embedding. Candidates are ranked with a softmax over the constrained candidate set.

Fourth, the two are trained in a loop. Starting from uniform query–rule assignments, each retriever is trained contrastively on the queries its rule received. The trained retrievers then execute all rules on the training queries, and the rank of the ground-truth evidence under each rule is converted into a reciprocal-rank reward. The router is updated to favor rules with better rewards — via an advantage weighted over all candidate rules for classifier routers, or via GRPO with eight sampled rollouts for the generation-based router. The updated router reassigns queries, providing fresh supervision for the next retriever round. All variants run three such alternate rounds.

At inference, the router picks the most probable rule, the corresponding retriever builds the constrained candidate set, and the top-ranked evidence is returned for downstream use. Evaluation uses Hits@1, Hits@5, Recall@20, and MRR.

Why This Matters

Impact on research. The paper reframes rule selection as latent-variable learning driven by retrieval outcomes rather than by constructed query–rule labels, and it connects the co-adaptation literature (agent-and-tool co-evolution) to structured retrieval. The reported rule-usage distributions suggest a path from hand-specified rule sets toward data-driven rule discovery, and the transfer results suggest routing policies can be shared across domains.

Real-world applications:

• Biomedical reasoning — a clinician following Symptom → Diagnosis → Treatment, or a pharmacology query following Anatomy → Disease → Drug (the PRIME dataset's setting).
• Scientific discovery — tracing intellectual lineage through Paper → Author → Paper, or searching by Author → Paper → Venue (the MAG setting).
• Enterprise knowledge management and e-commerce — product and customer-relation retrieval over enterprise schemas (the AMAZON setting).
• Cybersecurity incident analysis — tracing Alert → Host → Process → Vulnerability to find the root cause of a breach.

Industry relevance. The approach needs only query–evidence pairs, which most organizations already log, rather than curated reasoning-path annotations. Because routing and retrieval are separated, the retriever can be swapped or specialized per rule while a single router serves as a domain-agnostic entry point — an architecture that maps naturally onto enterprise search and assistant systems, and one the paper shows can be built on open models (Qwen2.5-3B-Instruct for routing, Qwen2.5-32B-Instruct for the foundation router) on four NVIDIA H200 GPUs.

Future Directions

• Scaling generation-based routing to large rule spaces. The eight-rollout sampling budget covers only a small fraction of PRIME's 65 rules; better credit assignment or cheaper rule-space exploration is an open problem the results directly motivate.

• Closing the loop from routing distributions to rule discovery. The paper shows marginalized routing yields interpretable rule-usage distributions and argues these can drive automatic rule refinement and discovery, but the mechanism for adding, pruning, or merging rules is left as future work.

• Generalizing to rules and schemas not enumerated in advance. Cross-dataset transfer is demonstrated when only the task description, graph schema, and candidate rules are swapped in; whether a router can handle rule sets it has never seen enumerated remains open.

• Extending beyond knowledge-graph retrieval. The rule-constrained retrieval formulation is presented generally (a candidate set plus a rule-conditioned scoring function) but is only instantiated for knowledge graphs, leaving other structured sources untested.

• Verifying downstream task impact. Evidence retrieval is evaluated with ranking metrics; how much the retrieved evidence improves end-task execution for scientific, biomedical, or enterprise workflows is not reported in the provided text.

Target Audience

Researchers and graduate students in information retrieval, knowledge-graph reasoning, and agentic/LLM systems; practitioners building enterprise or domain search over structured knowledge bases who lack annotated reasoning paths; and applied scientists interested in reinforcement-learning-driven co-adaptation between a planner and specialized retrieval tools. Readers should be comfortable with contrastive retrieval training, policy-gradient methods, and graph traversal concepts.

Authors’ abstract

Structured knowledge retrieval acquires task-relevant evidence from structured knowledge sources for downstream tasks. Existing methods rely on unconstrained retrieval or predefined query-rule supervision, such as author-paper rules for scientific discovery and anatomy-drug rules for precision medicine. Such supervision is rarely available in practice. We propose Co-Adaptive Rule Routing and Constraining Retrieval (CoRE), which jointly learns which reasoning rule to assign to each query and how to retrieve evidence under that rule, without predefined query-rule pairs. A routing agent selects a rule, and its corresponding retriever traverses the knowledge base and ranks evidence that satisfies the rule's constraints. Because routing and retrieval are interdependent, CoRE alternates between reinforcement learning for the routing policy and contrastive learning for the rule-constrained retriever. The selected rules provide supervision for retriever training, while retrieval performance rewards the routing policy. This closed-loop process enables both components to co-adapt without explicit query-rule supervision. Experiments on STaRK-QA datasets across scientific discovery, biomedical reasoning, and e-commerce knowledge management demonstrate CoRE's effectiveness. Ablations show the benefit of co-adapting the router and retriever. Generation-based routing transfers across datasets and benefits from multi-domain pretraining followed by target-specific fine-tuning. Analysis of rule usage across queries also informs automatic rule refinement and discovery.