Teaching a small model
to find the right document

Every retrieval method is a way of drawing that picture

Grep draws it with exact substrings — a query and a document either share a token or they don't. Everything upgrade from there is really just a better way of placing the query point and the document points so that relevant pairs end up close together and irrelevant ones end up far apart, under some distance you can compute fast.

The "small model" you train is the function that does the placing: text in, a fixed-length vector out. Train it well on your corpus's vocabulary and query style, and the geometry above does the retrieving — no LLM reasoning required at query time.

The bi-encoder (a.k.a. dense retriever, two-tower model)

One encoder, run twice: once over every document — offline, once, ahead of time — and once over the incoming query, online, per request. Each pass produces a vector; relevance is just the dot product (or cosine) between the two vectors. This is the architecture behind sentence-transformers, OpenAI/Cohere embedding endpoints, and models like bge-small, e5-small, gte-small.

The entire point of the two-tower shape is that the expensive half — encoding 2,000 documents — happens exactly once, whenever the corpus changes. At query time you run the encoder a single time and compare against a matrix you already have sitting in memory. This is why bi-encoders scale to query time, and it's the architecture you want if you're serving anything interactively.

The cross-encoder

Instead of embedding query and document separately, concatenate them into one input — [CLS] query [SEP] document [SEP] — and let the model attend across both at once, outputting a single relevance score directly. Because query tokens and document tokens interact inside every attention layer, cross-encoders are noticeably more accurate than bi-encoders at judging "does this document actually answer this query." The catch is in the name: there's no vector to precompute. You pay the full forward pass for every candidate, every time.

Retrieve with the bi-encoder, rerank with the cross-encoder

You don't have to choose. Nearly every production retrieval system is this exact pipeline: the cheap bi-encoder narrows 2,000 documents down to a shortlist — 20, 50 candidates — by cosine distance, and the expensive-but-accurate cross-encoder only has to score that shortlist. You get the cross-encoder's judgment quality at a fraction of its cost, because it never sees the 1,980 documents that were obviously irrelevant.

Contrastive fine-tuning

A pretrained embedding model already places text in a reasonable general-purpose space. Fine-tuning on your corpus nudges that space so that your queries land near your matching documents specifically — pulling the true match closer, pushing everything else away, one training step at a time. The loss almost every bi-encoder trainer uses is InfoNCE (in sentence-transformers, MultipleNegativesRankingLoss):

Read it as: given a query, its true document should score higher than every other document in the batch under the similarity function. Batches double as negative sampling — with a batch of 32 (query, doc) pairs, each query gets 31 free negatives for the price of one forward pass.

Where the (query, doc) pairs come from

With ~2,000 documents you almost certainly don't have enough real user queries to train on directly — which is exactly the gap synthetic data closes: prompt an LLM once per document ("write 2–3 realistic questions this document answers"), and you have a training set the size of your corpus × 2–3, generated in one batch job, for the cost of a few hundred completions. It's the same idea as the eval set already in this repo (eval/queries.json + eval/ground_truth.json) — except here the questions become training signal instead of scoring signal, so keep the two sets disjoint or your eval will just measure memorization.

Short answer: all of these, comfortably

At a few sentences per query and a few paragraphs per document, sequence lengths stay under 256 tokens — the memory-hungry part of a transformer (attention, which grows with sequence length squared) never gets a chance to matter. What dominates VRAM at this scale is just parameter count × Adam's bookkeeping. Rough order-of-magnitude, not benchmarked:

*mixed-precision fine-tuning, batch size 16–32, sequence length ≤ 256 — weights + gradients + Adam moments + activations. Leaves headroom on a 10 GB card even at the largest row.

The realistic constraint at your corpus size isn't compute — it's that a 335M-parameter model has far more capacity than 2,000 documents' worth of signal can steer. A small model (MiniLM/bge-small class) fine-tuned on good synthetic pairs will very likely outperform a large model fine-tuned on the same data, just because it overfits less. Start small; scale up only if the eval set says you're underfitting.

Skip the vector database

Approximate nearest-neighbor indexes (FAISS, HNSW, usearch) exist to make search over millions of vectors sub-linear. At 2,000 documents, brute-force cosine similarity — a single 2,000×384 matrix multiplied against one query vector — is a few hundred microseconds in NumPy, exact, and requires no index-building step at all. Reach for ANN only if the corpus grows past roughly 100k–1M documents; below that it's solving a problem you don't have yet.