EmbeddingGemma 2: a test bench for replacing BGE-M3 in my RAG

In my Quant project, the RAG relies on BGE-M3 for embeddings, with a reranker right behind it. Both run in the same container on a g6f.xlarge: one eighth of an NVIDIA L4, a little under 3 GB of GPU memory. Every megabyte counts.

On 6 October, Google DeepMind released EmbeddingGemma 2. On paper it ticks several boxes I care about: smaller, shorter vectors, and a licence that raises no questions. Before touching the ml-runner, I want to know how it does on my own documents, in French.

What just came out

According to the announcement and the official documentation:

• 744 million parameters in total, but the text part alone is 271 million. The image and audio encoders are optional and can be switched off at load time.
• 768 dimensions by default, trained with Matryoshka: you can keep only the first 512, 256 or 128 values.
• 8,192 tokens of context, more than 100 languages.
• Text, image, audio, video and code in the same vector space.
• Apache 2.0 licence, weights on Hugging Face (google/embeddinggemma-2).
• Supported by transformers and sentence-transformers (version 6.1.0 at least), as well as vLLM, llama.cpp and Ollama.

On the other side, BGE-M3 has about 568 million parameters and produces 1,024-dimension vectors, with the same 8,192-token context.

The fp16 trap

The documentation is clear: EmbeddingGemma 2 loads in bfloat16 or float32, and float16 can produce NaN or degraded embeddings without any error message.

That is the point I must not miss, because my ml-runner loads BGE-M3 in fp16. Copying the existing configuration would give a model that answers, with wrong vectors. The L4 supports bfloat16, so the fix is simple, as long as you think of it:

import torch
dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float32

The POC

The idea: a small set of hand-written questions, each linked to the passage that holds the answer, and both models evaluated on exactly the same questions. I look at recall@5 (is the right answer in the top 5?), MRR@10 (at what rank does it show up?), encoding time and GPU memory used.

Two input files:

• docs.jsonl: one passage per line, {"doc_id": "...", "text": "..."}, chunked the same way as the current RAG;
• eval.jsonl: one question per line, {"query": "...", "doc_id": "..."}.

"""Test bench: BGE-M3 against EmbeddingGemma class="tok-num">2 on my own documents.

pip install -U "sentence-transformers>=class="tok-num">6.1.class="tok-num">0"
"""
import json
import time

import torch
import torch.nn.functional as F
from sentence_transformers import SentenceTransformer


def load(path):
    with open(path, encoding="utf-class="tok-num">8") as f:
        return [json.loads(line) for line in f]


docs, evals = load("docs.jsonl"), load("eval.jsonl")
doc_ids = [d["doc_id"] for d in docs]
gold = torch.tensor([doc_ids.index(e["doc_id"]) for e in evals])
cuda = torch.cuda.is_available()
dtype = torch.bfloat16 if cuda and torch.cuda.is_bf16_supported() else torch.float32  # never fp16


def score(q, d, k=class="tok-num">10):
    top = (q @ d.T).topk(k, dim=class="tok-num">1).indices
    hit = top == gold[:, None]
    recall5 = hit[:, :class="tok-num">5].any(class="tok-num">1).float().mean().item()
    rank = hit.float().argmax(class="tok-num">1) + class="tok-num">1                      # rank of the first right passage
    mrr = torch.where(hit.any(class="tok-num">1), class="tok-num">1.0 / rank, torch.zeros(len(rank))).mean().item()
    return recall5, mrr


def bench(name, encode_query, encode_doc, dims):
    if cuda:
        torch.cuda.reset_peak_memory_stats()
    t0 = time.perf_counter()
    d = encode_doc([x["text"] for x in docs])
    t_docs = time.perf_counter() - t0
    q = encode_query([e["query"] for e in evals])
    peak = torch.cuda.max_memory_allocated() / class="tok-num">2**class="tok-num">20 if cuda else class="tok-num">0
    for dim in dims:
