--- library_name: sentence-transformers pipeline_tag: feature-extraction base_model: google/embeddinggemma-2 license: other license_name: apache-2.0-with-upstream-dataset-attribution license_link: https://huggingface.co/google/embeddinggemma-2 tags: - sentence-transformers - multi-vector - late-interaction - colbert - meanmaxsim - multimodal - text - image - audio - embeddinggemma-2 - matryoshka datasets: - sentence-transformers/msmarco-bm25 - mteb/nfcorpus - jxie/flickr8k - gijs/clotho - mteb/scifact --- # EmbeddingGemma 2 Multimodal Late Interaction (`B4-LI256`) **`JacobLinCool/embeddinggemma-2-multimodal-late-interaction`** adapts [`google/embeddinggemma-2`](https://huggingface.co/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`). ```bash pip install -U "transformers>=5.20.0" "sentence-transformers>=6.0.0" soundfile soxr pillow torch ``` ```python 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`](https://huggingface.co/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.