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
- 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). - 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.
- 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 + tokenizermlx_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.