erjigit17/embeddinggemma-300m-ane-coreml

<div align="center">

EmbeddingGemma-300M — Core ML / Apple Neural Engine build

    

768-dim sentence embeddings, running where a naive conversion of this model silently doesn't: the Apple Neural Engine.

</div>

A hand-converted Core ML build of `google/embeddinggemma-300m` that actually runs on the Apple Neural Engine (ANE) — not the CPU/GPU fallback most naive Core ML conversions of this architecture silently produce. 768-dimensional output, 8-bit linear-quantized weights, ~5.8 ms per embedding on an M4.

Quick facts

Good for fast, on-device RAG

Retrieval-augmented generation lives or dies on how much latency the embedding step adds to every query and every ingested chunk. This build embeds a passage in ~5.8 ms, entirely on-device, with no network round trip and no shared GPU/server queue to wait behind — for a RAG pipeline doing many small embedding calls per query (query embedding, re-ranking candidates, embedding freshly ingested chunks), that adds up to a pipeline where retrieval latency stops being the bottleneck.

It's built for exactly the asymmetric retrieval pattern RAG needs: a query task prefix for questions and a document task prefix for passages (see Usage prefixes) — the same mean-pooling architecture family as other strong open retrieval embedders, mapped onto hardware most RAG stacks never actually use for inference. If your RAG runs on Apple Silicon — a Mac server, a MacBook doing local-first retrieval, an on-device app — this is a way to get the embedding step off the CPU/GPU path entirely.

Jump to: Why this exists · Architecture notes · Usage · Code example · Confirming ANE residency · Limitations · Files · License

Why this exists

Converting this model to Core ML the naive way — tracing the model as-is and letting coremltools.convert() handle attention — produces a package that looks fine and loads without error, but silently falls back to CPU. The Apple Neural Engine compiler cannot compile the fused scaled_dot_product_attention op this architecture's attention module uses (ANECCompiler: ANECCompile() FAILED), and Core ML does not surface that failure anywhere you'd notice — MLModelConfiguration.computeUnits never errors, it just quietly runs slower. The only way to catch it is MLComputePlan, which reports the real per-operation compute device after ANE compilation.

This build reimplements attention as explicit matmul → +mask → softmax → matmul — mathematically identical to the model's published attention, but expressed in ops the ANE compiler accepts. The rewrite is verified bit-for-bit equivalent to the reference model by cosine similarity before conversion (see convert.py), not assumed. The result: ~5.8 ms per embedding on an M4, 98.8% of operations actually scheduled on the Neural Engine.

