I Built STAIR From Scratch: A Retriever That Returns Section Names, Not Chunk IDs

A full STAIR implementation (STructure Aware Information Retriever) rebuilt from the paper after its code went offline: the architecture, the trie that makes hallucination impossible, and what the numbers actually said

What I built: a retrieval system with no vector database. The index lives in a language model’s weights, and instead of returning chunk_0417 it returns "Brain Cells: Neurons", an actual section title from the book. A trie over the valid titles makes it structurally impossible for the model to cite a section that does not exist. Across every run I measured, the hallucination rate was exactly 0.00%.

The design comes from STAIR, short for STructure Aware Information Retriever, a 2026 generative retrieval paper. Its published code and SearchTome benchmark links are offline, so I reconstructed the whole of STAIR from the paper text. If you searched for a STAIR implementation and found nothing runnable, this post is what I wish had existed.

What this is not: an evaluation of STAIR. My training run was deliberately small, four epochs against the paper’s allowance of up to 200, on a different and smaller benchmark. The numbers describe what I built in two days, not what the method can do. Treat them as engineering notes.

What STAIR actually does

STAIR: STructure Aware Information Retriever is by Vineet Kumar, Meghanadh Pulivarthi, Vishwajeet Kumar, Jaydeep Sen, Riyaz Ahmad Bhat, and Sachindra Joshi at IBM Research.

STAIR’s starting point is a real problem. Standard retrieval-augmented generation (RAG) chops documents into fixed-length chunks. A 400-page textbook becomes 3,000 anonymous fragments, and the moment you do that you throw away something the author spent months building: the structure. Chapters, sections, subsections, a hand-made semantic map of the content, discarded because it does not fit the chunking pipeline.

STAIR keeps it. Instead of a vector database, the model is fine-tuned until it can generate the name of the section that answers a question:

Q: "What are the building blocks of the nervous system?"
A: "Brain Cells: Neurons"

That output is the retrieval result. There is no index to search, because the index lives in the model’s parameters. This lineage is called a Differentiable Search Index (DSI), and STAIR’s contribution is using the table of contents as the document identifier instead of an arbitrary number.

Their reported results on their SearchTome benchmark: 82.6% Recall@1, against 76.9% for a plain DSI, 68.7% for DPR, and 59.5% for BM25, with a hallucination rate under 0.05%.

Classic RAG vs STAIR, step by step

It is worth being precise about what actually changes here, because “generative retrieval” gets used loosely.

In classic RAG or vector search, you run a pipeline. Split the document into roughly 500-token chunks. Embed each chunk with a model like all-MiniLM or text-embedding-3. Store the vectors in FAISS, Pinecone, Chroma or pgvector. At query time, embed the question and return the nearest neighbours by cosine similarity. The unit you get back is chunk_0417, an offset into a file. It means nothing to a human, its boundaries were chosen by a character count, and a section that happens to straddle two chunks gets split by something that never read it.

In STAIR, there is no pipeline and no database at all. You take the book’s table of contents, fine-tune the model on the book’s text until the mapping from content to section name lives in its weights, and at query time the model generates the section name. The unit you get back is "Brain Cells: Neurons", a label a human author chose, with boundaries a human author decided.

Three differences matter in practice:

Classic RAGSTAIRWhere the index livesa vector databasethe model’s weightsWhat comes backan opaque chunk ida real section titleAdding a documentone insert, millisecondsretrain the model

That third row is the hidden cost, and I will come back to it.

I wanted to understand the method by building it. That is where the trouble started.

The STAIR artifacts are offline

The paper points its code and its SearchTome benchmark at an anonymous review-time URL.

$ curl -sI https://anonymous.4open.science/r/s_331/README.md
HTTP/1.1 410 Gone

410 Gone. No GitHub mirror, no HuggingFace dataset.

This is not a criticism of the authors. Anonymous review URLs expiring is a structural problem in how our field handles artifacts. It happens to good papers from good labs constantly, and the fix is institutional rather than individual. But it has one hard consequence: the original benchmark cannot be obtained, so nobody outside IBM can reproduce these numbers, including me.

The only surviving artifact is a community rebuild on HuggingFace: 6 books across 2 domains, against the original’s 18 across 6. A different, smaller benchmark. Every number below is measured on that one, and is not comparable to the paper’s.

Audit any benchmark you did not build

Before using the substitute, I checked something basic. Does every question’s gold answer actually exist in the corpus it is supposed to be retrieved from?

leaves = {r["leaf_id"] for r in load("leaf_docs.jsonl")}
orphans = [q for q in questions if q["leaf_id"] not in leaves]

Up to 21.8% of questions point at sections that are not in the corpus at all. No retrieval system can ever answer them.

This mattered more than it looked. Before I caught it, compgov appeared to be by far the hardest book at 47.5% R@1. On answerable questions only it scores 60.8%, which is mid-pack. The entire "difficulty" signal was an artifact of broken labels.

