STAIR: Table of Contents as a Generative Index for Long Document Retrieval

TL;DR

The paper studies long document retrieval for Retrieval Augmented Generation and argues that length-based chunking removes the semantic boundaries that make sections and subsections mutually exclusive. It introduces SearchTome, an 18 book benchmark across six domains whose retrieval targets are table of contents leaf entries, and STAIR, a Mistral based system that receives the whole table of contents and generates one valid leaf title. Across SearchTome, STAIR reaches average Recall at 1 of 82.6, Recall at 3 of 90.8 and nDCG at 3 of 87.5, while DSI, DPR, BM25 and out of the box Mistral receive 76.9, 68.7, 59.5 and 13.8 on Recall at 1 respectively. The ablations attribute much of the gain to reduced structural forgetting, with STAIR at a 0.05 percent nonleaf output rate.

Background positioning

This work sits at the intersection of a benchmark contribution and a method contribution. The benchmark angle is clear: SearchTome is presented as the first long-context retrieval benchmark with a clean table of contents and leaf-level gold identifiers. The method angle is also central: STAIR is not a new embedding architecture or a new reranker, but a structure-aware generative indexer. It refines the Differentiable Search Index paradigm by changing what the model sees and what it is allowed to output. The paper's most defensible claim is therefore not that long-context LLMs alone solve retrieval, but that a declared structural index, when exposed during training and inference, can make generative retrieval more stable and more accurate.

Problem and Motivation

Current retrievers commonly split long documents into fixed-length chunks because embedding models and sequence models prefer manageable passages. This is practical but semantically violent. A chapter boundary, a numbered section or a subsection title often carries more information than the surrounding text: it tells the reader which concept the passage belongs to and what it is not. When that signal is removed, chunks compete with each other on lexical similarity rather than on structural role. The paper notes that length-based chunks can lack coherent boundaries, which degrades retrieval quality in a way that cannot always be fixed by longer context, even though modern models can accept more tokens.

The second motivation is specific to model-based indexing. DSI, introduced by Tay and colleagues, compresses corpus knowledge into model parameters and generates document identifiers directly. This simplifies retrieval pipelines but places a heavy burden on training examples: if the model has to learn the corpus structure solely from query and identifier pairs, sections with few examples are harder to remember. The paper's intuition is that a table of contents is an already available externalization of that structure. A human looking at a table of contents can often identify the correct section without reading the whole book. STAIR asks the model to do the same thing: use the global structure as input and choose among valid leaf titles.

Core Chapter: SearchTome and STAIR as a Structure-Aware Retrieval Problem

Making the Table of Contents the Retrieval Target

The starting point is a formal distinction between the document and its structure. The paper considers a long document with a table of contents whose entries are organized as a hierarchy. A human can often answer a question by selecting a numbered subsection title before reading the body text. The conceptual figure in the paper shows this directly: for a question about plurality voting, a person can use the table of contents to pick the relevant section instead of scanning the entire book.

The benchmark target is the set of leaf entries in that hierarchy. The paper's notation is affected by PDF conversion, but the intended mathematical object is a collection of table of contents nodes and a leaf set. A table of contents node is any title entry in the structure, while a leaf node is an entry that has no child entry beneath it in the hierarchy.

Here $D$ denotes the long document being searched, ${\mathrm{ToC}}_D$ is the full table of contents for that document, $T_l$ is a candidate leaf title and $T_c$ is a potential child title. The relation $\mathrm{parent}$ captures the hierarchy among table of contents entries. The role of this set in the paper is to define the gold search space: SearchTome's retrieval outputs are not passages, not page ranges and not free text answers, but leaf titles that correspond to the relevant section content. If one removed the hierarchy and treated all titles as flat, the formula would collapse into a simple list, but then the model would lose the parent-child constraint that makes the leaf notion meaningful. If the table of contents is parsed incorrectly, however, the leaf set becomes the wrong output space, and the system can appear accurate on an ill-defined taxonomy while missing the actual section the user needs.

STAIR's Training Interface

SearchTome is constructed from opentextbooks across six domains: Education, Finance, Law, Medicine, Natural Sciences and Social Sciences, with three books per domain. For each book, the authors parse the PDF, extract the table of contents with pymupdf, and use Mixtral 8x7b to generate multiple questions per paragraph. Questions are split into train, development and test sets, and the gold retrieval target is the table of contents entry whose content answers the query. This construction is important because it creates a benchmark where the retrieval target is structural, not merely lexical.

