Sakura Gemma 4 E2B x EmbeddingGemma 2 - DualMode

One Gemma 4 E2B-it backbone for generation and EmbeddingGemma-2-distilled 768d retrieval. Research prototype. Community model. Not official from Google.

Version: Phase 3 / v0.2 (Phase 2 / v0.1 is kept in the git history of this repository).

In 30 seconds

What: one Gemma 4 E2B-it model that can chat and produce 768-dimensional search vectors for RAG, so no second embedding model has to be kept in memory.
How: the same backbone runs in two modes. Generate is the untouched Gemma 4. Embed uses bidirectional attention, a small LoRA (rank 16) and a 768d head: +3.56 M parameters (6.8 MB).
Quality: on external benchmarks the embed mode reaches 79 % (SciFact) and 71 % (NFCorpus) of the teacher's MRR; generation stays bitwise identical to plain Gemma 4.
Formats: this repository (Safetensors adapter + Python runtime) and a GGUF version (link becomes active when that repository is public) with a small llama.cpp patch.
Honest limit: research prototype. Not a drop-in for EmbeddingGemma 2 and not index-compatible with its vectors.

Positioning

Runtime backbone: google/gemma-4-E2B-it (the single shared model, loaded from the Hub).

Embedding teacher: google/embeddinggemma-2, used only for distillation. It is not merged into any weights and not required at runtime.

This is not a drop-in replacement for google/embeddinggemma-2 and is not index-compatible with vectors produced by it (see cross-space numbers below). Use it as a self-contained retriever: embed queries and documents with this model.

No claim of being the first of its kind is made. In the public implementations reviewed, no other Gemma 4 system was found that exposes both autoregressive generation and EmbeddingGemma-2-distilled dense retrieval from the same resident backbone.

What it is

Local RAG stacks usually keep two text models resident: a generator (Gemma 4) and a separate embedding model. DualMode reuses one Gemma 4 E2B-it backbone for both:

                        GENERATE MODE  (original Gemma 4 path, adapter fully bypassed)
Input -> Gemma 4 E2B-it backbone --+-- causal attention -> LM head
        (one resident model)       |
                                   +-- EMBED MODE
                                       bidirectional attention
                                       embedding-only LoRA (q_proj, o_proj, layers 15-33, rank 16)
                                       hidden state after layer 33  (hidden_states[34], before the final norm)
                                       mean pooling over valid tokens
                                       Linear 1536 -> 768 (no bias)
                                       L2 normalization   (MRL: 768 / 512 / 256 / 128)

Added parameters (Phase 3)

Component	Shape / scope	Parameters	BF16
Embedding head	Linear(1536 -> 768, bias=False)	1,179,648	2.25 MB
Embedding-only LoRA	rank 16, alpha 32, `q_proj` + `o_proj`, layers 15-33 (38 modules)	2,375,680	4.53 MB
Total overhead		3,555,328	6.78 MB

Relative to the stored Gemma 4 E2B-it checkpoint: +0.0694 % of 5.123 B total stored parameters, +0.1575 % of the 2.257 B dense text parameters.

Why only q_proj / o_proj and layers 15-33: Gemma 4 E2B shares K/V across its last 20 layers (layers 15-34 read the K/V of earlier layers), so k_proj/v_proj LoRA in those layers has no effect, and layer 34 lies after the pooled hidden state. The Phase 2 adapter (rank 8, q/k/v/o, layers 27-34) contained such dead modules; Phase 3 trains only modules that change the output.

Replacing a standalone EmbeddingGemma 2 text model (271.00 M parameters, 516.89 MB BF16) by this adapter saves about 267.4 M parameters / 510 MB of BF16 weights (744.37 M / 1,413 MB versus the full standalone checkpoint).

Results (Phase 3)

Everything below is teacher EmbeddingGemma 2 vs. this model on the same texts, queries with task: search result | query: , documents with title: <title> | text: . Both models see at most 128 tokens (the training length), so the teacher numbers are slightly below its published long-context numbers (SciFact teacher at 512 tokens: MRR 0.935).