To be clear, this is a defect in the community rebuild, not in IBM’s original SearchTome. But it is a good reminder that a benchmark you did not build is a dependency you have not tested.

The STAIR architecture

Here is the whole system on one page: the two training tasks, and what happens at query time.

Read it in two halves.

Training runs once per book. Task 1, ingestion, shows the model a passage and asks for its section title. This is what memorises the corpus into the parameters. Task 2, retrieval, shows it the entire table of contents plus a question and asks for the same kind of answer. Both tasks share one target, the section title, and the loss is masked so only that span is scored. Without the mask the model spends most of its gradient learning to echo a 1,000-token table of contents back at you.

Inference is a single generate call. The table of contents and the question go in. The output passes through a trie that admits only real section titles. That last box is the reason the hallucination rate is zero rather than merely low.

Building STAIR, piece by piece

Four pieces, reconstructed from the paper’s description.

1. The table of contents becomes the index. PyMuPDF gives you the embedded bookmark tree, and the leaves are the retrievable units.

def build_tree(entries):
    """pymupdf gives flat [level, title, page]; rebuild the hierarchy."""
    roots, stack = [], []
    for level, title, page in entries:
        node = TocNode(level=int(level), title=" ".join(title.split()), page=page)
        while stack and stack[-1].level >= node.level:
            stack.pop()
        (stack[-1].children if stack else roots).append(node)
        if stack:
            node.parent = stack[-1]
        stack.append(node)
    return roots

2. Constrained decoding. The model may only emit token sequences that spell a real section title. I built a trie over every valid title and used it as a hard mask at each decoding step.

class TokenTrie:
    def next_tokens(self, prefix):
        """Tokens that may legally follow. Empty means dead end."""
        node = self
        for tok in prefix:
            node = node.children.get(tok)
            if node is None:
                return []
        return list(node.children)

# Handed to HuggingFace generate() as a hard constraint
model.generate(**enc, prefix_allowed_tokens_fn=lambda b, ids:
               trie.next_tokens(ids[prompt_len:].tolist()))

The model cannot invent a section. Not “rarely does”, cannot. Across every run, every epoch, and both model sizes I tried, the hallucination rate was exactly 0.00%. This part of the paper reproduces perfectly, and it is the part I would actually use.

3. Two training tasks, as described above, with loss masked to the target span.

4. Statistical testing. The paper uses a randomization test. Building it was the best decision I made, for reasons I will get to.

Sizing the context before buying hardware

The paper allows a 14,000-token input, which sounds like it demands serious hardware. Before renting cloud GPUs, I measured what the table-of-contents prompt actually costs:

Every table of contents fits in 4k. The largest, a 489-section music theory textbook, is 22% of the paper’s budget. The 14,000 tokens is headroom for much larger books, not a floor.

Note the shape of it: the cost scales with the number of sections, not the page count. A 900-page reference with coarse chapter headings can have a smaller ToC prompt than a 200-page textbook with fine-grained subsections.

So I tried the paper’s actual model on consumer hardware. Mistral-7B-Instruct-v0.2, LoRA at r=16 and alpha=32, quantized to 4-bit with bitsandbytes:

It fits in 5 GB. If you have been assuming DSI-style generative retrieval needs a datacenter, for a single book it does not, and the binding constraint is the number of section titles the model has to memorise rather than the context length. That is worth knowing before anyone budgets for this.

What it scored

Engineering notes on a deliberately small training run, not a measurement of the method.

Here is what sat between my run and the paper’s setup, in case you are tempted to read these as comparable:

• Four epochs against an allowance of up to 200. Fifty times less training.
• My R@3 was still climbing when I stopped, 52 to 62 to 68 to 70, monotonic across epochs. The curve had not flattened.
• 790 tokens per ingestion example against the paper’s 14,000-token budget. That cap was mine, not theirs, and I only caught it afterwards.
• One book, not eighteen, on a smaller benchmark where up to a fifth of the questions are unanswerable.
• Going from a 0.5B model to 7B alone moved R@1 from 15.9 to 38.5, so scale was still paying off when I stopped buying it.

Any one of those could account for a large share of the difference. Together they make my numbers uninformative about the method.

The lesson that nearly got away

Early on I saw the no-ToC variant beating the ToC variant by 5 points on a 60-question dev set, and briefly concluded something interesting had happened.

Then I ran the randomization test. p = 0.25. Fifteen correct answers against twelve. Three questions.

On the full 610-query test set the gap vanished entirely, and at R@3 it reversed in the ToC’s favour (+5.57, p = 0.0025). I had been one blog post away from publishing noise as a finding about somebody else’s work.

If you take one practical thing from this post: a 5-point difference on a small evaluation set is noise, and you will believe it if you do not test it. That is true of your own results, and doubly true when you are about to say something about someone else’s.

An open question about STAIR

One observation I cannot explain, offered as a question rather than a result.

