JacobLinCool/embeddinggemma-2-multimodal-late-interaction

EmbeddingGemma 2 Multimodal Late Interaction (B4-LI256)

JacobLinCool/embeddinggemma-2-multimodal-late-interaction adapts google/embeddinggemma-2 into a native sentence_transformers.MultiVectorEncoder supporting Text, Image, and 16 kHz Audio token-level Late Interaction (Mean-MaxSim, 256d per token with native 128d Matryoshka prefix truncation) without requiring trust_remote_code or custom wrapper code.

---

1. Key Highlights (B4 G4 Scale-Up Checkpoint)

• Cross-Modal Text-to-Image (Flickr8k T→I): Reaches 87.40% R@1 (+3.00 pp over zero-shot B0-256 = 84.40%, +2.80 pp over L4 B3-LI256 = 84.60%, and surpasses the 3× wider 768d base model B0-768 = 86.80%). With 4× UPGMA Hierarchical Cosine Token Pooling (B4-LI256-Pool4x), accuracy improves further to 88.20% R@1 / 99.60% R@10 / 94.23% nDCG@10 while cutting index storage by 75% (65.0 MiB → 16.3 MiB).
• Cross-Modal Text-to-Audio (Clotho T→A, 15–30s 16 kHz clips): Reaches 14.33% R@1 (2.7× zero-shot B0-256 = 5.33%) and 53.67% R@10 (+22.00 pp over B0-256 = 31.67%, +17.33 pp over L4 B3-LI256 = 36.33%), while completely eliminating document length bias (Spearman r = +0.0482, p = 0.4056).
• Scientific Document Retrieval (BEIR SciFact T→T, 5,183 corpus docs): Reaches 81.45% nDCG@10 / 91.76% R@10 (+4.36 pp higher nDCG@10 than the exact same checkpoint's 256d single-vector pooling B4-D256 = 77.09%), and 81.94% nDCG@10 / 92.16% R@10 under 4× Token Pooling (B4-LI256-Pool4x).
• Multi-Dimensional Token-Level MRL (128d & 256d): Explicitly supervised at both 128d and 256d per token. Truncating each output token vector vec[:, :128] and L2-renormalizing + applying 4× Token Pooling (B4-LI128-Pool4x) reduces index storage by 87.5% (8× compression) while achieving 87.00% (T→I R@1), 14.67% / 53.67% (T→A R@1 / R@10), and **81.48% (SciFact nDCG@10)**—outperforming uncompressed B3-LI256 across all three modalities at 1/8th the index footprint.

---

2. Architecture & Training Methodology

• Backbone: google/embeddinggemma-2 (EmbeddingGemma2Model, Gemma 4 300M multimodal encoder with native SigLIP-style vision tower, Conformer-style audio tower, and 24-layer shared text Transformer + pretrained 512→768 Matryoshka projection language_model.embedding_projection).
• Native MultiVectorEncoder Pipeline (modules.json):
  1. 0_Transformer: Merged EmbeddingGemma2Model backbone (module_output_name="token_embeddings", outputting post-MRL 768d per-token states in bfloat16) with a modality-aware Jinja chat template that applies official query/document prompts only to text inputs while keeping image (~266 tokens) and 16 kHz audio (~25.1 tokens/sec) inputs unprompted.
  2. 1_Dense: Frozen prefix-selector Linear(768 -> 256, bias=False) initialized to W = [I_256 | 0_256x512], preserving the first 256 Matryoshka channels of each token without adding random projection weights.
  3. 2_MultiVectorMask: Masks padding tokens across variable-length batches.
  4. 3_Normalize: Applies token-wise L2 normalization (module_input_name="token_embeddings").
• Similarity Function (similarity_fn_name="meanmaxsim"): $$S_{\text{LI}}(q, x) = \frac{1}{|Z_{d}(q)|} \sum_{i=1}^{|Z_{d}(q)|} \max_{1 \le j \le |Z_{d}(x)|} \langle Z_{d}(q)i, Z{d}(x)_j \rangle, \quad d \in {128, 256}$$
• G4 Scale-Up Joint LoRA Adaptation (B4):
  • Full-network "full_shared" LoRA (r=32, alpha=64, dropout=0.05) attached to all 168 linear layers (q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj) across the 24 shared text Transformer layers (19,922,944 trainable parameters = 2.6041%), while keeping vision_tower, audio_tower, embed_vision, embed_audio, and embedding_projection frozen, then merged into base weights for zero-overhead inference.
  • Trained for 600 optimizer updates (200 updates/modality, batch_pairs=16, 2 denoised hard negatives per query -> Q=16, D=48 rectangular contrastive matrix per step) on an NVIDIA RTX PRO 6000 Blackwell Server Edition (96 GB GDDR7, bfloat16) over 79,405 multimodal query–asset pairs (20,000 text pairs from MS MARCO + SciFact + NFCorpus, 35,000 Flickr8k caption–image pairs, 24,405 Clotho caption–audio pairs; < 1 Epoch single-pass regime) with Multi-Dimensional Token-Level MRL and 768d Teacher KL Self-Distillation: $$\mathcal{L}{\text{G4}} = \frac{1}{3}\sum{d \in {768, 256, 128}} \mathcal{L}{\text{InfoNCE}}^{\text{dense-}d} + \frac{1}{2}\sum{d \in {256, 128}} \mathcal{L}{\text{InfoNCE}}^{\text{LI-}d} + 0.15 \cdot \tau^2 \operatorname{KL}\left(\operatorname{softmax}(S{\text{LI-256}}/\tau) ,|, \operatorname{softmax}(\operatorname{sg}(S_{\text{dense-768}})/\tau)\right)$$

---

3. Quickstart (Text, Image & 16 kHz Audio + 128d MRL Truncation)

Precision Requirement: Always run EmbeddingGemma 2 in bfloat16 or float32. Never cast the model to float16 (fp16), which causes activation overflow (NaN/Inf).

pip install -U "transformers>=5.20.0" "sentence-transformers>=6.0.0" soundfile soxr pillow torch

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from sentence_transformers import MultiVectorEncoder

REPO_ID = "JacobLinCool/embeddinggemma-2-multimodal-late-interaction"
model = MultiVectorEncoder(REPO_ID, trust_remote_code=False)

# 1. Encode a text query -> shape [L_q, 256]
query = model.encode_query("A dog barking outdoors")

# 2. Encode candidate documents across Text, Image, and 16 kHz mono Audio
text_doc = model.encode_document("A dog makes a barking noise in the park")
image_doc = model.encode_document(Image.new("RGB", (224, 224), "white"))
waveform = np.zeros(16000 * 3, dtype=np.float32)  # 3s 16kHz mono float32 audio
audio_doc = model.encode_document({"array": waveform, "sampling_rate": 16000})

# 3. Compute asymmetric 256d Mean-MaxSim similarity
for modality, doc in [("text", text_doc), ("image", image_doc), ("audio", audio_doc)]:
    score = model.similarity([query], [doc])
    print(f"256d T->{modality} MeanMaxSim score: {score.item():.4f} (doc tokens: {doc.shape[0]})")

# 4. Optional 128d Token-Level Matryoshka Prefix Truncation (2x smaller vectors)
q_128 = F.normalize(torch.as_tensor(query)[:, :128], p=2, dim=-1)
d_128 = F.normalize(torch.as_tensor(text_doc)[:, :128], p=2, dim=-1)
score_128 = (q_128 @ d_128.T).max(dim=-1).values.mean().item()
print(f"128d MRL T->text MeanMaxSim score: {score_128:.4f}")

---

4. Experimental Benchmark Results (seed=42, Held-Out Evaluation)

All models were evaluated on identical held-out candidate sets (0% overlap with training data):

• Flickr8k (T→I): 500 test queries (caption_0) against 500 test images (candidate SHA-256: e064acc86e2ce4fb...).
• Clotho (T→A): 300 evaluation queries (caption_1) against 300 full-duration 16 kHz mono audio clips (candidate SHA-256: cb520321dd6bb7e6...).
• BEIR SciFact (T→T): 300 test queries against the full 5,183-document scientific corpus (candidate SHA-256: 9948d896b35ae4e7...).

4.1 Full Progression (B0 Zero-Shot → B2/B3 L4 200 steps → B4 G4 600 steps)

Method	Flickr8k T→I R@1 (%)	Flickr8k T→I R@10 (%)	Clotho T→A R@1 (%)	Clotho T→A R@10 (%)	SciFact T→T nDCG@10 (%)	SciFact T→T R@10 (%)	T→I Δ vs B0-256 (pp)	T→A Δ vs B0-256 (pp)	T→T Δ vs B0-256 (pp)
B0-768	86.80	99.60	6.00	33.67	83.67	92.24	+2.40 [+0.80, +4.40]*	+0.67 [-1.00, +2.67]	+1.52 [+0.57, +2.52]*
B0-256	84.40	99.40	5.33	31.67	82.15	90.12	+0.00 [+0.00, +0.00]	+0.00 [+0.00, +0.00]	+0.00 [+0.00, +0.00]
B1-LI256	83.00	99.60	5.33	23.33	78.59	89.34	-1.40 [-4.20, +1.40]	+0.00 [-2.67, +2.67]	-3.56 [-6.76, -0.46]*
B2-D768	86.60	99.20	8.33	40.00	82.30	91.02	+2.20 [+0.20, +4.40]*	+3.00 [+0.00, +6.33]	+0.16 [-1.10, +1.39]
B2-D256	85.60	99.20	7.33	41.67	80.70	89.06	+1.20 [-0.80, +3.40]	+2.00 [-0.67, +5.00]	-1.44 [-2.75, -0.29]*
B3-D768	86.40	98.80	7.67	41.67	81.37	90.69	+2.00 [-0.20, +4.20]	+2.33 [-0.33, +5.33]	-0.78 [-2.29, +0.65]
B3-D256	85.20	99.00	7.00	42.33	80.33	89.06	+0.80 [-1.40, +3.00]	+1.67 [-1.00, +4.67]	-1.81 [-3.15, -0.48]*
B3-LI256	84.60	99.40	9.00	36.33	80.79	90.79	+0.20 [-2.20, +2.60]	+3.67 [+0.67, +6.67]*	-1.36 [-4.01, +1.47]
B3-Rerank50	84.60	99.40	9.00	36.33	80.18	90.12	+0.20 [-2.20, +2.60]	+3.67 [+0.67, +6.67]*	-1.97 [-4.42, +0.63]
B4-D768	86.80	99.20	15.33	52.00	79.29	90.24	+2.40 [+0.00, +4.80]	+10.00 [+6.00, +14.00]*	-2.86 [-5.41, -0.48]*
B4-D256	86.00	99.00	14.67	52.33	77.09	88.18	+1.60 [-0.40, +3.60]	+9.33 [+5.33, +13.33]*	-5.05 [-7.38, -2.83]*
B4-D128	84.40	99.00	14.67	52.33	74.69	86.12	+0.00 [-2.60, +2.60]	+9.33 [+5.99, +13.33]*	-7.46 [-10.32, -4.69]*
B4-LI256	87.40	99.40	14.33	53.67	81.45	91.76	+3.00 [+1.00, +5.20]*	+9.00 [+5.33, +13.00]*	-0.69 [-2.94, +1.82]
B4-LI128	87.20	99.40	14.00	53.33	81.05	91.59	+2.80 [+0.40, +5.40]*	+8.67 [+5.00, +12.67]*	-1.10 [-3.59, +1.52]
B4-Rerank50	87.40	99.40	14.33	53.67	79.69	89.76	+3.00 [+1.00, +5.20]*	+9.00 [+5.33, +13.00]*	-2.46 [-4.79, -0.14]*

4.2 UPGMA Hierarchical Cosine Document Token Pooling (Pool1x / Pool2x / Pool4x)

Applying parameter-free agglomerative cosine clustering to document token vectors at index time reduces storage by 50%–87.5% while acting as a spatial/lexical denoiser that merges redundant background image patches and stopwords into clean semantic centroids:

Method & Pool Factor	T→I R@1 (%)	T→I nDCG@10 (%)	T→I Index Size	T→A R@1 (%)	T→A R@10 (%)	T→A Index Size	T→T nDCG@10 (%)	T→T R@10 (%)	T→T Index Size
B4-LI256-Pool1x	87.40	93.90	65.02 MiB	14.33	53.00	83.20 MiB	81.45	91.76	593.65 MiB
B4-LI256-Pool2x	87.40	93.93	32.51 MiB	14.33	52.67	41.64 MiB	81.55	91.69	297.07 MiB
B4-LI256-Pool4x	88.20	94.23	16.30 MiB	13.33	52.00	20.85 MiB	81.94	92.16	148.77 MiB
B4-LI128-Pool1x	87.20	93.90	32.51 MiB	14.33	53.33	41.60 MiB	81.05	91.59	296.82 MiB
B4-LI128-Pool2x	87.40	93.87	16.26 MiB	15.33	53.33	20.82 MiB	81.34	91.66	148.53 MiB
B4-LI128-Pool4x	87.00	93.77	8.15 MiB	14.67	53.67	10.43 MiB	81.48	91.37	74.38 MiB

---

5. Scope Boundaries, Caveats & Licensing

1. Matryoshka Prefix Caveat: The 256d and 128d per-token representations select the leading channels of google/embeddinggemma-2's pretrained 768d MRL projection (language_model.embedding_projection) and are explicitly supervised at the token level via compute_g4_joint_mrl_loss.
2. Evaluated Modalities: Verified on Text→Image, Text→Audio (16 kHz mono), Text→Text, and diagnostic Image→Text / Audio→Text. Direct Image↔Audio and Video retrieval were not evaluated and are not claimed.
3. Licensing & Attribution:
   • Base model google/embeddinggemma-2 is licensed under Apache-2.0.
   • Fine-tuning data comprises sentence-transformers/msmarco-bm25 (16,000 pairs), mteb/scifact (917 train pairs), mteb/nfcorpus (3,083 train pairs), jxie/flickr8k (7,000 training images = 35,000 pairs), and gijs/clotho (4,881 development/validation audio clips = 24,405 pairs). Because upstream Flickr8k and Clotho datasets carry academic / non-commercial usage terms, users should respect upstream dataset licenses when deploying this checkpoint in commercial environments. No raw evaluation or training media files are redistributed in this repository.