FluidInference / embeddinggemma-2-coreml

EmbeddingGemma 2 (text + audio + image) — Core ML

The text, audio and vision encoders of google/embeddinggemma-2 converted to Core ML: text runs entirely on the Apple Neural Engine, audio and images on the GPU. Weights are the original ones (stored fp16 in the model; the token table is the original bf16). All three encoders are included (see Audio and Images); video uses the image encoder on frames.

Output: the same 768-d, L2-normalized embedding as SentenceTransformer("google/embeddinggemma-2").encode(...) (mean pooling over tokens, then normalization). Matryoshka truncation to 512/256/128 works as in the original: keep the first N values and re-normalize.

Audio (added): EmbeddingGemma2Audio.mlpackage turns 10 s of 16 kHz audio into 250 tokens in the text model's space; the text functions then embed them, so audio and text share one space (search audio with a text query).

Files

Functions

Every function has fixed shapes, which is what keeps it on the Neural Engine. Variable-length (enumerated) shapes put the whole graph on the CPU.

S ∈ {32, 48, 64, 128, 256, 512}. Inputs longer than 512 tokens must be truncated; on our test set capping at 512 tokens changed nothing measurable. At ≤ 512 tokens every sliding-window layer sees the whole sequence, so a padding mask is all the model needs.

pack_256 runs up to eight short texts in one call. Short calls are bound by streaming the weights from memory, so one 256-token call with eight texts costs about the same as two single calls.

Audio

EmbeddingGemma2Audio.mlpackage : waveform [1, 160160] fp32 (160 zero samples, then up to 10 s of 16 kHz mono, zero-padded) and frame_mask [1, 1000] (frame i valid when i*160 + 321 <= 160 + samples ) → audio_tokens [1, 250, 512]; token k is valid when frame_mask[4k] is. Build <bos> <|audio> tokens <audio|> <eos> (special tokens from embeddings.bf16 × √512, audio tokens as they are) and run embed_256 . Log-mel (Gemma 4 feature extractor: 20 ms Hann frames, 10 ms hop, 128 HTK mel bins, log(mag + 1e-3)) is computed inside the model.

Run it on the GPU ( cpuAndGPU ): the chunked local attention's 5-D blocks fall back to the CPU on the Neural Engine (46 ms per window there vs 13.6 ms on the GPU). fp16 drifts from fp32 on a few quiet tokens (window embedding cos ≥ 0.977, mean 0.9993); on a 1-hour earnings-call index, the top hit matched fp32 for 15/15 text queries and the top-5 overlap was 95%.

Images

EmbeddingGemma2Vision.mlpackage has one fixed-shape function per soft-token budget: vision_70 (630 patches), vision_140 (1,260) and vision_280 (2,520, the model's default). Resize the image keeping its aspect ratio so the sides are multiples of 48 px and it has at most 9 × budget patches of 16 px (HF Gemma4ImageProcessor ), scale RGB to [0, 1], cut it into patches row by row (each patch [16 rows][16 cols][RGB] ) and pad to the function's patch count. Inputs: patches [1, N, 768], positions [1, N, 2] (x, y; −1 for padding), position_embeddings [1, N, 768] ( table[0][x] + table[1][y] from position_embeddings.f16 ), valid [1, N] and pool [T, N] (1/9 for each patch in its 3 × 3 group). Output image_tokens [1, T, 512]; embed <bos> <|image> tokens <image|> <eos> with the text model.

Every op runs on the Neural Engine, but its time grows with the square of the patch count (full attention over all patches), so the GPU is faster at every budget. fp32 wrapper vs sentence-transformers: cosine 1.000000; fp16 token cosine ≥ 0.997 at 70/140 tokens.

Use

1. Prepend the task prefix (e.g. title: none | text: for documents, task: search result | query: for queries; all in config.json ).
2. Tokenize with tokenizer.json : <bos> (2), tokens, <eos> (1).
3. Look up each id's row in embeddings.bf16 (bf16 → float: shift the 16 bits left by 16), multiply by √512 ≈ 22.627417, pass as fp16.
4. Call the smallest embed_S that fits, or pack several texts into pack_256 .

A Swift implementation (tokenizer, table lookup, packing, downloads) is in FluidUse : EmbeddingGemma2Manager .

Requires macOS 15 / iOS 18 (multifunction models). The first load on a device compiles all seven functions for the Neural Engine, which takes several minutes once; Core ML caches the result.

Numbers (M5 Pro, macOS 27)

The model card warns against fp16: activations reach ~2,000, so squaring them inside RMSNorm overflows fp16. Every RMSNorm here divides by the row's max-abs value first, which keeps the math in range (plain fp16 RMSNorm: cosine 0.95 to the reference; this conversion: 0.9996). The final normalization is rescaled the same way.

Demos ( FluidUse , M5 Pro)

Each runs as a short show: Sources/<Demo>/demo.sh , then Play.

License: Apache 2.0, same as the source model.