Refactoring a retrieval pipeline through evaluation

How did I actually get in this mess, and background

For the past year my passion project has been developing Minne, which is a save-for-later and personal knowledge management system (PKMS). Since learning about the Zettelkasten method, I became quite enamored with the idea of having a personal knowledge database. I’ve tried a couple of different tools, but none really to my liking. While I truly believe a lot of the benefits of a PKMS comes from the manual process of creating links between knowledge and documents, by means of creating neural pathways and working through knowledge, I’ve learned that I’m not that enamored with PKMS. As a psychologist, I’m fascinated by how we offload memory to digital tools. So development began.

When designing Minne in 2024 I envisioned having content as a base layer, and then extracted knowledge entities as an abstraction of it. And that this higher level plane of knowledge could work well for searching and working on. That then these knowledge entities are based in and connected to the underlying documents/content that it was derived from. I think it came out quite nifty, and to some respect it mimics a human learning pattern.

  flowchart TB
subgraph Ingestion
A[Content] --> B[Chunks]
B --> C[Knowledge Entities]
end
subgraph Retrieval
C --> D[Graph Exploration]
C --> E[Vector + Text Search]
end

We humans learn these higher level abstractions connected to different entities which we then can reference and explore if we need more detail. That these second layer knowledge entities exists as a compressed version and that some fidelity loss occurs but that’s fine. Can’t overload the context anyway, or we lose quality.

Current retrieval pipeline

In the application we look through this knowledge entity database in two separate occurrences. For simplicity I call this slow search. As of writing, this retrieval pipeline consists of vector comparison, full text search, and exploration of the graph database.

  flowchart LR
Q[Query/Embedding] --> V[Vector Search]
Q --> L[Lexical Search]
V --> M[Fuse Scores]
L --> M
M --> D[Expand Graph]
D --> E[Attach Chunks]
E --> F[Rerank]
F --> G[Assemble]

• When we ingest new content, we run slow search so the LLM that creates new knowledge entities can link them to existing ones and work with a context-aware view of the graph.
• When we chat with an LLM in the application, it gets the result of slow search as context, which lets it answer in a knowledge-aware way and attach references.

While I like the current design, it’s not the fastest. We run several different database operations on quite a lot of data, using the knowledge entities as primary candidates. The project has started to see some real users, and I’d like to explore optimizing the retrieval pipeline. It should serve its purpose first, not just give me the pleasure of maintaining a complex retrieval pipeline.

Getting started

So to proceed I opted to set up a evaluation framework, measuring both speed and retrieval accuracy. At this point in time I’m not evaluating the actual LLM performance, but rather the accuracy of the slow search itself. Which means answering the question, is the answer to a question inside the data retrieved from our knowledge system?

Datasets

For datasets, I faced an interesting decision. What type of data should this “slow search” be optimized for. What are the users actually ingesting in the system? What am I ingesting? I’ve found I use it for an assortment of things, written notes, audio notes, links for reading later, some pdf documents. But creating a meaningful representative dataset would be hard. So even though it would have been preferred I opted to abandon creating my own dataset, and use established ones. I chose two:

• SQUAD v2, it acts as a simpler and lighter dataset. Shorter paragraphs.
• Natural Questions (dev). It’s a dataset on wikipedia articles. Some articles are really long and could prove to be challenging.

Initially, I stuck to these two. After tuning for these sets I tried adding hybrid search, which consequently lowered results. Which ran some bells about me not understanding something enough… So I eventually brought in some more tests in from the BEIR test suite. More on that those later.

Evaluation harness

For my current design, each evaluation means running each article through the ingestion pipeline to create the knowledge entities and relationships. This is a somewhat costly operation. To test retrieval precision effectively, I needed distractors. So for each question, I opted to ingest nine incorrect contexts along with the correct one. This meant that to answer the first 200 questions, I had to ingest a total of 2000 articles.

