EmbeddingGemma 2 : un banc d'essai pour remplacer BGE-M3 dans mon RAG

Google a publié EmbeddingGemma 2 le 6 octobre : 271 millions de paramètres pour la partie texte, des vecteurs de 768 dimensions qu'on peut couper à 256, licence Apache 2.0. Sur mon serveur, BGE-M3 partage un huitième de GPU avec un reranker. Le POC pour comparer les deux sur mes propres documents, et le piège du fp16.

Dans mon projet Quant, le RAG s'appuie sur BGE-M3 pour les embeddings, avec un reranker juste derrière. Les deux tournent dans le même conteneur, sur un g6f.xlarge : un huitième de NVIDIA L4, soit un peu moins de 3 Go de mémoire GPU. Chaque Mo compte.

Le 6 octobre, Google DeepMind a publié EmbeddingGemma 2. Sur le papier, il coche plusieurs cases qui m'intéressent : plus petit, des vecteurs plus courts, et une licence qui ne pose pas de question. Avant de toucher au ml-runner, je veux savoir ce qu'il donne sur mes documents à moi, en français.

Ce qui vient de sortir

D'après l'annonce et la documentation officielle :

  • 744 millions de paramètres en tout, mais la partie texte seule en fait 271 millions. Les encodeurs image et audio sont optionnels et se désactivent au chargement.
  • 768 dimensions par défaut, entraîné en Matryoshka : on peut garder seulement les 512, 256 ou 128 premières valeurs.
  • 8 192 tokens de contexte, plus de 100 langues.
  • Texte, image, audio, vidéo et code dans le même espace de vecteurs.
  • Licence Apache 2.0, poids sur Hugging Face (google/embeddinggemma-2).
  • Pris en charge par transformers et sentence-transformers (version 6.1.0 minimum), ainsi que vLLM, llama.cpp et Ollama.

En face, BGE-M3 fait environ 568 millions de paramètres et produit des vecteurs de 1 024 dimensions, avec le même contexte de 8 192 tokens.

Le piège du fp16

La documentation est claire : EmbeddingGemma 2 se charge en bfloat16 ou en float32, et le float16 peut produire des NaN ou des embeddings dégradés sans message d'erreur.

C'est le point à ne pas rater pour moi, parce que mon ml-runner charge BGE-M3 en fp16. Recopier la configuration existante donnerait un modèle qui répond, avec des vecteurs faux. La L4 gère le bfloat16, donc la correction est simple, à condition d'y penser :

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

Le POC

L'idée : un petit jeu de questions écrites à la main, chacune reliée au passage qui contient la réponse, et les deux modèles évalués sur exactement les mêmes questions. Je regarde le recall@5 (la bonne réponse est-elle dans les 5 premiers résultats ?), le MRR@10 (à quel rang arrive-t-elle ?), le temps d'encodage et la mémoire GPU consommée.

Deux fichiers en entrée :

  • docs.jsonl : un passage par ligne, {"doc_id": "...", "text": "..."}, découpé comme dans le RAG actuel ;
  • eval.jsonl : une question par ligne, {"query": "...", "doc_id": "..."}.
"""Banc d&#class="tok-num">39;essai : BGE-M3 contre EmbeddingGemma class="tok-num">2 sur mes propres 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  # jamais de 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                      # rang du premier bon 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 : on garde les class="tok-str">`dim` premières valeurs, puis on renormalise.
        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  pic GPU={peak:.0f} Mo")


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])                       # pas de Matryoshka : class="tok-num">1024 seulement
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},  # partie texte seule
)
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],
)

Quelques choix à expliquer :

  • encode_query et encode_document ajoutent les préfixes de tâche prévus par Google (task: search result | query: pour la question, title: none | text: pour le passage). BGE-M3 n'en utilise pas pour la recherche dense.
  • Les encodeurs image et audio sont désactivés au chargement. Pour un RAG texte, ce sont environ 470 millions de paramètres de moins en mémoire.
  • Après une troncature, il faut renormaliser les vecteurs. sentence-transformers coupe un vecteur déjà normalisé, et le morceau gardé n'est plus de norme 1.
  • BGE-M3 n'a pas été entraîné en Matryoshka : couper ses vecteurs n'aurait pas de sens, je le garde à 1 024.

Le calcul : ce que coûte une dimension

Dans pgvector, un vecteur vector(d) occupe 4 × d + 8 octets : 4 octets par valeur, plus un petit en-tête. Pour 100 000 passages, sans compter l'index :

Modèle Dimensions Octets par vecteur 100 000 passages
BGE-M3 1 024 4 104 ≈ 410 Mo
EmbeddingGemma 2 768 3 080 ≈ 308 Mo
EmbeddingGemma 2 256 1 032 ≈ 103 Mo

Passer de 1 024 à 256 dimensions divise donc le stockage par 4. L'index HNSW et le calcul des distances suivent à peu près la même proportion, puisque chaque comparaison parcourt d valeurs.

Côté qualité, Google publie ses propres moyennes sur le benchmark MTEB multilingue (v2) : 61,36 à 768 dimensions, 61,17 à 512, 60,41 à 256 et 57,89 à 128. Entre 768 et 256, on perd environ un point pour un stockage divisé par 3. En dessous, la chute s'accélère. Ces chiffres viennent de Google et ne disent rien de mes documents : c'est justement le rôle du banc d'essai.

Dernier point, qui vient de mon article sur les pipelines d'évaluation : un score mesuré sur n questions a un bruit d'écart type √(p(1 − p)/n). Avec 100 questions et un recall@5 autour de 0,8, cela fait 4 points d'écart type, donc environ ± 8 points d'intervalle. Un écart de 2 points entre les deux modèles ne voudra rien dire. D'où deux règles pour lire les résultats : comparer les modèles sur les mêmes questions, et regarder une par une les questions qui changent de verdict d'un modèle à l'autre.

Mes critères de décision

Pour que le remplacement vaille la peine dans mon cas, il faudrait :

  1. un recall@5 au moins équivalent à BGE-M3 à 256 dimensions, sur les mêmes questions ;
  2. aucune question importante qui passe de « trouvée » à « ratée » sans explication ;
  3. un pic de mémoire GPU nettement plus bas, pour laisser de la place au reranker sur le g6f.

Et un coût à ne pas oublier : changer de modèle d'embedding oblige à recalculer tous les vecteurs déjà stockés. Les vecteurs de BGE-M3 et ceux d'EmbeddingGemma 2 ne vivent pas dans le même espace, on ne peut pas les mélanger dans une même table.

Sources