janihal/embeddinggemma-300m-8bit-hi-mlx

embeddinggemma-300m — MLX 8-bit (accuracy-matched)

A mixed-precision MLX quantization of `google/embeddinggemma-300m`: the 24 transformer blocks are 8-bit (affine, group size 32), while the token-embedding table, both Dense projection heads, and all RMSNorms are kept in float precision (fp16). Runs the full EmbeddingGemma pipeline (mean-pool → 2× Dense → L2-normalize), 768-dim output, Matryoshka truncation to 512/256/128 supported.

Why this variant? Plain 8-bit MLX (the sibling) quantizes the token-embedding table too, which distorts the embedding space enough to cost ~0.2 pp STS / ~0.5 pp retrieval. Keeping that table (65 % of the parameters) in fp16 closes the gap entirely — this model is statistically indistinguishable from full precision on the benchmarks below, matching the llama.cpp ggml-org GGUF. The price is size: it is larger than plain 8-bit, though still smaller than fp16.

Usage

pip install mlx-embeddings

import mlx.core as mx
from mlx_embeddings.utils import load

model, tokenizer = load("janihal/embeddinggemma-300m-8bit-hi-mlx")

QUERY = "task: search result | query: "
DOC   = "title: none | text: "

texts = [QUERY + "What is the capital of France?",
         DOC   + "Paris is the capital and most populous city of France."]

enc = tokenizer.batch_encode_plus(texts, return_tensors="mlx", padding=True)
out = model(enc["input_ids"], attention_mask=enc["attention_mask"])
emb = out.text_embeds                    # (2, 768), mean-pooled + L2-normalized
sim = float((emb[0] * emb[1]).sum())     # cosine similarity

Prompt templates for other tasks (from config_sentence_transformers.json): Retrieval-query → task: search result | query: · Retrieval-document → title: none | text: · STS → task: sentence similarity | query: · Classification → task: classification | query: · Clustering → task: clustering | query: .

How it was produced

mlx_embeddings.convert quantizes every Linear + the Embedding. To quantize the transformer blocks only, a custom class_predicate was passed to mlx.nn.quantize:

import mlx.nn as nn
from mlx.utils import tree_flatten
from mlx_embeddings.utils import fetch_from_hub, save_weights, save_config

model, config, tok = fetch_from_hub(get_model_path("google/embeddinggemma-300m"), lazy=True)
model.load_weights([(k, v.astype(mx.float16)) for k, v in tree_flatten(model.parameters())])

def keep_fp(p, m):                       # True  -> quantize
    if not hasattr(m, "to_quantized"):   return False
    if "embed_tokens" in p:              return False   # keep token table fp16
    if p.startswith("dense"):            return False   # keep Dense heads fp16
    return m.weight.ndim >= 2 and m.weight.shape[-1] % 32 == 0

nn.quantize(model, group_size=32, bits=8, mode="affine", class_predicate=keep_fp)
save_weights("embeddinggemma-300m-q8-blk16", dict(tree_flatten(model.parameters())))

Evaluation

Controlled comparison on one Apple M5 Max. Reference = the same model converted to MLX float32. Candidates are scored against it on two public benchmarks and on raw embedding agreement.

• —STS: STS-Benchmark test, 1379 pairs → Spearman(cosine, gold 0–5).
• —Retrieval: NFCorpus (BEIR) test, 323 queries × 3633 docs → nDCG@10.
• —mean cos vs FP32: mean cosine to the FP32 reference's vectors, over all 6714 texts.
• —Δ confidence intervals are paired bootstrap, 2000 resamples.

Accuracy

† The unsloth GGUF is built from Google's QAT checkpoint and omits the two Dense projection layers. Its output lives in a different vector space (hence ~0 cosine to the reference); not a like-for-like point.

Read: both Δ CIs for this model straddle 0 on both benchmarks — it tracks full precision as tightly as the ggml-org GGUF does. A run with the embedding table in fp32 (939 MB) reaches mean cos vs FP32 = 0.99992; storing it fp16 trades a little raw-vector fidelity for half the size with no measurable task cost.

Weight-space quantization error

Param-weighted relative RMSE of the de-quantized block weights vs FP32 (the embedding + Dense are not quantized here, so 0 error there):

Performance (Apple M5 Max, both stacks on the Metal GPU)

· llama.cpp latency measured through llama-server and includes a localhost HTTP round-trip. Its model load is ~3× faster than MLX (no Python import).

Read: keeping the embedding table in fp16 rather than int8 costs essentially nothing at inference — throughput and latency match the compact sibling (an embedding lookup is a gather, not a matmul). On this machine MLX is ~1.8× faster for bulk embedding and ~1.3–1.5× faster per query than the GGUF, at about half the RAM. Single-run numbers, ±10 %.

Compared to other embedding models

Same benchmark (STS-B test, NFCorpus test) and same Apple M5 Max. MLX rows run on MLX/Metal, GGUF rows on llama.cpp/Metal. docs/s = wall time to embed the 3633-document NFCorpus corpus; query latency is single-text, warm.

Read: this model matches FP32 / the GGUF on both tasks. Among the others, only Qwen3-Embedding-4B clearly out-retrieves EmbeddingGemma-300m (nDCG 40.8 vs ~39) — at ~25× the embedding time and 4× the RAM. Qwen3-Embedding-0.6B is a stronger pure-similarity model (STS 91.3) but a weaker retriever (nDCG 36.7) and ~5× slower. Qwen3-VL-Embedding-2B (multimodal) is weaker than EmbeddingGemma on text on both axes. For retrieval / RAG at this size, EmbeddingGemma-300m is the best accuracy per byte and per second.

(Nemotron-3-Embed-1B GGUF did not produce usable embeddings through this llama.cpp build and is omitted.)

Caveats

• —"FP32 reference" is the MLX implementation; the GGUFs run in llama.cpp, so a small cross-framework gap (~0.1–0.2 %) is folded into their numbers.
• —One benchmark pair (STS + one retrieval set). Not a full MTEB run.
• —Performance measured on M5 Max; ratios shift with hardware, batch size, and text length. llama.cpp likely has some tuning headroom (-fa, threads, ubatch).
• —Loadable only with mlx-embeddings, not vanilla sentence-transformers.

License & attribution

Derived from [`google/embeddinggemma-300m`](https://huggingface.co/google/embeddinggemma-300m) (Google DeepMind) by post-training weight quantization only — no fine-tuning.

Use is governed by the [Gemma Terms of Use](https://ai.google.dev/gemma/terms) and the [Gemma Prohibited Use Policy](https://ai.google.dev/gemma/prohibited_use_policy). This is a modified version of EmbeddingGemma; the same terms and use restrictions apply to this model and its outputs.

Quantization tooling: `mlx-embeddings` · MLX. GGUF baselines: `ggml-org/embeddinggemma-300M-GGUF`, `unsloth/embeddinggemma-300M-GGUF`.