DEV Community

Multigrid
Multigrid

Posted on • Originally published at multigrid.ai

Running BGE Embedding Models Locally

BGE is the embedding family from the Beijing Academy of Artificial Intelligence, MIT licensed and free for commercial use. Getting vectors out of it takes four lines. Getting vectors that behave the way the model was trained to behave takes knowing two things the API does not tell you.

Install and run

pip install sentence-transformers

python - <<'PY'
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("BAAI/bge-base-en-v1.5")
vecs = model.encode(
    ["the cat sat on the mat", "a feline rested on a rug"],
    normalize_embeddings=True,
)
print(vecs.shape)          # (2, 768)
print(vecs @ vecs.T)
PY
Enter fullscreen mode Exit fullscreen mode

normalize_embeddings=True is not cosmetic. BGE was trained with a contrastive objective over normalised vectors, so cosine similarity is the intended metric; normalising once at encode time lets you use a plain dot product everywhere downstream, which is what every vector store is fastest at. Skip it and you will get inner products whose magnitude varies with text length.

The first call downloads to the Hugging Face cache. For an air-gapped machine, fetch once elsewhere, copy the cache directory, and set HF_HUB_OFFLINE=1 so a network attempt becomes a loud failure rather than a slow one.

Which size, and what each one costs

BAAI publishes three English v1.5 models, and the model cards give dimension and sequence length directly:

model                  dim    max seq   params   fp32 weights
bge-small-en-v1.5      384      512      33.4M     134 MB
bge-base-en-v1.5       768      512       110M     440 MB
bge-large-en-v1.5     1024      512       335M     1.34 GB
Enter fullscreen mode Exit fullscreen mode

The parameter counts are published; the weight column is the arithmetic, four bytes per parameter at fp32. The 512-token maximum sequence length is the one that catches people: it is inherited from the BERT backbone and it is a hard truncation, not an error. Feed a 3,000-token document and you silently embed its first 512 tokens. Chunking is not an optimisation here, it is a correctness requirement.

The dimension choice is a storage decision more than a quality one. Going from base to large multiplies both your index size and your query cost by 1.33 for a modest retrieval gain, and the RAM arithmetic decides it for you on a small machine.

BAAI also publishes Chinese (bge-*-zh-v1.5) and multilingual (bge-m3) members of the family, and they are separate models rather than configurations of these. The important consequence is that you cannot mix them in one index: vectors from bge-base-en and bge-base-zh occupy unrelated spaces even though both are 768 dimensions and both come from the same lab. If your corpus is mixed, you need one model that covers all of it, which is a different decision with its own arithmetic.

One further thing about the v1.5 suffix. The v1.5 models are not simply better-trained versions of v1; the release changed the similarity distribution deliberately, which is what the next-to-last section is about. A checkpoint pinned at v1 and a checkpoint at v1.5 therefore need different thresholds even at the same size, so record the exact revision alongside the index and not just the model name.

Verifying the pooling is CLS

A transformer produces one vector per token. Turning that into one vector per text is pooling, and different families made different choices: BGE uses the last hidden state of the first token, the [CLS] position. E5 and Nomic use a masked mean. If you reimplement encoding with raw transformers — which people do, to avoid the sentence-transformers dependency — and reach for mean pooling out of habit, you get vectors that look completely normal, have the right shape, and rank badly.

The check is to reproduce sentence-transformers by hand and compare:

import torch, torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel

tok = AutoTokenizer.from_pretrained("BAAI/bge-base-en-v1.5")
enc = AutoModel.from_pretrained("BAAI/bge-base-en-v1.5").eval()

batch = tok(["the cat sat on the mat"], padding=True, return_tensors="pt")
with torch.no_grad():
    out = enc(**batch).last_hidden_state

cls  = F.normalize(out[:, 0], dim=-1)
mean = F.normalize(out.mean(dim=1), dim=-1)

print(float(cls @ torch.tensor(vecs[0:1]).T))    # ~1.0
print(float(mean @ torch.tensor(vecs[0:1]).T))   # noticeably below 1.0
Enter fullscreen mode Exit fullscreen mode

The CLS vector matches what sentence-transformers gave you; the mean vector does not. That second number is the size of the mistake, and it is large enough to ruin retrieval while being small enough that nothing looks broken in a spot check.

The query instruction

BGE is asymmetric: it was trained so that a short query and a long passage that answers it land close together, which is a different objective from making two similar sentences land close together. To get the asymmetric behaviour, the query — and only the query — is prefixed with an instruction. The model card gives it verbatim: "Represent this sentence for searching relevant passages: "

INSTR = "Represent this sentence for searching relevant passages: "

doc_vecs   = model.encode(documents, normalize_embeddings=True)
query_vecs = model.encode([INSTR + q for q in queries],
                          normalize_embeddings=True)
Enter fullscreen mode Exit fullscreen mode

BAAI notes on the v1.5 cards that the instruction can be omitted with only slight degradation, and that it helps most for short queries against long documents. BAAI publishes the instruction and the guidance on the model card. What is not optional is consistency: index built without the prefix, queries sent with it, is the version of this that quietly halves your recall. Store the prefix policy next to the index, not in application code.

Why every score is above 0.6

Embed two unrelated sentences with BGE v1.5 and you will get a cosine similarity around 0.6 or 0.7. This surprises people who arrive with an intuition that 0.5 means “half similar”, and it is a documented property rather than a bug: the models were fine-tuned with a temperature of 0.01, which compresses the useful range of the score distribution into roughly 0.6 to 1.0.

BAAI states the consequence directly on the card — a similarity score greater than 0.5 does not indicate that two sentences are similar. The practical rules that follow:

  • Never hard-code a threshold from another model. A 0.75 cutoff that was sensible on a different family will pass almost everything here.
  • Rank, do not threshold, where you can. The ordering is what the model was trained to get right; the absolute value is an artefact of the temperature.
  • If you need a threshold, calibrate it. Embed a few hundred pairs you have judged yourself, plot the two distributions, and pick the cut. BAAI suggests trying 0.8, 0.85 and 0.9 as starting points, and that is a starting point rather than an answer.
  1. pip install sentence-transformers and encode with normalize_embeddings=True; confirm the shape is (n, 768) for the base model.
  2. Reproduce one vector by hand with CLS pooling and check it matches to ~1.0; check that mean pooling does not.
  3. Decide the prefix policy — instruction on queries only, or on neither — and record it with the index.
  4. Calibrate any similarity threshold against your own judged pairs before it reaches production.

Related

Top comments (0)