STAIR's training input makes this structural target explicit. For each example, the model receives a fixed prompt $P$, the user query $q$, and the complete table of contents for the book. The target $y$ is the correct leaf node.

In this expression, $x_{\mathrm{STAIR}}$ is the full prompt supplied to the language model, $P$ is a constant instruction prompt, $q$ is the synthetic query for a paragraph, ${\mathrm{ToC}}_D$ is the complete table of contents for the book, $T_l$ is one member of the leaf set, and $y$ is the supervised output title. Compared with DSI, the crucial difference is not that STAIR produces a shorter answer but that the structure is visible during training and inference. DSI must discover the semantic alignment between queries and section identifiers through the examples alone. STAIR receives that alignment as a structural prior. The paper's training configuration reinforces this design: it uses LoRA adapters on Mistral Instruct v0.2 with rank 16 and alpha 32, sets maximum input length to 14000 tokens for STAIR and 512 tokens for DSI, limits output length to 64 tokens, and early stops on development Recall at 1 with patience 20. The input length asymmetry is not incidental; the table of contents is long, so STAIR's training regime assumes that the model can absorb the full structure. This also means the reported result should be read as ToC-aware long-input generation rather than a pure parameter count difference.

The pipeline figure in the paper summarizes the system in three stages. Stage 1 prepares synthetic query and answer data from the book. Stage 2 fine tunes a language model to map a query plus table of contents to the correct leaf. Stage 3 performs constrained generation so that the model cannot emit arbitrary titles. The design rationale is therefore layered: the benchmark defines a leaf-level task, the model input exposes the global structure, and the decoder enforces structural validity at output time.

A useful way to compare STAIR with DSI is to ask where each system stores information. DSI stores corpus semantics in model weights and must generate identifiers without being shown the hierarchy directly. STAIR stores a mapping from query plus hierarchy to leaf title in the fine-tuned adapter and can exploit the hierarchy every time it sees a new query. If the query is about a topic whose leaf has many training examples, both systems may do well. The interesting regime, and the one the paper stresses, is the opposite case: leaf nodes with few training examples. Without the table of contents, the model has to remember that a rarely queried section exists. With the table of contents, the leaf is still present in the input, even if the training examples associated with it are sparse.

Constrained Decoding over Declared Leaves

The final mechanism is inference-time structure enforcement. The paper states that STAIR restricts the output vocabulary to valid table of contents leaves. This can be formalized as a constrained argmax over the leaf set, where ${\pi }_{h}eta$ is the conditional generation probability of the fine-tuned model.

In this formula, $\hat{T}$ is the predicted leaf title, $T$ ranges over the valid leaf set ${\mathrm{LN}}_D$, $heta$ denotes the parameters after LoRA fine-tuning, $q$ is the current query, ${\mathrm{ToC}}_D$ is the document structure, and $P$ is the instruction prompt. Its role is to implement the paper's hallucination control claim: instead of allowing the language model to generate any string that resembles a section title, decoding restricts the candidate set to the leaves extracted from the table of contents. This is not merely cosmetic. A generative retriever can produce a plausible but invalid identifier if it confuses subsection labels or blends two nearby titles. Constrained generation prevents many of those errors by construction. If the leaf set is empty or malformed, however, the constraint becomes harmful: the model may be forced to choose a wrong leaf because the correct one is absent. If the gold target is truly ambiguous across multiple sections, single-leaf generation can also be too restrictive, since the model is asked to commit to one title even when a better retrieval unit would be a set of related leaves.

Experimental Evidence

Main Retrieval Quality

The headline results are reported in Table 4. The average rows across the six domains are reproduced below. Haiku is described in the paper as Claude haiku 4.5, and DPR uses NV-Embed-v2 out of the box with passage and query lengths set to 512 and 256 respectively.

