google

/

embeddinggemma-300m

like 390

Follow

Google 28k

Hi, I'm having the same issue.
I'm testing this model on Korean retrieval benchmarks from MTEB (Ko-StrategyQA, AutoRAGRetrieval, PublicHealthQA, BelebeleRetrieval, MultiLongDocRetrieval, MIRACLRetrieval, MrTidyRetrieval) and this model shows very bad scores.

Also when evaluating especially on MultiLongDocRetrieval, MIRACLRetrieval, MrTidyRetrieval tasks (which contain long contexts), I see these message:

WARNING:mteb.evaluation.evaluators.RetrievalEvaluator:Found 33800 NaN values in the similarity scores. Replacing NaN values with -1.

I'm pretty much sure that I'm using the proper precision and model prompts. This is my code for model initialization:

model_prompts = {
    PromptType.query.value: "task: search result | query: ",
    PromptType.document.value: "title: none | text: ",
}

model = mteb.get_model(model_name, model_prompts=model_prompts, trust_remote_code=True, device=device)

And these logs are shown in the encoding step:

INFO:mteb.models.sentence_transformer_wrapper:Using prompt_name='query' for task=MultiLongDocRetrieval prompt_type=<PromptType.query: 'query'> with prompt='task: search result | query: '

INFO:mteb.models.sentence_transformer_wrapper:Using prompt_name='document' for task=MultiLongDocRetrieval prompt_type=<PromptType.document: 'document'> with prompt='title: none | text: '

Are there some problems?

• Youngjoon Jang

@LPN64 maybe your issue might be related to using float16 because this is stated in the README:

NOTE: EmbeddingGemma activations do not support float16. Please use float32 or bfloat16 as appropriate for your hardware.

@LPN64 maybe your issue might be related to using float16 because this is stated in the README:

NOTE: EmbeddingGemma activations do not support float16. Please use float32 or bfloat16 as appropriate for your hardware.

Well shit, I actually saw this sentence but I mixed up the var types, will update my post with new scores

I updated the code to :

    try:
        model = SentenceTransformer(
            embedding_model,
            trust_remote_code=True,
            device=device,
            model_kwargs={"dtype": torch.bfloat16},
        )
    except Exception:
        device = "cpu"
        model = SentenceTransformer(
            embedding_model,
            trust_remote_code=True,
            device=device,
            model_kwargs={"dtype": torch.bfloat16},
        )

but it doesn't help much

I updated the code to :

    try:
        model = SentenceTransformer(
            embedding_model,
            trust_remote_code=True,
            device=device,
            model_kwargs={"dtype": torch.bfloat16},
        )
    except Exception:
        device = "cpu"
        model = SentenceTransformer(
            embedding_model,
            trust_remote_code=True,
            device=device,
            model_kwargs={"dtype": torch.bfloat16},
        )

but it doesn't help much

1. Did you remove model.half()? Because even if you initialize the model with bf16, .half()would make it to fp16. reference:

