A local, from-scratch-trained multimodal embedding model covering all four modalities, text, image, audio, and video, in one shared 768-dim embedding space. Trained end-to-end on a single RTX 4060 laptop GPU (8GB VRAM).
Reproduces the architecture behind jina-embeddings-v5-omni ("GELATO": Geometry-preserving Embeddings via Locked Aligned TOwers, arXiv 2605.08384), built on public source encoders instead of Jina's exact adapted checkpoints. Full credit for the architecture and training recipe goes to Jina AI's GELATO paper. QuadEmbed is an independent, from-scratch reproduction of it, not a copy of Jina's released weights, and is not affiliated with or endorsed by Jina AI.
| Model weights | huggingface.co/Mithil-AI/quadembed-nano |
| Python package | pypi.org/project/quadembed — pip install quadembed |
| Write-up | I Built a Multimodal Embedding Model From Scratch on an RTX 4060 |
| Experiment log | EXPERIMENTS.md — seven rounds, what moved the needle and what didn't |
pip install quadembed # text + image + audio
pip install quadembed[video] # adds video supportfrom quadembed import QuadEmbed
from PIL import Image
model = QuadEmbed.from_pretrained() # pulls weights from the Hub
text = model.embed_text(["a dog running on the beach"])
image = model.embed_image([Image.open("photo.jpg").convert("RGB")])
similarity = text @ image.T # L2-normalized, so this is cosineThe rest of this README covers reproducing the training from source. The
packaged library lives in package/.
Cross-modal retrieval recall@k on held-out splits, text-query direction:
| Modality | R@1 | R@5 | R@10 | n | Peak VRAM |
|---|---|---|---|---|---|
| Image | 13.7% | 68.6% | 81.1% | 1024 | 1.56 GB |
| Audio | 67% | 97% | 100% | 33 | 2.69 GB |
| Video (held-out) | 40% | 86% | 94% | 50 | 0.89 GB |
Random-chance R@1 on the 1024-candidate image eval is ~0.1%. Every training run stayed under 2.7GB peak VRAM, well inside the 8GB budget.
Three frozen encoders, two small trainable projectors:
| Role | Model (frozen) | Output dim |
|---|---|---|
| text | jinaai/jina-embeddings-v5-text-nano (239M) |
768 (joint space, the anchor) |
| vision | google/siglip2-base-patch16-naflex |
768/patch, merged to 3072 |
| audio | openai/whisper-large-v3 encoder only |
1280/frame |
Trainable: VisionProjector.fc_vision_2 (3072→768, 2.36M params) and
AudioProjector.fc_audio (1280→768, 0.98M params), plus small per-modality
delimiter vectors. Everything else stays frozen and untouched, matching
GELATO's own ablation showing this outperforms unfreezing.
Video gets no separate encoder: quadembed/video.py samples frames from
an mp4 and runs each through the same trained vision projector, then
mean-pools over time.
Why these substitutions: the paper's vision/audio encoders (Qwen3.5's vision tower, Qwen2.5-Omni's audio tower) aren't distributed as standalone checkpoints and are described in the paper as adapted from SigLIP2 and Whisper-large-v3 respectively, so this reproduction uses those source models directly. The text encoder is an exact match.
All public, all pulled from the Hugging Face Hub at training time. No proprietary or manually collected data is used anywhere in this project.
| Dataset | Modality | Used for | Size used here | License |
|---|---|---|---|---|
jxie/flickr8k |
image-text | vision training + held-out eval | ~30,000 pairs (6k images × 5 captions) | CC BY 4.0 |
jxie/coco_captions |
image-text | vision training | ~124,560 pairs (40 of 182 shards) | CC BY 4.0 (COCO) |
google-research-datasets/conceptual_captions |
image-text | vision training (domain diversity) | 17,291 images | Google CC BY (see note) |
OpenSound/AudioCaps |
audio-text | audio training + eval | ~600 clip-caption pairs | MIT (AudioCaps: CC BY 4.0) |
VLM2Vec/MSR-VTT |
video-text | video training + held-out eval | 60 train / 50 test clips | research use |
Total vision training set at the final round: ~172,000 image-caption pairs across three datasets and two visual domains (curated photography from Flickr8k/COCO, plus web imagery from Conceptual Captions).
Notes on specific datasets:
- Conceptual Captions (CC3M) ships only
(image_url, caption)pairs, not image bytes.scripts/fetch_cc3m.pydownloads the actual images concurrently to local disk (32 workers, 6s timeout, skip-on-failure) so training never depends on live network calls. Expect ~58% success; dead links and expired stock-photo hosts account for the rest. It also filters degenerate images (min(w,h) < 32or aspect ratio beyond 6:1) — two one-pixel-wide "images" produced NaN in SigLIP2's patch-grid math and silently poisoned an entire training run before this guard existed. - AudioCaps is fetched by parquet shard path rather than streaming;
streaming the same data measured ~200x slower due to per-row-group HTTP
range requests against these audio-heavy files. Audio is decoded manually
via
soundfilerather than the defaulttorchcodecpath, which fails to load on Windows/CUDA (WinError 127). - COCO Captions is fetched in chunks of 4 shards and concatenated; a single bulk fetch stalled around 6 shards.
- Held-out evaluation always uses the Flickr8k test split for vision, regardless of what was added to training, so numbers stay comparable across all seven experiment rounds.
src/
quadembed/ importable library
encoders.py frozen text/vision/audio encoder wrappers
projectors.py the two trainable projectors + video pooling
losses.py InfoNCE + Matryoshka contrastive loss
data.py dataset loaders for all five datasets above
video.py frame sampling + video embedding
scripts/ command-line entrypoints
train.py train one modality's projector
eval.py recall@k for vision/audio
eval_video.py recall@k for video
fetch_cc3m.py one-time CC3M image downloader
smoke_test.py verify encoders load and shapes line up
hf_repo/ staging folder for the Hugging Face release
EXPERIMENTS.md full seven-round experiment log
python -m venv .venv
.venv\Scripts\Activate.ps1
pip install -r requirements.txt
cd srcVerify the frozen encoders load and tensor shapes line up before training:
python -m scripts.smoke_testSequential, one modality at a time (per GELATO's constrained-hardware
guidance: only one frozen encoder needs to be resident at a time). Run from
the src/ directory:
# vision, final configuration (~172k pairs)
python -m scripts.train --modality vision --steps 9000 --batch_size 128 --num_workers 4 --coco_shards 40 --cc3m_dir ../data_cache/cc3m
# audio
python -m scripts.train --modality audio --steps 500 --batch_size 16 --num_workers 2 --num_shards 6
# video (continues the vision projector on real sampled video frames)
python -m scripts.train --modality video --steps 300 --num_videos 60 --frames_per_video 3To use Conceptual Captions, download the images first:
python -m scripts.fetch_cc3m --num_urls 30000 --workers 32 --out_dir ../data_cache/cc3m--num_workers 4 matters for vision: SigLIP2's NaFlex preprocessing is
CPU-bound enough that num_workers=0 drops GPU utilization to near zero.
Continue any run from an existing checkpoint with --init_checkpoint.
python -m scripts.eval --modality vision --checkpoint ../checkpoints/vision_projector_FINAL_round7_best.pt
python -m scripts.eval --modality audio --checkpoint ../checkpoints/audio_projector_step500.pt
python -m scripts.eval_video --checkpoint ../checkpoints/video_projector_step300.pt --num_videos 50Reports text→modality and modality→text recall@{1,5,10} on a held-out split.
The short version: architecture was never the bottleneck, data volume was. Rounds 3 through 6 fixed a genuine spatial-merge bug, doubled batch size, added a genuinely different visual domain, and increased projector capacity. None of them moved R@1 off ~13%. Round 7 scaled the training set 3.2x and moved every metric at once.
Full round-by-round detail, including the negative results and the NaN incident, is in EXPERIMENTS.md.
- Batch size 16-128 vs. the paper's 256; in-batch negative count matters a lot for InfoNCE quality.
- ~172k image-caption pairs vs. the paper's enterprise-scale, multi-domain corpus. Confirmed empirically to be the actual bottleneck.
- No task-specific LoRA adapters (retrieval/classification/clustering) layered on top; base cross-modal alignment only.
Code in this repo is available for research and educational use. Note that
the frozen text encoder (jina-embeddings-v5-text-nano) is CC-BY-NC-4.0
(non-commercial), which carries through to any model trained on top of it.
SigLIP2 and Whisper-large-v3 are Apache-2.0. Check each dataset's own license
before redistribution.