llm-semantic-router / multi-modal-embed-small

multi-modal-embed-small

A compact multimodal embedding model that unifies text and image representations in a shared semantic space. Part of the MoM (Mixture of Models) family powering vLLM Semantic Router.

Model Description

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

• Text encoding via MiniLM-L6-v2 backbone
• Image encoding via SigLIP-base-patch16-512
• Cross-modal fusion via transformer attention
• 2DMSE : Two-Dimensional Matryoshka Sentence Embeddings for adaptive compute
• MRL : Matryoshka Representation Learning for flexible embedding dimensions

Key Features

Usage

Installation

pip install torch transformers pillow safetensors

Basic Usage

import torch
from PIL import Image
import requests
from io import BytesIO

# Load model
from transformers import AutoModel, AutoProcessor

# Or load from local checkpoint
import sys
sys.path.append("path/to/2DMSE-Multimodal-Embedder")
from src.models import MultimodalEmbedder

model = MultimodalEmbedder(
    text_encoder_name="sentence-transformers/all-MiniLM-L6-v2",
    image_encoder_name="google/siglip-base-patch16-512",
    output_dim=384,
    use_mobile_optimizations=True,
)
model.load_state_dict(torch.load("model.pt", map_location="cpu"))
model.eval()

Text Embedding

# Single text
text = "A photo of a cat sitting on a couch"
text_embedding = model.encode_text(text)  # 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
import torch.nn.functional as F
similarities = F.cosine_similarity(
    text_embeddings[0:1], 
    text_embeddings[1:], 
    dim=-1
)
print(f"Cat vs Dog similarity: {similarities[0]:.3f}")
print(f"Cat vs Car similarity: {similarities[1]:.3f}")

Image Embedding

from PIL import Image
import requests
from io import BytesIO

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

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

# Or from file
image = Image.open("my_image.jpg").convert('RGB')
image_embedding = model.encode_image(image)

Cross-Modal Retrieval

# 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",
    "A fluffy feline resting",
]
text_embs = model.encode_text(captions)

# Find most similar caption
similarities = F.cosine_similarity(image_emb, text_embs)
best_match_idx = similarities.argmax().item()
print(f"Best match: {captions[best_match_idx]}")
print(f"Similarity: {similarities[best_match_idx]:.3f}")

Matryoshka Dimension Reduction (MRL)

# Get full 384-dim embedding
full_emb = model.encode_text("Hello world")  # [1, 384]

# Truncate to smaller dimensions (MRL)
emb_256 = full_emb[:, :256]  # 256-dim, ~1.5x faster retrieval
emb_128 = full_emb[:, :128]  # 128-dim, ~3x faster retrieval
emb_64 = full_emb[:, :64]    # 64-dim, ~6x faster retrieval

# Normalize after truncation
emb_128_norm = F.normalize(emb_128, p=2, dim=-1)

2DMSE Adaptive Layer Exit

# Full model (all layers) - highest quality
full_emb = model.encode_text("Complex query", target_layer=None)

# Early exit at layer 3 (~50% compute) - faster
early_emb = model.encode_text("Simple query", target_layer=3)

# Even earlier exit (layer 1) - fastest
fastest_emb = model.encode_text("Quick lookup", target_layer=1)

Multimodal Fusion

# Combine text and image for richer representation
image = Image.open("cat.jpg").convert('RGB')
text = "A cute pet"

fused_embedding = model.encode_multimodal(
    texts=text,
    images=image
)  # Shape: [1, 384]

Training

Architecture

┌─────────────────────────────────────────────────────────────┐
│                   multi-modal-embed-small                        │
├─────────────────────────────────────────────────────────────┤
│  Text Encoder: MiniLM-L6-v2 (22M params)                   │
│  Image Encoder: SigLIP-base-patch16-512 (86M params)       │
│  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 Data

• LLaVA-CC3M : 595K image-caption pairs
• COCO Captions : Validation on 25K pairs

Training Configuration

• Hardware : 8x AMD MI300X GPUs
• Precision : BF16 mixed precision
• Batch Size : 256 per GPU (2048 effective)
• Optimizer : AdamW
• Learning Rate : 1e-4 with cosine decay
• Loss : InfoNCE contrastive + Matryoshka loss

Training Stages

1. Stage 1 (Frozen encoders): Align image-text space, 6 epochs
2. Stage 2 (Partial unfreeze): Fine-tune fusion + top encoder layers
3. Stage 4 (Full unfreeze): End-to-end fine-tuning

Evaluation

Image-Text Retrieval (COCO Validation)

Text Semantic Similarity

Cross-Modal Retrieval (Real-world test)

MRL Quality Retention (Matryoshka)

Limitations

• Optimized for English; multilingual performance may vary
• Image resolution fixed at 512x512
• Audio encoder included but not yet trained (see Roadmap)
• Best for semantic similarity, not generative tasks

Roadmap

Audio Modality Training (Planned)

The model architecture includes a Whisper audio encoder, but this release only trained on image-text data. Future releases will add audio-text alignment using:

This will enable:

• Audio-to-text retrieval
• Text-to-audio retrieval
• Audio-image-text multimodal fusion

Citation

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

License

Apache 2.0

Related Models

• mmbert-embed-32k-2d-matryoshka - Long context variant
• mmbert-embed-finance - Finance domain
• mmbert-embed-medical - Medical domain