Top-K Retrieval

What it is

Top-k retrieval returns the k best-scoring chunks for a query instead of only the single best one. If k is 5, the five most similar chunks are concatenated into the prompt. The idea is that the right answer is usually near the top of the ranking but not always exactly at the top, and that many questions need more than one chunk to answer.

k is a dial between recall and noise. Larger k means more chance the needed chunk is included, but also more irrelevant text in the prompt, higher token cost and more opportunity for the model to be distracted by a near-miss.

On this site the first endpoint used k = 3. When I built the chat endpoint I raised it to 5 because questions like 'what has he done with AWS?' touch several jobs at once, and three chunks kept missing one of them.

Why it's needed

A single chunk rarely contains the full answer to a multi-part question. 'Where has Abdur worked and what did he build there?' spans five employer paragraphs. At the same time, dumping the whole knowledge base into every prompt wastes tokens and buries the signal. Top-k is the practical middle: enough context to answer, little enough to stay focused.

Getting k right is also one of the cheapest quality improvements available. It is a one-line change that you can evaluate in minutes, unlike swapping the embedding model or re-chunking everything.

How it works

• Score every chunk against the query (cosine similarity), sort descending and keep the first k.
• Start with k between 3 and 5 for small corpora with clean chunks; go to 8 to 10 when chunks are small or questions are broad.
• Make k depend on chunk size: five 500-word chunks is 2500 words of context, five 100-word chunks is 500.
• Add a minimum-score cutoff only after inspecting real score distributions; a fixed threshold from a blog post will not match your model.
• Filter before ranking when you have metadata: restrict to the right document, date range or tenant, then take top k inside that subset.
• For larger corpora, retrieve a generous candidate set (k = 20) and rerank with a cross-encoder down to the 3 to 5 you actually send to the LLM.

Code example

import numpy as np

def top_k(index: np.ndarray, query_vec: np.ndarray, chunks: list[str],
          k: int = 5, min_score: float | None = None):
    """index: (N, d) unit vectors; query_vec: (d,) unit vector."""
    scores = index @ query_vec
    order = np.argsort(scores)[::-1][:k]
    results = []
    for i in order:
        if min_score is not None and scores[i] < min_score:
            break                      # ranking is sorted, so stop early
        results.append({"text": chunks[i], "score": float(scores[i])})
    return results

# Quick evaluation harness: does the expected chunk land in the top k?
def recall_at_k(index, encode, chunks, test_set, k):
    hits = 0
    for question, expected_substring in test_set:
        hits_for_q = top_k(index, encode(question), chunks, k)
        if any(expected_substring in r["text"] for r in hits_for_q):
            hits += 1
    return hits / len(test_set)

# Compare k values on the same questions before changing production:
# for k in (3, 5, 8): print(k, recall_at_k(index, enc, chunks, tests, k))

Where I used this

Use cases

• Every RAG prompt assembly step
• Search result pages that show the best N matches
• Candidate generation before a reranking stage
• Few-shot example selection: pick the k most similar labelled examples for the prompt

Common mistakes

• Picking k once and never measuring it again after the knowledge base grows or the chunking changes.
• Setting k so high that the prompt exceeds the model's useful attention. Models answer worse when the relevant chunk is buried in the middle of a long context.
• Treating k as the only knob. Chunk quality and the embedding model matter more than k = 3 versus k = 4.
• Returning top k even when every score is terrible, so the LLM confidently answers from unrelated text.
• Sending the chunks in random order. Put the highest-scoring chunk first; models weight the start of the context more.

Interview questions