💬 NLP

Baikal: Structured Search for Deep Research over Data Lakes

The paper introduces Baikal, a framework that addresses the limitations of iterative retrieval in deep research over data lakes by casting the task as a budgeted search problem, where evidence is clustered into semantic regions and adaptively explored to balance exploration and exploitation, resulting in significantly higher quality reports compared to existing baselines.

Original authors: Dhruv Agarwal, Rishitha Guttapalle Mohan, Aarti Kumari, Ashi Sinha, Athulya Anil, Kavitha Srinivas, Horst Samulowitz, Andrew McCallum

Published 2026-07-31

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Dhruv Agarwal, Rishitha Guttapalle Mohan, Aarti Kumari, Ashi Sinha, Athulya Anil, Kavitha Srinivas, Horst Samulowitz, Andrew McCallum

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a massive mystery, but instead of a few clue-filled notebooks, you are handed a warehouse the size of a city, packed with millions of scattered documents, spreadsheets, and handwritten notes. This is the world of "Deep Research" over "Data Lakes." In the real world, companies and scientists don't just have one file to read; they have thousands of tables of numbers and millions of paragraphs of text that don't always talk to each other. To solve a complex question, an Artificial Intelligence (AI) needs to act like a super-smart investigator: it has to find the right clues, figure out what they mean, and write a report that connects the dots. The tricky part is that the AI has a limited amount of time and energy (a "budget") to do this. If it gets too excited about one small, easy clue, it might miss the huge, important secrets hidden in a different corner of the warehouse. The big question scientists are asking is: How do we teach an AI to explore a giant, messy library without getting stuck in one aisle or missing the best stories?

This is exactly what the researchers behind Baikal set out to solve. They realized that the usual way AI searches for answers—just grabbing the next piece of text that looks interesting and moving on—is like a squirrel running in circles around the same tree. It finds a few nuts, but it never explores the whole forest. To fix this, the team invented a new method that treats the search like a strategic game of exploration.

First, Baikal doesn't just dump the whole warehouse in front of the AI. Instead, it acts like a master librarian who instantly organizes the chaos. It groups thousands of unrelated tables and documents into "semantic regions"—think of these as labeled bins or neighborhoods. One bin might hold everything about "Serbian athletes," while another holds "cricket venues." This happens before the AI even starts asking questions.

Once the library is organized, Baikal sends its AI agent on a mission with a strict budget of 50 steps (or "questions" it can ask). Instead of just picking a random document, the agent uses a smart strategy to decide which "neighborhood" to visit next. It's like playing a game where you have to choose which room to enter in a giant mansion. Some rooms might have already been explored and found empty, while others might be full of gold. Baikal uses a mathematical "guessing game" (called a bandit problem) to balance between visiting rooms it thinks are promising (exploitation) and trying out new, unexplored rooms to see if they hold something better (exploration).

Inside a chosen neighborhood, the AI generates specific, grounded questions. For example, if it's in the "Serbia" bin, it might ask, "Which sports did Serbian athletes win the most medals in?" It then digs through the specific tables and text in that bin to find the answer. A special "judge" AI scores how good that answer is based on whether it's true, useful, and different from what was found before. If the answer is great, the system remembers that neighborhood is valuable. If it's boring, it marks that area as "done" and moves on.

The results of this approach were impressive. The researchers tested Baikal on two massive data lakes: one with nearly 11,000 tables and 227,000 text passages (about Wikipedia), and another with over 2,700 tables and 13,000 financial reports. They gave the AI 15 complex research questions to answer. When they compared Baikal to the strongest existing methods, the difference was huge. Baikal's best setup improved the quality of the final research reports by 28% on the Wikipedia data and 36% on the financial data.

Crucially, the paper shows that the magic wasn't just in having the organized bins; it was in the search strategy. When they gave a powerful coding AI the same organized bins but let it search without Baikal's smart "neighborhood hopping" rules, it didn't get the same boost. This suggests that simply having the data organized isn't enough; you need a smart way to decide where to look next. By treating the search as a structured game of exploration rather than a random scavenger hunt, Baikal helps AI agents find more useful, distinct, and grounded facts, even when they are racing against a ticking clock. It turns a chaotic data dump into a systematic discovery journey, proving that sometimes, the best way to find the truth is to know exactly where not to look.

Technical Summary: Baikal – Structured Search for Deep Research over Data Lakes

1. Problem Formulation

The paper addresses the challenge of deep research over data lakes, defined as the task where an autonomous agent must synthesize a report by selectively investigating evidence distributed across thousands of heterogeneous tables and unstructured text passages.

The core difficulty lies in the budgeted search nature of the problem. Given a fixed budget of $B$ steps (where each step yields one finding), the agent must:

Retrieve query-relevant evidence from a massive, heterogeneous collection ( $L = T \cup P$ ).
Process evidence with reasoning.
Crucially: Identify a diverse set of salient aspects for the query without over-exploiting locally promising evidence.

Existing iterative retrieve-and-generate methods often fail here because they rely on accumulating context to shift model beliefs. This biases the agent toward locally promising evidence, leading to poor coverage of distinct semantic regions and low-quality reports that miss key aspects of the query. The authors formalize the objective as maximizing the average quality score of findings across the budget, where the quality score ( $\rho$ ) is a product of four dimensions: groundedness, relevance, distinctness, and utility.

2. Methodology: The Baikal Framework

Baikal reframes deep research as a budgeted multi-armed bandit problem over a structured semantic space. The framework operates in three main phases:

