A Hierarchical Bayesian Model for Selecting Relevant Schema Subgraphs,with an Application to Grounding Large Language Models

Abstract

Grounding a large language model on a relational database requires selecting the tables and joins relevant to a query. Existing schema-linking methods usually score tables or columns independently, although the target object is typically a connected subgraph of the foreign-key graph. We propose a hierarchical Bayesian model for schema linking as structured subset selection on small attributed graphs. Unary evidence from query-table and query-column features is combined with a pairwise autologistic coupling on foreign-key edges, allowing weakly-signalled tables to be selected when they connect other relevant tables. Database-level random effects model variation in the baseline inclusion rate and coupling strength, and the small size of the schema graphs permits exact evaluation of the conditional likelihood and inclusion probabilities, which we average over a Laplace approximation to the parameter posterior. On the benchmark datasets BIRD and Spider, positive graph coupling recovers tables missed by independent scoring but shifts the marginal inclusion probabilities upward. Jointly estimating the intercept and coupling restores calibration, and database-level partial pooling adapts the structural effect across schemas. The graph-coupled posterior assigns more probability to the exact relevant subgraph and produces informative predictive sets with near-nominal coverage, whereas predictive sets from the corresponding independent model undercover. The fitted model also reports a posterior over how strongly to weight the graph in each database. Stability and mean-shift results explain when coupling aids recovery and why the intercept must compensate for it. We therefore evaluate the method primarily as a posterior predictive distribution over table subsets rather than as a precision-maximizing selector.