STAIR (STructure Aware Information Retriever): A novel dataset and LLM based retriever for document structure augmentation

Abstract: Retrieval Augmented Generation (RAG) is a key component for generating accurate and hallucination free answers using LLMs. LLMs are improving at handling long context, but still suffer from "lost in the middle" problem. Thus, precise and accurate retrieval is important. Current retrievers chunk long context into length-based manageable chunks - in the process throwing away rich and informative semantic global structure in the corpus. We introduce a novel retrieval system STAIR that empowers an LLM to exploit global structure in a corpus such as a Table of Contents (ToC) to efficiently store and retrieve information from its model parameters. Our thorough and careful ablation studies with a finetuned Differentiable Search Index (DSI) system show that ToC helps build a low hallucination (less than 0.05%) generative Information Retrieval (IR) system and can generalize to examples where very few training samples are available. To further research in this novel direction of ToC based retrieval we release SearchTome - a diverse benchmark created from 18 books across 6 diverse domains to further research in this novel direction. STAIR achieves a high Recall@1 score of 82.6% on SearchTome as compared to DSI (76.9%), where the difference is found to be statistically significant. STAIR easily beats other strong baselines such as BM25 (59.5%), DPR (68.7%) and out-of-the-box Mistral (13.8%).

• The paper introduces STAIR, a retriever that uses a document's Table of Contents (ToC) for more accurate information retrieval by treating it as an explicit global representation of structured documents, enhancing systems like RAG and outperformance with regards to Recall@1 (82.6%).
• STAIR’s architecture and training pipeline, involving the fine-tuning of Mistral Instruct v0.2, showcases superior performance in retrieving relevant sections from a wide array of structured documents, particularly when training data is sparse.
• The research presents a new benchmark, including 18 open textbooks, illustrating that STAIR's structured retrieval effective significantly reduces hallucinations and improves recall metrics, providing a promising method for corporate and academic document retrieval.

Problem formulation and motivation

“STAIR (STructure Aware Information Retriever): A novel dataset and LLM based retriever for document structure augmentation” (2609.03874) addresses a specific weakness in long-document retrieval: conventional retrievers generally partition documents into length-based chunks, thereby discarding document-level organization and semantic boundaries. This problem is consequential for RAG systems because long-context LLMs do not reliably use all positions in an input; relevant evidence located in the middle of a context can be substantially less accessible than evidence near the beginning or end [2403.XXXX; (Li et al., 2024)]. STAIR instead treats a document’s Table of Contents (ToC) as an explicit global representation and uses its leaf sections as retrieval units.

The central claim is that a ToC provides information that is unavailable to a conventional chunking strategy. A ToC encodes hierarchical relations, topic boundaries, and semantically meaningful section titles. A query can therefore be aligned not only with local text, but also with the document’s organization. This design is particularly appropriate for textbooks, technical reports, enterprise documentation, and other corpora whose content is authored under an explicit hierarchy.

STAIR builds on model-based indexing, especially Differentiable Search Indexing (DSI), which fine-tunes a sequence-to-sequence model to map queries directly to document identifiers (Tay et al., 2022). The distinction is that DSI must infer the corpus structure indirectly from query–identifier training pairs, whereas STAIR receives the complete ToC at input time and generates a valid leaf-node title. The paper consequently frames retrieval as two coupled tasks: parametric ingestion of corpus knowledge and structured selection from the document’s ToC.

Benchmark construction

The paper introduces STAIR, a benchmark comprising 18 open textbooks distributed across six domains: education, finance, law, medicine, natural sciences, and social sciences. The books range from 179 to 1,327 pages and contain between 70 and 429 ToC leaves. Each book is parsed to extract its ToC and associate paragraphs with the corresponding leaf sections.

Questions are generated synthetically with Mixtral 8x7B (Jiang et al., 2024). For each paragraph, the model produces multiple questions intended to cover the important content of that paragraph. The resulting questions are divided into training, development, and test partitions, with the gold retrieval target defined as the relevant ToC leaf. The benchmark is therefore designed to evaluate whether a system can identify the section containing the answer rather than directly generate the answer itself.

This construction gives the benchmark broad topical coverage and provides an explicit structured retrieval target. It also imposes an important methodological qualification: the queries are synthetic and generated from the same books whose structure is used for retrieval. The benchmark therefore measures corpus-specific structured retrieval under controlled supervision more directly than it measures robustness to naturally occurring information needs.