The table shows a clean ordering among the strongest retrieval paradigms. Out of the box Mistral performs poorly at 13.8 Recall at 1, which is lower than lexical BM25 at 59.5. This is an important diagnostic result: simply using a large language model does not solve the task. The model needs both knowledge ingestion and task-specific training. DPR improves to 68.7 because it provides semantic matching, but it is not fine-tuned on SearchTome and therefore misses corpus-specific structure. DSI, which is fine-tuned and generates identifiers, reaches 76.9. STAIR, the only system in this comparison that sees the full table of contents, reaches 82.6. That is an absolute gain of 5.7 Recall at 1 points over DSI. The abstract describes the same comparison as a 7.4 percent improvement, which should be understood as the relative gap from DSI's 76.9 baseline rather than a percentage-point gain.

The paper reports Recall at 1, Recall at 3 and nDCG at 3 using BeIR, and the per-domain table in Table 4 shows the same ordering in every domain. In Education, STAIR reaches 83.3 against DSI's 76.2. In Finance, it reaches 82.8 against 78.1. In Law, it reaches 80.8 against 73.3. In Medicine, it reaches 86.1 against 82.1. In Natural Sciences, it reaches 80.6 against 76.2. In Social Sciences, it reaches 81.8 against 75.4. The authors also report a randomization test for retrieval systems and test the null hypothesis of no difference between DSI and STAIR at a 0.05 threshold for each of the six domains. The null is rejected for all domains, which strengthens the claim that the improvement is not a single-domain artifact.

Structure Reduces Hallucination and Helps Rare Leaves

The most interesting evidence is the ablation on training sparsity. The paper defines hallucination as generating a nonleaf node, which is an invalid document identifier. Figure 3 is described as showing that DSI's hallucination rate rises as the number of training examples per leaf falls, while STAIR's hallucination rate stays close to zero. The authors' interpretation is that STAIR has learned to generate outputs only from valid leaves, whereas DSI has to remember that a leaf exists and how to reach it. This is a low-data failure mode specific to generative indexing. If a section is rarely queried, DSI has fewer supervised signals that point to it. STAIR can still see it in the table of contents at inference time.

The paper also reports error rates in the qualitative analysis. Mistral's overall error rate is 86.20 percent, with 23.86 percent described as hallucinations and 26.81 percent of mistakes involving nonleaf predictions. DSI improves substantially to 24.31 percent error, with 3.25 percent nonleaf predictions, indicating that corpus-specific fine-tuning injects some structure into the model weights. STAIR has the lowest error rate at 18.67 percent and a nonleaf rate of 0.05 percent. These numbers are not all measuring the same denominator, and the Mistral split is somewhat ambiguous in the text, but the main comparison is clear: fine-tuning reduces invalid outputs, and constrained leaf generation reduces them dramatically.

Figure 4 makes the rare-leaf argument concrete. The paper states that for leaves with fewer training examples, the Recall at 1 difference between DSI and STAIR is much higher. As the number of training examples for a leaf increases, the gap narrows, but STAIR still achieves higher Recall at 1. This is exactly what the structural hypothesis predicts: when supervision is dense, a parametric index can learn the mapping from examples alone. When supervision is sparse, a visible structure acts as a memory aid. The benchmark is therefore not only a leaderboard artifact; it probes a regime where retrieval systems must cover the whole taxonomy, not only the popular topics.

Qualitative Evidence: Why the Structure Matters

Table 3 gives examples that clarify the ranking. For a query about how preschoolers react when caregivers belittle autonomous actions, the correct leaf is Initiative versus Guilt, which corresponds to the developmental stage named in the book. Mistral's output pattern suggests it may know the relevant idea but lacks the task constraint to choose a single valid leaf. BM25 selects a morally related section because of lexical overlap. DPR chooses a section about autonomy because the query contains the word autonomy, but it misses the specific developmental stage. DSI selects a plausible but incorrect title and, according to the paper's analysis, outputs a header that does not exist in the input table of contents. STAIR selects the gold leaf. The same partial ordering appears in other examples, including a nursing query about chronic illness and elderly activities of daily living, where STAIR recovers the section while BM25 and DPR pick more generic basic concepts.

These examples matter because they distinguish three different abilities. The first is knowledge: a model may know the answer. The second is task format: a model must produce one valid identifier. The third is structure: the identifier must correspond to a leaf that actually exists and maps to the right content. Mistral shows that raw knowledge is not enough. BM25 and DPR show that local similarity is not enough. DSI shows that fine-tuning is strong but can still invent or miss low-frequency headers. STAIR's design targets all three, especially the last.

Design Choices and What They Explain

