---
license: gemma
base_model: google/embeddinggemma-300m
language:
- en
pipeline_tag: feature-extraction
tags:
- embeddings
- code-retrieval
- code-search
- distillation
- quantization
- onnx
- mlx
- gemma3
---

# 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](https://arxiv.org/abs/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):

```python
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:

```python
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](https://arxiv.org/abs/2510.24256)
- Base model: [google/embeddinggemma-300m](https://huggingface.co/google/embeddinggemma-300m)
- Eval/fine-tune data: [CodeSearchNet](https://huggingface.co/datasets/code_search_net)

## 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](https://ai.google.dev/gemma/terms) and the
[Gemma Prohibited Use Policy](https://ai.google.dev/gemma/prohibited_use_policy),
which carry through to all downstream users. See the `NOTICE` file.