embeddinggemma-300m-code-8L-distill-int8

A distilled, pruned, int8 variant of google/embeddinggemma-300m for code retrieval (natural-language query → code). 8 of 24 transformer layers, 768-d, mean pooling. ~96% of the base model's code-search quality at ~3× the speed and ~1/6 the size.

Model Details

• Base model: google/embeddinggemma-300m
• Model type: bi-encoder sentence/code embedding
• Output: 768-dim, mean-pooled, L2-normalized
• Size: 235M params — 237 MB (ONNX int8) / 265 MB (MLX int8), vs ~1.6 GB base
• Languages: code (Python, Go, Java, JavaScript, PHP, Ruby) with English NL queries
• License: Gemma

Method

1. Curvature layer pruning — keep 8 of 24 transformer layers by diagonal-Fisher loss curvature (not first/last-N), following the loss-curvature memorization/reasoning spectrum (arXiv:2510.24256). Kept: [0,1,2,3,4,5,22,23]. This is the speedup (~3× fewer FLOPs/token).
2. int8 quantization — affine, group-wise (MLX) / dynamic (ONNX), on the transformer weights and the full token-embedding table (201M params, kept whole — it's a lookup, so int8 shrinks it 4× with no quality loss). This is the size reduction.
3. Distillation + fine-tune — self-distill the pruned backbone to the teacher's embeddings, then contrastive fine-tune on CodeSearchNet docstring↔code pairs (in-batch negatives). This recovers code-retrieval quality.

Usage

Do not add a task/instruction prefix to queries or documents. Unlike the base EmbeddingGemma (task: ... | query: ...), this model takes raw text. It is mean-pooled and has no Dense projection head, so a fixed prefix adds a constant vector that collapses query embeddings together and wrecks ranking (measured: CSN-python NDCG 0.76→0.42, ruby 0.42→0.05 with a prefix). Embed the bare query and the bare code.

MLX (Apple Metal) — uses the bundled mlx_encoder.py (bidirectional forward; required):

from mlx_encoder import load_encoder, embed_texts
from transformers import AutoTokenizer

model, layer_types, sw, hidden = load_encoder("mlx")
tok = AutoTokenizer.from_pretrained("mlx")
emb = embed_texts(model, layer_types, sw, hidden, tok, ["def add(a, b): return a + b"], max_len=256)

ONNX (CPU) — graph outputs token embeddings; mean-pool + normalize:

import onnxruntime as ort, numpy as np
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained(".")
sess = ort.InferenceSession("model.onnx")

def embed(texts, max_len=256):
    e = tok(texts, padding=True, truncation=True, max_length=max_len, return_tensors="np")
    h = sess.run(None, {"input_ids": e["input_ids"].astype(np.int64),
                        "attention_mask": e["attention_mask"].astype(np.int64)})[0]
    m = e["attention_mask"][:, :, None]
    v = (h * m).sum(1) / np.clip(m.sum(1), 1, None)
    return v / np.linalg.norm(v, axis=1, keepdims=True)

Evaluation

CodeSearchNet (Python test, 5000 corpus / 200 queries, gold = same index), Apple M5 Pro.

Fair same-backend comparison — this model vs the base, both int8, same backend, seq 256. The distillation lets the 8-layer model keep ~96% of NDCG while running ~3× faster:

backend	model	NDCG@10	MRR@10	R@1	docs/s	speedup
ONNX int8 CPU	this	0.902	0.881	0.830	83	3.0×
ONNX int8 CPU	embeddinggemma-300m	0.942	0.930	0.900	27	1×
MLX int8 Metal	this	0.910	0.890	0.845	485	2.8×
MLX int8 Metal	embeddinggemma-300m	0.946	0.935	0.905	175	1×

Per-language (CodeSearchNet, 5000/200 per language, NDCG@10):

	python	go	java	php	ruby	js	avg
this	0.958	0.939	0.822	0.843	0.763	0.759	0.847
base	0.972	0.926	0.935	0.935	0.878	0.723	0.895

Within 0.05 of the full base model on average; ahead on Go and JavaScript.

Limitations

• Code-search specialist. Weak on general text and instruction/QA retrieval; it has no Dense projection head or prompt support, so instruction prefixes degrade it (use none).
• Tuned on the 6 CodeSearchNet languages; other languages rely on the base model's coverage.
• int8 — small ranking differences vs the fp32 base.

Files

• model.onnx — int8 ONNX (CPU)
• mlx/ — int8 MLX weights + tokenizer
• mlx_encoder.py — bidirectional MLX forward (required for the MLX path)

References

• Loss-curvature memorization/reasoning spectrum (layer-pruning criterion): arXiv:2510.24256
• Base model: google/embeddinggemma-300m
• Eval/fine-tune data: CodeSearchNet

License

Gemma. This is a Model Derivative of google/embeddinggemma-300m.

Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms

Use and redistribution are governed by the Gemma Terms of Use and the Gemma Prohibited Use Policy, which carry through to all downstream users. See the NOTICE file.