I ran the model with no fine-tuning at all: base weights, table of contents in context, constrained decoding on.

In my runs the table of contents is worth 4.6x untrained, and nothing measurable once fine-tuned. One reading is that the two are redundant, that fine-tuning teaches the structure and showing it again adds little.

I want to be careful here. This comes from a model trained fifty times less than the paper specifies. A plausible alternative is that my model overfit its 475 examples, with training loss falling to 0.03, and learned to lean on memorised question-answer pairs rather than on the structure. That would be a property of my undertraining, not of the method. I genuinely do not know which it is, and settling it needs the full training run I did not do.

It did kill my own best idea, though. I had hypothesised that the expensive fine-tuning might be skippable, since constrained decoding gives the 0% hallucination rate for free. 11.50 against 43.50 says no. The training does the heavy lifting, exactly as the paper says.

The cost question worth asking

Setting my numbers aside, there is a design consideration worth naming for anyone evaluating DSI-style retrieval:

Because the index is the weights, adding one document means retraining. A vector database does an insert in milliseconds. For a corpus that changes, a product wiki, a docs site, a ticket queue, that is a serious operational constraint.

For a stable, high-stakes corpus, legal, regulatory or clinical reference, where you need the retriever to cite a real section and never invent one, the tradeoff looks very different.

What I would take into production

The constrained decoder, on its own.

The 0.00% hallucination rate is the most valuable property in this paper, and it comes from the trie rather than the fine-tuning. It is about 40 lines of code, it works with any causal language model, and it costs nothing.

You could put it on top of BM25, or on top of a vector retriever, and get a system that is structurally incapable of citing a section that does not exist. For legal, medical or regulatory search, that guarantee is often worth more than a few points of recall.

That combination is not what the paper is about. It is what I would build tomorrow, and I would not have thought of it without reading their work closely enough to rebuild it.

Takeaways

• Test your gaps before you believe them. A 5-point difference on 60 questions is noise. I nearly published it.
• Audit any benchmark you did not build. 21.8% of one book’s questions were unanswerable, masquerading as difficulty.
• Measure before you budget. The “14k token” requirement was about 1,200 tokens in practice, which moved this from a cloud project to a laptop project.
• A small replication says more about the replication than about the paper. Fifty times less training and eighteen times less text per example explains a lot.
• Constrained decoding is underrated. A guarantee beats an average.

The code

Everything described here is open source:

github.com/mithilai/stair-toc-retrieval

A complete, runnable STAIR implementation, which as far as I know is the only one available given the original artifacts are offline.

git clone https://github.com/mithilai/stair-toc-retrieval
cd stair-toc-retrieval

python -m venv .venv
.venv/Scripts/python.exe -m pip install torch==2.9.1+cu128     --extra-index-url https://download.pytorch.org/whl/cu128
.venv/Scripts/python.exe -m pip install -r requirements.txt

# Train with the table of contents in the prompt
python scripts/train_stair.py --book whole-child     --model mistralai/Mistral-7B-Instruct-v0.2 --four-bit --epochs 4

# Train the ablation: identical, minus the ToC
python scripts/train_stair.py --book whole-child --no-toc     --model mistralai/Mistral-7B-Instruct-v0.2 --four-bit --epochs 4

# Evaluate both against BM25, with significance testing
python scripts/evaluate_run.py --book whole-child     --model mistralai/Mistral-7B-Instruct-v0.2

What is in there:

• The full pipeline. PDF to table-of-contents tree to leaf sections to constrained generation, plus BM25, the metrics, and the randomization test.
• constrained.py, the trie. This is the piece worth stealing even if you never touch the rest. About 40 lines, works with any causal language model.
• Resumable training and sweeps. Checkpoints save every epoch, with the optimizer state and the adapter weights, because I learned that lesson the hard way when a battery died mid-run.
• 38 tests covering the logic where a bug produces plausible numbers rather than a crash, which is the failure mode that actually cost me time.
• Per-query predictions saved for every evaluation, so you can recompute comparisons without regenerating anything.

The benchmark books are open-licence textbooks, downloaded automatically, so you can reproduce every number in this post on your own machine.

If you run it at the paper’s full training budget, I would genuinely like to know what you get. Open an issue.

Credits

All credit for the STAIR method goes to Vineet Kumar, Meghanadh Pulivarthi, Vishwajeet Kumar, Jaydeep Sen, Riyaz Ahmad Bhat, and Sachindra Joshi at IBM Research, arXiv:2609.03874. The method is theirs and it is a good one. Everything described here as a shortfall is in my implementation, not in their work.

If the authors release the code and the SearchTome benchmark, I would happily re-run all of this properly and post an update.

Thanks also to the maintainer of the searchtome-subset community rebuild, without which there would have been no benchmark to measure against at all.

Built with PyTorch, HuggingFace Transformers, PEFT, bitsandbytes and PyMuPDF.