class=class="tok-str">"tok-com">        # Matryoshka: keep the first class="tok-str">`dim` values, then normalise again.
        qd = F.normalize(q[:, :dim].float().cpu(), dim=-class="tok-num">1)
        dd = F.normalize(d[:, :dim].float().cpu(), dim=-class="tok-num">1)
        r5, mrr = score(qd, dd)
        print(f"{name:<class="tok-num">17} dim={dim:<class="tok-num">5} recall@class="tok-num">5={r5:.3f}  mrr@class="tok-num">10={mrr:.3f}  "
              f"docs={t_docs:.1f}s  GPU peak={peak:.0f} MB")


bge = SentenceTransformer("BAAI/bge-m3", model_kwargs={"torch_dtype": dtype})
enc = lambda t: bge.encode(t, convert_to_tensor=True, normalize_embeddings=True, batch_size=class="tok-num">16)
bench("bge-m3", enc, enc, [class="tok-num">1024])                       # no Matryoshka: class="tok-num">1024 only
del bge
if cuda:
    torch.cuda.empty_cache()

eg2 = SentenceTransformer(
    "google/embeddinggemma-class="tok-num">2",
    model_kwargs={"torch_dtype": dtype},
    config_kwargs={"vision_config": None, "audio_config": None},  # text part only
)
bench(
    "embeddinggemma-class="tok-num">2",
    lambda t: eg2.encode_query(t, convert_to_tensor=True),
    lambda t: eg2.encode_document(t, convert_to_tensor=True, batch_size=class="tok-num">16),
    [class="tok-num">768, class="tok-num">512, class="tok-num">256, class="tok-num">128],
)

A few choices worth explaining:

• encode_query and encode_document add the task prefixes Google provides (task: search result | query: for the question, title: none | text: for the passage). BGE-M3 uses none for dense retrieval.
• The image and audio encoders are switched off at load time. For a text RAG, that is about 470 million fewer parameters in memory.
• After truncation, vectors must be normalised again. sentence-transformers cuts a vector that was already normalised, and the part you keep no longer has a norm of 1.
• BGE-M3 was not trained with Matryoshka: cutting its vectors would make no sense, so I keep it at 1,024.

The calculation: what a dimension costs

In pgvector, a vector(d) takes 4 × d + 8 bytes: 4 bytes per value, plus a small header. For 100,000 passages, index not included:

Going from 1,024 to 256 dimensions divides storage by 4. The HNSW index and distance computation follow roughly the same ratio, since each comparison runs over d values.

On quality, Google publishes its own averages on the multilingual MTEB benchmark (v2): 61.36 at 768 dimensions, 61.17 at 512, 60.41 at 256 and 57.89 at 128. Between 768 and 256, you lose about one point for a third of the storage. Below that, the drop speeds up. These numbers come from Google and say nothing about my documents: that is exactly what the test bench is for.

One last point, from my article on evaluation pipelines: a score measured on n questions has noise with a standard deviation of √(p(1 − p)/n). With 100 questions and a recall@5 around 0.8, that is 4 points of standard deviation, so an interval of about ± 8 points. A 2-point gap between the two models will mean nothing. Hence two rules for reading the results: compare the models on the same questions, and look one by one at the questions whose verdict changes from one model to the other.

My decision criteria

For the switch to be worth it in my case, I would need:

1. a recall@5 at least on par with BGE-M3 at 256 dimensions, on the same questions;
2. no important question going from "found" to "missed" without an explanation;
3. a clearly lower GPU memory peak, to leave room for the reranker on the g6f.

And one cost to keep in mind: changing embedding model means recomputing every stored vector. BGE-M3 vectors and EmbeddingGemma 2 vectors do not live in the same space, so they cannot be mixed in one table.

Sources

• EmbeddingGemma 2, Google announcement (6 October 2026)
• EmbeddingGemma 2 model card: task prefixes, dtypes, MTEB scores
• transformers documentation, EmbeddingGemma2: encode_query, encode_document, truncation