EmbeddingGemma 2 int4: text embedding model in 194 MiB

Text-only, 4-bit packed version of google/embeddinggemma-2. It keeps the text retrieval path only. The audio and vision towers are not included and were not evaluated.

Resident weights for the text path drop from 516.9 MiB (bf16) to 193.9 MiB (-62.5%). On the 1400-pair retrieval set below, NDCG@10 and Recall@10 stay within 0.08% and 0.20% of the bf16 model.

Scheme

Every Linear layer and the token embedding table is quantized to 4-bit codes with group size 16: per group of 16 input columns, an fp16 scale and an fp16 minimum, asymmetric min-max, no calibration data. Codes are bit-packed into uint8. This follows the group-wise setup in arXiv 2609.16391. RMSNorm weights, layer scalars and the other small tensors stay bf16.

Linears are dequantized one at a time during the forward pass, and the token table dequantizes only the rows it looks up, so the packed bytes are what stays in memory.

Usage

qpack.py is in this repo. It needs transformers==5.19.0 (the embedding_gemma2 architecture is not in earlier releases).

import torch, torch.nn.functional as F
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer
import sys

path = snapshot_download("ewin-reg/embeddinggemma-2-int4-text-embedding")
sys.path.insert(0, path)
from qpack import load_packed

tok = AutoTokenizer.from_pretrained(path)
model = load_packed(path, torch.bfloat16)

def embed(texts):
    enc = tok(texts, padding=True, truncation=True, max_length=1024, return_tensors="pt")
    with torch.inference_mode():
        h = model(input_ids=enc.input_ids, attention_mask=enc.attention_mask).last_hidden_state
    m = enc.attention_mask.unsqueeze(-1).to(h.dtype)
    return F.normalize(((h * m).sum(1) / m.sum(1)).float(), dim=-1)   # 768-dim, mean pooled

q = embed(["task: search result | query: how do plants make food"])
d = embed(["title: none | text: Photosynthesis converts light into chemical energy in plants.",
           "title: none | text: The stock market closed higher on Tuesday."])
print(q @ d.T)

Use the same prompt prefixes as the base model: task: search result | query: for queries and title: none | text: for passages.

load_packed builds the model on the meta device and attaches the packed tensors directly, so a full-size fp copy of the weights is never allocated.

Quality

Evaluation set: AtomicChat/embeddinggemma-2-GGUF-metrics (evalset/set.json), 1400 query/passage pairs (1000 general, 400 code). A passage counts as relevant when its doc id matches the query's. The reference is the original bf16 weights computed in fp32. Degradation is 1 - variant / reference; the interval is a paired bootstrap over queries (2000 resamples).

Other settings tried on the same set (fp32 compute): body INT4/g64 with table INT4/g64 reached 0.62% NDCG degradation (CI up to 0.99%), body INT4/g64 with table INT2/g64 reached 2.86% and fails a 1% budget. INT8 everywhere (g128) reached -0.04% but only cuts weights by 48%.

This is one retrieval set. It does not cover other tasks, languages beyond those in the set, or long documents beyond 1024 tokens.

Memory

Measured on a Kaggle CPU host (Linux, one fresh process per row, /proc/self/status, CUDA hidden). "Original" is the bf16 text weights built the standard way with transformers 5.19.0. Process RSS includes about 800 MB for Python and torch that no weight format changes, so the table also shows RSS minus that baseline.

Peak RSS (VmHWM): 2156 MB to 1808 MB at batch 8 x 512, and 1867 MB to 1454 MB at batch 1 x 2048.

Total process memory does not fall by 50%. The weights do (-62.5%), but the fixed runtime and the activation memory during the forward pass stay the same. Against AutoModel.from_pretrained on the full original checkpoint, which memory-maps the file and only touches the text pages, the RSS difference after a forward pass is small (1491 MB vs 1385 MB at batch 8 x 512). The saving here is mainly the size of the model on disk and in memory when the weights are held resident.

There is no KV cache to reduce. This is an encoder that processes the whole sequence in one pass. At 2048 tokens and batch 1, eager and SDPA attention used the same memory in the original model (1679 MB and 1684 MB RSS, measured in a separate earlier run).

Files

• model.safetensors: packed codes, fp16 scales and minimums, bf16 remainder (194 MiB)
• quant_config.json: bits, group sizes, embedding scale
• qpack.py: quantizer, packed save and the loader
• config.json, tokenizer files: copied from the base model

License

Apache-2.0, inherited from google/embeddinggemma-2. Check the base model card for its terms.

Inference optimization and efficiency techniques

What this package does to cut memory, what I measured, and what I did not try. Numbers come from the sections above.

Used and measured

Measured, no gain

• SDPA against eager attention at 2048 tokens and batch 1 used the same memory in the original model (1,684 MB and 1,679 MB RSS).
• There is no KV cache to shrink. This is an encoder that reads the whole sequence in one pass.

Not tried, so no claim

• A fused dequantize-and-multiply kernel. Each linear is expanded to bf16 before the matmul, so latency was not improved and I did not measure forward time for this card.
• torch.compile, ONNX or other runtime export.
• INT8 activations.

Process RSS after a forward pass fell by 19%, or 35% to 36% once the 800 MB Python and torch baseline is excluded, not by the 62.5% the weights shrank, because the runtime and the activations stay the same size.