Omni-Embed-Mini

What this is

Omni-Embed-Mini recasts cross-modal alignment as self-distillation through one shared frozen causal backbone. Every media sample is paired with a dense cascaded caption; the teacher target is the EOS-pooled embedding of that caption produced by the identical frozen backbone that processes the student input. Teacher and student therefore inhabit byte-identical geometry, and because no text-side parameter is ever updated, adding audio cannot degrade the inherited text and vision representations.

Only the audio projectors and small phased LoRA adapters on the audio encoders are trained. A Matryoshka SigLIP contrastive objective and an online hybrid hard-negative miner supply the contrastive signal. This variant keeps the Qwen3-VL native visual tower, so image, video, and document-page quality is inherited intact from the backbone while speech and audio are added on top.

Results

Six modalities, evaluated with the pipeline in the code repository.

v1 and v2 are tagged revisions of this same repository, so revision="v1.0" pins the numbers in the v1 column.

Quick start

pip install "transformers>=5.0" torch pillow numpy soundfile huggingface_hub

import torch
from transformers import AutoModel, AutoProcessor

REPO = "MBZUAI/Omni-Embed-Mini-2.3B"

model = AutoModel.from_pretrained(
    REPO, revision="v1.0", trust_remote_code=True, dtype=torch.bfloat16,
).cuda().eval()
processor = AutoProcessor.from_pretrained(REPO, revision="v1.0", trust_remote_code=True)

# OpenAI-style multimodal messages: text, audio, image, video, doc, or any composition.
messages = [{"role": "user", "content": [
    {"type": "audio", "audio": "path/to/clip.wav"},
    {"type": "text",  "text":  "rain on a tin roof at night"},
]}]

inputs = processor.apply_chat_template(
    messages, role="passage", tokenize=True, return_tensors="pt",
).to("cuda")
with torch.no_grad():
    doc = model(**inputs).pooler_output          # (1, 2048), already L2-normalised

# Query side. role="query" applies the retrieval instruction template.
# text_recipe="chat" puts a text query in the same subspace as media documents;
# omit it only for pure text-to-text retrieval, which uses the native recipe.
query = processor.apply_chat_template(
    [{"role": "user", "content": "rain on a tin roof at night"}],
    role="query", tokenize=True, return_tensors="pt", text_recipe="chat",
).to("cuda")
with torch.no_grad():
    q = model(**query).pooler_output

print("cosine:", float(q @ doc.T))               # normalised, so dot == cosine

Matryoshka (truncatable) embeddings

model(**inputs, truncate_dim=N).pooler_output returns a prefix-truncated, renormalised vector. Slicing pooler_output yourself is not equivalent unless you renormalise after. Supported N: 128, 256, 512, 1024, 2048. Use a smaller N to cut index size at a modest recall cost.

Architecture

Trained parameters: audio projectors plus phase-2 LoRA adapters on the audio encoders (48 Whisper and 24 Dasheng adapted tensors, merged into the released weights). The backbone is frozen at every stage and ships unmodified apart from a vocabulary resize that adds the media placeholder tokens. There is no vision_encoder.pt, which is expected for this variant, whose vision weights live inside backbone/.

Files

config.json                 modeling_omni_embed.py      processing_omni_embed.py
configuration_omni_embed.py chat_template.jinja         processor_config.json
tokenizer.json              tokenizer_config.json
backbone/                   # frozen Qwen3-VL backbone (vocab-resized), safetensors
whisper_encoder.pt          dasheng_encoder.pt
projector_weights.pt        # projectors; LoRA already merged into the encoders

The .pt files are PyTorch pickles, and the custom modeling code requires trust_remote_code=True. Load only from a source you trust.

Citation

@inproceedings{omniembedmini2026,
  title     = {Omni-Embed-Mini: Binding Modalities Without Forgetting via Dense Distillation},
  author    = {TBD},
  booktitle = {TBD},
  year      = {2026},
  note      = {Camera-ready in preparation}
}