External out-of-domain retrieval (never used for training or model selection)

Benchmark (500-doc subset)	Teacher MRR	This model MRR	MRR vs teacher	nDCG@10 teacher / this	nDCG vs teacher	Gate
MTEB SciFact test (300 queries)	0.8971	0.7097	79.1 %	0.9137 / 0.7413	81.1 %	STRONG (>= 70 %)
MTEB NFCorpus test (249 queries)	0.4609	0.3262	70.8 %	0.3099 / 0.1906	61.5 %	STRONG on MRR (nDCG 61.5 %)

SciFact detail (this model, student -> student): Recall@1 0.6233, Recall@5 0.8067, Recall@10 0.8733. Cross-space: student query -> teacher docs MRR 0.7190; teacher query -> student docs MRR 0.7517. NFCorpus: Recall@1 0.241, Recall@5 0.418, Recall@10 0.482.

For comparison, Phase 2 (v0.1): SciFact MRR 0.0380 (4.2 % of teacher). The earlier model was over-fitted to 400 training pairs; Phase 3 fixes this with a 12k-pair multi-domain distillation set.

The remaining gap is real: the model keeps roughly 71-79 % of the teacher's MRR, and nDCG@10 on NFCorpus is only about 61 % of the teacher. It is a research-grade retriever, not a teacher replacement.

Internal validation (1,280 held-out pairs from the 13 training sources, ranking against all 1,280 positives)

Direction	MRR	Recall@1	Recall@5
Teacher -> Teacher (upper bound)	0.9203	0.8891	0.9570
Student -> Student	0.8585	0.8109	0.9195
Student -> Teacher	0.8528	0.7953	0.9187
Teacher -> Student	0.8810	0.8391	0.9313

Generation is untouched

Regression run on this exact export (fresh unmodified google/gemma-4-E2B-it as reference, 6 chat prompts, greedy 24 tokens):

Check	Result
Token parity of `generate()` vs. unmodified model	6/6 identical
Max logit difference (last position, full vocabulary)	0.0
`embed -> generate -> embed` embedding difference	0.0
`generate -> embed -> generate` outputs identical	yes
18 alternating switching patterns (embed/generate)	0 embedding mismatches, 0 generation mismatches
Adapter save / reload	embeddings identical (max diff 0.0)
LoRA state after calls	disabled

The adapter is never merged into the backbone; generate() bypasses it completely.

Training (Phase 3)

Teacher: google/embeddinggemma-2 (revision 914f7f89142e33e77833254d9c9b90c3cef7303b), 768d embeddings cached.
Data: 11,954 train pairs + 1,280 validation pairs (query, positive) from public datasets, de-duplicated by query and passage; passages truncated to 900 characters, training length 128 tokens. Sources: Natural Questions, GooAQ, SQuAD, TriviaQA, MS MARCO (BM25 triplets), SPECTER, S2ORC (title-abstract), CodeSearchNet, StackExchange duplicates, Yahoo Answers, ELI5, GermanQuAD, WikiMatrix (German query -> English document). Domains: general knowledge, science, technical Q&A, code, German, cross-lingual, short queries, long documents. No SciFact, NFCorpus, ArguAna, FiQA or TREC-COVID data was used. Licenses follow the source datasets (some, e.g. MS MARCO and Yahoo Answers, carry research-use terms); check them before commercial use.
Stage 1 (epoch 1): representation distillation: cosine + smooth similarity-matrix (MSE) + MRL (512/256/128) losses.
Stage 2 (epochs 2-3): adds in-batch contrastive loss, student-query -> teacher-doc and teacher-query -> student-doc alignment. Source-homogeneous batches make in-batch negatives domain-near.
Model selection: by the internal validation set only. SciFact was logged as a monitor, NFCorpus was only evaluated at the end.
Candidates: rank 8 / layers 27-34 / q,k,v,o (SciFact MRR 0.684), rank 8 / layers 15-34 / q,o (0.699), rank 16 / layers 15-33 / q,o (0.710, chosen).
Diagnosis of v0.1: the 0.038 SciFact result was reproduced; no pipeline bug (masking, pooling, prefixes, padding, ids all verified). The cause was collapse of the embedding space (effective rank 32 vs. 140 for the teacher, mean pairwise cosine 0.91, positive-negative margin 0.027 vs. 0.196).