Considering the computational effort involved in converting the raw dataset to my knowledge design, I needed persistence. This allows me to decouple the ingestion cost from the retrieval experimentation. I can iterate on the retrieval pipeline against a static, pre-ingested set. If I need to change the ingestion logic, I can simply recreate the set. This setup also gives me the flexibility to modify parameters at runtime, like the number of snippets to pass, the total amount to retrieve from the initial search, or whether to enable re-ranking.

Initial findings

Initial runs on datasets showed the following results, which in my opinion are fairly respectable, except for latency. The majority of latency comes from the embedding and reranking, or so I thought! More detailed logging showed that for moderately sized datasets the majority of latency comes from the search and chunk loading. It worked fine for my smallish production dataset, but scaled really badly. Which makes it clear I need to do some heavy refactoring. So preliminary tests shows good hit@k1-5, but abhorrent latency.

Vector search

So, now we’ve got a decent enough evaluation pipeline. Let’s start by establishing a baseline.

Experimental setup

For all experiments, unless otherwise noted, we are using:

• Embeddings: nomic-ai/nomic-embed-text-v1.5
• Tokenizer: bert-base-cased
• Reranker: jinaai/jina-reranker-v1-turbo-en

Reranking surprise

I initially ran all my tests with a reranker enabled, assuming it would boost performance. However, when I finally ran a test without it, I was shocked.

The reranker was adding significant latency (~550ms per query) but actually lowered performance on the Natural Questions dataset.

Lesson Learned: Always verify your assumptions.

For the rest of this post, results are presented without reranking, unless otherwise specified.

Ablations

Here are the results for different vector chunk sizes and strategies.

Hybrid search

With the above results natural-questions-dev I was pleased with the details for tokenization, chunking & overlap. But I still wanted to explore hybrid retrieval. When I ran the first tests there was no clear improvement what so ever. Even some slightly degraded results. Which led me to reevaluate the choice of NQ-dev and Squad-v2 for my choice. It’s clear both these datasets respond favorably to dense queries. Whilst the user might ingest whatever content they wish to their deployment of Minne. I realized I needed to make my dataset more challenging. Some research and I came onto BEIR, and its datasets. I selected a sample of a series of datasets and wired them into my eval harness.

Results clearly demonstrated my pipeline were having significant issues with the more challenging datasets. Since I want a solution somewhat generalizable I did not want to tune for these sparse-favored datasets too much. I spent ample time running ablations tuning linear fusion parameters. Then I tried RRF and got improved results, with a increased simplicity. Easier to reason about, and fewer knobs.

I’m not at all particularly impressed by my current Hybrid search, it’s underperforming my expectations. And since I don’t want to tune for what I think might be more of an edge case, it will have to do for now.

Final comparisons

Clear improvements in latency compared to the initial implementation. Even with reranking off for the initial strategy, there would be over 1s latency for each query.

Fin

It’s been a long adventure, spanning several weeks. But it’s been mostly fun. I’ve been checked on my assumptions more than once, which I think is very healthy. The suggestion that I read more than once when researching this project was “Create your own dataset”. I get that now. Having my own dataset early on would have spared me significant time.

So to summarize. I went from this retrieval pipeline operating on chunks as well as a synthetic data layer.

  flowchart LR
Q[Query] --> V[Vector Search]
Q --> L[Lexical Search]
V --> M[Fuse w/ Linear Weights]
L --> M
M --> D[Graph Expansion]
D --> E[Attach Chunks]
E --> R[Reranker]
R --> G[Assemble]

To this:

  flowchart LR
Q[Query] --> V[Vector Search]
Q --> L[Lexical Search]
V --> R[RRF Fusion]
L --> R
R --> G[Result]

The new pipeline is easier to easier to reason about, easier to debug, having equal or better performance and significantly faster. I’m absolutely up for introducing more complexity again in attempts to boost results on sparse-favored datasets! Query rewriting, merging several queries, etc. But for now, and for Minne 1.0 I’m pleased.

Related posts

• Bundling ONNX Runtime in Rust with Nix, Docker and GitHub Actions - The dependency management challenges I faced when adding the reranker that this evaluation eventually removed.
• Using LLMs for psychological research - Another take on applying machine learning to real-world problems, this time analyzing student well-being.