The dataset contains substantial variation in corpus size and structural granularity. For example, Open Music Theory has 429 leaves, whereas several books have fewer than 100. This variation is useful for testing whether a model can distinguish among many structurally similar candidates and whether sparse supervision disproportionately harms underrepresented sections.

STAIR architecture and training procedure

STAIR fine-tunes Mistral Instruct v0.2 using LoRA with rank 16 and scaling parameter 32. During training, the model receives a prompt containing the book identity, the complete ToC, and a user query. It is trained to generate the gold leaf-node title. The maximum input length is 14,000 tokens, compared with 512 tokens for the DSI implementation.

The system pipeline has three stages. First, synthetic question–section pairs are prepared from the source book. Second, the LLM is fine-tuned to map a ToC-augmented query to the correct leaf node. Third, inference uses constrained generation so that the output is restricted to valid ToC leaves.

Figure 1: STAIR constructs synthetic question–section pairs, fine-tunes a model with the complete ToC as input, and performs constrained leaf-node generation.

The use of the ToC has two distinct effects. It supplies an explicit candidate space during decoding, and it exposes the model to a global semantic organization that would otherwise have to be memorized in the model parameters. The first effect directly addresses invalid identifier generation. The second should be most valuable when a leaf has few training examples, because the model can use the title and its position in the hierarchy as additional structural evidence.

This design differs from simply prepending metadata to a conventional retriever. The ToC is not merely an auxiliary feature used for ranking; it is the explicit input against which the model performs a constrained generative decision. Retrieval is consequently formulated as structured generation over a finite set of corpus-specific labels.

Experimental setting and baselines

The evaluation reports Recall@1, Recall@3, and nDCG@3 using the BEIR evaluation framework (Thakur et al., 2021). STAIR is compared with five systems:

• out-of-the-box Mistral Instruct v0.2, using beam search;
• BM25 over section content;
• NV-Embed-v2 as an out-of-the-box DPR-style dense retriever (Lee et al., 2024);
• fine-tuned DSI (Tay et al., 2022);
• Claude Haiku, included as an additional language-model baseline in the reported results.

The comparison is informative because the systems represent different retrieval paradigms. BM25 supplies a lexical baseline, DPR a dense semantic baseline, DSI a model-based index without explicit ToC input, and STAIR a model-based index with the global structure exposed during both training and inference.

Main retrieval results

STAIR achieves the strongest average performance on all three reported metrics. Its average Recall@1 is 82.6%, compared with 76.9% for DSI, 68.7% for DPR, 59.5% for BM25, 45.6% for Claude Haiku, and 13.8% for out-of-the-box Mistral.

System	Recall@1	Recall@3	nDCG@3
Mistral	13.8%	16.4%	15.4%
Claude Haiku	45.6%	64.1%	56.4%
BM25	59.5%	77.6%	70.1%
DPR	68.7%	85.4%	78.6%
DSI	76.9%	85.3%	81.9%
STAIR	82.6%	90.8%	87.5%

The 5.7-point absolute improvement over DSI corresponds to the paper’s reported 7.4% relative gain. The improvement is consistent across all six domains: STAIR obtains Recall@1 values between 80.6% and 86.1%, while DSI ranges from 73.3% to 82.1%. Medicine is the strongest domain for STAIR at 86.1% Recall@1, whereas natural sciences is the weakest at 80.6%.

The result supports the paper’s principal claim that explicit structural conditioning improves model-based retrieval. However, the comparison does not isolate all possible causes of the improvement. STAIR and DSI differ not only in whether the ToC is supplied but also in input length, output formulation, candidate constraints, and the amount of information available during inference. The observed gain is therefore evidence for the complete STAIR design, while attribution specifically to the ToC requires the ablation results.

Hallucination reduction and sparse supervision

The most pronounced qualitative advantage of STAIR concerns invalid generations. The paper defines a hallucination as generation of a non-leaf node, which is an invalid retrieval identifier under the benchmark’s target space. STAIR’s non-leaf prediction rate is 0.05%, compared with 3.25% for DSI. Its overall error rate is 18.67%, compared with 24.31% for DSI and 86.20% for out-of-the-box Mistral.

The result is notable because DSI has already undergone corpus-specific fine-tuning, yet it continues to produce structurally invalid identifiers. STAIR’s constrained, ToC-conditioned generation nearly eliminates this failure mode. The immediate implication is that structural validity can be enforced at the output layer rather than learned solely through statistical regularities in query–identifier pairs.

