Retrieval-augmented generation, visually

The first question anyone at work asked our docs bot was “how do I reset 2FA for a user?”, and it gave a clear, confident, well-formatted answer that described a settings page we removed in 2021. The answer was in the docs. It was in an old release note, which also counted as docs, and the retrieval step picked that over the current help article because the release note used the exact words “reset 2FA” and the help article said “remove a second factor”.

That’s retrieval-augmented generation in one story. The model did exactly what we asked: answer from the text we gave it. We gave it the wrong text.

RAG is everywhere this fall. OpenAI’s Assistants API, announced at DevDay a few weeks ago, has a retrieval tool built in, and every vendor I talk to has a “chat with your knowledge base” feature in beta. The idea has been around for a while (the name comes from a 2020 paper by Patrick Lewis and others at Facebook AI), but it’s only this year that the models got good enough at reading for it to be the default way to build anything that answers questions from your own content. This is how it works, drawn out, from building a small one for our internal docs and help center over the last couple of months.

Why not just ask the model?

A language model only knows what was in its training data, which stops at some date and never included your company’s runbooks. You can fine-tune a model on your docs, but fine-tuning is better at teaching style and format than at reliably teaching facts, it’s expensive to redo when the docs change, and the model still can’t tell you where an answer came from.

The other option is to paste the documents into the prompt. That works if the documents are small. Ours aren’t: about 1,400 help articles plus a few hundred internal pages, several million tokens in all. Even with the 128,000-token context GPT-4 Turbo announced this month, or the 200,000 tokens Anthropic announced for Claude 2.1 last week, it doesn’t fit, and if it did you’d be paying for all of it on every question.

So you find the few pieces of text that are most likely to contain the answer, and paste only those. That’s the whole trick. Everything else is about doing that finding well.

Chunking

First you cut the documents into pieces, because retrieval works on pieces, and because a 6,000-word article is a terrible thing to paste into a prompt to answer a one-line question.

How you cut matters more than anything else in this post, and nobody tells you. The simplest approach is fixed-size chunks, say 400 tokens each, with some overlap (we used about 50 tokens) so a sentence that straddles a boundary appears whole in at least one chunk. That works better than you’d expect, but it cuts tables in half, separates a heading from the steps under it, and produces chunks like “3. Click Save. 4. Repeat for each user.” that mean nothing on their own.

What worked better for us was cutting along the document’s own structure: split on headings, keep each section together if it’s under the size limit, split long sections on paragraphs, and then (this was the big one) prepend the article title and the heading path to every chunk. So a chunk starts with “Admin guide > Users > Removing a second factor” before its text. That one change made a bigger difference to answer quality than anything we did to the model or the prompt, because a chunk now carries the context it was cut out of.

Embeddings

Next, each chunk gets turned into a vector by an embedding model. We used OpenAI’s text-embedding-ada-002, which returns 1,536 numbers per input. Those numbers have no meaning individually. What matters is that texts with similar meaning end up pointing in similar directions in that 1,536-dimensional space. “Remove a second factor” and “reset 2FA” should land close together, which is the thing keyword search can’t do.

It’s cheap. Ada-002 costs $0.0001 per thousand tokens, and our roughly 12,000 chunks at around 400 tokens each came to about 4.8 million tokens, so embedding the entire corpus cost about 48 cents. Storing it is small too: 12,000 vectors × 1,536 dimensions × 4 bytes is about 74MB.

Similarity

At question time, you embed the question with the same model and look for the chunks whose vectors point in the most similar direction. The usual measure is cosine similarity, the cosine of the angle between two vectors:

cos(θ) = (a · b) / (|a| |b|)

which is the dot product divided by the product of the lengths. I spent a lot of my math degree on inner product spaces and it’s nice, for once, to see one doing honest work. It’s 1 when the vectors point the same way, 0 when they’re at right angles, and in practice the useful range for text embeddings is narrower than you’d think.

Here it is with three dimensions instead of 1,536. Say the question embeds to q = (0.9, 0.1, 0.3), chunk A is (0.8, 0.2, 0.4) and chunk B is (0.1, 0.9, 0.2).

• q · A = 0.72 + 0.02 + 0.12 = 0.86. |q| ≈ 0.954, |A| ≈ 0.917, so cos ≈ 0.86 / 0.874 ≈ 0.98.
• q · B = 0.09 + 0.09 + 0.06 = 0.24. |B| ≈ 0.927, so cos ≈ 0.24 / 0.885 ≈ 0.27.

A is the one you retrieve. OpenAI’s embeddings come normalized to length 1, so the lengths are all 1 and cosine similarity is just the dot product, which is one less thing to compute.

For 12,000 vectors you don’t need anything clever to search. Computing 12,000 dot products of length 1,536 is about 18 million multiply-adds, which a laptop does in a few milliseconds. We keep the vectors in a table in Postgres with the pgvector extension, mostly because we already run Postgres, and at our size an exact scan is fine. The approximate nearest-neighbor indexes (HNSW and friends) matter when you have millions of chunks, and they trade a little recall for a lot of speed.

Retrieval, and why it’s the hard part

You take the top k chunks (we use 6), put them in a prompt along with the question and instructions along the lines of “answer using only the context below, cite the article for each claim, and say you don’t know if the context doesn’t contain the answer”, and send that to the model.

The model part mostly just works. Where RAG systems fail is almost always before the model sees anything:

The right chunk isn’t retrieved. Embeddings are good at meaning and surprisingly bad at exact strings. Error codes, product SKUs, config keys and customer-specific jargon are where pure vector search falls down, because “ERR_4031” and “ERR_4013” embed almost identically. We now run a keyword search alongside the vector search (plain BM25, the same scoring I wrote about in the search engines post last year) and merge the two lists. That hybrid fixed most of our misses on identifiers.

The wrong chunk is retrieved, which is the 2FA story. Old release notes, deprecated pages and drafts are all “docs”. We now index only published, current help articles and runbooks marked as maintained, and store an updated date with each chunk so stale content can be filtered or at least shown as stale.

The right chunk is retrieved and ignored. There’s a paper from July, Lost in the Middle by Nelson Liu and others, showing that models use information at the start and end of a long context much better than information in the middle. With 6 chunks that’s a minor effect. It’s one reason I don’t think the new 128k and 200k context windows make retrieval obsolete, even setting cost aside: you can stuff in the whole manual, but that doesn’t mean the model reads all of it with equal care. And stuffing 128k tokens into GPT-4 Turbo at a cent per thousand input tokens costs about $1.28 a question, against a few cents for a RAG prompt.

There’s also a permissions problem, which is easy to miss in a demo. Our internal runbooks mention things customers shouldn’t see. If one index serves both an internal bot and a customer-facing one, the filter has to be applied at retrieval time, based on who’s asking, not by asking the model nicely not to repeat things. Anything you put in the context window, assume the user can get the model to show them. Retrieved text is also untrusted input from the model’s point of view, and a document with instructions in it is still a document with instructions in it, but that’s a topic for its own post.

What I’d tell someone starting

Spend your time on the documents and the retrieval, not the prompt. Build a set of real questions with known good answers before you build anything else (we have about 80 now), and look at what was retrieved for each one, not just the final answer. When the answer is wrong, the retrieved chunks tell you why nine times out of ten.

Our bot is still internal-only. It answers about two-thirds of our test questions well, cites the right article for most of those, and says “I don’t know” more often than it used to, which I count as an improvement. The release note about the old 2FA page is still in the help center, because somebody links to it from a forum post, but it’s out of the index.