The central design choice is the leaf unit. SearchTome asks the system to retrieve a leaf title rather than a chunk. This makes the benchmark directly sensitive to document organization. A length-based chunk might contain parts of two subsections, but a leaf title has a clear semantic boundary. The paper does not provide a formal proof that leaf titles are always the optimal retrieval granularity, but the benchmark design makes the assumption explicit. The second design choice is exposing the whole table of contents at inference. This is expensive in tokens, and the configuration makes that clear: STAIR uses a maximum input length of 14000 tokens. The method is therefore well matched to books and reports where the structure is compact relative to the body text but still rich enough to guide retrieval. The third design choice is constrained generation. This is the most direct defense against invalid identifiers. If the leaf set is correct, the constraint guarantees structural validity. The remaining question is semantic accuracy: the model may choose the wrong leaf, but it will at least choose a real one.

The comparison to RAPTOR is also revealing. RAPTOR builds a hierarchy through recursive embedding, clustering and summarization, then retrieves text at multiple abstraction levels using dense retrieval. STAIR instead treats the table of contents as an already available hierarchy and performs model-based indexing over leaf titles. The practical difference is where the structure comes from. RAPTOR creates a derived tree from the content; SearchTome and STAIR consume a parsed, authored tree. That makes the STAIR evaluation cleaner for studying the value of known structure, but it also limits applicability to corpora where the structure exists. For enterprise corpora or technical reports, the authored hierarchy may be present but noisy or incomplete. The paper acknowledges this by discussing dynamically induced table of contents structures as future work.

Evidence Quality and Limitations

The strongest evidence for the method is that STAIR improves over DSI while the only conceptual input difference is the table of contents, and the improvement appears across all six domains with domain-level significance testing. The low hallucination rate is also a mechanistic result, not just a performance number. If a system can generate only valid leaves, and the measured nonleaf rate is 0.05 percent, the paper has direct evidence that the structural output constraint is doing what it claims.

The weakest evidence concerns whether the comparison is a perfectly controlled ablation of ToC visibility alone. DSI is limited to 512 input tokens, while STAIR is allowed 14000 input tokens because it receives the whole table of contents. That is a design necessity, but it means the empirical result combines ToC visibility with a long-context input regime. A cleaner test would compare several variants: DSI with a compact learned identifier, STAIR with the full ToC, and an ablation where ToC information is injected into a similarly sized compact representation. The paper does not report such a length-controlled comparison. Still, the rare-leaf ablation and hallucination analysis support the interpretation that structure itself matters, because the gain is largest where the system must know that a section exists.

Another limitation is the gold data definition. SearchTome uses generated queries and leaf-level labels. The benchmark is therefore optimized for ToC retrieval, but it inherits whatever bias the generation process introduces. If Mixtral 8x7b produces questions that align well with section titles, the benchmark rewards systems that exploit those titles. If a query can be answered by multiple sections, forcing one leaf can make retrieval appear brittle even when a user would accept several. The paper's future-work note about iterative retrieval over leaf content is sensible for exactly this reason: real RAG often needs a set of relevant passages or multi-hop composition, not one label.

Finally, the current setting assumes that a global structure is already present. The authors state this limitation directly. The method is evaluated on books with parsed table of contents structures, not on arbitrary corpora where structure must be discovered. Extending STAIR to enterprise corpora with millions of URLs, to dynamically induced structures, or to zero-shot settings over unseen books is possible but nontrivial. In those settings, the hard part may no longer be choosing among valid leaves; it may be building the right leaf taxonomy in the first place.

Deeper Takeaway

The real contribution is a reframing of generative retrieval's output space. DSI asks a model to generate an identifier learned from data. STAIR asks the model to generate an identifier declared by the corpus structure. The latter changes the learning problem. Instead of discovering that a book has sections, the model can use the section inventory as a prior. Instead of avoiding malformed titles by hoping the language model has memorized them, the decoder can restrict the set of legal completions. The empirical payoff is a better average Recall at 1 than DSI and a much lower invalid-output rate, especially for low-frequency leaves. For engineers, this points toward a simple retrieval principle: when authoritative structure exists, expose it during training and enforce it during decoding. For researchers, SearchTome provides a testbed for whether model-based indexing can move beyond dense retrieval or sparse matching by treating the table of contents itself as an index.