Figure 2: DSI’s hallucination rate increases as the number of training examples for a leaf decreases, whereas STAIR remains close to zero.

The sparse-supervision ablation further supports the structural argument. For leaves with few training examples, the Recall@1 gap between STAIR and DSI is substantially larger. As the number of examples increases, the gap narrows, although STAIR remains superior. This pattern is consistent with the hypothesis that DSI must memorize both corpus content and the set of valid identifiers, while STAIR receives the candidate structure explicitly.

Figure 3: STAIR maintains higher per-leaf Recall@1 than DSI, particularly for leaves represented by few training examples.

The implication is not merely that STAIR produces better average scores. Its advantage is concentrated where parametric memorization is most fragile. In a long-tail retrieval distribution, explicit structural conditioning can compensate for low query coverage by making section identity recoverable from the ToC itself. This is a stronger claim than a general improvement from fine-tuning: the method changes the dependence of retrieval accuracy on the number of examples per section.

Qualitative error analysis

The qualitative examples illustrate complementary failure modes among the baselines. In the preschool-development example, BM25 selects a section based on lexical overlap with terms such as “autonomous,” while DPR retrieves a semantically adjacent section about autonomy and shame. DSI improves the topical match but generates a section title that does not exist in the supplied ToC, demonstrating the invalid-identifier problem. STAIR selects the correct leaf, “Initiative vs. Guilt (Preschool Years).”

The nursing example shows a different distinction. Several systems identify neighboring sections within the same chapter, but STAIR selects the section whose content explicitly discusses chronic illness, activities of daily living, and functional ability. In the comparative-politics example, STAIR retrieves the section containing the definition of irregular armed organizations, whereas lexical and dense baselines select nearby sections concerning insurgency, state-sponsored violence, or hybrid regimes.

These examples suggest that STAIR benefits from the interaction between semantic title information and corpus-specific section boundaries. Nevertheless, the examples are anecdotal and do not establish causal explanations for individual predictions. The aggregate ablations provide stronger evidence for the proposed mechanism than the qualitative cases alone.

Limitations and open questions

The evaluation is restricted to documents with an existing global structure. This assumption is central to the method: STAIR requires a ToC, a hierarchy, and a set of valid leaf nodes. Consequently, the results do not establish performance on arbitrary web corpora, unstructured document collections, conversational archives, or datasets where section boundaries are noisy or inconsistent.

The benchmark also relies on synthetic questions generated from the source paragraphs. Although this enables broad coverage and precise gold labels, it may produce queries that are more closely aligned with the document’s wording and organizational vocabulary than naturally occurring user queries. The paper does not report evaluation on independently authored questions, adversarial paraphrases, cross-book transfer, or unseen-document retrieval.

A further open issue is scalability. The experiments use books with at most 429 leaves and input sequences capped at 14,000 tokens. The paper proposes enterprise settings involving potentially millions of URLs, but it does not demonstrate how the complete ToC would be represented, selected, or constrained at that scale. Hierarchical decomposition, multi-stage routing, and dynamic structure construction remain untested.

Finally, the paper describes constrained generation as restricting outputs to valid leaves, while the reported hallucination rate is nonzero at 0.05%. This discrepancy should be clarified: either the constraint is implemented approximately, applied only during part of inference, or the reported metric includes other forms of invalid output. Resolving this detail is important for reproducing the claimed structural-validity guarantee.

Conclusion

STAIR presents a structured model-based retriever that conditions an LLM on a document’s complete ToC and generates the relevant leaf section. Across 18 books and six domains, it achieves 82.6% Recall@1, exceeding fine-tuned DSI by 5.7 absolute points and outperforming BM25, DPR, Claude Haiku, and out-of-the-box Mistral. Its strongest evidence concerns low-supervision leaves and invalid-generation control: STAIR maintains nearly zero non-leaf predictions and degrades less severely than DSI when training coverage is sparse.

The paper’s contribution is therefore both methodological and infrastructural. It identifies document structure as an explicit retrieval signal, operationalizes retrieval as constrained generation over semantic section units, and releases a benchmark for evaluating this setting. The principal unresolved question is whether the same gains persist when structures are automatically induced, queries are naturally authored, and the candidate hierarchy is substantially larger than that of the evaluated books.