>>> model = SentenceTransformer(model_name, model_kwargs={"torch_dtype": torch.bfloat16})
>>> model.dtype
torch.bfloat16
>>> model.half()
SentenceTransformer(
  (0): Transformer({'max_seq_length': 2048, 'do_lower_case': False, 'architecture': 'Gemma3TextModel'})
  (1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
  (2): Dense({'in_features': 768, 'out_features': 3072, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity'})
  (3): Dense({'in_features': 3072, 'out_features': 768, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity'})
  (4): Normalize()
)
>>> model.dtype
torch.float16
>>> 

2. Are you using the prompts for retrieval ?

• "task: search result | query: " should be used for encoding queries
• "title: none | text: " should be used for encoding documents

I removed .half() and moved away from encode() to use encode_query and encode_document from sentence-transformers

By prompts, that's what you meant, right ?

New results :

It's better but still very far from bge-m3

Do you have any ideas, @tomaarsen ?
This is my test code using mteb:

"""Benchmarking all datasets constituting the MTEB Korean leaderboard & average scores"""
from __future__ import annotations

import os
import logging
from multiprocessing import Process, current_process
import torch
import hashlib

from sentence_transformers import SentenceTransformer
from sentence_transformers.models import StaticEmbedding

import mteb
from mteb import MTEB, get_tasks
from mteb.encoder_interface import PromptType
from mteb.models.sentence_transformer_wrapper import SentenceTransformerWrapper
from mteb.models.instruct_wrapper import instruct_wrapper

import argparse
from dotenv import load_dotenv
from setproctitle import setproctitle
import traceback
import logging

load_dotenv() # for OPENAI

parser = argparse.ArgumentParser(description="Extract contexts")
parser.add_argument('--quantize', default=False, type=bool, help='quantize embeddings')
args = parser.parse_args()

logging.basicConfig(level=logging.INFO)

logger = logging.getLogger("main")

TASK_LIST_RETRIEVAL_GPU_MAPPING = {
    0: [
        "Ko-StrategyQA",
        "AutoRAGRetrieval",
        "PublicHealthQA",
        "BelebeleRetrieval",
        "MultiLongDocRetrieval",
    ],
    1: ["MIRACLRetrieval"],
    2: ["MrTidyRetrieval"],
}

model_names = [
    "google/embeddinggemma-300m"
]

save_path = "./RESULTS_DEV"

def evaluate_model(model_name, gpu_id, tasks):
    import torch
    try:
        device = torch.device(f"cuda:{str(gpu_id)}") 
        torch.cuda.set_device(device)
        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

        
        model = None
        if not os.path.exists(model_name):
            if "m2v" in model_name: 
                static_embedding = StaticEmbedding.from_model2vec(model_name)
                model = SentenceTransformer(modules=[static_embedding], model_kwargs={"attn_implementation": "sdpa"}, device=device)
            else:
                if model_name == "nlpai-lab/KoE5" or model_name == "KU-HIAI-ONTHEIT/ontheit-large-v1_1" or "KUKE" in model_name:
                    model_prompts = {
                        PromptType.query.value: "query: ",
                        PromptType.passage.value: "passage: ",
                    }
                    model = SentenceTransformerWrapper(model=model_name, model_prompts=model_prompts, model_kwargs={"attn_implementation": "sdpa"}, device=device)
                elif "snowflake" in model_name.lower():
                    model_prompts = {
                        PromptType.query.value: "query: ",
                    }
                    model = SentenceTransformerWrapper(model=model_name, model_prompts=model_prompts, device=device)
                elif "Qwen3" in model_name:
                    model = mteb.get_model(model_name, model_kwargs={"attn_implementation": "flash_attention_2", "torch_dtype": torch.bfloat16}, device=device, tokenizer_kwargs={"padding_side": "left"},)
                elif "frony" in model_name:
                    model_prompts = {
                        # PromptType.query.value: "Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery:",
                        PromptType.query.value: "<Q>",
                        PromptType.passage.value: "<P>",
                    }
                    model = mteb.get_model(model_name, model_prompts=model_prompts, model_kwargs={"attn_implementation": "sdpa"}, device=device)
                elif "gte-multilingual" in model_name or "nomic-embed" in model_name:
                    model = mteb.get_model(model_name, trust_remote_code=True, device=device)
                elif "gemma" in model_name:
                    model_prompts = {
                        PromptType.query.value: "task: search result | query: ",
                        PromptType.document.value: "title: none | text: ",
                    }

                    model = mteb.get_model(model_name, model_prompts=model_prompts, trust_remote_code=True, device=device)
                else:
                    model = mteb.get_model(model_name, device=device)
        else: # 직접 학습한 모델의 경우
            file_name = os.path.join(model_name, "model.safetensors")
            if os.path.exists(file_name):
                if "m2v" in model_name: # model2vec의 경우: 모델명에 m2v를 포함시켜주어야 model2vec 모델로 인식합니다.
                    static_embedding = StaticEmbedding.from_model2vec(model_name)
                    model = SentenceTransformer(modules=[static_embedding], model_kwargs={"attn_implementation": "sdpa"}, device=device)
                else:
                    model = mteb.get_model(model_name, model_kwargs={"attn_implementation": "sdpa"}, device=device)

        if model:
            output_folder_name = os.path.basename(model_name)
            if os.path.isdir(model_name) and len(output_folder_name) > 100:
                model_hash = hashlib.md5(model_name.encode()).hexdigest()[:6]
                output_folder_name = f"{output_folder_name[:93]}_{model_hash}"

            if os.path.isdir(model_name):
                try:
                    model.model_meta.name = output_folder_name
                except AttributeError:
                    logger.warning("Could not override model_meta.name. Path might still be too long.")
            
            setproctitle(f"{output_folder_name}-{gpu_id}")
            print(f"Running tasks: {tasks} / {model_name} on GPU {gpu_id} in process {current_process().name}")
            evaluation = MTEB(
                tasks=get_tasks(tasks=tasks, languages=["kor-Kore", "kor-Hang", "kor_Hang"])
            )

            if "multilingual-e5" in model_name or "KoE5" in model_name or "ontheit" in model_name or "KUKE" in model_name:
                batch_size = 512
            elif "jina" in model_name:
                batch_size = 8
            elif "bge-m3" in model_name or "Snowflake" in model_name:
                batch_size = 64
            elif "gemma2" in model_name:
                batch_size = 256 
            elif "Salesforce" in model_name:
                batch_size = 8
            else:
                batch_size = 256

            if args.quantize: # quantized model의 경우
                evaluation.run(
                    model,
                    output_folder=f"{save_path}/{output_folder_name}-quantized",
                    encode_kwargs={"batch_size": batch_size, "precision": "binary"},
                )
            else:
                evaluation.run(
                    model,
                    output_folder=f"{save_path}/{output_folder_name}",
                    encode_kwargs={"batch_size": batch_size},
                )
    except Exception as ex:
        print(ex)
        traceback.print_exc()

if __name__ == "__main__":
    torch.multiprocessing.set_start_method('spawn')
    
    for model_name in model_names:
        print(f"Starting evaluation for model: {model_name}")
        processes = []
        
        for gpu_id, tasks in TASK_LIST_RETRIEVAL_GPU_MAPPING.items():
            p = Process(target=evaluate_model, args=(model_name, gpu_id, tasks))
            p.start()
            processes.append(p)
        
        for p in processes:
            p.join()
        
        print(f"Completed evaluation for model: {model_name}")

Back at the office (a little more horsepower here) , here is gemma compared to others embeddings models on this Rag task :

@yjoonjang you should be able to simply load the model for MTEB evaluation with model = mteb.get_model("google/embeddinggemma-300m"), it's defined here: https://github.com/embeddings-benchmark/mteb/blob/729f20adbbcda328ac38528d27c7136d79f946da/mteb/models/google_models.py#L241-L259, and the model itself already stores the prompts for retrieval: https://huggingface.co/google/embeddinggemma-300m/blob/main/config_sentence_transformers.json#L19-L20

@LPN64 The encode_query and encode_document that you started using are indeed recommended, those take care of the prompts for you. I still think it's rather weird that this model ranks after a literal static embedding model though. In my tests on e.g. MIRIAD, the non-finetuned EmbeddingGemma was the #2 for <500M parameters after only snowflake-arctic-embed-m-v2.0: https://huggingface.co/blog/embeddinggemma#finetuned-evaluation

I also just ran some tests with https://huggingface.co/datasets/dwzhu/LongEmbed (i.e. needle tests) via MTEB, and results look strong there: #4 and #5 on LEMBSummScreenFDRetrieval and LEMBQMSumRetrieval out of all <500M models.

We did run into some potential issues in the implementation, so some questions that might help narrow it down:

• Are you using FA2? i.e. is flash_attn installed? pip show flash_attn?
• Are the chunk sizes indeed 512, never larger?

• Tom Aarsen

Thanks for your reply, @tomaarsen .

I've changed my code to mteb.get_model(model_name, device=device, model_kwargs={"attn_implementation": "flash_attention_2", "torch_dtype": torch.bfloat16}), and I confirmed that FA2 is installed.

One good news is that I do not get any warnings about NaN values anymore. However, I still have some retrieval performance problems.
To make it short, I'll report the nDCG@10 result for Korean task in BelebeleRetrieval (which is multilingual, human-translated benchmark)

Other benchmark results also show similar patterns (low retrieval scores for embeddinggemma on Korean retrieval)

• Youngjoon Jang

Sorry for the delayed answer, I wanted to came back with mode results (see at the bottom)

Are the chunk sizes indeed 512, never larger?

Technically it's 512 + needle size (50 max) at very max but yes.

Chunk sizes are in characters not tokens, it's (laziness and) to ensure it's "fair" between two different models tokenizers.

flash attention is indeed installed.

my benchmark, as described in my first post doesn't seem really .... hard, it's a simple needle in the haystack challenge.

But it includes multiple challenges described here (+sample)

The "subtle" dataset might seems really hard, so hard it may even sound stupid, but if any of the embeddings manages to correctly rank the needle, it means all others models are wrong when they fail to do so.

From memory bge-m3 is one of the very rare model to have a very good "resilience" to the difference size of query->doc chunk size difference

For the following experiment I ran it in all those chunk sizes : chunks_size = [512, 1024, 2048, 4096]

Full table :

Thank you for the excellent details! Given @yjoonjang 's results without fa2, I think it might not matter, but can you try without flash attention 2 @LPN64 ? We're investigating a bug here: https://github.com/huggingface/transformers/pull/40700

@yjoonjang , Google themselves have also evaluated on Belebele, you can find all results in: https://github.com/embeddings-benchmark/results/tree/main/results/google__embeddinggemma-300m/64614b0b8b64f0c6c1e52b07e4e9a4e8fe4d2da2

Here's Belebele: https://raw.githubusercontent.com/RyanMullins/mteb-results/c7359540dab61156ed875185c99df95a29e42e3f/results/google__embeddinggemma-300m/64614b0b8b64f0c6c1e52b07e4e9a4e8fe4d2da2/BelebeleRetrieval.json
Their reported results:

• kor_Hang-kor_Hang: 0.94142
• kor_Hang-eng_Latn: 0.95134
• eng_Latn-kor_Hang: 0.92925

If your scores is only kor_Hang-kor_Hang, then Google reports that EmbeddingGemma outperforms bge-m3 and snowflake-arctic-embed-l-v2.0 despite being much smaller. A finetuned EmbeddingGemma might then also outperform KURE-v1.

Somewhere, something is a bit off, it seems.

• Tom Aarsen

Oh, are you two on the correct transformers version? See https://github.com/huggingface/transformers/releases/tag/v4.56.0-Embedding-Gemma-preview

pip install git+https://github.com/huggingface/[email protected]

If you're on older versions of transformers, the bidirectional configuration options in this model will be ignored, and it will use causal attention instead of bidirectional attention like what it was trained for. Using the latest transformers GitHub, with

pip install -U git+https://github.com/huggingface/transformers/

I get:

• BelebeleRetrieval, kor_Hang-kor_Hang: 0.94242
• BelebeleRetrieval, kor_Hang-eng_Latn: 0.95205
• BelebeleRetrieval, eng_Latn-kor_Hang: 0.92919

And with the potentially patched version, with

pip install -U git+https://github.com/vasqu/transformers@fix-gemma-embedding-fa

I get:

• BelebeleRetrieval, kor_Hang-kor_Hang: 0.94228
• BelebeleRetrieval, kor_Hang-eng_Latn: 0.95205
• BelebeleRetrieval, eng_Latn-kor_Hang: 0.92919

Using this CLI command:

mteb run -m "google/embeddinggemma-300m" -t BelebeleRetrieval --verbosity 3 --batch_size 8 --languages kor

So those already match Google's results very closely. I think you two just need to upgrade transformers to the latest (or already to the upcoming patch. Note: the patch fixes some issues with FA2 and the sliding window on all attention implementations for inputs of 256+ tokens).

• Tom Aarsen

The transformers version was indeed the issue. Now I'm able to reproduce the results successfully.
Thank you so much @tomaarsen !

It would be helpful if the Google authors could specify the recommended transformers and torch versions in their documentation.

Great to see embeddinggemma performing so well !

Glad it worked! The https://huggingface.co/blog/embeddinggemma blogpost is a bit clearer in this regard: https://huggingface.co/blog/embeddinggemma#sentence-transformers

For future readers: Weird performance? Try to install transformers from source (pip install git+https://github.com/huggingface/transformers), you might be on an incompatible version.

• Tom Aarsen

Same conclusion, updating transformers boosts the scores a lot :

On the rag/needle in haystack benchmark it's "good", but not SOTA

I would consider the issue fixed as it's performances makes more sense now (bfloat + transformers branch) but still like a LOT of embeddings model, bad crosslanguage performances.