Architecture notes (why a naive conversion doesn't work here)

EmbeddingGemma-300M is Gemma 3's decoder architecture adapted into a bidirectional encoder — the attention rewrite has to account for all of the following, not just the missing-SDPA-support issue:

• —Bidirectional, not causal (use_bidirectional_attention: true) — no autoregressive mask.
• —Grouped-query attention: 3 query heads share 1 key/value head; each head is 256-wide even though the model's hidden size is 768 (head width is independent of hidden size in Gemma 3).
• —Alternating attention window: 5 of every 6 layers are sliding-window (radius 257, rope_local_base_freq: 10000.0); every 6th layer is full attention (rope_theta: 1000000.0) — two different RoPE frequencies depending on layer type.
• —QK-norm: RMSNorm applied per-head, after the head split, before RoPE.
• —Sandwich normalization: four RMSNorms per layer (input_layernorm → attn → post_attention_layernorm → +residual, then pre_feedforward_layernorm → MLP → post_feedforward_layernorm → +residual), not the usual two.
• —Embedding scale: token embeddings are multiplied by sqrt(hidden_size) (≈27.71) — already applied inside Gemma3TextScaledWordEmbedding; applying it again (an easy mistake) silently wrecks the output (cosine similarity ~0.18 against reference in an early, broken draft of this conversion).
• —Two extra Dense layers after mean pooling (768→3072→768, no bias, no activation between them) — part of the published model (sentence_transformers calls them 2_Dense/3_Dense), not an add-on.

convert.py is the exact, runnable script that produces this package from the original weights — read it before trusting any of the above, don't take the claims on faith.

Usage prefixes

The model uses different task prefixes for queries vs. documents — this is the model's own documented convention, not specific to this conversion, but getting it backwards silently hurts retrieval ranking rather than erroring:

• —Query: "task: search result | query: " + your text
• —Document: "title: none | text: " + your text

Inputs / outputs

• —input_ids: int32[1, 128] — token IDs, padded/truncated to exactly 128 (see tokenizer.json; configure padding to length 128 with pad id 0 and truncation to 128 — the raw file has neither set by default).
• —attention_mask: int32[1, 128] — 1 for real tokens, 0 for padding.
• —embedding: float32[1, 768] — L2-normalized sentence embedding.

import coremltools as ct
import numpy as np
from tokenizers import Tokenizer

tokenizer = Tokenizer.from_file("tokenizer.json")
tokenizer.enable_padding(length=128, pad_id=0, pad_token="<pad>")
tokenizer.enable_truncation(max_length=128)

model = ct.models.MLModel(
    "embeddinggemma-300m-8bit.mlpackage",
    compute_units=ct.ComputeUnit.CPU_AND_NE,
)

encoding = tokenizer.encode("task: search result | query: how do I use claude code")
out = model.predict({
    "input_ids": np.array([encoding.ids], dtype=np.int32),
    "attention_mask": np.array([encoding.attention_mask], dtype=np.int32),
})
embedding = out["embedding"][0]

Confirming it's actually running on the Neural Engine

MLModelConfiguration(computeUnits: .cpuAndNeuralEngine) never errors even when Core ML can't use the ANE for a given op — it silently runs that op on CPU instead. To confirm real ANE residency, compile the package and inspect it with MLComputePlan:

let compiledURL = try await MLModel.compileModel(at: packageURL)
let plan = try await MLComputePlan.load(contentsOf: compiledURL, configuration: config)
// walk plan.modelStructure's operations; plan.deviceUsage(for:).preferred per op

This package measures 2001/2025 operations (98.8%) actually scheduled on .neuralEngine.

Limitations

• —128-token fixed context. Long passages are truncated, not chunked. google/embeddinggemma-300m's native window is 512 tokens; this build trades window size for a smaller, faster fixed shape. The attention mask logic already generalizes to any sequence length (the alternating sliding-window mask isn't hardcoded to 128), so SEQ_LEN in config.py can be raised and convert.py re-run — but measure first: at 512 this build's median latency was ~29 ms, roughly 5x slower than at 128, because the full-attention layers (1 in every 6) scale quadratically with sequence length. This is a real trade-off, not a free option.
• —Apple Silicon only. Core ML has no Linux/CUDA equivalent. convert.py's attention rewrite is plain PyTorch until the final ct.convert() call, so the same module can run on CUDA directly if you need a portable version — only the last few lines would need to change.
• —8-bit only. 4-bit and 6-bit palettization (coremltools.optimize.coreml.palettize_weights, mode="kmeans", no calibration data) were tried and rejected: 4-bit produced cosine similarity ~0.41 against the 8-bit output (effectively random), 6-bit ~0.92 (still a real quality regression), and neither was faster — palettization here saves package size, not Core ML inference time, since Core ML still unpacks the palette table for the matmul.

Files

• —embeddinggemma-300m-8bit.mlpackage — the Core ML package.
• —tokenizer.json — the exact tokenizer snapshot used to verify this conversion.
• —convert.py — the full, runnable conversion script (rebuilds this package from unsloth/embeddinggemma-300m from scratch, including the cosine-similarity verification gate for both the query and document task prefixes).
• —config.py — SEQ_LEN and the two task prefixes, imported by convert.py. If you're also running a server built on this package, both must import from the same place — a mismatch between the compiled shape and what's actually sent fails silently.

License

Inherited from google/embeddinggemma-300m: the Gemma license. Read it before commercial use — it carries use restrictions the base Apache-2.0-style license on many other embedding models does not.