A. Adding Structure (Pre-computation)

Instead of searching raw tables and passages, Baikal first organizes the data lake into semantic regions:

Embedding: A single encoder (e.g., Qwen3-Embedding) embeds both table metadata (title, schema) and text passages into a shared vector space.
Clustering: Using BERTopic (UMAP for dimensionality reduction + HDBSCAN for clustering), the system groups thousands of items into $|R|$ semantic regions.
Region Definition: Each region $c$ contains a subset of tables $T_c$ and passages $P_c$ , along with a description. Regions are designed to fit within an agent's context window while remaining topically coherent. This reduces the search space from $|L|$ items to $|R|$ "arms."

B. Search and Investigation (Runtime)

For a given query $q$ , Baikal performs a sequential search over the candidate regions $C(q)$ :

Candidate Selection: The query is embedded, and top- $k$ tables/passages are retrieved. Any region overlapping with these top- $k$ items is retained as a candidate. This ensures that retrieving one item brings its entire semantic region into scope.
Region Selection Policy: A policy $\pi$ $π$ selects a region $c \in C(q)$ $c \in C (q)$ to investigate. The paper evaluates several policies:
- Random: Uniform sampling.
- LLM-guided: An LLM selects based on region descriptions and history.
- Bayesian $\epsilon$ -greedy: Exploits the region with the highest estimated mean reward, with probability $\epsilon$ exploring via softmax over posterior means.
- Bayes-UCB: Selects the region with the highest posterior upper quantile, balancing exploration and exploitation.
- Note: Bayesian policies are initialized with LLM-elicited priors to avoid wasting the budget on uninformed exploration.
Subquestion Generation & Selection: Within the selected region, an LLM generates up to $K$ analytical subquestions (e.g., counts, comparisons) grounded in the region's content. One subquestion is selected (randomly or via an LLM policy) to be answered.
Execution: A coding agent (e.g., OpenCode) executes the subquestion using region-scoped SQL against a SQLite database of the region's tables and reads its passages. The agent may also perform region expansion (using grep to fetch additional relevant passages) if the initial region is insufficient.
Reward Feedback: An LLM judge scores the resulting finding on the four rubric dimensions. This score serves as the reward $r_t$ to update the policy's value estimate for that region (using a pseudo-count scheme for Bayesian updates).

C. Report Synthesis

Once the budget is exhausted or no candidates remain, the collected findings are synthesized into a final report, with claims attributed to specific evidence sources.

3. Key Contributions

Formalization: Casting deep research over data lakes as a budgeted search problem over structured semantic regions, explicitly addressing the trade-off between exploration and exploitation.
Baikal Framework: A system that clusters heterogeneous data into semantic regions and uses sequential decision-making (bandit policies) to guide evidence gathering, rather than relying solely on iterative context accumulation.
New Testbed & Evaluation: The construction of deep-research queries over HybridQA (10,993 tables) and TAT-QA (2,757 tables) data lakes, accompanied by a practical rubric for evaluating groundedness, relevance, distinctness, and utility using LLM judges.
Empirical Validation: Demonstrating that the gains come from the search procedure (sequential allocation of budget) rather than the structure alone. Providing a strong coding agent with the same regions does not replicate Baikal's performance.

4. Results

The authors evaluated Baikal against strong baselines, including an OpenCode agent (with and without retrieval/clustering artifacts) and DeepSearcher, on 15 queries per data lake with a budget of 50 steps.

Performance: The best Baikal configuration (Bayes-UCB on HybridQA; Random/Bayes $\epsilon$ -greedy on TAT-QA) improved report scores by 28% on HybridQA and 36% on TAT-QA over the strongest baselines.
Quality Dimensions: Baikal significantly improved groundedness (0.92 vs. ~0.72 for baselines) and distinctness (0.86 vs. ~0.62), which are critical for avoiding repetitive or hallucinated content. While baselines sometimes scored higher on relevance, their lack of groundedness resulted in lower overall scores due to the multiplicative nature of the rubric.
Search Efficiency: Baikal policies separated from baselines within the first 10 steps. While random selection performed surprisingly well (due to the structure itself), Bayesian policies showed clear advantages as the budget increased (e.g., Bayes-UCB gained 8.8% over random at budget 200).
Cost Analysis: Baikal incurred higher total costs per query because it produced more valid findings (44.6 vs. 30.9 for OpenCode). However, the cost per valid finding was comparable to baselines, indicating the premium buys usable evidence rather than expensive failed steps.

5. Significance and Claims

The paper claims that structured semantic exploration is essential for systematic research over heterogeneous data lakes. The primary contribution is not just the clustering of data, but the explicit management of the search process under a fixed budget.

Structure vs. Search: The authors emphasize that providing a strong agent with the same structured regions (retrieval + clustering) does not transfer the performance gains. The improvement stems from the sequential decision-making policy that allocates the budget across regions to maximize coverage and distinctness.
Scalability: The approach handles the "long-horizon" nature of deep research by reducing the search space to semantic regions, allowing the agent to cover distinct parts of the lake that iterative context-based methods often miss.
Reproducibility: The evaluation rubric (groundedness, relevance, distinctness, utility) was validated for inter-LLM agreement, showing high consistency in distinctness and groundedness scoring.

In conclusion, Baikal demonstrates that treating a data lake as a structured space to be searched, rather than a corpus to be retrieved from, offers a robust foundation for deep research tasks, particularly when evidence is distributed across thousands of tables and passages.