Important behavior notes

No BOS token in embed mode. The tokenizer in this repository does not prepend <bos> and the adapter was trained that way. (Runtimes that always prepend BOS, such as llama.cpp, must mask it in embed mode.)
128-token training length. Longer inputs work but are untested beyond 512 tokens; truncating to 128 tokens reproduces the reported numbers.
Prefixes: queries task: search result | query: , documents title: none | text: (or title: <title> | text: ). embed() adds the prefix automatically if missing.
bf16 numerics matter: a few short texts differ by cos ~0.97 between bf16 and fp16 reference runs. This is rounding noise of the bf16 backbone, not a bug.

Files

File	SHA256
`embedding_head.safetensors`	`A008DA414E71FBCD41F8D9244D9F633DAD0DF89861B727EE594F832286335B00`
`embedding_lora.safetensors`	`ED890947BB440AC4F0E86D6D099270BF8A52C6E258F9BBBBD53810E98E671FF0`

dual_mode_config.json, runtime/ (Python runtime: dual_mode_model.py, embedding_lora.py), tokenizer and config files. The Gemma 4 E2B-it base weights are loaded from google/gemma-4-E2B-it (pinned base commit 3e22461f65e89153144f8adb70e3b8c2cc9845a7); they are not duplicated here.

Quickstart

import torch
from runtime.dual_mode_model import Gemma4DualModeForConditionalGeneration

# download this repository (runtime/ + adapter) and keep `runtime` importable
model = Gemma4DualModeForConditionalGeneration.from_pretrained(
    "webmp3/Sakura-Gemma4-E2B-EmbeddingGemma2-DualMode", torch_dtype=torch.bfloat16, device_map="cuda")

# GENERATE (original Gemma 4 path)
msgs = [{"role": "user", "content": "Explain rotary position embeddings in one sentence."}]
inp = model.tokenizer.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", return_dict=True).to("cuda")
out = model.generate(**inp, max_new_tokens=60, do_sample=False)
print(model.tokenizer.decode(out[0], skip_special_tokens=True))

# EMBED (bidirectional, LoRA on, 768d, L2 normalized)
q = model.embed("How do dense vector representations work?", task="SearchQuery")
d = model.embed("Dense embeddings map text to points in a continuous vector space.", task="Document")
print(float(q @ d))
# Matryoshka: model.embed(text, truncate_dim=256)

Limitations

Trained on about 12k pairs for three epochs; broader or longer training should help. The OOD gap to the teacher (21-29 % MRR) is real.
Not index-compatible with EmbeddingGemma 2 vectors; cross-space retrieval is clearly worse than native student -> student retrieval for NFCorpus-like data.
Text only. The vision and audio towers of Gemma 4 are not part of the embedding path.
Memory numbers from v0.1 (AMD Radeon 8060S UMA, PyTorch ROCm: 516 MB VRAM/UMA and 251 MB RSS saved versus a standalone EmbeddingGemma 2 text model) were measured with the smaller Phase 2 adapter; the Phase 3 adapter is 3.4 MB larger in BF16, which is below that measurement's noise.
The GGUF version needs a small llama.cpp patch; stock llama.cpp cannot run embed mode (see the GGUF repository).

Citation

@misc{sakura_gemma4_dualmode_2026,
  author = {webmp},
  title = {Sakura Gemma 4 x EmbeddingGemma 2: DualMode for Generation and Dense Retrieval},
  year = {2026},
  publisher = {Hugging Face},
  howpublished = {\url{https://huggingface.co/webmp3/Sakura-Gemma4-E2B-EmbeddingGemma2-DualMode}}
}

Downloads last month: 36

Model tree for webmp3/Sakura-Gemma4-E2B-EmbeddingGemma2-DualMode

Base model

google/embeddinggemma-2

Finetuned

(27)

this model