README.md

---
license: apache-2.0
language:
- en
library_name: transformers
tags:
- sentence-transformers
- multimodal
- embeddings
- image-text
- audio-text
- retrieval
- 2DMSE
- matryoshka
pipeline_tag: sentence-similarity
model-index:
- name: multi-modal-embed-small
  results:
  - task:
      type: image-text-retrieval
    dataset:
      name: COCO
      type: coco
    metrics:
    - name: Image-to-Text R@1
      type: recall_at_1
      value: 41.88
    - name: Image-to-Text R@5
      type: recall_at_5
      value: 71.64
    - name: Image-to-Text R@10
      type: recall_at_10
      value: 82.16
  - task:
      type: audio-text-retrieval
    dataset:
      name: LibriSpeech
      type: librispeech
    metrics:
    - name: Audio-to-Text R@1
      type: recall_at_1
      value: 36.38
    - name: Audio-to-Text R@5
      type: recall_at_5
      value: 68.22
    - name: Audio-to-Text R@10
      type: recall_at_10
      value: 79.52
---

# multi-modal-embed-small

A compact multimodal embedding model that unifies text, image, and audio representations in a shared semantic space. Part of the [MoM (Mixture of Models)](https://huggingface.co/llm-semantic-router) family.

## Model Description

**multi-modal-embed-small** is a lightweight multimodal encoder (~120M parameters) supporting:

- **Text encoding** via MiniLM-L6-v2 (22M params)
- **Image encoding** via SigLIP-base-patch16-512 (86M params)
- **Audio encoding** via Whisper-tiny encoder (8M params)
- **Cross-modal fusion** via 2-layer transformer attention
- **2DMSE**: Two-Dimensional Matryoshka Sentence Embeddings for adaptive compute
- **MRL**: Matryoshka Representation Learning for flexible embedding dimensions

### Key Features

| Feature | Description |
|---------|-------------|
| **Embedding Dimension** | 384 (supports MRL truncation to 32, 64, 128, 256) |
| **Image Resolution** | 512×512 |
| **Audio Input** | Up to 30s, 16kHz (Whisper Mel spectrogram) |
| **Modalities** | Text, Image, Audio, Multimodal fusion |
| **2DMSE Support** | Early exit at any encoder layer |
| **Languages** | English |

## Installation

```bash
pip install torch transformers pillow safetensors
```

## Usage

### Load Model

Two checkpoint formats are available:
- `model.pt` (932 MB) - PyTorch format
- `model.safetensors` (1.35 GB) - SafeTensors format

```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer, SiglipModel, SiglipProcessor, WhisperModel, WhisperFeatureExtractor
from huggingface_hub import hf_hub_download

class MultiModalEmbedder(nn.Module):
    """Standalone multimodal embedder - no external dependencies."""
    
    def __init__(self):
        super().__init__()
        # Text encoder (384d, no projection needed)
        self.text_tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
        self.text_encoder = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
        
        # Image encoder (768d -> 384d projection)
        self.image_processor = SiglipProcessor.from_pretrained("google/siglip-base-patch16-512")
        self.image_encoder = SiglipModel.from_pretrained("google/siglip-base-patch16-512").vision_model
        self.image_proj = nn.Linear(768, 384)
        
        # Audio encoder (384d, no projection needed)
        self.audio_processor = WhisperFeatureExtractor.from_pretrained("openai/whisper-tiny")
        self.audio_encoder = WhisperModel.from_pretrained("openai/whisper-tiny").encoder
    
    def encode_text(self, texts):
        if isinstance(texts, str):
            texts = [texts]
        inputs = self.text_tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
        inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()}
        outputs = self.text_encoder(**inputs)
        embeddings = outputs.last_hidden_state.mean(dim=1)  # Mean pooling
        return F.normalize(embeddings, p=2, dim=-1)
    
    def encode_image(self, images):
        inputs = self.image_processor(images=images, return_tensors="pt")
        inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()}
        outputs = self.image_encoder(**inputs)
        embeddings = outputs.pooler_output
        embeddings = self.image_proj(embeddings)  # 768 -> 384
        return F.normalize(embeddings, p=2, dim=-1)
    
    def encode_audio(self, waveform):
        # waveform: numpy array or tensor at 16kHz
        if isinstance(waveform, torch.Tensor):
            waveform = waveform.squeeze().numpy()
        inputs = self.audio_processor(waveform, sampling_rate=16000, return_tensors="pt")
        inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()}
        outputs = self.audio_encoder(**inputs)
        embeddings = outputs.last_hidden_state.mean(dim=1)  # Mean pooling
        return F.normalize(embeddings, p=2, dim=-1)

# Load model
model = MultiModalEmbedder()

# Download and load trained weights
checkpoint_path = hf_hub_download(
    repo_id="llm-semantic-router/multi-modal-embed-small",
    filename="model.pt"
)
state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=False)

# Load text encoder weights
model.text_encoder.load_state_dict({
    k.replace("text_encoder.encoder.", ""): v 
    for k, v in state_dict.items() 
    if k.startswith("text_encoder.encoder.")
})

# Load image encoder and projection weights
model.image_encoder.load_state_dict({
    k.replace("image_encoder.vision_encoder.", ""): v 
    for k, v in state_dict.items() 
    if k.startswith("image_encoder.vision_encoder.")
})
model.image_proj.load_state_dict({
    k.replace("image_encoder.projection.", ""): v 
    for k, v in state_dict.items() 
    if k.startswith("image_encoder.projection.")
})

# Load audio encoder weights
model.audio_encoder.load_state_dict({
    k.replace("audio_encoder.encoder.", ""): v 
    for k, v in state_dict.items() 
    if k.startswith("audio_encoder.encoder.")
})

model.eval()
print("Model loaded successfully!")
```

### Text Embedding

```python
import torch.nn.functional as F

# Single text
text_embedding = model.encode_text("A photo of a cat")  # Shape: [1, 384]

# Batch of texts
texts = ["A fluffy orange cat", "A golden retriever dog", "A red sports car"]
text_embeddings = model.encode_text(texts)  # Shape: [3, 384]

# Compute similarity
similarities = F.cosine_similarity(text_embeddings[0:1], text_embeddings[1:], dim=-1)
print(f"Cat vs Dog: {similarities[0]:.3f}")
print(f"Cat vs Car: {similarities[1]:.3f}")
```

### Image Embedding

```python
from PIL import Image
import requests
from io import BytesIO

# Load image
url = "https://upload.wikimedia.org/wikipedia/commons/thumb/3/3a/Cat03.jpg/1200px-Cat03.jpg"
image = Image.open(BytesIO(requests.get(url).content)).convert('RGB')

# Get embedding
image_embedding = model.encode_image(image)  # Shape: [1, 384]
```

### Audio Embedding

```python
import torchaudio

# Load audio (16kHz)
waveform, sample_rate = torchaudio.load("speech.wav")
if sample_rate != 16000:
    waveform = torchaudio.functional.resample(waveform, sample_rate, 16000)

# Get embedding
audio_embedding = model.encode_audio(waveform)  # Shape: [1, 384]
```

### Cross-Modal Retrieval

```python
# Image-to-text retrieval
image = Image.open("cat.jpg").convert('RGB')
image_emb = model.encode_image(image)

captions = [
    "A cat sleeping on a bed",
    "A dog playing in the park",
    "A car driving on the highway",
]
text_embs = model.encode_text(captions)

similarities = F.cosine_similarity(image_emb, text_embs)
best_idx = similarities.argmax().item()
print(f"Best match: {captions[best_idx]} ({similarities[best_idx]:.3f})")
```

### Matryoshka Dimension Reduction

```python
# Full 384-dim embedding
full_emb = model.encode_text("Hello world")  # [1, 384]

# Truncate to smaller dimensions
emb_256 = F.normalize(full_emb[:, :256], p=2, dim=-1)  # 1.5x faster retrieval
emb_128 = F.normalize(full_emb[:, :128], p=2, dim=-1)  # 3x faster retrieval
emb_64 = F.normalize(full_emb[:, :64], p=2, dim=-1)    # 6x faster retrieval
```

## Architecture

```
┌──────────────────────────────────────────────────────────────┐
│                  multi-modal-embed-small                     │
├──────────────────────────────────────────────────────────────┤
│  Text Encoder:  MiniLM-L6-v2           (22M params, 6 layers)│
│  Image Encoder: SigLIP-base-patch16-512 (86M params)         │
│  Audio Encoder: Whisper-tiny encoder    (8M params, 4 layers) │
│  Fusion:        2-layer Transformer                          │
├──────────────────────────────────────────────────────────────┤
│  Output: 384-dim normalized embeddings                       │
│  2DMSE:  Layer 0-5 early exit support                        │
│  MRL:    32, 64, 128, 256, 384 dim truncation                │
└──────────────────────────────────────────────────────────────┘
```

## Training

### Training Data

| Modality | Dataset | Samples | Purpose |
|----------|---------|---------|---------|
| Image-Text | LLaVA-CC3M | 595K | Image-text alignment |
| Image-Text | COCO Captions | 25K | Validation |
| Audio-Text | LibriSpeech | 105K | Audio-text alignment |

### Training Stages

| Stage | Description | Trainable | Epochs |
|-------|-------------|-----------|--------|
| 1 | Initial alignment | Projection layers only | 6 |
| 2 | Partial unfreeze | Top encoder layers + projections | 3 |
| 4 | Full image-text | All image/text parameters | 3 |
| 5 | Audio alignment | Audio encoder (text/image frozen) | 5 |

### Training Configuration

- **Hardware**: 8× AMD MI300X GPUs
- **Precision**: BF16 mixed precision
- **Batch Size**: 64 per GPU (512 effective)
- **Optimizer**: AdamW
- **Learning Rate**: 1e-4 → 5e-5 (stage dependent)
- **Loss**: InfoNCE contrastive + Matryoshka loss

## Evaluation

### Image-Text Retrieval (COCO Validation)

| Metric | Image→Text | Text→Image |
|--------|------------|------------|
| R@1    | 41.88%     | 39.21%     |
| R@5    | 71.64%     | 69.15%     |
| R@10   | 82.16%     | 80.02%     |

### Audio-Text Retrieval (LibriSpeech)

| Metric | Audio→Text |
|--------|------------|
| R@1    | 36.38%     |
| R@5    | 68.22%     |
| R@10   | 79.52%     |

### MRL Quality Retention

| Dimension | Compression | Quality |
|-----------|-------------|---------|
| 384 (full)| 1×          | 100%    |
| 256       | 1.5×        | ~98%    |
| 128       | 3×          | ~95%    |
| 64        | 6×          | ~90%    |

## Limitations

- English language only
- Image resolution fixed at 512×512
- Audio limited to 30 seconds
- Best for retrieval/similarity, not generation

## Citation

```bibtex
@misc{multi-modal-embed-small-2026,
  title={multi-modal-embed-small: Compact Multimodal Embeddings with 2DMSE},
  author={Semantic Router Team},
  year={2026},
  url={https://huggingface.co/llm-semantic-router/multi-modal-embed-small}
}
```

## License

Apache 2.0